Conflict-Averse Gradient Descent for Multi-task Learning

In plain words

Training one network on several goals at once is common: for example, label every pixel of a photo and also estimate how far each pixel is from the camera. The simplest approach averages the learning signals from all goals and takes one step. That step can help one goal while quietly hurting another, and the damage is often worst for the goal that is hardest to learn.

The method described here, CAGrad, keeps the average as its starting point but does not follow it blindly. It looks for a nearby step that makes the currently worst-off goal improve as much as possible, and it limits how far that step may stray from the average. A single knob sets that limit. At zero the method is plain averaging. Larger values let the step lean more toward protecting the weakest goal. The authors show that, with a suitable knob setting, the method still settles at a good point for the average goal.

The paper reports better results than several earlier ways of combining these signals, on supervised vision, robot-control and semi-supervised benchmarks. These are the authors' own numbers on their own setups.

TL;DR

• Changes how the gradients of several losses are combined into one update. It is a drop-in replacement for averaging §3.1.
• It is plain gradient descent when c = 0, and approximates the MGDA direction for large c §3.1.
• For c < 1 the paper proves convergence to a stationary point of the average loss §3.2. The experiments do not use the theorem's plain step; see Pitfalls.
• Costs one backward pass per task each step, plus a small solve over the task weights §3.3.
• Official MIT-licensed code exists and implements the same rule in three variants cagrad.py.

Problem

With several losses L_1 ... L_K on shared parameters θ, the usual goal is the average loss L_0 = (1/K) Σ_i L_i §2.1 eq. 1. Following its gradient can still hurt individual tasks, because the task gradients may point in different directions or have very different sizes §1 §2.1. Earlier remedies either re-weight the losses by heuristics, which gives no convergence guarantee, or change the update direction, which may converge to any point on the Pareto set and so ignores L_0 §2.2. The paper wants a method that keeps the average loss as the target and still protects each task §1 §3.1.

Method

Setup. Let g_i = ∇L_i(θ) and g_0 = ∇L_0(θ) = (1/K) Σ_i g_i §2.1. For a step θ' = θ − α·d, the paper measures conflict with the worst local improvement across tasks, R(θ, d), which is approximately −min_i ⟨g_i, d⟩ for small α §3.1 eq. 2. A direction d that keeps every task decreasing has R < 0 §3.1.

CAGrad direction. CAGrad chooses d to maximise the worst per-task decrease, subject to staying within a ball of radius c‖g_0‖ around g_0 §3.1 eq. 3:

\max_{d}\; \min_{i \in [K]} \langle g_i, d \rangle \quad \text{s.t.} \quad \lVert d - g_0 \rVert \le c\,\lVert g_0 \rVert eq. 3

Here c ∈ [0, 1) is a fixed hyperparameter that controls the convergence rate §3.1. The paper states that the ball is centred on the average gradient because the average loss is the chosen main objective §3.1.

Dual form used in practice. The primal in eq. 3 has one variable per parameter, which is too large for a network. Let φ = c²‖g_0‖² and let g_w = Σ_i w_i g_i, where w lies on the probability simplex W. The paper derives a problem in the K weights only §3.1:

w^* = \arg\min_{w \in \mathcal{W}} \; g_w^\top g_0 + \sqrt{\phi}\,\lVert g_w \rVert

The optimal direction is then d* = g_0 + (√φ / ‖g_{w*}‖) · g_{w*}, equivalent to d* = g_0 + g_{w*}/λ* with λ* = ‖g_{w*}‖ / √φ §3.1. The derivation is in Lagrangian form and uses strong duality §3.1 app. A.3. Because √φ = c‖g_0‖, the update adds a term of length c‖g_0‖ pointing along the normalised blended gradient.

Algorithm. The paper's CAGrad loop is shown in alg. 1. The pseudocode below is a restatement of it, with the citations kept.

# One CAGrad update for K task losses  [alg. 1]
# Input:  parameters theta, K task losses, c in [0, 1), step size alpha
# Output: updated parameters

def cagrad_step(theta, task_losses, c, alpha):
    # Gradient of each task loss at theta  [alg. 1]
    grads = [gradient(loss, theta) for loss in task_losses]
    # Average gradient, the gradient of the main objective  [§2.1] [alg. 1]
    g0 = sum(grads) / len(grads)
    # Radius of the search ball, sqrt(phi) = c * ||g0||  [§3.1] [alg. 1]
    radius = c * norm(g0)
    # Task weights on the simplex minimising g_w.g0 + radius*||g_w||  [§3.1]
    w = solve_simplex_problem(grads, g0, radius)
    # Blended gradient g_w = sum_i w_i * g_i  [§3.1] [alg. 1]
    g_w = sum(w_i * g_i for w_i, g_i in zip(w, grads))
    # Update direction d = g0 + radius * g_w / ||g_w||  [§3.1] [alg. 1]
    d = g0 + radius * g_w / norm(g_w)
    # Plain step with size alpha, as written in the algorithm  [alg. 1]
    return theta - alpha * d

