← All tutorials

PCGrad (Projecting Conflicting Gradients)

PCGrad removes the conflicting part of each task’s gradient before the gradients are summed, and leaves the rest unchanged.

PCGrad (projecting conflicting gradients), introduced by Yu, Kumar, Gupta, Levine, Hausman and Finn (2020), changes one line of a multitask training step: the sum of the task gradients. Two task gradients conflict when they point more than 90 degrees apart. For every conflicting pair, PCGrad removes from one gradient the part that points against the other, then sums the edited gradients. It runs every training step on that step’s batch and carries no state between steps. Its cost is one backward pass per task over the shared parameters, plus a dot product for each pair of tasks. It learns nothing and adds no hyperparameters. The weighting methods, GradNorm and Uncertainty Weighting, change each loss’s weight; PCGrad changes the gradients’ directions.

On this page

1 Removing the conflicting parts

When two task gradients conflict, adding them cancels part of each, and the sum can raise one task’s loss. PCGrad does not rescale the gradients. It removes only the part of each gradient that points against the other task’s gradient, and keeps every part that does not, so an edited gradient gets shorter only by the part removed, and the sum can come out longer than the plain sum. The paper also calls this gradient surgery.

2 The PCGrad algorithm

1. Compute each task’s gradient g_i over the shared parameters, one backward pass per task.
2. For each task i, visit every other task j in random order. If g_i conflicts with task j’s original gradient g_j, subtract from g_i its projection onto g_j (Section 3).
3. Sum the edited gradients and give the result to the optimizer as the shared parameters’ gradient.

3 Conflict test and projection

The test for a conflict is the dot product of two gradients, which equals their lengths times the cosine of the angle between them. Its sign is the sign of that cosine, called the cosine similarity: positive below 90 degrees, zero at 90 degrees, negative above. A negative dot product means a conflict. The part of g_i that points along g_j is its projection onto g_j; in a conflict that projection points against g_j, and PCGrad subtracts it. What remains is orthogonal to g_j (at 90 degrees). Figure 11 of Introduction to Multitask Learning lets you drag two gradients and watch the projection.

The core of PCGrad in PyTorch is two functions. project is the eq card in Section 3. pcgrad applies it to every pair in random order and sums the results; G holds one flattened task gradient per row.

import random
import torch

def project(g, g_j):
    dot = g @ g_j
    if dot < 0:  # conflict
        coef = dot / (g_j @ g_j)
        g = g - coef * g_j
    return g

def pcgrad(G):  # G: (K, P)
    K = len(G)
    out = []
    for i in range(K):
        g = G[i].clone()
        others = list(range(K))
        others.remove(i)
        random.shuffle(others)
        for j in others:
            g = project(g, G[j])
        out.append(g)
    return torch.stack(out).sum(0)

1. 1The conflict test and the projection, lines 2 and 1 of the eq card in Section 3.
2. 2Every other task in a new random order each call, step 2 in Section 2. The authors’ TensorFlow code calls a shuffle whose result is discarded, so its order is fixed in practice; LibMTL shuffles. With two tasks the order does not matter.
3. 3Tests and projects against each task’s original G[j], as the paper’s pseudocode does, then sums, step 3 in Section 2. A mean instead of the sum would divide the step by K.

4 Why this projection

For a small step, only the projection changes task j’s loss, so it is the smallest edit that removes the conflict.

1. Why edit both gradients? Editing only one leaves the other still conflicting with it. With two tasks, editing both leaves nothing in the sum that conflicts with either task. With three or more, a later projection can bring part of a conflict back.
2. Why a random order? With three or more tasks the result depends on the order of the projections. A new random order each step spreads that effect over the tasks.

5 Conflicting and non-conflicting pairs

The paper’s Figure 2. arXiv:2001.06782

Panels (a) and (b).

Panels (c) and (d). From the paper’s Figure 2, arXiv:2001.06782.

The original caption reads “… PCGrad projects task i's gradient onto the normal vector of task j's gradient, and vice versa. …”

Figure 3. (a) conflict; (b) g_i minus its part along g_j; (c) the reverse; (d) no conflict.

6 Both projections, step by step

7 The tragic triad

Joint training can do worse than training each task alone. The paper attributes this to conflicting gradients, and argues the damage is large when three conditions hold together, which it calls the tragic triad.

8 Implementation in PyTorch

PCGrad replaces the usual single loss.backward(). The step below differentiates each task’s batch-mean loss separately over the shared parameters, edits the results with pcgrad and project from Section 3, and writes the edited sum into the shared parameters’ gradients.

import torch
from torch.autograd import grad

RETAIN = dict(retain_graph=True)

def flat_grad(loss, ps):
    gs = grad(loss, ps, **RETAIN)
    flat = [g.ravel() for g in gs]
    return torch.cat(flat)

def fill_grads(params, vec):
    i = 0
    for p in params:
        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())

def pcgrad_step(losses):
    rows = []
    for L in losses:
        g_k = flat_grad(L, body)
        rows.append(g_k)
    g = pcgrad(torch.stack(rows))
    opt.zero_grad()
    fill_grads(body, g)
    total = sum(losses)
    total.backward(inputs=head_ps)
    opt.step()

1. 1flat_grad returns one task’s gradient over a parameter list ps as one flat vector, keeping the graph for the next task.
2. 2fill_grads writes a flat vector back into each parameter’s .grad, using numel and view_as.
3. 3K backward passes over the shared parameters body and K flattened gradient copies, step 1 in Section 2.
4. 4PCGrad replaces the plain sum, at K(K−1) dot products; the result becomes the shared gradient, step 3 in Section 2. Mixed precision multiplies every loss by one scale, so divide the gradients by it before projecting (Section 11 of Gradient Balancing Inside the Training Loop), and clip after fill_grads.
5. 5The heads get their ordinary gradients, since each head receives only its own task’s loss (Section 3 of Gradient Balancing Inside the Training Loop shows this write-back style), then one optimizer step for all parameters.

The official code is tianheyu927/PCGrad, a TensorFlow 1 optimizer wrapper used as PCGrad(tf.train.AdamOptimizer()) followed by optimizer.minimize(losses) with a list of task losses. PCGrad adds no hyperparameters; the paper keeps the base method’s settings. LibMTL ships a PyTorch version (--weighting PCGrad) that edits only the shared encoder’s gradients and shuffles the task order.

Loss scale and momentum still matter. A loss 10 times larger gives a gradient 10 times longer, and momentum can bring back a removed direction; see Sections 5 and 8.1 of Gradient Balancing Inside the Training Loop.

9 Comparison with other methods

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:

PCGrad leaves each edited gradient orthogonal to the one it conflicted with, a cosine similarity of zero. GradVac asks whether zero is the right target for every pair of tasks.

10 References

1. Yu, T., et al. (2020). Gradient Surgery for Multi-Task Learning. NeurIPS. arXiv:2001.06782.