Skip to content

← Flax step 65 of 100

Medium Primitives

Multi-Task Two-Head Loss

Jump to the task

Why this matters

A multi-task model can share a feature extractor while using separate output heads for different predictions. A detector, for example, may predict both object classes and box coordinates. Sharing can reduce duplicated computation, but the tasks can also compete. A shared trunk does not guarantee that every task improves or that their gradients are balanced.

This exercise uses a shared Dense/Linear layer with eight outputs, followed by ReLU and two separate scalar-output heads. Both heads receive the same hidden features, but have separate weights and biases.

Combining the losses

For fixed alpha in [0, 1]:

loss_a = mean((prediction_a - target_a) ** 2)
loss_b = mean((prediction_b - target_b) ** 2)
combined = alpha * loss_a + (1 - alpha) * loss_b

For example, losses 4 and 1 with alpha 0.25 give a combined loss of 1.75. Alpha zero ignores task A in the combined objective; alpha one ignores B. Equal weights do not guarantee equal gradient contributions. Target units, error scales and the model's sensitivities all matter.

For fixed weights, the shared parameters receive the weighted sum of the two loss gradients. Normalizing loss values alone does not equalize gradient magnitudes. Learning an unconstrained mixing weight can simply favor the easier task; methods for adaptive weighting require their own objectives.

The model

Use this structure and construction order so initialization matches the exercise. Keep the default initializers and biases.

import jax
import jax.numpy as jnp
import flax.linen as nn

class TwoHead(nn.Module):
    hidden: int = 8

    @nn.compact
    def __call__(self, x):
        h = nn.Dense(self.hidden)(x)
        h = nn.relu(h)
        out_a = nn.Dense(1)(h)
        out_b = nn.Dense(1)(h)
        return out_a, out_b

Inside the compact method, construct the layers in the stated order. Linen assigns them the default names Dense_0, Dense_1 and Dense_2. Names and construction order affect the initialized parameters, so changing them can change seeded fixture values even when the architecture looks equivalent.

Problem

Implement multi_task_loss(seed, x, y_a, y_b, alpha):

  1. Construct TwoHead(hidden=8) and initialize it with jax.random.PRNGKey(int(seed)) and x.
  2. Apply the model using the returned variables.
  3. Reshape each (N, 1) prediction to (N,) to match its target.
  4. Compute the two MSEs and the weighted combined loss.
  5. Return jnp.stack([combined, loss_a, loss_b]).

Inputs are float32 x of shape (N, D), float32 targets each of shape (N,), a nonnegative integer-valued seed, and alpha in [0, 1]. N and D are positive. Return a float32 array of shape (3,). There is no training step: all three losses use the same newly initialized model.

Common pitfalls

  • Match prediction and target shapes before subtraction. (N, 1) - (N,) broadcasts to (N, N), which compares each prediction with every target.
  • Use separate Dense instances for the two heads. Reusing a layer instance can intentionally tie parameters, but that is not this architecture.
  • The result of model.init is a variables dictionary containing a params collection. Pass the complete dictionary to model.apply.
  • Preserve JAX values when returning losses. jnp.stack supports gradients; converting each loss with Python float does not.
  • If wrapping this setup function in jax.jit, mark seed static because it is converted to a Python integer. Alpha can remain dynamic.