Add weight_decay_mask to dadapt_adamw, and fix it for prodigy.

optax/contrib/_dadapt_adamw.py

@@ -18,11 +18,14 @@
 D-Adaptation" (https://arxiv.org/abs/2301.07733) by Aaron Defazio and Konstantin
 Mishchenko (ICML 2023 Outstanding Paper award).
 """
-from typing import NamedTuple, Optional
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Union
+
 import jax
 import jax.numpy as jnp
 from optax._src import base
 from optax._src import numerics
+from optax.transforms import _adding
 import optax.tree
 
 
@@ -45,6 +48,9 @@ def dadapt_adamw(
     eps: jax.typing.ArrayLike = 1e-8,
     estim_lr0: jax.typing.ArrayLike = 1e-6,
     weight_decay: jax.typing.ArrayLike = 0.0,
+    weight_decay_mask: Optional[
+        Union[Any, Callable[[base.Params], Any]]
+    ] = None,
 ) -> base.GradientTransformationExtraArgs:
   """Learning rate free AdamW by D-Adaptation.
 
@@ -62,6 +68,12 @@ def dadapt_adamw(
     estim_lr0: Initial (under-)estimate of the learning rate.
     weight_decay: AdamW style weight-decay. To use Regular Adam decay, chain
       with add_decayed_weights.
+    weight_decay_mask: A tree with same structure as (or a prefix of) the params
+      PyTree, or a Callable that returns such a pytree given the params/updates.
+      The leaves should be booleans, ``True`` for leaves/subtrees you want to
+      apply the weight decay to, and ``False`` for those you want to skip. The
+      mask must be static for the gradient transformation to be jit-compilable.
+      Passed through to :func:`optax.add_decayed_weights`.
 
   Returns:
     The corresponding :class:`optax.GradientTransformation`.
@@ -130,12 +142,17 @@ def update_fn(
     d_estimate = numerator_weighted / ((1 - sb2) * grad_sum_l1)
     estim_lr = jnp.maximum(state.estim_lr, d_estimate)
     p_update = jax.tree.map(
-        # pyrefly: ignore[unsupported-operation]
-        lambda ea, eas, p: -weight_decay * dlr * p - ea / (jnp.sqrt(eas) + eps),
+        lambda ea, eas: -ea / (jnp.sqrt(eas) + eps),
         exp_avg,
         exp_avg_sq,
-        params,
     )
+    # Unlike in prodigy, `dlr` cannot be factored out of `p_update` (`exp_avg`
+    # already accumulates it), so the decay scale is dynamic. It is negated
+    # because `p_update` is a descent direction.
+    decay_tx = _adding.add_decayed_weights(
+        -jnp.asarray(weight_decay * dlr), weight_decay_mask
+    )
+    p_update, _ = decay_tx.update(p_update, decay_tx.init(params), params)
     new_state = DAdaptAdamWState(
         exp_avg,
         exp_avg_sq,

optax/contrib/_prodigy.py

@@ -26,6 +26,7 @@
 import jax.numpy as jnp
 from optax._src import base
 from optax._src import numerics
+from optax.transforms import _adding
 import optax.tree
 
 
@@ -81,8 +82,9 @@ def prodigy(
     weight_decay_mask: A tree with same structure as (or a prefix of) the params
       PyTree, or a Callable that returns such a pytree given the params/updates.
       The leaves should be booleans, ``True`` for leaves/subtrees you want to
-      apply the weight decay to, and ``False`` for those you want to skip. Note
-      that the Adam gradient transformations are applied to all parameters.
+      apply the weight decay to, and ``False`` for those you want to skip. The
+      mask must be static for the gradient transformation to be jit-compilable.
+      Passed through to :func:`optax.add_decayed_weights`.
 
   Returns:
     A :class:`optax.GradientTransformation` object.
@@ -94,6 +96,10 @@ def prodigy(
   beta1, beta2 = betas
   if beta3 is None:
     beta3 = beta2**0.5
+  # Applies (masked) decoupled weight decay. This is stateless, so it can be
+  # applied inline within `update_fn`.
+  # pyrefly: ignore[bad-argument-type]
+  decay_tx = _adding.add_decayed_weights(weight_decay, weight_decay_mask)
 
   def init_fn(params: base.Params) -> ProdigyState:
     # Define state parameters with the lowest dtype of the parameters to avoid
@@ -163,36 +169,15 @@ def update_fn(
     lr_estimate = estim_lr_coef * numerator_weighted / denominator
     estim_lr = jnp.maximum(state.estim_lr, lr_estimate)
 
-    p_update = jax.tree.map(
-        lambda ea, eas: -dlr * ea / (jnp.sqrt(eas) + estim_lr * eps),
+    # Factor out `-dlr` so that the decay scale is the plain `weight_decay`,
+    # letting us reuse `add_decayed_weights` (which also handles the masking).
+    ascent_dir = jax.tree.map(
+        lambda ea, eas: ea / (jnp.sqrt(eas) + estim_lr * eps),
         exp_avg,
         exp_avg_sq,
     )
-
-    # Resolve weight decay mask.
-    if weight_decay_mask is not None:
-      # pyrefly: ignore[not-callable]
-      mask_tree = (
-          weight_decay_mask(params)
-          if callable(weight_decay_mask)
-          else weight_decay_mask
-      )
-      p_update = jax.tree.map(
-          lambda u, p, m: jnp.where(
-              # pyrefly: ignore[unsupported-operation]
-              m, u - weight_decay * dlr * p, u
-          ),
-          p_update,
-          params,
-          mask_tree,
-      )
-    else:
-      p_update = jax.tree.map(
-          # pyrefly: ignore[unsupported-operation]
-          lambda u, p: u - weight_decay * dlr * p,
-          p_update,
-          params,
-      )
+    ascent_dir, _ = decay_tx.update(ascent_dir, decay_tx.init(params), params)
+    p_update = optax.tree.scale(-dlr, ascent_dir)
 
     new_state = ProdigyState(
         exp_avg,