What the loss, the optimizer, clipping and logging do to a balanced gradient.
In multitask training, each task has its own loss and gradient. A balancing method (a balancer) reweights the task losses or edits the task gradients so no task takes over the shared update. This page follows the balanced gradient through clipping, the optimizer and logging.
On this page
A balancer acts once per training step, on that step’s batch. Loss-weighting methods (GradNorm, Uncertainty Weighting) change the task losses before the backward pass; gradient methods (PCGrad, GradVac, CAGrad, MGDA) replace the summed gradient after it (Figure 1).
A balancer can do exactly what it promises and the step still come out unbalanced, because the stages around it shape its input and its output.
The shared layers, the trunk, are the part every task uses; each task also has its own head. K is the number of tasks.
Figure 1. Where each kind of method acts in one training step.
The 2D arrows on these pages are pictures. A real task gradient has one number per shared weight, all shared layers flattened into one vector, often millions long; lengths, dot products and projections work on it as on the arrows.
Each loss is a batch mean (PyTorch’s default reduction="mean"), so gradients do not grow with batch size. One torch.autograd.grad call per task over the shared parameters returns its gradient without touching .grad; retain_graph=True keeps the graph for the next task. Every step costs K backward passes and K copies of the shared gradient.
| Method | Gradients taken over | How the result reaches the weights |
|---|---|---|
| PCGrad | all shared parameters | written into the shared .grad |
| GradVac | shared parameters, per layer in the paper | written into the shared .grad |
| CAGrad | all shared parameters | written into the shared .grad |
| MGDA | all shared parameters, or only the trunk’s output (MGDA-UB) | official code re-weights the losses and backpropagates again |
| GradNorm | only the sizes, at the last shared layer | learned loss weights; the weighted loss is backpropagated |
| Uncertainty Weighting | none | learned loss weights; one ordinary backward |
Writing into the shared .grad leaves each head its own plain gradient, the PCGrad paper’s rule and how LibMTL, a widely used library of these methods, writes back PCGrad and GradVac. Backpropagating weighted losses scales the heads’ gradients too.
This step writes into .grad, with a plain sum as the combine; vector_to_parameters would overwrite the weights, not their gradients.
import torch
from torch.autograd import grad
import torch.nn.utils as nnu
clip = nnu.clip_grad_norm_
KEEP = dict(retain_graph=True)
MAX_NORM = 1.0 # example clip cap
def flat_grad(loss, body):
gs = grad(loss, body, **KEEP)
flat = [g.ravel() for g in gs]
return torch.cat(flat) # (P,)
def fill_grads(params, vec):
i = 0
for p in params: # unflatten
n = p.numel()
part = vec[i:i + n]
p.grad = part.view_as(p)
i += n
body = list(trunk.parameters())
head_ps = list(heads.parameters())
feats = trunk(x)
losses = []
for k in range(K): # batch means
pred = heads[k](feats)
L = crits[k](pred, ys[k])
losses.append(L)
rows = []
for L in losses: # K backwards
g_k = flat_grad(L, body)
rows.append(g_k)
G = torch.stack(rows) # (K, P)
g = G.sum(dim=0) # combine here
opt.zero_grad()
fill_grads(body, g)
total = sum(losses)
total.backward(inputs=head_ps)
params = body + head_ps
pre_clip = clip(params, MAX_NORM)
opt.step()flat_grad flattens one task’s gradient; fill_grads writes a flat vector back into each .grad.Loss weights w act before the backward passes and change a task gradient’s length, never its direction; gradient methods act after and can turn it. The question each one answers:
| Method | The question it answers |
|---|---|
| GradNorm | How much weight should each task's loss get so all tasks train at a similar rate? (learned during training from how fast each loss falls) |
| Uncertainty Weighting | How much weight should each task's loss get, given how noisy that task's targets are? (the noise level is learned along with the model) |
| PCGrad | How should a task gradient be edited when it conflicts with another task's gradient? |
| GradVac | How should each pair of task gradients be moved toward a target cosine, a running average of that pair's past cosines? |
| CAGrad | How can a step raise the smallest per-task improvement while staying near the average gradient? |
| MGDA | Which convex combination of task gradients lowers every task's loss at once, and when does no such combination exist? |
A loss's scale, from its reduction (mean or sum over the batch) or a constant factor, changes the gradient's length, never its direction, and length is what balancers compare.
These two loss functions (PyTorch calls them criteria) differ only in reduction; the sum’s gradient is n times longer for n elements.
crit_mean = nn.MSELoss() # mean
kw = dict(reduction="sum")
crit_sum = nn.MSELoss(**kw)Loss-weighting methods undo a constant scale once their weights adapt. Gradient methods feel it directly.
Figure 4. A larger task 1 loss turns PCGrad and MGDA but not GradNorm.
Plain SGD (stochastic gradient descent) steps against the gradient times the learning rate. Adam keeps two running averages for each coordinate of the gradient; each coordinate is one weight.
The ten-step average smooths noise but lags: after a 90 degree turn the momentum takes 7 steps to turn 45 degrees, while Adam's step can turn in 1, onto a coordinate whose gradient had been small.
At β1 = 0 the momentum keeps only this step; at 0.5 it spans about 2 steps, at 0.9 about 10 and at 0.99 about 100.
Figure 5. The noise-free line turns at the tick; a larger β1 delays the momentum. this frame is 19 steps later.
The per-coordinate division cancels some of the balancer's changes, delays others, and tilts every step.
A change in the task mix is not a common factor: Adam’s step turns with it, then turns back as each coordinate’s scale catches up. Below, task 1’s gradient triples, with fixed gradients (3, 1) and (−1, 2) summed and Adam at PyTorch’s defaults.
| Optimizer, time after the change | Turn of the step |
|---|---|
| SGD, at once and afterwards | about 24 degrees |
| Adam, 10 steps later | about 18 degrees |
| Adam, 500 steps later | about 4 degrees |
| Adam, 1000 steps later | about 2 degrees |
In a two-bowl example, two tasks share two weights; each loss is a bowl, lowest at (3, 0) and (0, 3), ten times steeper along one weight. PCGrad balances them; SGD and Adam run at learning rate 0.01 for 20,000 steps. The Pareto set holds the points where neither loss can improve without the other getting worse.
| SGD | Adam | |
|---|---|---|
| Angle of the step off the balanced direction (median, far from the Pareto set) | 0 degrees | about 8 degrees, almost all from the per-coordinate scale |
| Task losses at step 20,000 | 4.03 and 4.15 | 3.66 and 4.61 |
PCGrad's and MGDA's guarantees cover the direction they hand over, for one batch; momentum and the per-coordinate scale change the step after.
Momentum, here as in PyTorch's SGD (Optimizers from SGD to Adam), carries the past into every step.
Training still converges: for two smooth, strongly bowl-shaped losses with exact gradients and tuned step sizes, the PCGrad paper proves (Theorem 4, Appendix A.4) that PCGrad with heavy-ball momentum reaches either the optimum or a point where the two gradients point exactly opposite.
Near the Pareto front, the task losses of the Pareto set from Section 7.2 plotted against each other, MGDA’s combined gradient stops shrinking once it reaches the size of the batch-to-batch noise in the gradients, instead of going to zero. Below, the step size over the last 1000 of 6000 steps as a share of the first, in the two-bowl example with gradient noise (5 runs, learning rate 0.02).
| Optimizer | Step size near the front |
|---|---|
| SGD | 2 to 5 percent |
| Adam | 13 to 19 percent |
Watch the per-task losses to decide when to stop, not a gradient size.
The balancer's own learned numbers, GradNorm's task weights and Uncertainty Weighting's noise levels σ, meet the optimizer too.
The GradNorm paper renormalizes after every step “to decouple gradient normalization from the global learning rate”; without it the weights’ sum drifts and acts as a learning-rate change. The weights get their own optimizer opt_w, in the paper Adam at 0.025 against 2×10−5 for the network, about 1250 times larger. In the fragment, Lgrad trains only w; the last two lines rescale w to sum to K inside no_grad, so the rescale is not trained.
# GradNorm's own loss L_grad,
# backpropagated into w only
L_grad.backward(inputs=[w])
opt_w.step()
opt_w.zero_grad()
# rescale so the weights sum to K
with torch.no_grad():
w.mul_(K / w.sum())Weight decay shrinks the weights a little each step inside opt.step(), after the balancer wrote the gradient, so the balancer never sees it; a larger wd shrinks them faster. Adam adds wd times the weights to the gradient before m and v; AdamW shrinks the weights directly, by the learning rate times wd, outside m and v.
import torch.optim as optim
kw = dict(lr=lr, weight_decay=wd)
# coupled L2: wd * weights joins
# the gradient before m and v
opt = optim.Adam(params, **kw)
# decoupled: shrinks the weights
opt = optim.AdamW(params, **kw)An L2 penalty written into a task's loss does pass through the balancer. Decay on Uncertainty Weighting's log σ pulls each task's weight toward its value at σ = 1, which is 0.5 for regression and 1 for classification.
The Uncertainty Weighting paper (three street-scene tasks, SGD with momentum) saw all three uncertainties fall and lowered the learning rate on a schedule. Under Adam, the per-coordinate scale on a task's head largely absorbs a slow weight change.
A smaller τ clips more steps and shortens them more; a very large τ never clips. PyTorch has no default; it is set per model. The three blocks clip globally, per task and per value, reusing clip, nnu, G and params from Section 3; only the first scales the whole gradient by one factor.
# tau: the max norm
# global: direction kept
pre_clip = clip(params, tau)
# per task, before the combine
for k in range(K):
f = tau / G[k].norm()
G[k] *= f.clamp(max=1.0)
# per value: direction turns
nnu.clip_grad_value_(params, cv)Mixed precision runs most arithmetic in 16-bit floats; to stop small gradients rounding to zero, it multiplies every loss by one scale s. Every task carries the same s, so no balancer's direction changes.
When a scaled gradient overflows, the scaler skips the step; a balancer with its own state (GradNorm’s weights, GradVac’s targets) should update it only on steps not skipped. Any size compared with a fixed number, such as the max norm or a logged size, is first divided by s; unscale_ does that for .grad. Below is the Section 3 step with a GradScaler.
scaler = torch.amp.GradScaler()
rows = []
for L in losses: # s times big
big = scaler.scale(L)
g_k = flat_grad(big, body)
rows.append(g_k)
G = torch.stack(rows)
s = scaler.get_scale()
sizes = G.norm(dim=1) / s # true
opt.zero_grad()
fill_grads(body, G.sum(dim=0))
total = scaler.scale(sum(losses))
total.backward(inputs=head_ps)
scaler.unscale_(opt) # .grad / s
pre_clip = clip(params, MAX_NORM)
scaler.step(opt) # skips inf, NaN
scaler.update()
skipped = scaler.get_scale() < sskipped for steps the scaler dropped.A batch too big for memory is split into micro-batches, balanced in one of two orders: average the micro-batch gradients and balance once, or balance each and average. two_orders computes both; bal is any balancer mapping a (K, P) matrix of task gradients to one vector, such as pcgrad on the PCGrad page.
def two_orders(bal, parts):
# parts: (K, P) per part
stack = torch.stack(parts)
g_once = bal(stack.mean(0))
each = [bal(G) for G in parts]
g_each = torch.stack(each)
return g_once, g_each.mean(0)Balancing is not linear, so with noise the two orders give steps about 7 degrees apart for PCGrad and 19 for MGDA (median, two-bowl example, 4 micro-batches, gradient noise about as large as the gradient). Averaging first and balancing once gives the balancer the same gradients a full batch would.
The weighted total loss hides all of this, so this block logs its parts every step. G, s, pre_clip and skipped come from the mixed precision step; writer is a torch.utils.tensorboard..
add = writer.add_scalar
true_G = G / s # undo AMP scale
for k in range(K):
t = f"task{k}"
raw = raw_losses[k].item()
add(f"loss/{t}", raw, step)
size = true_G[k].norm().item()
add(f"size/{t}", size, step)
add(f"w/{t}", w[k], step)
g1, g2 = true_G[0], true_G[1]
den = g1.norm() * g2.norm()
cos = (g1 @ g2 / den).item()
add("cos/task0_task1", cos, step)
add("pre_clip", pre_clip, step)
add("amp_skipped", skipped, step)Compare each raw loss with that task trained alone, and log how much each number varies, not only its average.