Introduction to Multitask Learning

Training one model on several tasks, the three common ways it goes wrong, and which methods fix each.

1One model, several tasks

Multitask learning trains one model on several tasks at once. A shared trunk of layers feeds one small head per task, each task has a loss (a number that scores how badly the model does that task), and training lowers the sum of the losses. Training on that sum can go wrong in three common ways. The fixes fall into two families: methods that learn a weight for each task’s loss, and methods that combine the tasks’ gradients (Section 3.1) afresh at every step (Figure 1).

2Why sharing one trunk helps

Caruana (1997) tested the extra-signal benefit with a robot’s camera: a small network looked at a horizontal strip of pixels from each door image to find the doorknob and to tell single from double doors, once alone and once with eight extra tasks that locate and measure the doorway, the door jambs and the door edges. Error is root-mean-squared error (RMSE) on test images, compared against the best of three single-task networks.

Scene understanding is a standard use: one network predicts segmentation (a label for every pixel), depth and surface normals (the direction each surface faces) (Vandenhende et al., 2022).

3Shared and task-specific parameters

The trunk-and-heads design is called hard parameter sharing. Ruder (2017) calls it the most common approach to multitask learning in neural networks and says it greatly reduces the risk of overfitting, the third benefit in Section 2.

The trunk’s parameters, the weights training adjusts, are θ_sh. Task i’s head has its own parameters θ_i, and K is the number of tasks. Each step, task i’s loss L_i is averaged over a batch, the small set of examples the model sees in that step. θ stands for all parameters, trunk and heads together.

3.1A task gradient has one entry per shared weight

Task i’s gradient g_i lists how fast its loss L_i changes with each shared weight: it points where L_i rises fastest, and its length says how fast. With one number per shared weight, a trunk with millions of weights gives a vector with millions of entries. Every arrow on this page is a 2D drawing of such vectors; two vectors always lie in one plane, so the lengths and angles drawn are exact. How each g_i is computed in PyTorch is on Gradient Balancing Inside the Training Loop.

4Train on the weighted sum of the losses

Take NYUv2, a set of indoor photos, and one network learning two tasks on it: semantic segmentation (seg, a class for every pixel) and depth (meters to every pixel). With w_seg = w_depth = 0.5 the weighted total is half of each loss, so the step follows half of each task’s gradient. For two of the shared weights, illustrative gradients are g_seg = (2, 1) and g_depth = (−6, 2).

Training several tasks at once adds their losses into one total and walks all the parameters θ downhill on that single number.

L(θ) = Σ_i w_i L_i(θ)

The total loss L adds each task’s loss L_i times its loss weight w_i; Σ_i means add over the tasks i = 1 to K. Raising w_i makes task i count for more in every step, lowering it shrinks the task’s share of each step, at 0 the task drops out of training, and every weight at 1 gives the plain sum.

θ ← θ − η∇L(θ)

Each step of gradient descent replaces the parameters θ (← reads “becomes”) by themselves minus η times ∇L, the gradient of the total over all parameters; the gradient points uphill, so subtracting it, a step against ∇L, moves downhill. The step size η (the learning rate) sets how far each step goes: larger moves faster but can overshoot and make the loss jump, smaller is steadier but slower; PyTorch’s SGD defaults to 0.001 (SGD, Momentum and Nesterov). The methods in this series keep this step and change the weights inside L or the direction the step takes.

So the whole recipe is one weighted total and one downhill step on it: with w = 0.5 each, the gradient of the total on the two shared weights is half of (2, 1) plus half of (−6, 2), which is (−2, 1.5), and the step goes against it.

5Three common causes of negative transfer

When tasks are trained together on the plain sum, a task can end up worse than a separate model trained on that task alone; this is called negative transfer. Every task’s gradient changes the same trunk weights, so one task’s update can undo another’s. Three common causes, each with a different fix, are shown in Figure 1.

5.1Losses on different scales

