← All tutorials

GradNorm (Gradient Normalization)

GradNorm learns one loss weight per task so that every task’s gradient has a similar size and every task learns at a similar pace.

On this page

The introduction lists three causes of negative transfer (training tasks together making one of them worse than training it alone) when training on the plain sum of task losses. GradNorm, introduced by Chen, Badrinarayanan, Lee and Rabinovich (2018), fixes the first two, scale and speed. At every training step it measures each task’s gradient norm (the gradient’s length) and how far each task has progressed. It then adjusts one loss weight per task so that each gradient norm moves toward a target, and the target is higher for a task that is behind. Between steps it keeps the K loss weights, their optimizer state and each task’s starting loss. Each step costs one extra gradient per task at the last shared layer, about 5 percent more training time in the paper.

1 Unequal gradient norms

Several tasks share one trunk, the part of the model every task uses; each task also has its own small head. Task losses often sit on very different scales, for example a depth loss near 800 and a classification loss near 2. In the plain sum the larger loss dominates. Even with similar losses, one task can have much larger gradients on the shared parameters, so the trunk moves mostly for that task. Hand-tuned loss weights also go out of date as training progresses.

A task dominates through the size of its gradient, so GradNorm measures and evens out the gradient norm, the gradient’s length. GradNorm is a loss-weighting method: it changes each task’s weight, never the gradients’ directions. Gradient methods such as PCGrad edit directions instead.

2 One loss weight per task

With K tasks, GradNorm gives each task i a loss weight w_i that multiplies its loss. Every step it adjusts the weights toward two goals. First, every task’s weighted gradient norm should be about the same. Second, a task whose loss is falling more slowly than the others should get a larger gradient norm until it catches up. One hyperparameter, α, sets how strongly the second goal acts.

3 The GradNorm update

Each step has four parts: measure each task’s gradient norm, set each task’s target norm, move the weights toward the targets, and rescale the weights so they sum to K. The weights start at 1. Sections 4 and 5 define each part, and Figure 3 then steps through one update.

4 The three measurements

The gradient norm says how strongly each task pushes on the shared weights. The loss ratio says how much of its starting loss each task still has. The relative inverse training rate compares that ratio with the other tasks. The norms are measured only at the last shared layer, whose weights are W; the paper chose it to keep the extra cost small.

4.1 Gradient norm

4.2 Loss ratio

Dividing each loss by its value at the start of training gives a number with no units: 1 at the start, smaller as the task learns. The paper calls it the inverse training rate, because a small value means the task learned fast. For a classification task whose starting loss depends strongly on the random initialization, the paper suggests the theoretical starting value log C for C classes instead.

4.3 Relative inverse training rate

Each task’s loss ratio is then divided by the average ratio over all tasks. A value above 1 means the task is behind the others; below 1 means it is ahead.

5 Target norm and gradient loss

The average fixes the scale problem. The factor fixes the speed problem.

GradNorm then defines a loss on the weights that measures how far each norm is from its target, and lowers it with an ordinary optimizer.

Two rules keep the weight update well behaved.

1. The target is held fixed. The target is treated as a constant when the weights are updated. Otherwise shrinking every weight would also shrink the target and lower the loss without balancing anything; the paper does this “to prevent loss weights from spuriously drifting towards zero”.
2. The weights are rescaled to sum to K after every update. The paper does this “to decouple gradient normalization from the global learning rate”.

6 Simulation with two tasks

The simulation is a schematic model, not a real network. Each step, each task’s loss shrinks by about 2 percent times its share of the total gradient norm.

7 Results on two toy systems

The paper built two small artificial systems, one with 2 tasks and one with 10. Each task’s target outputs are multiplied by a different scale σ, so its squared loss grows with σ². Here σ is only that loss scale, not the noise level of the same name on the Uncertainty Weighting page. The tasks differ only in loss scale, so evening out the gradient norms does all the work. The task-normalized loss divides each task’s loss by its starting loss and sums them, so large and small scales count equally.

Figure 5. The toy networks and their loss scales σ. From the paper’s Figure 2.

Figure 6. Small-σ tasks end with larger weights. From the paper’s Figure 2.

Figure 7. GradNorm ends below equal weights. From the paper’s Figure 2.

