Skip to content

← Flax step 74 of 100

Hard Primitives

Masked L2 Regularization with Adam

Jump to the task

Choose which parameters receive a penalty

A parameter mask lets you apply regularization selectively. This exercise adds an L2 penalty to a Dense layer's kernel while leaving its bias out of that penalty. The bias still learns from the data loss. Choosing which parameters to regularize is a training choice, not a universal rule based only on their shape.

The order of optimizer transformations matters. Here the kernel's incoming gradient becomes g_kernel + wd * kernel, and then Adam processes it. This is coupled L2 regularization. The equivalent training objective is:

MSE + (wd / 2) * sum(kernel ** 2)

The factor of one half makes the penalty's derivative wd * kernel. The value returned by this exercise is the post-update MSE alone, without adding the regularization penalty to that reported value.

Problem

Implement wd_mask_step(seed, x, y, lr, wd):

  1. Initialize a standalone nn.Dense(4) on x, using jax.random.PRNGKey(int(seed)) without splitting the key. Keep its default initializer and bias. Extract variables["params"].
  2. Build the mask {"kernel": True, "bias": False}.
  3. Construct the optimizer in this exact order: optax.chain(optax.masked(optax.add_decayed_weights(wd), mask), optax.adam(lr)). Use Adam's default settings and a fresh optimizer state.
  4. Differentiate the mean squared error between flattened predictions and flattened targets. The optimizer adds the kernel penalty, so do not also add it to the loss used here.
  5. Apply one update and compute the MSE again. Return [final_loss] as a JAX array of shape (1,).

Inputs: x has shape (N, D) with positive N and D. y has shape (N * 4,), in row-major output order. seed is an integer-valued scalar that may arrive as a float. lr and wd are nonnegative scalar values. Output: the post-step MSE, averaged over all N * 4 elements.

What the mask changes

A masked transformation processes the selected updates and passes the others through. With this mask, the bias gradient reaches Adam unchanged; it is not zeroed or frozen. This small example isolates that behavior:

import jax.numpy as jnp
import optax

params = {"kernel": jnp.array([[2.0]]), "bias": jnp.array([1.0])}
grads = {"kernel": jnp.array([[3.0]]), "bias": jnp.array([-2.0])}
mask = {"kernel": True, "bias": False}
penalty = optax.masked(optax.add_decayed_weights(0.1), mask)
adjusted, _ = penalty.update(grads, penalty.init(params), params)
assert bool(jnp.allclose(adjusted["kernel"], 3.2))
assert bool(jnp.allclose(adjusted["bias"], -2.0))

For a nested model, jax.tree_util.tree_map_with_path can select leaves by dictionary keys. Check the model's parameter names rather than assuming every matrix is named kernel.

Why this chain is not AdamW

Adam maintains moving averages from the updates it receives. In this task, those updates already include the kernel penalty, so Adam's adaptive scaling also acts on that penalty.

AdamW instead adds decay after Adam's moment-based scaling, before the learning-rate scaling. Its decay contribution is -lr * wd * parameter on selected leaves. In Optax, use optax.adamw(lr, weight_decay=wd, mask=mask) for that behavior. It is a different optimizer and will not generally match this exercise's answer. See the Optax AdamW reference.

With wd=0, the requested chain reduces to ordinary Adam. With lr=0, its parameters stay unchanged. Neither regularization nor an optimizer step guarantees that the data loss decreases for arbitrary inputs.

Keep seed, lr, and wd static when JIT-compiling the supplied implementation, which converts them to Python scalars during setup. Return the loss with jnp.reshape(final_loss, (1,)) to keep it as a JAX value.

The diagram shows an illustrative nested parameter tree. This task uses just one Dense layer's kernel and bias.

Loading visualization…