Multiplying a loss by a number multiplies its gradient by the same number, so a loss measured on a larger scale, such as a loss near 600 next to one near 0.4, usually has a far longer gradient and so decides most of the step. It also makes up almost all of the total loss (Figure 6).

Hand-tuned weights w_i can fix this, but the tuning has to be redone whenever the tasks change.

5.2Tasks learn at different speeds

When one task is nearly solved early and another keeps improving, weights that were fair at the start are unfair later, so the balance has to change during training.

5.3Task gradients can point in conflicting directions

The cosine of the angle between two task gradients measures how much they point the same way. Past a right angle, part of each gradient works against the other; when that opposition is strong enough, a step against their sum raises one task’s loss (Section 7 gives the exact condition). Loss weights can pick a step that helps both tasks at one moment. A fixed weight keeps the same factor for all of training while the gradients’ lengths and angle change, so a setting that helps both tasks early can hurt one of them later. A combination chosen afresh at every step can be picked to help both whenever such a step exists.

6Four vector operations on gradients

Two of the failures in Section 5, scale and direction, come down to the lengths and angles of task gradients.

6.1The norm is a gradient’s length

In the pair of Section 4, depth’s gradient (−6, 2) is much longer than seg’s (2, 1).

A gradient’s length says how hard its task is pushing the shared weights at this step.

‖v‖ = √v₁² + v₂² + … + v_n²

The task gradients measured here are the parts of the weighted-sum step of Section 4. For a vector v with entries v₁ to v_n, square each entry, add, and take the square root: Pythagoras’s theorem over all n entries, where for a gradient n is the number of shared weights.

So the norm turns each task’s push into one length to compare: seg’s (2, 1) has length √5 ≈ 2.236 and depth’s (−6, 2) has √40 ≈ 6.325, almost three times longer.

6.2The dot product and the cosine

Recap The dot product and the cosine say whether two gradients point the same way, the measure behind the direction failure of Section 5.3. Taught in Vectors and Matrices, Section 3.

In the same pair, seg’s (2, 1) points right and up while depth’s (−6, 2) points mostly left, so the two lean apart.

The dot product says whether two gradients pull the same way, and the cosine says it with their lengths taken out.

⟨u, v⟩ = Σ_j u_j v_j

Multiply matching entries u_j and v_j and add over all j.

⟨u, v⟩ = ‖u‖ ‖v‖ cos ψ

The same number equals both lengths, the norms of Section 6.1, times cos ψ, where ψ (psi) is the angle between the vectors. Lengths are positive, so the dot product’s sign is the sign of cos ψ: positive when the vectors lean the same way, zero at a right angle, negative when they lean apart.

cos ψ = ⟨u, v⟩‖u‖ ‖v‖

Dividing by both lengths removes scale: a long and a short gradient pointing the same way have cosine 1, which is why the cosine is the standard check for conflicting directions and the norm the check for scale.

So a negative dot product is the sign that one shared step cannot fully please both tasks: for seg and depth it is 2 × (−6) + 1 × 2 = −10, and the cosine is −10 over 2.236 × 6.325, which is −0.707, an angle of 135 degrees.

Figure 10 shows the dot product turning negative once two gradients lean apart past a right angle.

6.3Projecting one gradient onto another

Recap Projection splits one gradient into the part along another gradient and the part at a right angle to it. Taught in Vectors and Matrices, Section 4.

Because seg and depth lean apart, the part of seg’s (2, 1) that lies along depth’s (−6, 2) points against depth.

Projecting keeps the part of one gradient that lies along another, and what is left over sits at right angles to it.

proj = ⟨g₁, g₂⟩‖g₂‖² g₂

The top of the fraction is the dot product of Section 6.2. The fraction is one number: how many copies of g₂ make up g₁’s part along it. It is negative when the two point apart, and then the projection points against g₂.

g₁ − proj ⟂ g₂

Subtracting the projection leaves the remainder; ⟂ means “at a right angle to”, so the remainder’s dot product with g₂ is 0.