Special cases. With c = 0 the radius is zero and the rule reduces to d = g_0, plain gradient descent on L_0 §3.1 alg. 1. As c grows without bound, minimising the blended-gradient norm becomes the MGDA rule, which the paper calls its generalisation of MGDA. It notes that MGDA can settle at any Pareto point, whereas CAGrad aims at the average loss §3.1 app. A.1.

Practical speedup for many tasks. Each step needs one backward pass per task. The paper therefore proposes sampling a subset S of tasks, computing their gradients and the average g_0, and solving the same problem with the unsampled part replaced by its average §3.3 eq. 4:

\max_{d} \; \min \Big( \Big\langle \frac{K g_0 - \sum_{i \in S} g_i}{K - |S|},\, d \Big\rangle,\; \min_{i \in S} \langle g_i, d \rangle \Big) \quad \text{s.t.} \quad \lVert d - g_0 \rVert \le c\,\lVert g_0 \rVert eq. 4

The constraint is unchanged, so the convergence statement still applies according to the paper §3.3 app. A.3. The stated cost is O((|S|N + T)), where N is one backpropagation pass and T the optimisation time. The paper says T is much smaller than N when K < 10 §3.3.

Convergence (theory). Assume each L_i is differentiable with H-Lipschitz gradient and that L_0 is bounded below §3.2. With a fixed step α ≤ 1/H and 0 ≤ c < 1, the paper's Theorem 1 bounds the accumulated squared gradient norm of L_0 over T steps by 2(L_0(θ_0) − L_0*) / (α(1 − c²)) §3.2. This means the iterates approach a stationary point of L_0 §3.2. For c ≥ 1 the theorem only states that fixed points are Pareto-stationary §3.2. The proof is in app. A.3.

Hyperparameters

Results

Author-reported throughout. Setups differ between experiments, so the numbers do not compare across tables.

Multi-task supervised learning, NYU-v2 (SegNet with MTAN attention, 3 tasks). Δm is the average per-task relative change from single-task training; lower is better. Values are means over 3 seeds tab. 1 §5.2.

The caption marks baseline entries that differ significantly from CAGrad at p = 0.1, using asterisks in the table tab. 1. Under the paper's text, CAGrad matches the other methods on segmentation and depth, and is better on surface normals, which the other methods leave behind §5.2.

Multi-task supervised learning, CityScapes (2 tasks). Values are means over 3 seeds tab. 2.

Appendix sweep over c on the same setup tab. 5 tab. 6 (Δm %):

Multi-task reinforcement learning, Meta-World (SAC; success rate; mean ± standard error over 10 runs) tab. 3 §5.3.

Per-step training time on the same benchmarks, in seconds, as reported tab. 7 app. B.3: PCGrad 9.7 (MT10) and 59.8 (MT50); CAGrad 10.3 and 27.8; CAGrad-Fast 4.8 and 11.4. The paper states the CAGrad-Fast speedups as about 2× (MT10) and 5× (MT50) over PCGrad §5.3.

Semi-supervised learning on CIFAR-10 (ARML auxiliary tasks; average test accuracy over 3 runs) tab. 4 §5.4.

Toy and MNIST-pair experiments §5.1. On a two-dimensional toy problem, the paper shows that plain GD gets stuck in 2 of 5 starting points, and that CAGrad with c = 0 recovers GD while c = 10 approximates MGDA §5.1 fig. 3. On Multi-Fashion+MNIST, the paper reports final training losses over three runs only in a figure, with CAGrad at c = 0.2 lowest in average loss §5.1 fig. 4. No numbers are given in text or table.

Paper vs code

Official code: https://github.com/Cranial-XIX/CAGrad at commit dc3d48152b6196945cfd56144879b9d42353b095, MIT licence (Copyright Bo Liu, 2022) LICENSE. The code repository covers the toy example, supervised NYU-v2 and CityScapes, and the Meta-World RL agent. It contains no code for the MNIST-pair or semi-supervised experiments, which the paper says use other codebases app. B.1 app. B.4.

