Repository: huggingface/transformers

Issue:
Flax text-classification multi-optimizer incorrect
The below code is incorrect:

https://github.com/huggingface/transformers/blob/8d171628fe84bdf92ee40b5375d7265278180f14/examples/flax/text-classification/run_flax_glue.py#L173-L188

- The `dict` from `traverse_util.flatten_dict` has keys which are tuples of strings, not one long string with the path separated by periods.
- `optax.masked` applies the transformation wherever the mask is True, so the masks are flipped.
- Flax's LayerNorm calls the scale parameter `scale` not `weight`

I believe following code would be correct:

```py
# We use Optax's "masking" functionality to create a multi-optimizer, one
# with weight decay and the other without. Each sub-optimizer will be applied
# wherever the mask is True
decay_path = lambda p: p[-1] != "bias" and p[-2:] != ("LayerNorm", "scale")  # noqa: E731

tx = optax.chain(
    optax.masked(adamw(0.0), mask=traverse(lambda path, _: not decay_path(path))),
    optax.masked(adamw(weight_decay), mask=traverse(lambda path, _: decay_path(path))),
)
```

But, for the networks in that example, the parameters that shouldn't be decayed all have 1 dimension (biases and layernorm scales). L173-L188 can therefore be simplified to:

```py
# We use Optax's "masking" functionality to create a multi-optimizer, one
# with weight decay and the other without. Each sub-optimizer will be applied
# wherever the mask is True. The bias parameters and LayerNorm scale
# parameters should not be decayed, and these are the only parameters
# with 1 dimension.
tx = optax.chain(
    optax.masked(adamw(0.0), mask=partial(jax.tree_map, lambda p: p.ndim == 1)),
    optax.masked(adamw(weight_decay), mask=partial(jax.tree_map, lambda p: p.ndim != 1)),
)
```

Though the latter solution is simpler, perhaps the first one is better since it illustrates a more general way for users to construct a multi-optimizer. I'd be happy to open a PR with either.

cc: @marcvanzee @patrickvonplaten 


Discussion:
Actually `optax.adamw` can take a `mask` directly to mask weight decay (since that's pretty common):

```py
tx = optax.adamw(learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay, 
                 mask=partial(jax.tree_map, lambda p: p.ndim != 1))
```
Hey @n2cholas,

Thanks for diving into this! 

I think a nice solution would be:

```python
tx = optax.adamw(learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay, 
                 mask=partial(jax.tree_map, lambda p: p[-1] != "bias" and p[-2:] != ("LayerNorm", "scale")))
```

This way, it's consistent, but we also show the user clearly that `bias` and `LayerNorm` are not included in the weight decay. It would be great if you could open a PR for this :-) I can re-run the experiments afterward

Great, will do!

Resolving patch:
diff --git a/examples/flax/text-classification/run_flax_glue.py b/examples/flax/text-classification/run_flax_glue.py
--- a/examples/flax/text-classification/run_flax_glue.py
+++ b/examples/flax/text-classification/run_flax_glue.py
@@ -165,25 +165,17 @@ class TrainState(train_state.TrainState):
         logits_fn: Callable = struct.field(pytree_node=False)
         loss_fn: Callable = struct.field(pytree_node=False)
 
-    # Creates a multi-optimizer consisting of two "Adam with weight decay" optimizers.
-    def adamw(decay):
-        return optax.adamw(learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=decay)
-
-    def traverse(fn):
-        def mask(data):
-            flat = traverse_util.flatten_dict(data)
-            return traverse_util.unflatten_dict({k: fn(k, v) for k, v in flat.items()})
-
-        return mask
-
-    # We use Optax's "masking" functionality to create a multi-optimizer, one
-    # with weight decay and the other without. Note masking means the optimizer
-    # will ignore these paths.
-    decay_path = lambda p: not any(x in p for x in ["bias", "LayerNorm.weight"])  # noqa: E731
-
-    tx = optax.chain(
-        optax.masked(adamw(0.0), mask=traverse(lambda path, _: decay_path(path))),
-        optax.masked(adamw(weight_decay), mask=traverse(lambda path, _: not decay_path(path))),
+    # We use Optax's "masking" functionality to not apply weight decay
+    # to bias and LayerNorm scale parameters. decay_mask_fn returns a
+    # mask boolean with the same structure as the parameters.
+    # The mask is True for parameters that should be decayed.
+    def decay_mask_fn(params):
+        flat_params = traverse_util.flatten_dict(params)
+        flat_mask = {path: (path[-1] != "bias" and path[-2:] != ("LayerNorm", "scale")) for path in flat_params}
+        return traverse_util.unflatten_dict(flat_mask)
+
+    tx = optax.adamw(
+        learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay, mask=decay_mask_fn
     )
 
     if is_regression:


Test patch: