Skip to content

← Optax step 15 of 25

Hard Primitives

multi_transform per-Param Group

Jump to the task

Apply different SGD learning rates to two groups of parameters. The first half of the input array uses lr_a; the second half uses lr_b.

For params = [1, 2, 3, 4], grads = [0.1, 0.1, 0.1, 0.1], lr_a = 0.1 and lr_b = 1.0, return approximately [0.99, 1.99, 2.9, 3.9]. The second half takes a larger step because its learning rate is larger. A zero rate leaves that group's values unchanged in this plain SGD exercise.

Give each group a label

Optax treats an array as one leaf in a parameter tree. To assign separate transforms to its halves, first split both parameters and gradients into dictionaries with matching keys, such as a and b.

Use optax.multi_transform with:

  • transforms {'a': optax.sgd(lr_a), 'b': optax.sgd(lr_b)};
  • labels {'a': 'a', 'b': 'b'}.

Each label selects a transform for that array leaf. An array of integer labels does not assign optimizers to individual elements of a single leaf. More generally, Optax also accepts labels for whole subtrees; matching one label to each leaf is sufficient here.

Initialize state from the parameter dictionary, update using the gradient dictionary, and apply the updates to the parameter dictionary. Finally, concatenate the updated a and b arrays in their original order.

Inputs

Implement multi_transform_step(params, grads, lr_a, lr_b) using the partitioned optimizer above. Perform one step with plain SGD and no momentum.

  • params: a 1-D floating-point JAX array of positive, even length.
  • grads: a floating-point JAX array with the same shape.
  • lr_a: the scalar learning rate for indices before len(params) // 2.
  • lr_b: the scalar learning rate for the remaining indices.

Output

Return one JAX array with the original shape and element order. Leave the input arrays unchanged; do not return the intermediate dictionaries.