Skip to content

← Flax NNX step 93 of 100

Hard Primitives

NNX Surgery Freeze

Jump to the task

Why this matters

Freezing a layer means keeping its parameters unchanged while updating another part of the model. This exercise uses one step of plain SGD with no momentum or weight decay. For this optimizer, replacing a parameter's gradient with zero makes its update zero.

This is not a general rule for every optimizer. Momentum from earlier steps can move a parameter even when its current gradient is zero; weight decay can change it too. For those cases, exclude frozen parameters from updates or ensure their final optimizer updates are zero. Masking incoming gradients and masking final updates are different operations.

Mask the intended layer

Match the first dictionary key in each gradient path exactly. A prefix string check for l1 would also catch a layer named l11.

import jax
import jax.numpy as jnp

def zero_first_layer_grads(path, leaf):
    is_first_layer = (
        bool(path)
        and isinstance(path[0], jax.tree_util.DictKey)
        and path[0].key == "l1"
    )
    return jnp.zeros_like(leaf) if is_first_layer else leaf

grads = {
    "l1": {"kernel": jnp.ones((2, 3)), "bias": jnp.ones(3)},
    "l11": {"kernel": jnp.ones((3, 1))},
    "l2": {"bias": jnp.array([2.0])},
}
masked = jax.tree_util.tree_map_with_path(zero_first_layer_grads, grads)
assert jnp.all(masked["l1"]["kernel"] == 0)
assert jnp.all(masked["l1"]["bias"] == 0)
assert jnp.all(masked["l11"]["kernel"] == 1)
assert jnp.all(masked["l2"]["bias"] == 2)

The same predicate works for the nested gradient state of the two-layer NNX model below. The map reaches each array inside its Variable wrapper. Both the kernel and bias under the top-level l1 entry must be masked.

Verify values, not just norms

This exercise reports the first kernel's norm before and after the step. That is a useful summary, but equal norms do not prove equal parameters: [3, 4] and [0, 5] both have norm 5. A kernel norm also says nothing about the layer's bias. When checking a freeze locally, compare every frozen array with its saved pre-step values and verify the intended trainable updates.

Likewise, restoring parameters after an optimizer step restores their values but does not rewind optimizer state. Decide how to handle that state when changing which parameters are trainable. Zeroing gradients is not a universal remedy for stale momentum.

Problem

Implement freeze_first_dense(seed, x, y, lr):

  1. Construct TwoLayer(nnx.Module) with l1 mapping x.shape[-1] to 4 and l2 mapping 4 to y.shape[-1]. Create l1 first, sharing one nnx.Rngs(int(seed)) between the constructors. Keep default initializers and biases. Forward is l2(jax.nn.relu(l1(x))).
  2. Create nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param).
  3. Save jnp.linalg.norm(model.l1.kernel[...]) before the update.
  4. Compute loss and gradients with nnx.value_and_grad. The loss is jnp.mean((model(x) - y) ** 2), averaged over every batch and output entry.
  5. Use jax.tree_util.tree_map_with_path to zero all gradient leaves under the exact top-level key l1, leaving l2 gradients unchanged.
  6. Call optimizer.update(model, masked_grads) once.
  7. Measure the first kernel's norm again and return jnp.stack([loss, norm_after, norm_before]).

Return the pre-update loss produced by value_and_grad; do not recompute it after the step. Keep loss, norms and learning rate as JAX values, rather than converting them to Python floats. The seed configures initialization and must be static if JIT-compiling the whole entrypoint.

This function constructs a model and optimizer for a single-step exercise. A training loop should create them once and reuse them across steps. The output tests check only loss and kernel norms; they cannot prove that you preserved the bias or actually updated the second layer. Follow the entire update procedure, and use full-array comparisons when debugging it.

Inputs: x has shape (N, D_in) and y has shape (N, D_out), with matching positive batch size and positive widths. Both are finite float32 arrays. seed is a nonnegative integer-valued scalar supplied as a float; cast it to int. lr is a finite nonnegative learning rate. Inputs are small enough to keep the loss, gradients and updates finite.

Output: shape (3,), containing [loss, norm_after, norm_before]. The frozen kernel values, and therefore its norm, remain unchanged under the specified plain-SGD update.