1. Normalisation of the update. The paper's update has no extra scale alg. 1. Every code path divides the update by a constant:
   • supervised, both datasets: 1 + c² utils.py utils.py utils.py;
   • Meta-World agent: 1 + c² in the default path [code: mtrl/mtrl_files/cagrad.py#L254-L256]; 1 + c in the fast path cagrad.py; 1 + λ in the exact path cagrad.py;
   • toy example: 1 + c toy.py. The code's scaling is not described in the paper. At c = 0.4 the 1 + c² divisor shrinks the step by about 14%.
2. Optimiser. The paper's algorithm is a plain step θ − αd alg. 1. The code writes the CAGrad direction into the gradients, then calls an Adam optimiser in the supervised trainer utils.py model_segnet_mt.py, and the toy example uses Adam toy.py. The Meta-World agent writes the vector into .grad and leaves stepping to the configured SAC optimisers cagrad.py.
3. Sampled average for CAGrad-Fast. Eq. 4 averages the unsampled tasks over K − |S| tasks §3.3. The code sums the K − n tasks after the first n − 1 sampled ones and divides by K − n + 1. The sampled task at position n − 1 is therefore dropped, and the average is scaled [code: mtrl/mtrl_files/cagrad.py#L173-L180]. This is a deviation from the paper's formula.
4. Task-weight solver. The paper leaves the solver open §3.1. The Meta-World default uses softmax-parameterised weights, 20 SGD steps, and keeps the best iterate cagrad.py. The exact variant uses SLSQP cagrad.py. The supervised trainer uses bounded SLSQP utils.py. The CityScapes trainer uses a one-dimensional bounded search, because it has two tasks utils.py.
5. Toy example formula. The paper's f_2 contains −tanh(−θ_2) §5.1. The code adds the term: +torch.tanh(-x2) toy.py. The paper's figures may come from either version; the card cannot say which.
6. Toy starting points. The paper uses five starts §5.1 fig. 3. The code runs three: (-8.5, 7.5), (-8.5, -5), (9, 9) toy.py. The Figure 1 caption says three runs fig. 1. The starts differ from the paper's list in the second and third positions.
7. CityScapes c. The text says c = 0.2 §5.2. The appendix says c = 0.4 app. B.2, and so does the reported table row tab. 2 tab. 6. The code runs c = 0.4 and c = 0.6 run.sh. Nothing in the repository shows c = 0.2 for CityScapes.
8. Degenerate weight norm. The update divides by ‖g_w‖, which the paper does not address §3.1. The code adds small constants: 1e-8 in the supervised trainer utils.py, 1e-4 in the Meta-World agent cagrad.py, and 1e-4 in the toy example toy.py.
9. Default config. The Meta-World YAML sets cagrad_method: "exact" cagrad_state_sac.yaml. The agent's assertion accepts only cagrad, cagrad_exact and cagrad_fast<k> cagrad.py, so the default fails unless overridden. The run script overrides it run.sh.

Pitfalls

• The theorem does not cover the experiments. Theorem 1 is for the plain step θ − αd with α ≤ 1/H §3.2 alg. 1. The supervised, toy and RL runs feed the direction into Adam or SAC optimisers. The step-size condition and the convergence bound therefore do not carry over automatically.
• Two statements of the theorem disagree. The main text says fixed points are Pareto-stationary for c ≥ 1 §3.2. The appendix says any c ≥ 0 app. A.3. Use the appendix version only with its proof in hand.
• c is picked on training loss, but the appendix sweep favours other values. The supervised text selects c by average training performance app. B.2. In the appendix sweep for CityScapes, c = 0.8 gives a lower Δm than the c = 0.4 row reported in the main table tab. 6.
• Significance is weak. The test uses p = 0.1 and three seeds tab. 1 tab. 2. Several baseline differences are marked, but the margins in segmentation are small: CAGrad 75.16 against MTAN 75.18 mIoU on CityScapes tab. 2.
• Best-epoch reporting in RL. The RL setup reports the highest average test performance seen during training app. B.3. That is an optimistic estimate, and the paper does not report a held-out selection.
• Per-step cost. Each step needs K backward passes §3.3 [code: mtrl/mtrl_files/cagrad.py#L194-L205]. The task-weight solve adds up to 20 inner steps in the RL setup app. B.3. Timings in tab. 7 are per update step and the hardware is not stated in the table.
• Degenerate directions. If g_w = 0 the update is undefined. The paper says nothing about this case; the code adds small constants instead.
• Ambiguous c for CityScapes. Three sources give 0.2, 0.4, or both. Use 0.4 only with the table and code, and say so in any reproduction.
• Small, unseen components. The MNIST-pair setup and the semi-supervised setup have no code in this repository. Reproducing them relies on the appendix descriptions and on other codebases app. B.1 app. B.4.
• Unstated weights. The semi-supervised auxiliary weights w_1 and w_2 are "the same as ARML", and their values are not given in this paper §5.4.
• Toy optimiser and length. The paper does not state the toy optimiser or number of steps. Figure 1 says GD and the other methods run with Adam fig. 1; the code uses 100,000 Adam steps toy.py.
• Compute scaling. The paper says T ≪ N usually holds for K < 10 §3.3, which is the regime of its supervised experiments. The 10- and 50-task RL speedups rest on the timing table alone tab. 7, and the paper gives no scaling analysis beyond the O((|S|N + T)) expression §3.3.

Lineage

• Builds on: gradient descent, which it recovers at c = 0 §3.1; multiple gradient descent (MGDA), which it approximates as c grows §3.1 app. A.1.
• Compared against: MGDA, PCGrad, GradDrop, Cross-Stitch, MTAN, GradNorm, Soft Modularization, CARE, and multi-task SAC variants, as listed in tab. 1–tab. 4.
• Superseded by: not known from wiki content.