Simon Pezold spezold
spezold
/ coupled_weight_decay.py
Last active
July 20, 2026 13:03
Coupled weight decay mixin (AdamC, MuonC) for PyTorch
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """ | |
| 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 |
OlderNewer