I have a model with a few outputs, each output with shape:
Shape: (batch_size, labels_1)->
Shape: (batch_size, labels_2)->
Shape: (batch_size, labels_3)->
Now this model is used for a classification problem, so I'm applying a softmax activation to each and using the categorical_crossentropy (or just "crossentropy") loss.
Basically, everything is working reasonably well, but I noticed the losses are unbalanced. (Not because of class unbalance, but because of a different number of neurons)
An output with more neurons tend to have greater losses because their initial predictions for the correct label (the only label considered by this loss) is closer to zero.
I could try to balance them with a simple
loss / number_of_neurons, but I'm not confident that this will indeed balance the losses (since softmax is not linear), it sounds like unbalancing in a different way.
Is there a mathematically correct way to balance these losses to compensate for their number of neurons?