Skip to content

← Flax NNX step 59 of 100

Medium Primitives

NNX Multi-Task 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
from flax import nnx

class TwoHead(nnx.Module):
    def __init__(self, d_in, hidden, rngs):
        self.trunk = nnx.Linear(d_in, hidden, rngs=rngs)
        self.head_a = nnx.Linear(hidden, 1, rngs=rngs)
        self.head_b = nnx.Linear(hidden, 1, rngs=rngs)

    def __call__(self, x):
        h = jax.nn.relu(self.trunk(x))
        return self.head_a(h), self.head_b(h)

Create the trunk, head_a and head_b in that order, passing the same nnx.Rngs(int(seed)) container through the constructors. The container advances its stream as each layer requests keys; do not create a new identically seeded container for each head.

Problem

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

  1. Construct TwoHead(x.shape[-1], 8, rngs=nnx.Rngs(int(seed))).
  2. Create nnx.Optimizer(model, optax.sgd(0.05), wrt=nnx.Param).
  3. Define a scalar loss function with the weighted MSE formula above.
  4. Obtain gradients with nnx.value_and_grad(loss_fn)(model, x, y_a, y_b) and perform exactly one optimizer.update(model, grads).
  5. Run the updated model again. Recompute both task losses and their weighted combination using the same alpha.
  6. Return jnp.stack([combined, loss_a_final, loss_b_final]).

Inputs are float32 x of shape (N, D), float32 targets each of shape (N, 1), a nonnegative integer-valued seed, and alpha in [0, 1]. N and D are positive. The float32 output has shape (3,) and reports losses after the step. This differs from the Linen forward-only exercise.

Common pitfalls

  • The loss passed to value_and_grad must be scalar. Combine both losses within the differentiated computation; helper functions are fine. A tuple of model outputs is not itself a scalar loss.
  • Keep predictions and targets at (N, 1). Changing just one to (N,) creates an unintended (N, N) broadcast for batches larger than one.
  • With alpha zero, head A has no direct loss gradient, but its predictions may still change because the shared trunk updates through task B.
  • One SGD step is not guaranteed to lower both task losses, or even the combined loss for an arbitrary input and fixed learning rate.
  • Keep alpha and reported losses as JAX values. Python float conversions are unnecessary and can break tracing or differentiation.
  • Initialize once, update once, then evaluate that same model. Recreating it before reporting would discard the update.