Skip to content

← Flax step 95 of 100

Medium Primitives

Param Surgery - Freeze First Dense

Jump to the task

A parameter tree and its gradient tree have matching paths. To keep a layer fixed during plain SGD, replace that layer's gradient arrays with zeros while leaving its parameter arrays intact. This exercise practices selecting those paths in a two-layer Flax Linen model.

Change gradients, not weights

The model is Dense(8), ReLU, then Dense(1). Its automatically named layers are Dense_0 and Dense_1. Freeze both Dense_0/kernel and Dense_0/bias; train both arrays under Dense_1.

Replacing a kernel with zeros changes the model immediately. Replacing its gradient with zeros instead preserves the learned values during the plain SGD step used here. Keep the original parameters available while constructing the modified gradient tree and the updated parameter tree.

Select leaves by dictionary key

jax.tree_util.tree_map_with_path passes a path and an array to your callback. For the first kernel, the path contains dictionary keys Dense_0 and kernel. Select the layer key and return jnp.zeros_like(grad) for its leaves:

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

Apply this callback to the gradient tree. Matching Dense_0 includes the bias as well as the kernel. Other leaves keep their original gradients. zeros_like preserves the selected array's shape and dtype.

Scope of the freezing rule

Plain SGD without momentum or weight decay uses update = -lr * gradient. That is why gradient zeroing freezes this layer. With stored momentum or weight decay, a zero current gradient may still produce a nonzero update. General freezing requires zero final updates for the selected parameters; how optimizer statistics are retained is a separate decision.

Verify the actual arrays

The requested output includes the first kernel's norm before and after the step. jnp.linalg.norm on that matrix is its Frobenius norm: the square root of the sum of squared entries. The two reported norms should agree.

Equal norms do not prove that a kernel is unchanged. For example, exchanging two entries preserves the norm. When checking your implementation, compare the complete first-layer kernel and bias before and after the step. Check that the second layer receives its SGD update too; skipping all updates is not the requested operation.

Problem

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

  1. Define a Linen MLP with a compact __call__ that applies unnamed nn.Dense(8), nn.relu, then unnamed nn.Dense(1). Use default Dense settings and biases with no final activation.
  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. Compute full-batch mean squared error and its gradients with jax.value_and_grad. Reshape predictions from (N, 1) to (N,) before subtracting targets. Do not include a one-half factor in the loss.
  4. Use tree_map_with_path to zero the first layer's gradient leaves. Initialize optax.sgd(float(lr)) with default settings and apply one update using the modified gradients.
  5. 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 both inputs unchanged. With lr=0, both layers retain their initialized values. Report the loss computed before the update.