How to implement layer-wise learning rate decay?

[andsteing]

(originally asked by @debidatta)

How can I implement an Optax optimizer that uses different learning rates for different layers?

[andsteing]

(originally answered by @levskaya)

Adapted from #1453

This can easily be done with optax.multi_transform.

For Flax it can be very handy to use flax.traverse_util.ModelParamTraversal to create the second parameter:

def flattened_traversal(fn):
  """Returns function that is called with `(path, param)` instead of pytree."""
  def mask(tree):
    flat = flax.traverse_util.flatten_dict(tree)
    return flax.traverse_util.unflatten_dict(
        {k: fn(k, v) for k, v in flat.items()})
  return mask

# Specify layer-wise learning rate.
lrs = {'Dense_0': 0.1, 'Dense_1': 0.2, 'head': 0.3}
label_fn = flattened_traversal(lambda path, _: path[0])

tx = optax.multi_transform(
    {name: optax.sgd(lr) for name, lr in lrs.items()}, label_fn)

Full example

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax


class Model(nn.Module):
  num_layers: int
  depth: int
  @nn.compact
  def __call__(self, x):
    for i in range(self.num_layers - 1):
      x = nn.relu(nn.Dense(self.depth)(x))
    return nn.Dense(self.depth, name='head')(x)


model = Model(num_layers=3, depth=10)
x = jnp.zeros([1, 10])
params = model.init(jax.random.PRNGKey(0), x)['params']

jax.tree_map(jnp.shape, params)

FrozenDict({
    Dense_0: {
        bias: (10,),
        kernel: (10, 10),
    },
    Dense_1: {
        bias: (10,),
        kernel: (10, 10),
    },
    head: {
        bias: (10,),
        kernel: (10, 10),
    },
})

def flattened_traversal(fn):
  """Returns function that is called with `(path, param)` instead of pytree."""
  def mask(tree):
    flat = flax.traverse_util.flatten_dict(tree)
    return flax.traverse_util.unflatten_dict(
        {k: fn(k, v) for k, v in flat.items()})
  return mask

# Specify layer-wise learning rate.
lrs = {'Dense_0': 0.1, 'Dense_1': 0.2, 'head': 0.3}
label_fn = flattened_traversal(lambda path, _: path[0])

tx = optax.multi_transform(
    {name: optax.sgd(lr) for name, lr in lrs.items()}, label_fn)

fake_grads = jax.tree_map(jnp.ones_like, params.unfreeze())
opt_state = tx.init(params.unfreeze())
updates, opt_state = tx.update(fake_grads, opt_state)

jax.tree_map(lambda x: jnp.sum(jnp.abs(x)), updates)

{'Dense_0': {'bias': DeviceArray(1., dtype=float32),
  'kernel': DeviceArray(9.999999, dtype=float32)},
 'Dense_1': {'bias': DeviceArray(2., dtype=float32),
  'kernel': DeviceArray(19.999998, dtype=float32)},
 'head': {'bias': DeviceArray(3., dtype=float32),
  'kernel': DeviceArray(29.999994, dtype=float32)}}

[dwromero]

In addition to @andsteing 's answer and after a few discussions with him, we found a glitch in the previous answer while using optax.multi_transform together with a flax.train_state.TrainState.

The previous solution leads to a ValueError:

[google3/third_party/py/jax/_src/tree_util.py](http://google3/third_party/py/jax/_src/tree_util.py) in <listcomp>(.0)
    204   """
    205   leaves, treedef = tree_flatten(tree, is_leaf)
--> 206   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    207   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
    208 

ValueError: Expected dict, got FrozenDict({ ...

This error is raised because flattened_traversal constructs a mapping that returns a dict and not a FrozenDict. To solve this issue, it is sufficient to return a FrozenDict in flattened_traversal. That is, replace the previous answer by:

def flattened_traversal(fn):
  """Returns function that is called with `(path, param)` instead of pytree."""
  def mask(tree):
    flat = flax.traverse_util.flatten_dict(tree)
    return flax.core.freeze(flax.traverse_util.unflatten_dict(
        {k: fn(k, v) for k, v in flat.items()}))
  return mask

[cgarciae]

A slightly modified version of @andsteing's flattened_traversal function is now available as flax.traverse_util.path_aware_map, you can find how to use it in the new Transfer Learning guide on the Freeze Layers section.