GradNorm: Loss Weighting for Multi-Task Learning (2D)

Loss Weighting Strategy

Task Setup

Task hardness spread 0.60

GradNorm asymmetry α 1.50

Per-Task Loss Weight w

Per-Task Gradient Norm |G|

Training step

Mean grad norm Ḡ

0.00

Drag inside the diagram to orbit the task heads around the trunk; scroll/pinch to zoom. The right-hand scope plots each task's loss on a log axis.

How it works

Three task heads share one trunk. Every step each task's loss shrinks toward zero at a rate set by its own difficulty and by how much weight the shared gradient gives it: L_i(t+1) = L_i(t)·exp(−w_i·k_i·η), and its gradient-norm contribution to the trunk is G_i(t) = w_i·k_i·L_i(t).

Equal:        w_i = 1  for every task

Uncertainty (Kendall et al. 2018):
  total = Σ_i L_i·exp(−2·s_i) + s_i     (s_i = log σ_i, learned)
  w_i   = exp(−2·s_i) = 1/σ_i²
  s_i  -= lr_s · (1 − 2·L_i·exp(−2·s_i))

GradNorm (Chen et al. 2018):
  r_i(t) = [L_i(t)/L_i(0)] / mean_j[L_j(t)/L_j(0)]      (relative pace)
  target_i = Ḡ(t) · r_i(t)^α
  w_i  += lr_w · sign(target_i − G_i)
  then renormalize:  w_i *= N / Σ_j w_j

Equal weighting — the fast, easy task dominates the shared trunk; a hard task can stall (its bar barely moves).
Uncertainty weighting — a task with high loss (= high learned "uncertainty" σ) is automatically down-weighted, a task that is confident and low-loss gets amplified.
GradNorm — directly targets equal training pace across tasks: whichever task is falling behind (large r_i) gets pushed harder, up to the α exponent's strength, then all weights are rescaled to keep the total gradient magnitude constant.
Particles flowing head → trunk are colored per task; their size and speed track that task's current gradient norm |G_i| — a visibly bigger, faster stream is pulling the shared weights harder.

Real-world relevance: this is exactly why a production multi-task model (e.g. one CNN doing detection + segmentation, or one BERT doing sentiment + NER) needs a weighting scheme at all — without one, whichever task has the largest raw loss silently starves the others of shared capacity.

drag = orbit · wheel = zoom