Skip to content

← Optax step 22 of 25

Hard Primitives

Train Step with Frozen Params (Mask)

Jump to the task

Freeze selected parameter values while applying one SGD step to the others. Here, frozen_mask = 1 means keep the original value, and frozen_mask = 0 means update it.

For params = [2, -3], grads = [10, -2], frozen_mask = [1, 0] and lr = 0.1, return [2, -2.8]. The first parameter stays at 2, not zero, despite its nonzero gradient.

Zero the update, not the parameter

Use jnp.where(frozen_mask >= 0.5, 0.0, grads) to replace the frozen positions' gradients with zero. Apply plain SGD to that gradient array:

new_params = params - lr * masked_grads

Do not multiply the updated parameters by a mask or replace frozen values with zero. Freezing means retaining their original values.

Zeroing gradients works here because SGD has no momentum or weight decay. With those additions, a parameter may move even when its current gradient is zero. Freezing then requires ensuring that its final update is zero.

For a parameter tree, a partition can route frozen leaves through optax.set_to_zero(). optax.masked alone does not freeze unselected leaves: it passes their incoming updates through unchanged.

Inputs

Implement train_with_freeze(params, frozen_mask, grads, lr) using Optax SGD without momentum or weight decay. Initialize state and perform one update with the masked gradients.

  • params and grads: 1-D floating-point JAX arrays with the same shape.
  • frozen_mask: a JAX array of that shape containing only 0.0 and 1.0. A value of 1.0 marks a frozen position.
  • lr: a scalar learning rate.

Output

Return a JAX array with the original shape. Frozen values must equal their original values; other positions take the SGD step. Leave all input arrays unchanged.