Repository: google-deepmind/optax

Issue:
Feature Proposal: multi-optimizer / multi-transform
Currently, to apply different optimizers to different sets of parameters, you need to construct multiple masks and chain them. For example, to optionally apply weight decay to some parameters and not to others:

```py
tx = optax.chain(
    optax.masked(optax.adamw(0.01, weight_decay=0.0), 
                 mask=partial(jax.tree_map, lambda p: p.ndim == 1)),
    optax.masked(optax.adamw(0.01, weight_decay=1e-3), 
                 mask=partial(jax.tree_map, lambda p: p.ndim != 1)),
)
```

However, this has a few issues:
1. When you have more than two groups, creating mask functions for them all becomes verbose
2. The burden of verifying the masked groups are mutually exclusive falls on the user (which can be bothersome for groups > 2)
3. This solution isn't obvious to a new user

I'd like to propose a `multi_transform` (or `multi_transformation`, as always naming is hard) that generalizes `optax.masked`:

```py
def multi_transform(transforms: Sequence[GradientTransformation], 
                    partition: Union[PyTree, Callable[[base.Params], PyTree]]):
  ... 
```
Where `partition` is a pytree where the leaves contain an index (or a function that returns such a pytree given the parameters). The index corresponds to a `GradientTransformation` in `transforms`.

With this, our above example would be:
```py
tx = optax.multi_transform(
    [optax.adamw(0.01, weight_decay=1e-3), optax.adamw(0.01, weight_decay=0.0)],
    partition=partial(jax.tree_map, lambda p: int(p.ndim == 1)))  # 0 for weight decay, 1 for no weight decay
```

