Skip to content

← Optax step 18 of 25

Hard Primitives

masked: Apply WD Only To Certain Params

Jump to the task

Apply weight decay to selected elements while applying the supplied loss gradients to every element. A zero in this mask skips decay; it does not freeze the parameter.

Add only the selected penalty terms

For each position, use:

combined_grad = grad + wd * param    # mask is 1
combined_grad = grad                 # mask is 0
new_param = param - lr * combined_grad

For params = [2, 2], grads = [0.5, 0.5], mask = [1, 0], wd = 0.1 and lr = 0.1, the combined gradients are [0.7, 0.5]. Return [1.93, 1.95]. Both parameters change, but only the first receives the extra decay term.

This task adds the penalty before plain SGD. For this optimizer, it gives the same values as subtracting lr * wd * params separately at selected positions, using the original parameters. The equivalence does not generally hold for adaptive optimizers such as Adam.

Element masks and tree masks

Use jnp.where(mask > 0.5, wd * params, 0.0) to build the penalty array, then add it to grads and take one step with optax.sgd(lr).

optax.masked serves a related but different interface: its Boolean mask selects array leaves or subtrees, not individual elements within one array. Unselected updates pass through unchanged. This exercise therefore uses jnp.where for element selection rather than calling optax.masked.

Inputs

Implement mask_specific_params(params, grads, lr, wd, mask).

  • params and grads: 1-D floating-point JAX arrays with the same shape.
  • mask: a JAX array of that shape containing only 0.0 and 1.0. A value of 1.0 selects the decay term.
  • lr: a scalar SGD learning rate.
  • wd: a scalar weight decay coefficient.

Use plain SGD without momentum. Initialize its state and perform one update.

Output

Return the updated JAX array with the original shape. Leave all input arrays unchanged. Positions with mask zero still receive their ordinary SGD update.