Repository: google-deepmind/optax

Issue:
Option to not decay bias with additive_weight_decay
Currently, `additive_weight_decay` will decay all the parameters. [Jia et al. 2018](https://arxiv.org/abs/1807.11205) show it is beneficial not to decay the bias parameters (and only decay the weights). Many training examples implement this too, such as [Flax's Imagenet example](https://github.com/google/flax/blob/master/examples/imagenet/train.py#L119-L121). 

One way to support this is to add a `decay_bias: bool` argument to `additive_weight_decay`, then within the decay update:

```python
updates = jax.tree_multimap(
    lambda g, p: g + weight_decay * p if p.ndim > 1 or decay_bias else g, 
    updates, params)
```

As far as I'm aware for common NN layers, the bias always has one dimension, so checking if `p.ndim > 1` is sufficient (please correct me if I'm wrong).

I'd be happy to contribute this if it sounds reasonable.

EDIT: I realized batch norm scale parameters also have only one dimension, so this filter would wrongly include those. However, many training pipelines do not regularize batchnorm scale/bias (e.g. the [reference imagenet implementation from mlperf](https://github.com/mlcommons/training/blob/4b0a874e2eb4c4d7c5e9ff299e99c4324d97acd8/image_classification/tensorflow/official/resnet/resnet_run_loop.py#L314-L329)). 


Discussion:
This seems great, maybe to be more clear you could name the argument `decay_scalars: bool`
I don't think `decay_scalars` is the right term, since we want to not decay parameters with only one dimension (i.e. vectors). `decay_one_dim_params: bool` would be accurate, but is a mouthful 😅. Maybe `decay_vectors: bool`? (though that isn't entirely clear).
How about we just allow to provide a mask instead?

The new constructor argument could be an (optional) py-treee of the same structure as params and grads,
that specifies for each element in the tree whether or not to decay weights.

In the doc string we could also show, as an example, creating the mask by checking `p.ndim > 1`
I like that! What do you think of making it a [wrapper](https://github.com/deepmind/optax/blob/master/optax/_src/wrappers.py) instead so it's more generic? The signature could be:

```python
def masked_apply(
    inner: transform.GradientTransformation,
    mask: PyTree
) -> transform.GradientTransformation:
```

Another possibility is to make an `apply_if` wrapper, which accepts a predicate. So the signature could be:
```python
def apply_if(
    inner: transform.GradientTransformation,
    predicate: Callable[[jnp.array, jnp.array], bool]  # (gradient, param) -> bool
) -> transform.GradientTransformation:
```
And potential usage:
```python
apply_if(additive_weight_decay, lambda grad, param: param.ndim > 1)
```
I think making it a wrapper sounds great.

I prefer your first proposed API as it allows the user more freedom in **how* the mask is constructed.
For instance in many neural networks libraries in JAX the params structure include some form of scoping, and it is common to filter based on the names instead of the shapes alone.

Maybe name it just `mask` instead of `masked_apply`
Right, that makes sense. I also agree the name `mask` is cleaner. I'll create a PR for this soon!

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -91,6 +91,8 @@
 from optax._src.wrappers import lookahead
 from optax._src.wrappers import LookaheadParams
 from optax._src.wrappers import LookaheadState
+from optax._src.wrappers import masked
+from optax._src.wrappers import MaskedState
 from optax._src.wrappers import MultiSteps
 from optax._src.wrappers import MultiStepsState
 
@@ -135,6 +137,8 @@
     "LookaheadParams",
     "LookaheadState",
     "linear_onecycle_schedule",
+    "masked",
+    "MaskedState",
     "measure_valued_jacobians",
     "moving_avg_baseline",
     "multi_normal",
diff --git a/optax/_src/wrappers.py b/optax/_src/wrappers.py
--- a/optax/_src/wrappers.py
+++ b/optax/_src/wrappers.py
@@ -431,3 +431,67 @@ def _lookahead_update(
       last_difference)
 
   return LookaheadParams(fast=fast_updates, slow=slow_updates)
+
+
+class MaskedState(NamedTuple):
+  """Maintains inner transform state for masked transformations."""
+  inner_state: Any
+
+
+def masked(inner: transform.GradientTransformation,
+           mask: Any) -> transform.GradientTransformation:
+  """Mask updates so only a subset of them are computed.
+
+  For example, it is common to skip weight decay for BatchNorm scale and all
+  bias parameters. In many networks, these are the only parameters with only
+  one dimension. So, you may mask these out as follows:
+
+  ```
+  mask = jax.tree_util.tree_map(lambda x: x.ndim != 1, params)
+  custom_weight_decay = optax.masked(optax.add_decayed_weights(0.001), mask)
+  ```
+
+  For the `inner` transform, state will only be stored for the parameters that
+  have a mask value of `True`.
+
+  Args:
+    inner: Inner transformation to mask.
+    mask: A PyTree with the same structure as the parameters or is a prefix of
+      the parameter PyTree. The leaves should be booleans which are `True` for
+      leaves/subtrees you want to apply the transformation to, and `False` for
+      those you want to skip.
+
+  Returns:
+    New GradientTransformation wrapping `inner`.
+  """
+  flat_mask, treedef = tree_flatten(mask)
+
+  def init_fn(params):
+    flat_params = treedef.flatten_up_to(params)
+    masked_params = [p for p, m in zip(flat_params, flat_mask) if m]
+    return MaskedState(inner_state=inner.init(masked_params))
+
+  def update_fn(updates, state, params=None):
+    # Flatten then filter out updates/params not in the mask:
+    flat_updates = treedef.flatten_up_to(updates)
+    masked_updates = [g for g, m in zip(flat_updates, flat_mask) if m]
+
+    if params:
+      flat_params = treedef.flatten_up_to(params)
+      masked_params = [p for p, m in zip(flat_params, flat_mask) if m]
+    else:
+      masked_params = None
+
+    # Compute new updates
+    new_masked_updates, new_inner_state = inner.update(
+        masked_updates, state.inner_state, masked_params)
+
+    # Incorporate new_masked_updates into flat_updates, then unflatten
+    new_masked_updates = iter(new_masked_updates)
+    for i, m in enumerate(flat_mask):
+      if m: flat_updates[i] = next(new_masked_updates)
+
+    new_updates = treedef.unflatten(flat_updates)
+    return new_updates, MaskedState(inner_state=new_inner_state)
+
+  return transform.GradientTransformation(init_fn, update_fn)


Test patch:
diff --git a/optax/_src/wrappers_test.py b/optax/_src/wrappers_test.py
--- a/optax/_src/wrappers_test.py
+++ b/optax/_src/wrappers_test.py
@@ -330,5 +330,118 @@ def test_lookahead_edge_cases(self, sync_period, slow_step_size,
     chex.assert_tree_all_close(final_params.slow, correct_result)
 
 
+class MaskedTest(chex.TestCase):
+  """Tests for the masked wrapper."""
+
+  @chex.all_variants
+  @parameterized.parameters(
+      lambda: transform.scale(-1.),  # stateless test
+      _build_sgd  # stateful test
+  )
+  def test_masked(self, opt_builder):
+    mask = {'a': True,
+            'b': [False, True],
+            'c': {'d': True, 'e': (False, True)}}
+    params = {'a': 1, 'b': [2, 3], 'c': {'d': 4, 'e': (5, 6)}}
+    input_updates = jax.tree_util.tree_map(lambda x: x/10., params)
+
+    # Negate the updates wherever the mask is True
+    def masked_negate(updates):
+      return jax.tree_util.tree_multimap(
+          lambda upd, m: -upd if m else upd, updates, mask)
+    correct_updates = masked_negate(input_updates)
+
+    init_fn, update_fn = wrappers.masked(opt_builder(), mask)
+    update_fn = self.variant(update_fn)
+    state = init_fn(params)
+    updates, state = update_fn(input_updates, state, params)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+    # Check repeated application, this time with no params.
+    correct_updates = masked_negate(correct_updates)
+    updates, state = update_fn(updates, state)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+  @chex.all_variants
+  @parameterized.parameters(
+      lambda: transform.scale(-1.),  # stateless test
+      _build_sgd  # stateful test
+  )
+  def test_prefix_mask(self, opt_builder):
+    """Test when the mask is a prefix of the updates PyTree."""
+    mask = {'a': True, 'b': False, 'c': {'d': False, 'e': True}}
+    params = {'a': 1, 'b': {'f': 2}, 'c': {'d': 3, 'e': ([4, 5], 6)}}
+    input_updates = jax.tree_util.tree_map(lambda x: x/10., params)
+
+    # Negate the updates wherever the mask (or mask parent) is True
+    def _masked_sgd_on_updates(m, upd):
+      return jax.tree_util.tree_map(lambda x: -x, upd) if m else upd
+    correct_updates = jax.tree_util.tree_multimap(
+        _masked_sgd_on_updates, mask, input_updates)
+
+    init_fn, update_fn = wrappers.masked(opt_builder(), mask)
+    update_fn = self.variant(update_fn)
+    state = init_fn(params)
+    updates, state = update_fn(input_updates, state, params)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+    # Check repeated application, this time with no params.
+    correct_updates = jax.tree_util.tree_multimap(
+        _masked_sgd_on_updates, mask, correct_updates)
+    updates, state = update_fn(updates, state)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+  @chex.all_variants
+  def test_update_requires_params(self):
+    weight_decay = 0.1
+    mask = {'a': True,
+            'b': [False, True],
+            'c': {'d': True, 'e': (False, True)}}
+    params = {'a': 1, 'b': [2, 3], 'c': {'d': 4, 'e': (5, 6)}}
+    input_updates = jax.tree_util.tree_map(lambda x: x/10., params)
+
+    correct_updates = jax.tree_util.tree_multimap(
+        lambda m, u, p: u + weight_decay * p if m else u,
+        mask, input_updates, params)
+
+    init_fn, update_fn = wrappers.masked(
+        transform.additive_weight_decay(weight_decay), mask)
+    update_fn = self.variant(update_fn)
+
+    state = init_fn(params)
+    updates, state = update_fn(input_updates, state, params)
+    chex.assert_tree_all_close(updates, correct_updates)
+
+    params = update.apply_updates(params, updates)
+
+    # Test repeated application
+    new_correct_updates = jax.tree_util.tree_multimap(
+        lambda m, u, p: u + weight_decay * p if m else u,
+        mask, correct_updates, params)
+    updates, state = update_fn(correct_updates, state, params)
+    chex.assert_tree_all_close(updates, new_correct_updates)
+
+  @parameterized.parameters(list, tuple, dict)
+  def test_empty(self, container):
+    init_fn, update_fn = wrappers.masked(_build_sgd(), container())
+    update_fn(container(), init_fn(container()))
+
+  @parameterized.parameters(True, False)
+  def test_tree_mismatch_fails(self, extra_key_in_mask):
+    mask = {'a': True,
+            'b': [False, True],
+            'c': {'d': True, 'e': (False, True)}}
+    params = {'a': 1, 'b': [2, 3], 'c': {'d': 4, 'e': (5, 6)}}
+
+    if extra_key_in_mask:
+      mask['c']['extra'] = True
+    else:
+      params['c']['extra'] = 7
+
+    init_fn = wrappers.masked(_build_sgd(), mask)[0]
+    with self.assertRaises(ValueError):
+      init_fn(params)
+
+
 if __name__ == '__main__':
   absltest.main()