From the paper's Figure 2 (arXiv:1711.02257): “Gradient Normalization on a toy 2-task (top) and 10-task (bottom) system. …”

8 Implementation in PyTorch

GradNorm runs once every training step, on that step’s batch. Its state is the K weights, their optimizer and each task’s starting loss. Its extra cost is one gradient per task down to the last shared layer, plus a second-order term through that layer; the paper measured about 5 percent more training time on NYUv2. The weighted loss is backpropagated as usual, so the heads also receive gradients scaled by w_i (Section 3 of Gradient Balancing Inside the Training Loop compares this with methods that write into the shared gradients).

This step follows the paper’s Algorithm 1. W is the last shared layer’s weight tensor, params holds all network parameters, losses holds the K batch-mean losses and L0 the losses at step 0. GRAPH passes create_graph=True to grad, which makes each norm G differentiable with respect to w; RETAIN passes retain_graph=True, which keeps the graph for the network step; detach() holds the target T fixed.

import torch
from torch.autograd import grad
from torch.optim import Adam

ALPHA = 1.5  # NYUv2 in the paper
RETAIN = dict(retain_graph=True)
GRAPH = dict(create_graph=True)

w = torch.ones(K).requires_grad_()
opt_w = Adam([w], lr=0.025)
opt = Adam(params, lr=2e-5)

def gradnorm_step(losses, L0, W):
    L = torch.stack(losses)
    G = []
    for k in range(K):
        wL = w[k] * L[k]
        g_k, = grad(wL, W, **GRAPH)
        G.append(g_k.norm())
    G = torch.stack(G)
    ratio = L.detach() / L0
    r = ratio / ratio.mean()
    T = G.mean() * r ** ALPHA
    T = T.detach()  # fixed target
    L_grad = (G - T).abs().sum()
    gw, = grad(L_grad, w, **RETAIN)

    opt.zero_grad()
    total = (w.detach() * L).sum()
    total.backward()
    opt.step()
    w.grad = gw
    opt_w.step()
    with torch.no_grad():
        w.mul_(K / w.sum())

1. 1State kept across steps: the weights w start at 1, with their own Adam optimizer.
2. 2Gradient norm of each weighted loss at W, the eq card in Section 4.1. Keep W one shared layer; a larger W costs more, and the paper does not report its effect on accuracy.
3. 3Loss ratio, relative inverse training rate r and the fixed target T, the eq cards in Section 4.2, Section 4.3 and Section 5. detach() matters: a target that can move lets every weight shrink together.
4. 4L_grad and its gradient with respect to the weights only, part 3 in Figure 2.
5. 5The ordinary network step on the weighted sum; heads and trunk both get it.
6. 6Weight step, then rescaling so the weights sum to K, part 4 in Figure 2. Rescale after every weight step, or the weights change the effective learning rate. A large weight learning rate can push a weight below zero; the simulation in Figure 4 keeps them above 0.02.

The paper does not say how raising or lowering the weight learning rate changes the result. There is no official code. LibMTL ships a version (--weighting GradNorm) that differs from the paper: by default it measures norms over the whole shared encoder instead of the last shared layer, it keeps the weights summing to K through a softmax instead of rescaling, and it trains the heads on the unweighted losses.

Gradient Balancing Inside the Training Loop covers how the weights interact with the optimizer in Sections 9.1 and 7.1.

9 When to use GradNorm

The hyperparameter α sets how strongly a task that is behind is pushed to catch up. Raising it pushes the weights further apart and cuts fast-learning tasks harder; lowering it toward 0 leaves only the equal-norm goal. The more the tasks differ, the larger the best α.

Reach for it when tasks have very different loss scales or learn at very different speeds. The classic case is dense prediction, one prediction per pixel, such as segmentation, depth and surface normals in one network.

It costs one extra gradient per task down to the last shared layer, about 5 percent more training time on NYUv2.

It does not come with a convergence proof, and it ignores gradient direction. Two gradients with equal norms can still point against each other.

GradNorm sets the weights from training speed. Uncertainty Weighting sets them from how noisy each task is.

Two methods learn each task’s loss weight (GradNorm and Uncertainty Weighting). Four methods edit the task gradients (PCGrad, GradVac, CAGrad and MGDA). The question each one answers:

10 References

1. Chen, Z., et al. (2018). GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks. ICML. arXiv:1711.02257.