So the projection splits g₁ into a part that fights g₂ and a part g₂ never feels; PCGrad keeps that remainder. For seg on depth the fraction is −10 over 40 = −0.25, so the projection is (1.5, −0.5) and the remainder (0.5, 1.5).

6.4A convex combination is a weighted average

Recap A convex combination mixes the task gradients with weights that are zero or more and add to one; the plain average gives every task 1/K.

The weighted sum of Section 4 with w = 0.5 each steps against half of seg’s (2, 1) plus half of depth’s (−6, 2): one step that mixes both tasks.

A convex combination is a weighted average of the task gradients, so the step always lies between them.

d = Σ_i w_i g_i

The projection remainder of Section 6.3 changes one gradient; a blend keeps every gradient whole and only weights it. d is the combined direction. Each g_i is the gradient of its own loss, so d is the gradient of the weighted sum of Section 4 with these weights: choosing a combination is choosing loss weights for one step.

w_i ≥ 0, Σ_i w_i = 1

The weights are shares of a whole, so for two tasks the tip of d lies on the segment joining the two gradient tips.

So the plain average, w = 0.5 each, is the midpoint (−2, 1.5) of the segment from (2, 1) to (−6, 2), and other weights pick other points.

7Whether a step helps each task

Each method in this series steps against some direction d over the shared weights; whether that step helps task i depends on the dot product of g_i with d.

In the illustrative pair of Section 4, the average step (−2, 1.5) has dot product −2.5 with seg’s gradient (2, 1) and 15 with depth’s (−6, 2), so at a step size of 0.1 seg’s loss rises by about 0.25 while depth’s falls.

A step helps a task exactly when the step and that task’s gradient point the same way.

s_i = ⟨g_i, d⟩

This is the dot product of Section 6.2, taken between a task gradient and the step. Task i’s score s_i is the dot product of its gradient with d, a measure of how much d agrees with g_i. It is positive when d leans the same way as g_i, and then a step against d lowers task i’s loss.

L_i(θ_sh − η d) ≈ L_i(θ_sh) − η s_i

For a small step size η, stepping the shared weights against d, with the heads held fixed, changes task i’s loss by about η times its score: down when s_i is positive, up when it is negative. ≈ means “close to, for small η”.

d = 12(g₁ + g₂)

This is the plain average of the two task gradients, the direction that equal loss weights give.

s₂ = 12⟨g₁, g₂⟩ + 12‖g₂‖²

Put the average into task 2’s score: half its dot product with g₁ plus half its own squared length. The squared length is never negative, so only a negative dot product can pull the score below zero.

s₂ < 0 ⇔ ⟨g₁, g₂⟩ < −‖g₂‖²

⇔ means “exactly when”: the average hurts task 2 only when the dot product is more negative than task 2’s own squared length. Task 1 works the same way with ‖g₁‖², so at a right angle, where the dot product is 0, both tasks gain.

So the test reads straight off the gradients whether one shared step hurts a task. With d the average of Section 6.4, s_seg = −2.5 and s_depth = 15. With seg as task 2 the last line reads −10 < −5, true, so seg is hurt; with depth as task 2 it reads −10 < −40, false.

8Which method fixes which failure

Each failure from Section 5 leaves a different signal in the logged gradients or losses.

Failure, and what the logs show	Methods that fix it
1 Losses on different scales (Section 5.1)One task’s gradient norm is far larger than the others’ from the start.	GradNorm, Uncertainty Weighting
2 Tasks learn at different speeds (Section 5.2)The losses, each as a fraction of its start, drift apart during training.	GradNorm
3 Conflicting directions (Section 5.3)The cosine between two task gradients is negative.	PCGrad, GradVac, CAGrad, 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	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?

Whichever method you use, the loss scale still sets the gradients going in, and the optimizer and gradient clipping (capping a gradient’s length) still change the combined gradient coming out; Gradient Balancing Inside the Training Loop shows how.

Which of these failures does your own run have? The Multitask Field Guide reads it off the logged norms and cosines.