Multi-Task Learning: One Backbone, Many Objectives

An autonomous vehicle needs a lot from a single camera frame: where the other road users are, which pixels are drivable, how far away everything is, where the lane boundaries run. Training a separate network per question is the simplest thing that works, and it is also the thing you cannot afford — the compute budget on the vehicle is fixed, and every model you add spends part of it re-extracting features that the other models already computed.

Multi-task learning takes the other route: one shared backbone, several lightweight heads. Beyond the obvious latency win, the tasks act as regularizers on each other. A representation good enough to predict depth is usually a better detection feature than one trained on boxes alone.

The catch is that you now have one number to descend on, and it is a sum of things that have no business being added together.

The combined loss

The usual formulation is a weighted sum over tasks:

Everything hard about multi-task learning lives in $w_i$.

The first problem is scale. A focal loss for classification settles around $10^{-1}$; an L1 depth loss in meters may start in the tens. Left alone, depth owns the gradient and classification is noise. So you tune the weights — and now you have $T-1$ extra hyperparameters that interact, on a model that takes a day to train.

The second problem is conflict, and it is the one that does not go away with tuning. Two tasks can want the same parameter to move in opposite directions. Writing $g_i = \nabla_\theta \mathcal{L}_i$, the tasks conflict when

No choice of positive scalar weights fixes a negative inner product. Scaling changes the magnitude of the disagreement, not its sign.

Ways to set the weights

Uncertainty weighting (Kendall, Gal & Cipolla) learns them, by treating each task’s weight as a homoscedastic noise term:

The $\log\sigma_i$ term is what keeps it honest — without it the optimizer would drive every $\sigma_i \to \infty$ and declare victory. In practice you parameterize $s_i = \log \sigma_i^2$ for stability:

class UncertaintyWeighting(nn.Module):
    def __init__(self, num_tasks):
        super().__init__()
        # s = log(sigma^2), learned jointly with the network
        self.s = nn.Parameter(torch.zeros(num_tasks))

    def forward(self, losses):
        total = 0.0
        for i, loss in enumerate(losses):
            total = total + 0.5 * torch.exp(-self.s[i]) * loss + 0.5 * self.s[i]
        return total

This handles scale well and costs almost nothing. It does not address gradient conflict.

GradNorm targets the training dynamics instead, tuning weights so that each task’s gradient magnitude at the shared layer stays proportional to how fast that task is still improving. Tasks that have converged get quieter; stragglers get louder.

PCGrad attacks conflict directly. When $g_i \cdot g_j < 0$, it projects one gradient onto the normal plane of the other before summing:

def pcgrad(grads):
    """grads: list of flattened per-task gradients. Returns the summed update."""
    projected = []
    for i, g_i in enumerate(grads):
        g = g_i.clone()
        # project away the conflicting component of every other task, in random
        # order -- the result depends on order, so do not fix it
        for j in torch.randperm(len(grads)).tolist():
            if j == i:
                continue
            overlap = torch.dot(g, grads[j])
            if overlap < 0:                       # conflict, and only then
                g -= overlap / grads[j].norm().pow(2) * grads[j]
        projected.append(g)
    return torch.stack(projected).sum(dim=0)

It is more expensive — you need per-task gradients, not just per-task losses, which means $T$ backward passes or a batched trick — but it is the only one of the three that changes the direction of the update rather than its scale.

What actually bites in practice

Negative transfer is real, and it is not symmetric. Adding a task can make another one worse. Semantic segmentation and depth tend to help each other; a task with a very different label distribution often just steals capacity. Before assuming a task belongs in the model, train it alone and keep that number as the baseline you have to beat.

Watch per-task metrics, never the aggregate. Total loss going down is compatible with one head quietly collapsing. Plot each task separately and set a per-task regression gate in CI.

Datasets are rarely co-labelled. Most frames have boxes or masks, not both. You end up sampling per-batch by task, which means the effective weight on a task is w_i × (its share of batches) — two knobs pointing at the same thing. Fix the sampling ratio first, then tune weights.

Give heads their own learning rates. A head training from scratch on a backbone that is being fine-tuned wants a higher LR than the backbone does. This is a one-line change that often does more than a week of weight search.

Where I would start

For a new stack: uncertainty weighting, per-task learning rates, per-task metrics on the dashboard from day one. That combination gets you most of the way with very little machinery. Reach for GradNorm or PCGrad only once you can show a conflict — a per-task metric that is worse in the joint model than it was standalone. Otherwise you are adding complexity to solve a problem you have not confirmed you have.