Repository: google-deepmind/optax

Issue:
Stateless Transformation
Right now, there's some boilerplate for defining simple gradient transformations. If a user wants to implement custom weight decay, clipping, constraints, etc without any state, they still have to define a nested function with an `init` that does nothing, and handle the empty state.

I think it'd be convenient to provide a `stateless` transformation that accepts a function to apply to the updates and params. We can also do the `jax.tree_multimap` for the user by default.

```python
weight_decay = optax.stateless(lambda g, p: g + 0.1 * p)
```

If the user wants to define a function that does the tree_multimap themselves:

```python
def my_function(updates, params):
    return jax.tree_multimap(lambda g, p: ..., updates, params)

optim = optax.stateless(my_function, on_leaves=False)
```

In my view, this is a very clean way to implement simple stateless transformations. I think the `on_leaves` argument could use a better name, though. I'd be happy to implement this if it sounds reasonable.


Discussion:
This is an interesting suggestion. Lets discuss our options here a bit more.

Some initial thoughts:

One concern for me is that `optax.stateless` suggests the handling of the empty state,
but it does not immediately suggest that it would do the tree mapping for the user as well.

While `on_leaves` allows me to disable it,
it's weird to have to disable something that it's not clear was happening in the first place.

Not the end of the world ofc, maybe we just need a different name than `stateless`?

An alternative would be to disentangle the two aspects: statelessness and tree_mapping,
but maybe `optax.stateless(optax.map(fn)` is a bit two cumbersome.

In any case we would probably want to make this (or these) decorators so that you can do things like
```
@optax.stateless
@optax.tree_map  # if we don't incorporate it in stateless directly
def user_fn(g, p):
  return g - 0.1*p
```



Good point @mtthss, I agree it's more transparent to make the `tree_multimap` opt-in instead of the default. I think making `on_leaves=False` by default is a good way to achieve this.

