Coupled weight decay mixin (AdamC, MuonC) for PyTorch

coupled_weight_decay.py

"""
Provide a mixin for coupling weight decay, e.g. AdamC, as proposed in A. Defazio, “Why Gradients Rapidly Increase
Near the End of Training,” (arXiv:2506.02285).

Use e.g. as `class AdamC(CoupledWeightDecayMixin, optim.AdamW): pass` (see below); then provide `{"normalized": True}`
for parameter groups where the coupling/correction factor should be applied (norm layers usually).
"""

from typing import Any, Mapping
from torch import no_grad, optim


class CoupledWeightDecayMixin:

    def _get(self, group: Mapping[str, Any], param: str, on_missing: type[BaseException] | Any = ValueError) -> Any:
        # Try group → try `self.defaults` → raise `on_missing` (if it is an exception type) or return it as fallback
        if param in group:
            return group[param]
        if param in self.defaults:
            return self.defaults[param]
        if isinstance(on_missing, type) and issubclass(on_missing, BaseException):
            raise on_missing(param)
        return on_missing

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Store reference lr for normalized groups; store None for others
        param_groups = self.param_groups
        self._ref_lrs = [(self._get(g, "lr") if self._get(g, "normalized", False) else None) for g in param_groups]
        self._pre_wds = []

    @no_grad()
    def _correct_wd_values(self):
        # Keep preset weight decay, then apply correction factor
        for group, ref_lr in zip(self.param_groups, self._ref_lrs):
            self._pre_wds.append(group_wd := self._get(group, "weight_decay"))
            if ref_lr is not None:
                group["weight_decay"] = group_wd * (self._get(group, "lr") / ref_lr)

    @no_grad()
    def _reset_wd_values(self):
        # Reset weight decay to preset value
        for group, ref_lr, pre_wd in zip(self.param_groups, self._ref_lrs, self._pre_wds):
            if ref_lr is not None:
                group["weight_decay"] = pre_wd
        self._pre_wds.clear()

    def step(self, closure=None):
        self._correct_wd_values()
        try:
            result = super().step(closure)
        finally:
            self._reset_wd_values()
        return result

    def state_dict(self):
        return super().state_dict() | {"mixin_ref_lrs": self._ref_lrs}

    def load_state_dict(self, state_dict):
        self._ref_lrs = state_dict.pop("mixin_ref_lrs")
        return super().load_state_dict(state_dict)


class AdamC(CoupledWeightDecayMixin, optim.AdamW):
    pass


class MuonC(CoupledWeightDecayMixin, optim.Muon):
    pass