Skip to content

← Flax step 73 of 100

Hard Primitives

Pretrain, Then Train Only the Head

Jump to the task

Reuse a trained feature extractor

This exercise has two stages. First, update every parameter of a small network on one dataset. Then reuse those parameters on a second dataset, updating only the final layer. The first layer, called the trunk, stays fixed during the second stage. The final layer is the head.

You keep the pretrained head too; this task does not replace or reinitialize it. Three training steps are enough to practice the state transition, but do not establish that the network has learned useful transferable features.

Freeze updates explicitly

Optax's multi_transform routes each parameter's update to a labeled transformation. Send the trunk to set_to_zero and the head to SGD. The following standalone example shows that even nonzero trunk gradients produce zero trunk updates:

import jax
import jax.numpy as jnp
import optax

params = {
    "Dense_0": {"kernel": jnp.ones((2, 8)), "bias": jnp.zeros(8)},
    "Dense_1": {"kernel": jnp.ones((8, 1)), "bias": jnp.zeros(1)},
}
labels = jax.tree_util.tree_map_with_path(
    lambda path, _: "frozen" if any(
        getattr(entry, "key", None) == "Dense_0" for entry in path
    ) else "trainable",
    params,
)
tx = optax.multi_transform(
    {"frozen": optax.set_to_zero(), "trainable": optax.sgd(0.1)}, labels
)
grads = jax.tree.map(jnp.ones_like, params)
updates, _ = tx.update(grads, tx.init(params), params)
assert bool(jnp.all(updates["Dense_0"]["kernel"] == 0))
assert bool(jnp.all(updates["Dense_0"]["bias"] == 0))
assert bool(jnp.allclose(updates["Dense_1"]["kernel"], -0.1))
assert bool(jnp.allclose(updates["Dense_1"]["bias"], -0.1))

This label rule fits this two-layer model. For a larger model, inspect the parameter paths before deciding what to freeze.

A plain optax.masked wrapper is not enough to freeze its unselected leaves: it passes their incoming updates through. Zeroing gradients alone can also leave momentum or weight-decay updates active. Zeroing the final updates for the frozen group prevents that group from moving.

The reference computes gradients for the whole parameter tree before zeroing trunk updates. Freezing therefore does not automatically skip trunk differentiation or guarantee a speedup. Restricting differentiation to the head is a separate optimization.

Problem

Implement pretrain_finetune_loss(seed, x_pre, y_pre, x_fine, y_fine, lr_pre, lr_fine, finetune_steps) using this exact architecture:

import flax.linen as nn

class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(8)(x)
        x = nn.relu(x)
        return nn.Dense(1)(x)

Keep both Dense layers' default initialization and biases. Their generated names are Dense_0 and Dense_1; these names also affect Flax's random initialization. Initialize the model on x_pre with jax.random.PRNGKey(int(seed)), without splitting the key.

  1. Pretrain: take exactly three full-batch SGD steps on (x_pre, y_pre) at learning rate lr_pre. Update both layers, with no momentum or weight decay.
  2. Switch optimizers: keep the resulting parameter tree and create a fresh optimizer state for the second stage. Freeze both the kernel and bias of Dense_0. Only Dense_1 may change.
  3. Train the head: take int(finetune_steps) full-batch SGD steps on (x_fine, y_fine) at learning rate lr_fine.
  4. Report: evaluate the final model's MSE on (x_fine, y_fine) and return it as a JAX array of shape (1,).

In both stages, flatten the predictions from (N, 1) to (N,) before subtracting the targets. MSE is the mean squared error over the examples. Recompute gradients at the current parameters for every step.

Inputs:

  • x_pre: float array (N_p, D); y_pre: float array (N_p,).
  • x_fine: float array (N_f, D); y_fine: float array (N_f,).
  • seed: integer-valued random seed, possibly supplied as a float.
  • lr_pre, lr_fine: nonnegative scalar learning rates.
  • finetune_steps: nonnegative integer-valued scalar, possibly a float.

The batches are nonempty and have the same positive feature width D. Output: [final_finetune_loss], measured after all requested updates. This is a loss on the second-stage training data, not a held-out estimate of generalization.

Check the stage boundary

With zero fine-tuning steps, return the pretrained model's loss on the fine-tuning data. Do not return its pretraining-data loss. With lr_fine=0, the second stage also leaves the pretrained parameters fixed. Otherwise, compare the trunk immediately before and after fine-tuning; it must be unchanged, even though it changed during pretraining.

Keep seed, both learning rates, and finetune_steps static when JIT-compiling the supplied implementation, which converts them to Python scalars during setup. Keep the returned loss as a JAX value rather than calling Python float on it.