Dynamic Loss Weight

Dynamic weight that allows for gradient balancing in multi task models

Multi task models can share certain modules the whole architecture. Regardless of which modules are shared, separate loss functions are utilized for the separate tasks.

To simplify this, I will give an example.

Let's say we want to detect cars & motorcycles in traffic. There is a lot of information which can be interchanged in its feature exctraction. Tyres, headlights, etc. Sharing the backbone & encoder between both tasks would be logical.

The decoder will however be separate, for the sake of this example. The decoder will then have its separate loss per task. The loss function between both tasks can be drastically different, leading to losses which are very imbalanced. Where the car task loss might be 50x as large as that of the motorcycle loss, and vice versa.

Now how will the gradient update react to this?

The loss would look as follows:

And the gradient update as follows:

Now if is significantly larger than , the gradient update will be largely influenced toward the motocycle task if and are both the same value. We need to update the weights so that there is a gradient balance between both tasks.

We want want to ensure that our car_over_motorcycle_weighted is = 1

car_over_motorcycle_weighted : the ratio between the weighted loss of both tasks

With that goal in mind, let's get into the specifics on how we achieve that.

For our shared modules we want to measure our gradient magnitudes.

Now we take the L2 norm over all flattened per-parameter gradients:

Now we have the magniute of each gradient pushing in the shared module for both respective tasks.

From linearity (between the loss and gradients) the total shared gradient is:

So the balance we are trying to achieve i sbetween the two weighted magnitudes:

We want with target (each task pulls the shared module equally hard). We hold fixed and solve for the that would balance this step:

The relationship is inverse: if moto pushes harder (), then rises to give the car task more say, and vice versa.

We don't snap straight to our calculated however, as our per-step gradient normas are noisy and spiky. As the training data that is processed can vary significantly. So a raw target would make the weight jitter. Instead we ease toward the target with EMA (Exponential Moving Average), done in log space (weights are positive and multiplicaative, so log space keeps the steps symmetric).

For example: Let's say our current weight is . Two cases, both a factor of 2 off:

Both are equally off by a factor of 2, however in our linear space it says one is twice as far off as the other. So an averaging step would overreact to the doubling and underreact to the halving.

In log space those gaps are equal and opposite:

So our equation will look as follows:

this is simply a smooth multiplicative nudge with rate e.g. ():

And at last we want to clamp it so a bad spike can't send the weight to zero or infinity, and we can do that as follows:

we then feed back into our $L_{\text{total}} = w_{\text{car\_new}} \cdot L_{\text{car}} + w_{moto} \cdot L_{\text{moto}} $

Now our training will always seek to balance the gradient magnitudes between both tasks.