DL0070 Multi-Task Loss Balancing

Start from the failure mode: with a naive sum , each task contributes gradient in proportion to its loss scale, so a depth-regression MSE sitting near 200 drowns out a segmentation cross-entropy sitting near 1, and the shared trunk optimizes the loud task while the quiet ones stall. The fix comes in three tiers. First, put the losses on a common scale by dividing each by a fixed or running estimate of its magnitude, such as its initial value. Second, learn the weights: Kendall et al.’s uncertainty weighting attaches a trainable per task and minimizes , derived from Gaussian and categorical likelihoods, so noisy tasks are down-weighted automatically while the term stops from blowing up. Third, when scales are balanced but gradient directions still conflict, move to gradient-level methods: GradNorm tunes the task weights so each task’s gradient norm on the shared trunk approaches a common scale adjusted by its relative inverse training rate, so tasks that are learning slowly get pushed harder, and PCGrad projects away mutually conflicting components.

(1) Scale Equals Loudness: task ‘s gradient share scales with the units of , so the unweighted sum is an implicit weighting set by arbitrary unit choices; expressing depth in millimeters instead of meters multiplies its MSE by and hands it the entire gradient budget.
(2) Uncertainty Weighting: treat as task ‘s observation noise; the weight falls as noise grows while the penalty (the likelihood’s normalizing constant) rises, so the optimum is a genuine trade-off learned by gradient descent alongside the network weights.
(3) Magnitude Is Not Direction: scale balancing fixes how loudly tasks speak, not whether they agree; when task gradients point in opposing directions (negative transfer), gradient-space methods such as GradNorm or PCGrad are the right lever.

Figure 1: Illustrative three-task trunk, gradient norms on a log scale: unweighted, the depth MSE contributes over 100x the others; once each loss carries a learned weight, the three contribute comparably.

The uncertainty objective is not ad hoc. Modeling regression noise as gives the negative log-likelihood per task, up to constants; classification slots in through a scaled softmax likelihood whose approximation yields the analogous weight (without the factor 2) alongside the same penalty. Two properties matter in practice: is learned by the same optimizer as the network, so balancing needs no manual grid search; and the term keeps the objective honest, because without it every would grow without bound, all task weights would collapse to zero, and nothing would be learned. Most implementations optimize and compute for numerical stability.

Figure 2: The learned trade-off: as a task’s noise grows, its weight collapses while the penalty rises, so the optimizer cannot silence a noisy task for free.

Mathematical Formulation:

Where:

• is the loss of task (cross-entropy for segmentation, MSE for depth) and the number of tasks sharing the trunk.
• is a learned per-task scalar modeling homoscedastic (task-level, input-independent) observation noise, initialized at 1 and trained by the same optimizer as the weights.
• is the derived task weight and the likelihood’s normalizing term that penalizes inflating ; in code, parameterize and optimize instead.