Skip to content

← Flax step 76 of 100

Hard Primitives

Param Freezing via Grad Zeroing

Jump to the task

Freezing a layer means keeping its parameters unchanged while training other parts of the model. This exercise freezes the first layer of a two-layer MLP by replacing its gradients with zeros before one plain SGD update.

Why zero gradients work here

For SGD without momentum or weight decay, each parameter update is -lr * gradient. A zero gradient therefore produces a zero update. The first layer keeps both its kernel and bias; the second layer trains normally.

This reasoning does not apply to every optimizer. Momentum from earlier steps can move parameters even when the current gradient is zero. Weight decay can also change parameters independently of that gradient. To freeze parameters with such an optimizer, ensure their final updates are zero, and decide separately how their optimizer state should behave when training resumes.

Select gradient leaves by path

Gradients have the same nested structure as the parameters. For this model, Dense_0 is the first layer and Dense_1 is the second. Use jax.tree_util.tree_map_with_path to visit each gradient array together with its path. A dictionary leaf such as grads["Dense_0"]["kernel"] has path entries for both dictionary keys.

def zero_frozen(path, grad):
    frozen = any(getattr(entry, "key", None) == "Dense_0" for entry in path)
    return jnp.zeros_like(grad) if frozen else grad

grads = jax.tree_util.tree_map_with_path(zero_frozen, grads)

This zeros both first-layer gradient leaves and preserves their shapes and dtypes. Match the dictionary key, not a substring in a printed path. In this model the layer names are unique, so the test identifies the intended subtree.

Check values, not just norms

The returned norms summarize the first kernel's size before and after the step. For a matrix, jnp.linalg.norm(kernel) uses the square root of the sum of squared entries, also called the Frobenius norm.

Equal norms are a useful check but do not prove equal parameters: [3, 4] and [4, 3] have the same norm and different values. When debugging, compare every first-layer array before and after the update, including the bias. Also check that the second layer receives its ordinary SGD update. Merely skipping all training would not perform the requested operation.

Problem

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

  1. Define a Linen MLP module with a compact __call__: unnamed Dense(8), ReLU, then unnamed Dense(1). Use default Dense settings and biases, with no final activation. These automatic names determine the seeded parameters.
  2. Initialize with jax.random.PRNGKey(int(seed)) and x, without splitting the key. Extract the params collection and record the first kernel's norm.
  3. Use jax.value_and_grad to compute full-batch mean squared error and its gradients. Reshape predictions from (N, 1) to (N,) before subtracting targets. There is no one-half factor in the loss.
  4. Zero the Dense_0 gradient leaves using tree_map_with_path.
  5. Initialize optax.sgd(float(lr)) with default settings, compute updates from the modified gradients and apply those updates once.
  6. Return a floating-point JAX array of shape (3,) containing [loss_before, frozen_kernel_norm_after, frozen_kernel_norm_before].

x is a finite float32 array of shape (N, D) and y is a finite float32 array of shape (N,), with N >= 1 and D >= 1. seed is a nonnegative integer-valued scalar that may arrive as a float; lr is finite and nonnegative. Leave the input arrays unchanged. With lr=0, both layers keep their initial values. Return the pre-update loss, not a newly evaluated post-update loss.