Letting the user do the `tree_multimap` themselves could work too (not `tree_map`, since we're tree-mapping over both updates and parameters). I don't think the following is _too_ bad

```python
tx = optax.stateless(optax.tree_multimap(lambda g, p: g - 0.1*p))
```

Though perhaps a bit clunky, like you said. It's not much better than the user just doing that in native JAX:

```python
from functools import partial
tx = optax.stateless(partial(jax.tree_multimap, lambda g, p: g - 0.1*p))
```

Also, I don't think `optax.stateless` can be a decorator, since it returns a `GradientTransformation` and not a callable.

<hr>

A bit tangential, but I think it's considering alternative use cases for `optax.tree_multimap` before creating the API. In particular, it may be useful for optax to provide a specialized multimap for transformation `update_fn`s that lets you write it for each parameter instead of for the pytree. For example, the `update_fn` for [`optax.trace`](https://github.com/deepmind/optax/blob/master/optax/_src/transform.py#L65-L72) becomes much cleaner:

```python
  @optax.multimap_update  # different name so users don't confuse it with a normal multimap
  def update_fn(update, state, param=None):
    del params
    new_trace = update + decay * state.trace
    new_update = update + decay * new_trace if nesterov else new_trace
    new_trace = new_trace.astype(accumulator_dtype)
    return updates, TraceState(trace=new_trace)
```
This `multimap_update` would have to do a few things:

- Not map over `params` if it is None
- `tree_transpose` the state so it is a pytree of `OptStates` (instead of an `OptState` of pytrees).
- `tree_transpose` the new returned state back into an `OptState` of pytrees.

I can open a new issue for this if you want to further discuss it.
I really like the `optax.stateless` suggestion! As you said, it would make defining simple transformations much easier.

An additional consideration would be how hyperparameters should be handled, e.g. for the `weight_decay` parameter in `add_decayed_weights`. We would need a function that creates a gradient transformation so that we can pass it e.g. to `inject_hyperparams`.

With your current suggestion I think this would be:

```
def add_decayed_weights(weight_decay):
  return optax.stateless(lambda g, p: g + weight_decay * p, on_leaves=True)
  
opt = add_decayed_weights(weight_decay=0.1)
```

If we make optax.stateless a decorater that directly returns a function that creates a gradient transformation (passing all the hyperparameters through) it would be:

```
add_decayed_weights = optax.stateless(lambda g, p, weight_decay: g + weight_decay * p, on_leaves=True)
opt = add_decayed_weights(weight_decay=0.1)
```

(We'd probably want a different name for `stateless` in this case.)

The latter is a bit shorter and would work as a simple decorator, but I think the former (your suggestion) is easier to understand and more readable.

So yes, I think `optax.stateless` as you proposed with `on_leaves=False` by default sounds great to me!

Awesome, thanks @mkunesch! I'll get started on a PR.
Great, thanks @n2cholas! And apologies for the slow turnaround on this issue.
Thanks a lot @n2cholas! Also, I didn't get back to your separate point about `@optax.multimap_update`. Thanks a lot for the suggestion!

I agree that this would simplify quite a few optimizers, but I think there are two potential downsides:

1. it could cause confusion when update functions across optax or user code sometimes operate on trees and sometimes on leaves.
2. for an arbitrary state I think the mapping could become confusing (as it differs between updates and state).

In my opinion, for stateless transforms the benefits of being able to define an optimizer in just one line outweigh the potential downside of the confusion (especially since there is no state so point 1 above doesn't apply). However, in the general case for optimizers that have non-trivial state where we have already added all the boiler plate of init and update functions I am not sure whether the benefit of being able to define the update functions on leaves outweigh the downside.

In short, I agree that it would be great to have the `on_leaves` option for stateless transforms, but I'm not sure we should also have `@optax.multimap_update` in general. What do you think?
+1 to @mkunesch analysis
Hi @n2cholas just checking if there is any further action needed on our side to proceed with the PR?
No more action needed on your side @mtthss, I'll work on this over the weekend!

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -37,6 +37,8 @@
 from optax._src.base import Params
 from optax._src.base import Schedule
 from optax._src.base import set_to_zero
+from optax._src.base import stateless
+from optax._src.base import stateless_with_tree_map
 from optax._src.base import TransformInitFn
 from optax._src.base import TransformUpdateFn
 from optax._src.base import Updates
@@ -277,6 +279,8 @@
     "sm3",
     "smooth_labels",
     "softmax_cross_entropy",
+    "stateless",
+    "stateless_with_tree_map",
     "trace",
     "TraceState",
     "TransformInitFn",
diff --git a/optax/_src/base.py b/optax/_src/base.py
--- a/optax/_src/base.py
+++ b/optax/_src/base.py
@@ -138,3 +138,61 @@ def update_fn(updates, state, params=None):
     return jax.tree_map(jnp.zeros_like, updates), state
 
   return GradientTransformation(init_fn, update_fn)
+
+
+def stateless(
+    f: Callable[[Updates, Optional[Params]], Updates],
+) -> GradientTransformation:
+  """Creates a stateless transformation from an update-like function.
+
+  This wrapper eliminates the boilerplate needed to create a transformation that
+  does not require saved state between iterations.
+
+  Args:
+    f: Update function that takes in updates (e.g. gradients) and parameters
+      and returns updates. The parameters may be `None`.
+
+  Returns:
+    An `optax.GradientTransformation`.
+  """
+
+  def init_fn(_):
+    return EmptyState()
+
+  def update_fn(updates, state, params=None):
+    del state
+    return f(updates, params), EmptyState()
+
+  return GradientTransformation(init_fn, update_fn)
+
+
+def stateless_with_tree_map(
+    f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
+) -> GradientTransformation:
+  """Creates a stateless transformation from an update-like function for arrays.
+
+  This wrapper eliminates the boilerplate needed to create a transformation that
+  does not require saved state between iterations, just like optax.stateless.
+  In addition, this function will apply the tree_map over update/params for you.
+
+  Args:
+    f: Update function that takes in an update array (e.g. gradients) and
+      parameter array and returns an update array. The parameter array may be
+      `None`.
+
+  Returns:
+    An `optax.GradientTransformation`.
+  """
+
+  def init_fn(_):
+    return EmptyState()
+
+  def update_fn(updates, state, params=None):
+    del state
+    if params is not None:
+      return jax.tree_map(f, updates, params), EmptyState()
+    else:
+      f_ = lambda u: f(u, None)
+      return jax.tree_map(f_, updates), EmptyState()
+
+  return GradientTransformation(init_fn, update_fn)


Test patch:
diff --git a/optax/_src/base_test.py b/optax/_src/base_test.py
--- a/optax/_src/base_test.py
+++ b/optax/_src/base_test.py
@@ -15,7 +15,10 @@
 """Tests for base.py."""
 
 from absl.testing import absltest
+
 import chex
+import jax
+import jax.numpy as jnp
 import numpy as np
 
 from optax._src import base
@@ -59,5 +62,78 @@ def test_set_to_zero_is_stateless(self):
         self.variant(base.set_to_zero().init)(params=None), base.EmptyState())
 
 
+class StatelessTest(chex.TestCase):
+  """Tests for the stateless transformation."""
+
+  @chex.all_variants
+  def test_stateless(self):
+    params = {'a': jnp.zeros((1, 2)), 'b': jnp.ones((1,))}
+    updates = {'a': jnp.ones((1, 2)), 'b': jnp.full((1,), 2.0)}
+
+    @base.stateless
+    def opt(g, p):
+      return jax.tree_map(lambda g_, p_: g_ + 0.1 * p_, g, p)
+
+    state = opt.init(params)
+    update_fn = self.variant(opt.update)
+    new_updates, _ = update_fn(updates, state, params)
+    expected_updates = {'a': jnp.ones((1, 2)), 'b': jnp.array([2.1])}
+    chex.assert_trees_all_close(new_updates, expected_updates)
+
+  @chex.all_variants
+  def test_stateless_no_params(self):
+    updates = {'linear': jnp.full((5, 3), 3.0)}
+
+    @base.stateless
+    def opt(g, _):
+      return jax.tree_map(lambda g_: g_ * 2, g)
+
+    state = opt.init(None)
+    update_fn = self.variant(opt.update)
+    new_updates, _ = update_fn(updates, state)
+    expected_updates = {'linear': jnp.full((5, 3), 6.0)}
+    chex.assert_trees_all_close(new_updates, expected_updates)
+
+  def test_init_returns_emptystate(self):
+    def weight_decay(g, p):
+      return jax.tree_map(lambda g_, p_: g_ + 0.1 * p_, g, p)
+
+    opt = base.stateless(weight_decay)
+    state = opt.init(None)
+    self.assertIsInstance(state, base.EmptyState)
+
+
+class StatelessWithTreeMapTest(chex.TestCase):
+  """Tests for the stateless_with_tree_map transformation."""
+
+  @chex.all_variants
+  def test_stateless_with_tree_map(self):
+    params = {'a': jnp.zeros((1, 2)), 'b': jnp.ones((1,))}
+    updates = {'a': jnp.ones((1, 2)), 'b': jnp.full((1,), 2.0)}
+
+    opt = base.stateless_with_tree_map(lambda g, p: g + 0.1 * p)
+    state = opt.init(params)
+    update_fn = self.variant(opt.update)
+    new_updates, _ = update_fn(updates, state, params)
+    expected_updates = {'a': jnp.ones((1, 2)), 'b': jnp.array([2.1])}
+    chex.assert_trees_all_close(new_updates, expected_updates)
+
+  @chex.all_variants
+  def test_stateless_with_tree_map_no_params(self):
+    updates = {'linear': jnp.full((5, 3), 3.0)}
+
+    opt = base.stateless_with_tree_map(lambda g, _: g * 2.0)
+    state = opt.init(None)
+    update_fn = self.variant(opt.update)
+    new_updates, _ = update_fn(updates, state)
+    expected_updates = {'linear': jnp.full((5, 3), 6.0)}
+    chex.assert_trees_all_close(new_updates, expected_updates)
+
+  def test_init_returns_emptystate(self):
+    opt = base.stateless_with_tree_map(lambda g, p: g + 0.1 * p)
+    state = opt.init(None)
+    self.assertIsInstance(state, base.EmptyState)
+
+
 if __name__ == '__main__':
   absltest.main()