(a) Negative cosine. PCGrad lefttop, GradVac rightbottom.
(b) Positive cosine below the target.
Figure 2. Paper’s Figure 5, arXiv:2010.05874; a2 is our a: “Comparing PCGrad (left) [top] with GradVac (right) [bottom] in two cases. …”
GradVac keeps each pair of task gradients about as aligned as that pair usually is.
On this page
GradVac (Gradient Vaccine; Wang et al., 2021) edits the task gradients before they are added into one update. When two tasks’ gradients agree less than usual, it tilts one toward the other. Agreement is the cosine similarity, the dot product divided by both lengths: 1 for the same direction, 0 at a right angle, −1 for opposite. Each task pair keeps a target cosine, the running average of its past cosines. Below the target, GradVac adds the coefficient a times one gradient to the other, with a chosen so the cosine equals the target.
GradVac runs every training step on that step’s minibatch and costs one backward pass per task, K for K tasks. Its state is one target per task pair per parameter group, a block of shared weights; one layer per group was best in the paper.
PCGrad (projecting conflicting gradients) edits a pair only when its cosine is negative, projecting one gradient so the cosine becomes 0: every pair has target 0. A pair that usually sits at 0.6 can drop to 0.2, and PCGrad does nothing. GradVac gives each pair its own target, so it acts here too; with every target at 0 it is exactly PCGrad. Like a vaccine, it acts before the cosine turns negative.
(a) Negative cosine. PCGrad lefttop, GradVac rightbottom.
(b) Positive cosine below the target.
Figure 2. Paper’s Figure 5, arXiv:2010.05874; a2 is our a: “Comparing PCGrad (left) [top] with GradVac (right) [bottom] in two cases. …”
The authors trained two translation models, one from every language into English and one from English into every language. Each language pair, such as English to French, is one task. Every 1000 steps they measured the cosine between every pair of task gradients, for analysis only.
| Observation | What the paper found |
|---|---|
| Related languages | Languages of one family have higher cosines and cluster. |
| Alignment and quality | Trained on English to French plus one other pair, a higher cosine went with better English to French. |
| Layers and training time | The cosine differs by layer and over training. |
| Sign of the cosine | Mostly positive, so PCGrad’s projection acted on few steps. |
Each pair has its own usual cosine; a fixed target of 0 ignores it.
For tasks i and j with gradients gi and gj (subscript ij marks the pair), GradVac tracks the pair’s usual cosine with an exponential moving average (EMA).
An illustrative pair of translation tasks, English to French and English to Hindi, at averaging rate 0.01. The walkthrough in Section 5 draws the same pair.
| value | |
|---|---|
| old target φij | 0.45 |
| this batch’s cosine | 0.30 |
| new target, 0.99 × 0.45 + 0.01 × 0.30 | 0.4485 |
Raising β follows recent steps faster but lets one noisy batch move the target; lowering it averages over about the last 1/β steps. At β = 0 the target stays at its start, 0, and GradVac is PCGrad; at β = 1 it is the last batch’s cosine.
The same illustrative pair, with gi = (2.2, 0.6), batch cosine 0.30 and target 0.45.
| value | |
|---|---|
| lengths of gi and gj | 2.280 and 1.8 |
| current and target angle | 72.5° and 63.3° |
| sine of their difference | 0.161 |
| coefficient a | 0.229 |
| cosine after the edit | 0.45 |
Figure 3 draws its own illustrative pair, with g1 as gi and g2 as gj.
Figure 3. A target of 0 gives PCGrad; one at or below the cosine gives no edit.
Figure 4. One pair, g1 as gi and g2 as gj, with only task 1 edited.
Translation is scored in BLEU, the overlap with reference translations; tagging in F1, which balances precision and recall.
| Benchmark | Compared with | Result |
|---|---|---|
| WMT translation (the Workshop on Machine Translation data), 4 languages: French and Czech with much data, Hindi and Turkish with little | Plain model, GradNorm, MGDA, PCGrad | GradVac best on average and on every language; no other method beat the plain model on every language. |
| In-house translation, 25 languages crawled from the web | Plain model, oversampling low-resource languages (the paper’s default) | Largest gains on the high-resource languages. |
| Same data | Plain model, sampling by data size | Gains at every resource level. |
| Multilingual BERT, named-entity tagging (marking names of people, places and organizations) per language | PCGrad | Average F1 83.8 against 83.1, not on every language. |
The Field Guide lists the pairwise cosines among the values to log.
Two listings, setup then method, make one GradVac step for a shared trunk with task heads. As in LibMTL, only the trunk is edited; the paper’s model shares every parameter.
import math
import random
import torch
from torch.autograd import grad
BETA = 0.01 # paper's best
EPS = 1e-8
KEEP = dict(retain_graph=True)
cos = torch.cosine_similarity
def flat_grad(loss, body):
gs = grad(loss, body, **KEEP)
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
kids = trunk.children()
layers = [m for m in kids if list(m.parameters())]
slices, i = [], 0
body = []
for m in layers:
ps = list(m.parameters())
n = sum(p.numel() for p in ps)
slices.append(slice(i, i + n))
i += n
body += ps
C = len(slices)
head_ps = list(heads.parameters())
phi_hat = torch.zeros(K, K, C)
EDIT = range(K)def coef(g_i, g_j, phi, tgt):
s = math.sqrt(1 - phi * phi)
t = math.sqrt(1 - tgt * tgt)
n_j = g_j.norm() + EPS
r = g_i.norm() / n_j
num = tgt * s - phi * t
return r * num / (t + EPS)
def edit(g_i, g_j, ijk):
phi = cos(g_i, g_j, 0).item()
tgt = phi_hat[ijk].item()
ema = (1 - BETA) * tgt
new = ema + BETA * phi
phi_hat[ijk] = new
if phi >= tgt:
return g_i
a = coef(g_i, g_j, phi, tgt)
return g_i + a * g_j
def edit_ij(out, G, i, j):
for k in range(C):
sl = slices[k]
g_i = out[i, sl]
g_j = G[j, sl]
ijk = (i, j, k)
new = edit(g_i, g_j, ijk)
out[i, sl] = new
def gradvac(G, edit_tasks):
out = G.clone()
for i in edit_tasks:
others = list(range(K))
others.remove(i)
random.shuffle(others)
for j in others:
edit_ij(out, G, i, j)
return out.sum(0)
def gradvac_step(losses):
rows = []
for L in losses:
g_k = flat_grad(L, body)
rows.append(g_k)
G = torch.stack(rows)
g = gradvac(G, EDIT)
opt.zero_grad()
fill_grads(body, g)
total = sum(losses)
total.backward(inputs=head_ps)
opt.step()out.| Setting | Paper value | Effect |
|---|---|---|
| Averaging rate β | 0.01, best of 0.1, 0.01, 0.001 | About 100 steps of memory; raising and lowering as in Section 3. |
| Parameter groups | One per layer, best of whole model, encoder and decoder, per layer, per weight matrix | Finer groups track each layer’s cosine; “neither too coarse nor too fine-grained perform the best”. |
| Tasks edited | Low-resource only, best of all, low-resource and high-resource | Leaves high-resource gradients untilted. EDIT in the listing. |
| Rescaling the edited gradient | None; “sufficient to ignore this step” | The listing does not rescale. |
| Target start | 0 | Targets need about 100 steps to catch up, so early steps act like PCGrad. |
No official code was found. LibMTL (github.com/median-research-group/LibMTL) defaults to β = 0.5 and the whole model as one group. Each step costs K backward passes, which dominate with many tasks, and K copies of the shared gradient; the K × K targets per layer are small.
EPS to both.GradVac replaces loss.backward().
Momentum and Adam’s per-coordinate scaling can still turn the step (Sections 6 and 7.2 of Gradient Balancing Inside the Training Loop).
Two methods learn each task’s loss weight (GradNorm and Uncertainty Weighting). Four methods edit the task gradients (PCGrad, GradVac, CAGrad and MGDA).
| 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 (this page) | 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? |
GradVac makes no promise about the smallest per-task improvement; CAGrad raises it.