Skip to content

← Flax step 75 of 100

Hard Primitives

Separate Learning Rates for Kernel and Bias

Jump to the task

Different step sizes for different parameters

A Dense layer has a kernel matrix and a bias vector. In this exercise, each gets its own learning rate. One SGD step updates them as:

  • kernel_new = kernel - lr_kernel * kernel_gradient
  • bias_new = bias - lr_bias * bias_gradient

Both gradients come from the same loss evaluated at the original parameters. There is no momentum or weight decay in this task. The inputs are two actual learning rates, not multipliers of an additional base rate.

Grouping parameters is useful when you want different update sizes, such as a smaller step for a pretrained component and a larger step for a new component. The appropriate rates depend on the model and data. A larger rate does not guarantee faster learning or a lower loss.

Problem

Implement multi_lr_step(seed, x, y, lr_kernel, lr_bias):

  1. Initialize a standalone nn.Dense(4) with jax.random.PRNGKey(int(seed)) and x. Keep its default kernel initializer and zero bias. Use the "params" collection.
  2. Label the kernel leaf "kernel" and the bias leaf "bias".
  3. Use optax.multi_transform to route kernel updates to optax.sgd(lr_kernel) and bias updates to optax.sgd(lr_bias).
  4. Compute the MSE between flattened predictions and flattened y, then take exactly one optimizer step from a fresh optimizer state.
  5. Evaluate the MSE again with the updated parameters. Return it as a JAX array of shape (1,).

Inputs: x is a float array of shape (N, D), with positive N and D. The target y has shape (N * 4,) and lists outputs in row-major order. seed is an integer-valued scalar, possibly a float. lr_kernel and lr_bias are nonnegative scalar rates. Output: [loss_after_step]. Return the loss after the update, not the loss used to calculate the gradients.

The mean is over all N * 4 output elements. Summing output errors per example and then averaging only over N would scale the gradients by four and change the requested update.

A labels tree routes updates

The labels tree identifies which transformation handles each parameter. Here the parameter dictionary is flat, so it can be written directly:

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])}
labels = {"kernel": "kernel", "bias": "bias"}
tx = optax.multi_transform(
    {"kernel": optax.sgd(0.1), "bias": optax.sgd(0.5)}, labels
)
updates, _ = tx.update(grads, tx.init(params), params)
new_params = optax.apply_updates(params, updates)
# kernel: 2 - 0.1 * 3 = 1.7; bias: 1 - 0.5 * (-2) = 2
assert bool(jnp.allclose(new_params["kernel"], 1.7))
assert bool(jnp.allclose(new_params["bias"], 2.0))

For a nested model, jax.tree_util.tree_map_with_path can generate labels from parameter paths. Inspect the actual tree before assigning groups: treating every non-kernel leaf as a bias is only safe for a model whose leaves really are kernels and biases. Every label used must have a matching entry in the transform dictionary.

The diagram shows a larger illustrative tree. The exercise itself has only the two leaves of one standalone Dense layer.

Loading visualization…

Useful checks

If both rates are zero, the parameters and loss stay unchanged. If only one rate is zero, that parameter group stays fixed. Equal rates should give the same update as one ordinary SGD optimizer. A sufficiently large rate can increase the loss.

Keep seed, lr_kernel, and lr_bias static when JIT-compiling the supplied implementation, which converts them to Python scalars during setup. Keep the returned loss as a JAX array instead of converting it to a Python float.