Cautious Weight Decay

TL;DR

Addresses the mismatch between gradient-driven updates and weight decay by proposing cautious weight decay, a sign-aligned, optimizer-agnostic modification. Theoretical contributions include Lyapunov-based stability showing unbiased optimization toward the stationary manifold and a sliding-mode interpretation that yields a locally Pareto-optimal frontier; discrete-time analysis provides a standard $O(T^{-1/2})$ convergence rate for Adam with CWD. Empirically, CWD improves final loss and zero-shot accuracy on large-scale language models and ImageNet across multiple optimizers without additional tuning. The method is simple to implement and broadly applicable, offering practical gains for large-scale training with minimal overhead.

Abstract

We introduce Cautious Weight Decay (CWD), a one-line, optimizer-agnostic modification that applies weight decay only to parameter coordinates whose signs align with the optimizer update. Unlike standard decoupled decay, which implicitly optimizes a regularized or constrained objective, CWD preserves the original loss and admits a bilevel interpretation: it induces sliding-mode behavior upon reaching the stationary manifold, allowing it to search for locally Pareto-optimal stationary points of the unmodified objective. In practice, CWD is a drop-in change for optimizers such as AdamW, Lion, and Muon, requiring no new hyperparameters or additional tuning. For language model pre-training and ImageNet classification, CWD consistently improves final loss and accuracy at million- to billion-parameter scales.

Figures (12)

• Figure 1: Final validation loss vs. weight decay coefficient $\lambda$ for 338M models trained on C4 under Chinchilla scaling. Our approach (red) achieves lower final loss than standard weight decay (blue) while preserving the optimizer-specific optimum in $\lambda$. For each optimizer (AdamW, Lion, Muon), both methods use the same hyperparameters.
• Figure 2: Trajectories of Adam, AdamW, and Adam + CWD on a toy example. Adam halts at a minimizer, while AdamW minimizes the objective within a constrained region (green). In contrast, Adam + CWD exhibits sliding mode dynamics within the minimizer manifold.
• Figure 3: Toy objectives and trajectories. Left: $f(x,y)=((y-3)^2-(x-3)^2-1)^2$. Right: $f(x,y)=(y-3-(x-3)^2)^2$. We compare Adam, AdamW, and Adam + CWD; AdamW and CWD use the same weight decay $\lambda$, and all other hyperparameters $(\eta,\beta_1,\beta_2,\epsilon)$ are identical. For both objectives, Adam converges to a generic point on the minimizer manifold, whereas AdamW converges to a solution of the box-constrained problem $\min_{x,y} f(x,y)$ subject to $\max\{x,y\}\leq\frac{1}{\lambda}$. In contrast, Adam + CWD converges to the Pareto front of the minimizer manifold.
• Figure 4: Evaluation loss across scales. 3$\times$3 grid for 338M, 986M, and 2B Transformer models trained with AdamW, Lion, and Muon on C4 dataset. All panels show a zoom into the final $\sim$40% of training steps to highlight late-stage behavior. Baseline curves (dashed blue) use standard weight decay with tuned hyperparameters (learning rate schedule, $\beta$’s, weight decay, etc.; see Appendix \ref{['app:configs']}). Our method (solid red) follows Algorithm \ref{['alg:cwd']} and reuses the baseline hyperparameters without additional tuning. Full (non-zoomed) curves are in Figures \ref{['fig:muon_loss']}, \ref{['fig:adam_loss']} and \ref{['fig:lion_loss']} in Appendix \ref{['ssec::addtitional_results']}.
• Figure 5: Training loss of OLMo 1B on 100B tokens. Left: AdamW. Right: Muon.

Theorems & Definitions (22)

• Theorem 1
• proof : Proof sketch
• Remark 1
• Remark 2
• Definition 1: $L$-smoothness
• Definition 2: Coerciveness
• Remark 3
• Remark 4
• Lemma 1
• proof