This solves our above problems:
1. You only need one pytree (function)
2. The groups are mutually exclusive by design, since there's only one pytree with indices
3. The name matches what users would expect from such a feature (due to similar naming in other libraries, e.g. [Flax](https://flax.readthedocs.io/en/latest/flax.optim.html#flax.optim.MultiOptimizer)).

Minor details:
- `partition` could just be `partition_fn: Callable[[base.Params], PyTree]` (I left it as `partition` to mirror the signature of `masked`).
- `partition` could have a different name, perhaps `indices`.

cc: @mtthss @jheek @andsteing


Discussion:
I agree that it's hard to make sure that masks don't overlap thus I agree with the motivation here.

But I am concerned about the proposed API because it means you have to keep track of integer indices in the first argument to `multi_transform`. And if someone adds another optimizer somewhere in the list (not at the end) then things will silently fail.

Can there be an API that looks somehow more like the old flax MultiOptim? Where you directly couple a "mask" to an Optax optimizer, rather than going through indices?

```py
optax.multi_transform(
  partition=jax.tree_map, lambda p: optax.adamw(0.01, weight_decay=0.0) if p.ndim == 1) else 
  
Good point @avital, I like that idea. I think the example in my post is a bit too simple, here's something more realistic that illustrates the issue and your proposal for new readers (assuming we use Flax Conv and BatchNorm layers):

```py
def map_params(fn):  # Convenience function to construct mask_fns
    def map_fn(params):
        flat_params = flax.traverse_util.flatten_dict(params)
        flat_mask = {path: fn(path) for path in flat_params.keys()}
        return flax.traverse_util.unflatten_dict(flat_mask)
    return map_fn

# Suppose we want different learning rates for kernels, biases, and all Batchnorm params
# Below implementation has a silent bug: both (b) and (c) will be applied to batchnorm biases
tx = optax.chain(
    optax.masked(optax.sgd(0.1), map_params(lambda path: path[-1] == 'kernel')),  # (a)
    optax.masked(optax.sgd(0.2), map_params(lambda path: path[-1] == 'bias')),  # (b)
    optax.masked(optax.sgd(0.3), map_params(lambda path: path[-2].startswith('BatchNorm'))),  # (c)
)

# Original Proposal:
# It is _impossible_ to have overlap by design, but using indices is brittle
def partition_fn(path):
    if path[-2].startswith('BatchNorm'): return 0
    if path[-1] == 'kernel': return 1
    if path[-1] == 'bias': return 2

tx = optax.multi_transform([optax.sgd(0.3), optax.sgd(0.1), optax.sgd(0.2)], 
                           map_params(partition_fn))

# Avital's Proposal:
# This will throw an error, since multi_transform will check for overlap.
tx = optax.multi_transform(
    (optax.sgd(0.1), map_params(lambda path: path[-1] == 'kernel')),  # (a)
    (optax.sgd(0.2), map_params(lambda path: path[-1] == 'bias')),  # (b)
    (optax.sgd(0.3), map_params(lambda path: path[-2] == 'BatchNorm')),  # (c)
)
# User would correct to something like this
tx = optax.multi_transform(
    (optax.sgd(0.1), map_params(lambda path: path[-1] == 'kernel'),  # (a)
    (optax.sgd(0.2), map_params(lambda path: path[-1] == 'bias' and not path[-2].startswith('BatchNorm')),  # (b)
    (optax.sgd(0.3), map_params(lambda path: path[-2].startswith('BatchNorm')),  # (c)
)
```

Should the `multi_transform` check that users provide a transform for all parameters? i.e., check that the union of all the masks covers all parameters. This protects users against the opposite of the overlap bug: not having a transform for some parameters. If users truly want to skip optimizing some parameters, they can always use `optax.masked(optax.multi_transform(...))`.
Another simple modification to your original proposal would be to simply use `dict` instead of `Sequence` when providing the individual gradient transformations, and then to refer to them via their key, i.e.

```python
# Original Proposal - modified:
# It is _impossible_ to have overlap by design, but using indices is brittle
def partition_fn(path):
    if path[-2].startswith('BatchNorm'): return 'batch_norm'
    if path[-1] == 'kernel': return 'kernel'
    if path[-1] == 'bias': return 'bias'

tx = optax.multi_transform(dict(batch_norm=optax.sgd(0.3),
                                kernel=optax.sgd(0.1),
                                bias=optax.sgd(0.2)), 
                           map_params(partition_fn))
```
Thanks @n2cholas for kicking off the discussion!
I like the version by @andsteing where the set of transforms is a dict
and the partition function returns the keys

In particular I like that:
1) it makes it clear that the order doesn't have a meaning
2) keys feel less brittle than indices
I like @andsteing's version as well! 

Question on naming: what's a better name than `partition` for the second argument? My issue with partition is in other Python libraries, partition is a function that takes one collection and returns two (e.g. `itertools.partition`, [`haiku.data_structures.partition`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.data_structures.partition)). In our case, it's a function that simply "labels" each parameter (or subtree), so that may be confusing to a new user.

Perhaps a better name is `label_fn` or `label`? Though that also has a more common meaning in ML. Any ideas?

I also think the argument type should just be `Callable[[Params], PyTree]` instead of `Union[PyTree, Callable[[Params], PyTree]]`. (and perhaps the `optax.masked` should eventually deprecate passing in a pytree and only take callables, for consistency).

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -44,6 +44,8 @@
 from optax._src.clipping import ClipByGlobalNormState
 from optax._src.clipping import ClipState
 from optax._src.combine import chain
+from optax._src.combine import multi_transform
+from optax._src.combine import MultiTransformState
 from optax._src.constrain import keep_params_nonnegative
 from optax._src.constrain import NonNegativeParamsState
 from optax._src.constrain import zero_nans
@@ -202,6 +204,8 @@
     "matrix_inverse_pth_root",
     "masked",
     "MaskedState",
+    "multi_transform",
+    "MultiTransformState",
     "measure_valued_jacobians",
     "moving_avg_baseline",
     "multi_normal",
diff --git a/optax/_src/combine.py b/optax/_src/combine.py
--- a/optax/_src/combine.py
+++ b/optax/_src/combine.py
@@ -14,7 +14,14 @@
 # ==============================================================================
 """Flexibly compose gradient transformations."""
 
+from typing import Callable, NamedTuple, Union, Mapping, Hashable
+
+import jax
+
 from optax._src import base
+from optax._src import wrappers
+
+PyTree = wrappers.PyTree
 
 
 def chain(
@@ -51,3 +58,95 @@ def update_fn(updates, state, params=None):
     return updates, new_state
 
   return base.GradientTransformation(init_fn, update_fn)
+
+
+class MultiTransformState(NamedTuple):
+  inner_states: Mapping[Hashable, NamedTuple]
+
+
+def multi_transform(
+    transforms: Mapping[Hashable, base.GradientTransformation],
+    param_labels: Union[PyTree, Callable[[PyTree], PyTree]]
+) -> base.GradientTransformation:
+  """Partitions params and applies a different transformation to each subset.
+
+  Below is an example where we apply Adam to the weights and SGD to the biases
+  of a 2-layer neural network::
+
+    import optax
+    import jax
+    import jax.numpy as jnp
+
+    def map_nested_fn(fn):
+      '''Recursively apply `fn` to the key-value pairs of a nested dict'''
+      def map_fn(nested_dict):
+        return {k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
+                for k, v in nested_dict.items()}
+      return map_fn
+
+    params = {'linear_1': {'w': jnp.zeros((5, 6)), 'b': jnp.zeros(5)},
+              'linear_2': {'w': jnp.zeros((6, 1)), 'b': jnp.zeros(1)}}
+    gradients = jax.tree_map(jnp.ones_like, params)  # dummy gradients
+
+    label_fn = map_nested_fn(lambda k, _: k)
+    tx = optax.multi_transform({'w': optax.adam(1.0), 'b': optax.sgd(1.0)},
+                               label_fn)
+    state = tx.init(params)
+    updates, new_state = tx.update(gradients, state, params)
+    new_params = optax.apply_updates(params, updates)
+
+  Instead of providing a ``label_fn``, you may provide a PyTree of labels
+  directly.  Also, this PyTree may be a prefix of the parameters PyTree. This
+  is demonstrated in the GAN pseudocode below::
+
+    generator_params = ...
+    discriminator_params = ...
+    all_params = (generator_params, discriminator_params)
+    param_labels = ('generator', 'discriminator')
+
+    tx = optax.multi_transform(
+        {'generator': optax.adam(0.1), 'discriminator': optax.adam(0.5)},
+        param_labels)
+
+  If you would like to not optimize some parameters, you may wrap
+  ``optax.multi_transform`` with :func:`optax.masked`.
+
+  Args:
+    transforms: A mapping from labels to transformations. Each transformation
+      will be only be applied to parameters with the same label.
+    param_labels: A PyTree that is the same shape or a prefix of the
+      parameters/updates (or a function that returns one given the parameters as
+      input). The leaves of this PyTree correspond to the keys of the transforms
+      (therefore the values at the leaves must be a subset of the keys).
+
+  Returns:
+    An ``optax.GradientTransformation``.
+  """
+  def make_mask(labels, group):
+    return jax.tree_map(lambda label: label == group, labels)
+
+  def init_fn(params):
+    labels = param_labels(params) if callable(param_labels) else param_labels
+
+    label_set = set(jax.tree_leaves(labels))
+    if not label_set.issubset(transforms.keys()):
+      raise ValueError('Some parameters have no corresponding transformation.\n'
+                       f'Parameter labels: {list(sorted(label_set))} \n'
+                       f'Transforms keys: {list(sorted(transforms.keys()))} \n')
+
+    inner_states = {
+      group: wrappers.masked(tx, make_mask(labels, group)).init(params)
+      for group, tx in transforms.items()
+    }
+    return MultiTransformState(inner_states)
+
+  def update_fn(updates, state, params=None):
+    labels = param_labels(updates) if callable(param_labels) else param_labels
+    new_inner_state = {}
+    for group, tx in transforms.items():
+      masked_tx = wrappers.masked(tx, make_mask(labels, group))
+      updates, new_inner_state[group] = masked_tx.update(
+          updates, state.inner_states[group], params)
+    return updates, MultiTransformState(new_inner_state)
+
+  return base.GradientTransformation(init_fn, update_fn)
diff --git a/optax/_src/wrappers.py b/optax/_src/wrappers.py
--- a/optax/_src/wrappers.py
+++ b/optax/_src/wrappers.py
@@ -389,4 +389,3 @@ def reject_update(_):
     return updates, MaybeUpdateState(new_inner_state, state.step + 1)
 
   return base.GradientTransformation(init_fn, update_fn)
-


Test patch:
diff --git a/optax/_src/combine_test.py b/optax/_src/combine_test.py
--- a/optax/_src/combine_test.py
+++ b/optax/_src/combine_test.py
@@ -15,10 +15,13 @@
 """Tests for `combine.py`."""
 
 from absl.testing import absltest
+from absl.testing import parameterized
 
 import chex
+import jax
 import jax.numpy as jnp
 
+from optax._src import alias
 from optax._src import combine
 from optax._src import transform
 from optax._src import update
@@ -71,5 +74,77 @@ def update_fn(updates, state):
     chex.assert_tree_all_close(manual_params, chain_params, rtol=1e-4)
 
 
+def _map_keys_fn(fn):
+  def map_fn(nested_dict):
+    return {k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
+            for k, v in nested_dict.items()}
+  return map_fn
+
+
+class MultiTransformTest(chex.TestCase):
+  """Tests for the multi_transform wrapper."""
+
+  @chex.all_variants
+  @parameterized.parameters(True, False)
+  def test_multi_transform(self, use_fn):
+    params = {'a1': 1., 'b1': 2., 'z1': {'a2': 3., 'z2': {'c1': 4.}}}
+    params = jax.tree_map(jnp.asarray, params)
+    input_updates = jax.tree_map(lambda x: x / 10.0, params)
+    tx_dict = {'a': transform.scale(-1.0),
+               'b': transform.ema(0.0),  # stateful
+               'c': transform.scale(2.0)}
+    param_labels = _map_keys_fn(lambda k, _: k[0])
+    if not use_fn:
+      param_labels = param_labels(params)
+    tx = combine.multi_transform(tx_dict, param_labels)
+    update_fn = self.variant(tx.update)
+    state = self.variant(tx.init)(params)
+
+    correct_update_fn = _map_keys_fn(
+        lambda k, v: {'a': -v, 'b': v, 'c': 2.0*v}[k[0]])
+
+    updates, state = update_fn(input_updates, state, params)
+    correct_updates = correct_update_fn(input_updates)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+    # Check repeated application, this time with no params.
+    correct_updates = correct_update_fn(correct_updates)
+    updates, state = update_fn(updates, state)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+  @parameterized.parameters(list, tuple, dict)
+  def test_empty(self, container):
+    init_fn, update_fn = combine.multi_transform({0: alias.sgd(1.)},
+                                                  lambda _: 0)
+    updates, _ = update_fn(container(), init_fn(container()))
+    self.assertEqual(updates, container())
+
+  @chex.all_variants
+  @parameterized.parameters(
+    (False, False), (False, True), (True, False), (True, True))
+  def test_labels_mismatch(self, use_extra_label, use_fn):
+    # The labels from label_fn must be a subet of the keys for the tx.
+    params = {'a': 1., 'b': [2., 3.], 'c': {'d': 4., 'e': (5., 6.)}}
+    params = jax.tree_map(jnp.asarray, params)
+    label_tree = {'a': 0, 'b': [1, 0], 'c': 1}  # prefix of params
+
+    if use_extra_label:
+      label_tree['a'] = 3
+
+    transforms = {0: alias.sgd(1.),
+                  1: alias.adam(1., b1=0., b2=0.),
+                  2: transform.trace(1.0)}
+    init_fn, update_fn = combine.multi_transform(
+      transforms, (lambda _: label_tree) if use_fn else label_tree)
+
+    if use_extra_label:
+      with self.assertRaises(ValueError):
+        self.variant(init_fn)(params)
+    else:
+      state = self.variant(init_fn)(params)
+      updates = jax.tree_map(lambda x: x / 10.0, params)
+      self.variant(update_fn)(updates, state)
+
+
 if __name__ == '__main__':
   absltest.main()