# Using a multi-headed neural network, how should I approach the regression head loss

I have a multi-headed NN where one head performs multi-label classification and the other a regression task on a set of images. The classification head outputs a one-hot vector where each value in the vector represents a probability for the given class being present in the image. The regression head should provide a prediction of the percentage of the image that is covered by each class class.

So a dummy example of the target classes and network predictions would look like:

classes = ['tree', 'car', 'ground', 'sky']
class_prob_preds = np.array([0.03, -0.17, 0.58, 0.73])
area_preds = np.array([0.05, 0.01, 0.42, 0.50]) # regression values don't need to add to 1
class_preds = np.where(class_prob_preds > 0.5) # gives the index of the predicted classes


My question is, how should I pass to the loss function for the regression part? Should I feed the entire vector of area predictions to the regression loss function, or should I use the predicted classes of the classification head to select only the values in the regression vector that correspond to the predicted classes and pass these to the loss function?