How to use per-parameter learning rates?

[cgarciae]

Here's a strategy for defining per-parameter learning rates using optax.multi_transform and leveraging Variable metadata.

import jax
from flax import nnx

model = nnx.Linear(
  2,
  5,
  kernel_init=nnx.with_metadata(
    nnx.initializers.lecun_normal(),
    learning_rate=0.2, # add learning rate metadata
  ),
  bias_init=nnx.with_metadata(
    nnx.initializers.zeros,
    learning_rate=0.1, # add learning rate metadata
  ),
  rngs=nnx.Rngs(42),
)

state = nnx.state(model, nnx.Param)
optimizers = {}

def leaf_optimizer(path, value):
  # here we both populate `optimizer` and return the corresponding tag which is the learning rate
  optimizers.setdefault(value.learning_rate, optax.adamw(value.learning_rate))
  return value.learning_rate

state_optimizer = nnx.map_state(leaf_optimizer, state)

optimizer = nnx.Optimizer(
  model,
  tx=optax.multi_transform(optimizers, state_optimizer),
  wrt=nnx.Param,
)

@nnx.jit
def train_step(model, optimizer, x, y):
  def loss_fn(model):
    predictions = model(x)
    loss = jnp.mean((predictions - y) ** 2)
    return loss

  loss, grads = nnx.value_and_grad(loss_fn)(model)
  optimizer = optimizer.update(grads)
  return loss

x = jax.random.normal(jax.random.key(42), (32, 2))
y = jax.random.normal(jax.random.key(43), (32, 5))

losses = []
for _ in range(50):
  loss = train_step(model, optimizer, x, y)
  losses.append(loss)

Using nnx.with_metadata is not necesary, the leaf_optimizer could also just decide the learning rate for each value on its own. If you control the creation of the Variables is easier to just assign the metadata directly e.g:

class Linear(nnx.Module):
  def __init__(self, din, dout, *, rngs):
    self.w = nnx.Param(nnx.initializers.lecun_normal(rngs.params(), (din, dout)), learning_rate=0.2)
    ...