From d47bc8d8974267c9a35ec7597fc1aa24e96f6210 Mon Sep 17 00:00:00 2001 From: James Martens Date: Sun, 30 Aug 2026 07:44:42 -0700 Subject: [PATCH] Add `weight_decay_mask` to `dadapt_adamw`, and fix it for `prodigy`. `prodigy`'s `weight_decay_mask` previously required a mask whose tree structure exactly matched the params. Masks that were a prefix of the params tree, or a plain `True`/`False`, raised a pytree structure error even though the docstring documented prefix masks as supported. Both optimizers now route the masking through `optax.add_decayed_weights`. PiperOrigin-RevId: 973454596 --- optax/contrib/_dadapt_adamw.py | 25 ++++++++++++++++---- optax/contrib/_prodigy.py | 43 +++++++++++----------------------- 2 files changed, 35 insertions(+), 33 deletions(-) diff --git a/optax/contrib/_dadapt_adamw.py b/optax/contrib/_dadapt_adamw.py index 2fc90c40e..a24e89d0f 100644 --- a/optax/contrib/_dadapt_adamw.py +++ b/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, diff --git a/optax/contrib/_prodigy.py b/optax/contrib/_prodigy.py index 94482f369..98e2a3ba9 100644 --- a/optax/contrib/_prodigy.py +++ b/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,