Cautious Weight Decay
Authors: Lizhang Chen
Jonathan Li
Kaizhao Liang
Baiyu Su
Cong Xie
Chen Liang
Ni Lao
qiang liu
Venue: ICLR 2026 Poster (2026)
Keywords: optimization
regularization
weight decay
decoupled
lyapunov
training
deep learning
Primary area: optimization

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.

Published as a conference paper at ICLR 2026
CAUTIOUS WEIGHT DECAY
Lizhang Chen 1,2
Jonathan Li
Kaizhao Liang'
Baiyu Su
Cong Xie
Nuo Wang Pierse2
Chen Liang
Ni Lao
Qiang Liu'
Thee University of Texas at Austin
Google LLC, Mountain View California, USA
ABSTRACT
We introduce Cautious Weight Decay (CWD), a one-line, optimizer-agnostic modi- fication that applies weight decay only to parameter coordinates whose signs align with the optimizer update. Unlike standard decoupled decay, which implicitly op- timizes a regularized or constrained objective, CWD preserves the original loss and admits bilevel interpretation: it induces aeddng-mode behavior upon reaching the stationary manifold, allowing it to search for locally Pareto Pareto-optimal stationary points of unmodified objective. mizers such as ADAMW, LION, and MUON, requiring no new hyperparameters additional tuning. For language model pre-training and ImageNet classification, CWD consistently improves final loss and accuracy at million- billion billion-paramete parameter scales.
practice, CWD a drop-in change for opti-
INTRODUCTION
Algorithm
Cautious Weight Decay (CWD)
given parameters xx, optimizer update ue learning rates > 0, weight decay coefficient À 2 0
X++1
Uu XT[[u,,; 2 0)xe
b entrywise multiplication
Optimization algorithms lie at the core of modern deep learning, shaping not only convergence speed but also training stability and generalization ability across domains such natural language processing and computer vision. As models and datasets scale, traditional methods such as stochastic gradient descent (SGD) and SGD with momentum (Sutskever et al., 2013) encounter limitations, including slow convergence in non-convex landscapes, sensitiviry to learning rate schedules, and poor robustness sparse or noisy gradients (Scaman & Malherbe, 2020. Zhao et al.] (025).. In
response, wide range of alternatives have emerged, including adaptive gradient methods (Duchi Kingma & Ba 2015), approximate second-order approaches (Martens Grosse 2015
al. 2011
Gupta et al. 2018 Yao et al. 2021 [iu et al. 20244] Nguyen et al., 2024 Wen et al. 2025) and
specialized algorithms for extreme training regimes (Liang ct al. 2024b Luo et al. 2024 Xie et al,
2024
Huang et al. 2025 Zhang et al. 2025).
Among these advances, decoupled weight decay (Loshchilov Hutter 2019) has proven especially influential. In its general form, decoupled weight decay augments any optimizer update u, with decay term applied directly to the parameters, i.c.
X++1 X - (u, + Ax;),
u, OptimiizerUpdaa((xx.))
This technique improves training stability and generalization by preventing the adaptive learning rates from interfering with regularization, as exemplified by the success of ADAMW in large model
training (Brown al. 2020 Dosovitskiy al. 202 Touvron et al. 2023) and the subsequent
development of statc-of-the. the optimizers such as LION (Chen et al. 2023). LION-K (Chen et al.) 2024), and MUON (Jordan al. 2024. Liu et al. 2025).
However, decoupled weight decay remains agnostic to the directional aligument between the opti- mizer update and the parameters, which may hurt performance when they conflict. Intuitively, when
Equual contribution by I.C. and JL. Correspondence: Izchen, j1 1....c... wtwxaa........ edu

---

Published as conference paper at ICLR 2026
Adam W
Muon
Lion
3,050
3.020
Ours
Validattion
3,100
3.000
3.000
3,000
0.00
0.10
0.20
0.30
0.40
0.00
0.10
0.20
0.30
0.40
2 Weight Decay
3
Weight Decay
Weight Decay
Figure 1: Final validation loss weight decay coefficient À 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-speciife optimum in A. For each optimizer (ADAMW, LION, MUON), both methods use the same hyperparameters.
the update u, and parameters X; point in the same direction for a given dimension, weight decay acts as a regularizer that improves stability; however, when their directions differ, applying decay actively resists beneficial movement toward the optimum. Furthermore, decoupled weight decay has been shown to implicitly impose regularization terms on the objective function (Chen et al. 2024 Xie & 2024). which corresponds parameter norm constraints for ADAMW, LION, and MUON.
In light of these limitations, we propose simple re- finement: cautious weight decay (CWD), in which de- cay is applied only in dimensions where the update and parameter signs align (Algorithm D contributions are as follows.
Feasible Region
Our main
Adam (wd=0)
AdamW
Ours
We introduce cautious weight decay, sign-selective extension of decoupled decay that applies weight de- cay only when the parameters and update align. Our technique can be implemented one-line modifica- tion without introducing additional hyperparameters compared standard decoupled decay.
y
(3, 3)
We use Lyapunov analysis to show that standard optimizers (SGD(M), LION-K, Adam) with cautious weight decay are asymptotically stable and unbi- ased, in the sense that they optimize the original loss rather than regularized surrogate. The regulariza- tion effect of cautious weight decay instead becomes bilevel objective of finding locally Pareto-optimal points within the stationary manifold (Figure 2). Fur- thermore, we show a convergence rate for discrete- time ADAM with cautious weight decay in the smooth nonconvex setting under additional assumptions.
X
Figure 2: Trajectories of ADAM, ADAMW, and ADAM + CWD on a toy example. ADAM halts at minimizer, while ADAMW mini- mizes the objective within a constrained re- gion (green). In contrast, ADAM CWD exhibits sliding mode dynamics within the minimizer manifold.
In language modeling (LMo et al. 2025 Kamath
et al. 2025) and ImageNet classification (Deng al. 2009) observe that cautious weight de- cay generally accelerates convergence and lowers final validation loss for ADAMW, LION, and MUON (e.g., Figure [T). These improvements translate into higher zero-shot accuracy standard benchmarks from 338M to 2B parameters and across architectures without retuning baseline settings (0:20,000 NVIDIA H100 HBM3--0GB GPU hours for all experiments).
2 BACKGROUND AND MOTIVATION
2. DECOUPLED WEIGHT DECAY
Gradient-baan based optimizers with decoupled weight decay can be characterized by the update rule
X = A)xx uu
(1)
where u, := L((x,, .i..... gt is an adaptive, often sign-normalizer update vector constructed from first and recond-moment estimates (e.g., momentum buffers, diagonal preconditioners), m 0 the learning rate, and > is the decoupled weight decay coefficient. This framework encapsu- lates wide range of standard optimizers for machine learning, including ADAMW and LION-K.
2

---

Published as a conference paper at ICLR 2026
m,, ADAMW. The update vector is given by u, D,' where D, is diagonal preconditioner and m is bias-corrected first-moment estimate. Explicitly,
m. 3 -,m--1 - B1) 1 81
V 82V1 + (1 3e)g2
D; diag (V ((++
where and 82 are momentum coefficients and is numerical stability constant.
LION-K. Given convex function K, the update vector u, is a momentum-filerre step that is preconditioned using a subgradient, i.e.
m. -2m,-1 (1 B2) m -1m--1 -1- B.)
-VK(m))
where 31 and 32 momentum coefficients and VK is subgradient of K. Examples include LION when = ||-1|, and MUON when = II- where ||-I1, denotes the nuclear norm when the parameters treated as matrix.
2.2 IMPLICIT REGULARIZATION EFFECTS OF WEIGHT DECAY
In general, the application of decoupled weight decay imposes certain regularization or constraint effect on the objective function, where the specific effect depends on the choice of u,. For example, SGD with decoupled weight decay exactly SGD on an fzrregularized objective. To see the equiva- lence, let f Rdd - IR be differentiable and consider the regularized variant f(xx = f(x) + ||x|
A single SGD step on f with learning rate 1/: yields the update
X++1 x 7o (xx) )x,) (1
n:A)xx 7e f(()),
which is precisely the decoupled weight decay update given by (T.
Given convex function K. with subgradient VK and convex conjugate K suppose the iterates of LION-K converge to a fixed point (x*, mn, m²). Then moment estimators stabilize so that m m" fx(*)), and the fixed-point condition yields - -VK(- V Ax* O. Rearranging
and using the identity (VK)- VK", we obtain ff(xx) VA (Ax*) = 0, where the left-hand
side the gradient of the function
=xx = f(x) xxx).
This suggests that LION-K optimizes the regularized objective f. an observation made by Chen al. 2024). In the special cases of LION and MUON, XC is the 0-c0 indicator function of dual norm ball, corresponding the constrained optimization problems
min (x)
||x||
and
min X X
S(X)
S.C.
X
lop S
XE2
X
respectively, where I1-Ilap is the spectral norm when the parameters are treated matrix.
similar analysis for ADAMW suggests that it solves the box-constrained problem of minimizing f (xx such that ||x|| oo < 1.4 but convergence cannot be established due to the lack of a Lyapunov function. For more discussion, see Appendix Cadd Xie & Li (2024).
While ADAMW and LION-K are practically strong, they implicitly optimize regularized surrogate that dependent on the weight decay coefficient X. This motivates the development of mecha- nism that maintains the beneficial effects of decoupled weight decay (e.g. regularization, training acceleration) while optimizing the original objective.
3 CAUTIOUS WEIGHT DECAY
Cautious weight decay (CWD) modifies the update rule I as
X++1
(u, AT(u x,
x2);

---

Published as a conference paper at ICLR 2026
where @ denotes entrywise multiplication! | As one-line modification, cautious weight decay implementation trivial and universally compatible with gradient- based optimization algorithms. Theoretically, cautious weight decay also exhibits the following behavior.
Unbiased optimization. in the sense that every accumulation point x* of the trajectory satisfies fx(x)) 0 under the same convergence conditions required of the base optimizer without weight decay. In over-paranen aver-parameterized deep models, the set of stationary points is typically a union of con- nected submanifolds rather than isolated points. Consequently, the w-limit set of the trajectory is contained in some stationary manifold, and the iterates eventually remain arbitrarily close to it.
Sliding mode dynamics within the stationary manifold, where cautious weight decay allows the trajectory to traverse along the manifold until it cannot decrease the parameter magnitudes in every coordinate. In other words, cautious weight decay steers the trajectory towards local Pareto front of the stationary manifold under the ordering that prioritizes smaller parameter magnitudes.
3.1 CONVERGENCE TO THE STATIONARY MANIFOLD
We construct Lyapunov functions for the continuous-time limits several standard optimizers equipped with cautious weight decay. A Lyapunov function is lower bounded function with non- positive derivative that is used to certify the stability of systems of differential equations.
Consider the continuous-ting time dynamics of SGD with cautious weight decay
x = (xx)) - XI (xx X² 2 0)x;
This ODE has the Lyapunov function H(xx) = fxx), since H is lower bounded and (llxx))--(VVVffaaa XI(V 0)xx V fxx) (x))(x) II $ 0,
dPp - - dt
where (.)+ == (0, -). LaSalle's invariance principle (LaSalle 1960) states that the accumulation points of any trajectory lie within the union of trajectories z, that satisfy 4 (z,) O for all 0. Consequently, we conclude that SGD with cautious weight decay produces trajectories that approach the stationary set {x fx) 0} of the original loss. This holds because cautious weight decay is applied only in secondary fashion and is automatically deactivated whenever it conflicts with the main objective, thereby ensuring that the loss landscape remains unbiased.
Beyond the simple case of SGD, the same Lyapunov-type argument can be extended to more sophis- ticated algorithms such as SGDM, LION-K, and ADAM. In each case, cautious weight decay still minimizes the original objective without introducing explicit bias, but key difficulty lies in con- structing appropriate Lyapunov functions. Table summarizes the Lyapunov functions of several major optimizers with cautious weight decay, and detailed derivations provided in Appendix [ By applying LaSalle's invariance principle, we can show that the momentum- based algorithms in Table converge the stationary of the original objective, together with vanishing momentum:
((x, ()) fx(x) = 0, m = =0}.
3.2 SLIDING MODE DYNAMICS
Although both standard optimization (with no weight decay) and cautious weight decay are unbiased with respect the original objective, their behaviors diverge within the stationary manifold. In the former, the dynamics halt as the momentum m decays to zero, while, in contrast, the cautious weight decay dynamics induce sliding mode, continuing move along the manifold while reducing the parameter magnitudes as much as possible. Consequently, the algorithm converges to a subset of the stationary manifold where further simultaneous reduction of all coordinates of x is no longer possible. Equivalently, it converges to locally Pareto-optimal stationary point under a preference for smaller parameter magnitudes.
To provide mathematical background, consider possibly time-varying discontinuous ODE
L f((z))
z,
E ddd
Thhooughout the paper, when it is clear from context, also drop and write for simplicity.
4

---

Published as a conference paper at ICLR 2026
Table I: Comparison of the continuous- dynamics of different optimizers. SGDM represents SGD with momentum. LION-K includes LION (K I1-11)) and MUON (K = Ill) as special cases. f: R +
R is assumed to be differentiable and lower bounded by f*
Optimizer
Continuous-time dynamics
Lyapunov function
SGD + CRD
k = V (xx))) - AI(V f(x))x 2 0)x:
H(x) f(x)
SGDM CWD
x m. AI(mx 0)x
H(x, m) B3(()) [ (mx)) ||
IiM ((v)(xx) me)
LION-IC + CWD
& V.C(m,) - AI(mxx 0)xx
H(x, m) = =f ( ((m)) + A (-mx)+ ||
m = -aV f(x.) mm
ADAM + CRD
* = 0: - AI(m 2 0)x,
H,(x,m,h) = ax((x)+
a.m²
+ (mxx*
h:
2h
t = a ( me)
(V/ (xx)) v.)
Notation. We drop @ for simplicity, = (1 eex((=a()) ,7 (1 exp( xxp((-())) "h == / tc.
Due to the discontinuity of Se. the solution may not well defined in the classical or Carathéodory sense, especially across switching surfaces. We therefore interpret solutions in the Filippov sense (Filippov] 1988). where a discontinuous ODE is formally differential inclusion that specifies that Z; belongs to the closed convex envelope of the discontinuous vector field, i.c.
è, E Flf(((z) = n n 830 pS))
co (f:( 8) S)),
where y denotes the Lebesgue measure, (z, 8) is the 6-ball centered at z, and to denotes the closed convex envelope. This construction captures possible limiting directions of the vector field near discontinuities, ensuring well-defined dynamics even when ft is not continuous. The key idea is that the values of 2, must be determined by the behavior of ft in a neighborhood around z,, rather than at the point itself. The inclusion, therefore, defines range of admissible velocities consistent with nearby values of the vector field.
In particular, whenever ft contains coordinatewise indicators such as T(g(2)) 2 0), the Filippov set replaces them by selectors $; E [0, 1]dt on the switching set [[g = 0}:
(1)
[g (+(z)); ; > 0,
[sc]
{0}
[g(2,)): }) < 0,
[0. [g(2))]; = 0.
Recalling the Lyapunov analysis in Section 3.1 continuous-time dynamics of standard opti- mizers with cautious weight decay converge the stationary manifold M == fx f 0). with the momentum m, also decaying to for momentum- methods. Consequently, once the trajectory enters the stationary manifold, the residual dynamics reduce to
As; X
St E [0.
Moreover, since the Lyapunov function confines the dynamics to the stationary set, the selectors St must be chosen such that the trajectory remains within the manifold. Differentiating the stationarity condition yields
di =
0, Se [0, 1J
This relation allows to solve for admissible choices of that guarantee invariance of the manifold. In general, the solution for s, need not be unique, and the actual value realized practice may be implicitly determined by the discretization scheme employed.
5

---

Published as a conference paper at ICLR 2026
Adam A
Ours
start
Pareto Set
Feasible Region
fxx, y) = ((x 3)2 (y- 3)² - 1)2
f(x, y) (y 3 (x - 33)
1-
1
X
X
Figure 3: Toy objectives and trajectories. Left: [(x, M) = (ly 3)2 + (r - 3)2 - 1)2 Right: [(r, y) (y - (x 3)²)² We compare ADAM, ADAMW, and ADAM CWD; ADAMW and CWD use the same weight decay X, and all other hyperparameters (17, 31. 32. are identical. For both objectives. ADAM converges
generie point on the minimizer manifold, whereas ADAMW converges to solution of the box-constrained problem minz,y fzz y) subject to max (z. v) S x In contrast, ADAM CWD converges the Pareto front of the minimizer manifold.
Effectively, cautious weight decay decreases parameter magnitudes along each coordinate while staying within the stationary manifold, pushing x toward the local Pareto front of the manifold
P :== {x E M 36 Vy E (B(x, 6) n M) \ (xx,, |y| Σ |x]},
where the tangent space no longer allows nonzero 8, in (2). In other words, stationary point locally Pareto-optimal it has a neighborhood in the stationary manifold that contains no other point with a smaller or equal magnitude in every coordinate.
This argument shows that cautious weight decay dynamics converge to P. Since IP may not be singleton, the exact limit point depends intricately on initialization and the discretization of the continuous-timg dynamics. Figure illustrates this behavior on two toy problems.
3.3 DISCRETE-TIME ANALYSIS
Leveraging the Lyapunov functions in Table we can extend our continuous-time analysis obtain convergence guarantees for the discrete-timg dynamics of various optimizers with cautious weight decay. As a concrete example, we provide in Appendix [E an explicit convergence rate for discrete- time ADAM with cautious weight decay.
4 EXPERIMENTS
Overview. We evaluate CWD against three standard optimizens-AAAMW. LION, and MUON- autoregressive language modeling and ImageNet classification. For Transformer models with simi- lar architecture to Gemma (Kamath et al.] 2025) with 338M, 986M, and 2B parameters in the Sim- ply (Liang al. 2025) codebase, follow the Chinchilla compute-optimal scaling rule-20 tokens per parameter (TPP) (Hoffmann et al. (002) and train on C4 (Raffel et al.) 2020). For each size, we grid-search batch size, learning rate, weight decay, warmup ratio, and optinizer-ppecffic hyperpa- rameters for the baselines (ADAMW W. LION, MUON); we then reuse the selected baseline settings for CWD without retuning (details in Appendix IT. Under matched settings, CWD lowers final validation loss and improves shot accuracy. On the OLMo codebase (OLMo et al., (025). we further study an over-training regime- OLMo- 1B trained on 100B tokens (100 TPP) from Dolma (Soldaini al.) 2024). Under matched settings, CWD lowers final validation loss and improves zero-shot accuracy (Table 4.. We also observe similar gains on ImageNet (Deng et al. 2009) across ViT (Dosovitskiy al. 2021) and ResNet (He et al. 2016).
Ablations of weight decay. Figure sweeps the deight-decay coefficient À for a 338M model on C4: A [0, 0.4] for MUON and ADAMW, and À E [0, 3.0] for LION. Two patterns are consistent across runs: (i) at a fixed X, CWD attains lower final loss than the corresponding baseline with decoupled weight decay; (ii) the minimizing value X is essentially unchanged when replacing the baseline with CWD. In practice, one can swap in CWD at an already tuned À and obtain improvements without additional sweeps.

---

Published as a conference paper at ICLR 2026
Eval Loss (338M)
Eval Loss (986M)
Eval Loss (2B)
Ours
Ours
Ours
3.15
AdamW
AdamW
AdamW
Loss
Loss
3.10
3.05
Step
Step
Step
Eval Loss (338M)
Eval Loss (986M)
Eval Loss (2B)
3.20
Ours
Ours
Ours
3.15
Lion
Lion
Lion
Loss
Loss
3.00
18000
21000
16000
15000
20000
Step
Step
Step
Eval Loss (338M)
Eval Loss (986M)
Eval Loss (2B)
5.20
2:50
Ours
Ours
Ours
Muon
Muon
Muon
3.30
Loss
Loss
3.05
2.70
12000
15000
18000
21000
12000
14000
18000
12000
14000
360C0
18200
20000
Step
Step
Step
Figure 4: Evaluation loss across scales. 3x3 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 ~40% of training steps to highlight late-stage behavior. Baseline curves (dashed blue) use standard weight decay with tuned byperparameters (learning rate schedule, f's, weight decay, etc.; see Appendix F). Our method (solid red) follows Algorithm [ and reuses the baseline hyperparameters without addi- tional tuning. Full (non-zoomed) curves are in Figures 8 9 and 10 in Appendix G
Ablation study of selective weight decay strategies on OLMo-IB (100B tokens). We com- pare our momentum based selection against alternative masking approaches. Baseline: standard weight decay (A tuned). Ours: update-based mask I(ux 2 0) using baseline's À without retuning. Random: time-varying Bernoulli mask matching our method's sparsity ratio (see Figure 7 in Ap-
Table 2:
pendix Gb. Gradient: uses (gg 2 0) instead. No WD: X = 0. Lower validation loss is better.
Weight Decay Active
Ablated Masks
Disabled
Optimizer
Baseline
Ours
Random
Gradient
No WD
ADAMW
2.65
2.56
2.82
2.75
2.70
MUON
2.51
2.42
2.73
2.74
2.62
Ablations on masking. teble whether the benefits arise from the amount of decay applied or from CWD's structure. Replacing our mask with time-matched Bernoulli "random mask" substan- tially degrades performance (e.g.. 2.56 - 2.82 for ADAMW, 2.42 - 2.73 for MUON), showing that simply reducing the frequency of decay insufficient. Substituting the indicator with the gradient- based 1l(gx 2 0) also underperforms. Finally, X = 0 remains worse than tuned decay, illustrating that explicit regularization is helpful and CWD leverages more effectively. We further verify the dif- ference between CWD and SPD (Tia et al., 2024) with additional instruction-folloiing fine-tuning experiments. We fine-tune on the Alpaca GPT-4 dataset (Peng et al. 2023). which contains 52,000 instruction-responss pairs generated by GPT-4 (OpenAl 2023), 3), and evaluate on three benchmarks: MMLU (Hendrycks et al. 2021), comprising 57 tasks and 14,079 questions covering a broad range of world knowledge; AGlEval (Zhong et al. 2024). human-centrr benchmark with 9,316 instances targeting general reasoning and croblem--solvin skills; and WinoGrande (Sakaguchi et al.) 202I), large-scale commonsense reasoning dataset with 44,000 instances. We consider two base models,

---

Published as a conference paper at ICLR 2026
Table ImageNet validation accuracy (%) across architectures and optimizers. All models train for 300 epochs with standard augmentation. Base: optimizer with tuned weight decay. Ours: cautious weight decay using the same coefficient as baseline (no retuning).
ADAMW
LION
MUON
Model
Params
Base
Ours
Base
Ours
Base
Ours
ViT-S/16
22.05M
78.84
79.45
79.29
79.82
79.35
79.91
ResNet-50
25.56M
76.30
76.68
76.41
76.75
76.47
76.83
ViT-B/16
86.57M
80.15
80.71
80.76
80.92
80.83
81.04
Optimizer
Hellaswag acc_norm
ARC-Easy acc_norm
ARC-C1 acc.norm acc_norm
PIQA f MMLU
ComQA
acc
acc
ADAMW
0.38
0,50
0.25
0.67
0.23
0.29
ADAMW+CKD
0.40
0.53
0.27
0.69
0.25
0.31
MUON
0.39
0.51
0.26
0.68
0.24
0.30
MUON+CWD 0.41
0.51
0.28
0.71
0.26
0.33
Table 4: Downstream accuracy across diverse reasoning benchmarks. All runs use the OLMo code- base with 1B IB-parameter models trained for 100B tokens under an over-training regime. Here ARC- C=ARC-Challenge and Com@A=Commonsense(A Figure shows the corresponding loss curves.
TamyLlama Zhang et al. 2024) and Mistral-7B (Jiang et al. 2023), and compare LoRA (Hu et al, 2022). SPD (Tian et al. 2024), layerwise "inner-produe variant of CWD using (( x) 2 0), and our proposed "elementwise" CRD. Table reports accuracy on MMLU, AGIEval, and WinoGrande
for Tiama and Mistral-7B.
Training dynamics. On IB models trained for 100B tokens, we observe that CWD tends to improve the loss trajectory relative to tuned ADAMW and MUON, rather than only the final value (Figure 5 similar pattern appears at 986M: Figure I1 in Appendix showss evaluation/training loss and RMS parameter over time. CWD generally achieves lower loss while ending with intermediate norm. In contrast, removing decay entirely (A = 0) descends faster mid-training but plateaus earlier, finishing higher loss and the largest norm; tuned ADAM with À > O yields the smallest norm. Overall, these results suggest that the gains come from more selective application of regularization rather than from disabling it.
CWD outperforms standard decay across optimizers and scales. Under the common setup across 338M, 986M, and 2B parameters, CWD consistently lowers eval loss for ADAMW, LION, and MUON (see Figure 4 and Figures 8 in Appendix G) and increases downstream accuracy (Table
CWD yields lower gradient norms than standard decay. Across model sizes, CWD produces lower RMS-normalized gradient norms than the corresponding baselines (see Figure 12 in Appendix G This coincides with the lower end-of-training loss Figure the accuracy gains Table
Model
Method
MMLU (5-shot) t
AGlEval (3-shot) 1
WinoGrande (5-shot)
TinyLJlma
LoRA (haseline)
25.81 + 0.07
19.82 + 0.11
61.33 + 0.09
LinyLlama
SPD
26.14 + 0.08
20.21 + 0.10
61.92 + 0.08
TinyLlama
Inner-product CMD
26.02 1 0.08
19.80 1 0.10
61.70 1
TinyLlama
Elementwise CWD
26.42 + 0.09
20.12 1 0.09
62.18 1 0.08
Mistral-7B
LoRA (baseline)
61.78 0.09
27.56 1 0.07
78.85 1 0.11
Mistral-7B
SPD
62.05 1 0.08
27.98 1 0.06
78.81 1 0.10
Mistral-7B
Inner-product CWD
61.76 + 0.09
27.90 + 0.07
78.83 + 0.10
Mistral-7B
Elementwise CWD
62.13 + 0.07
28.31 + 0.06
78.92 + 0.09
Table Accuracy on MMLU, AGIEval, and WinoGrande for TinyL. and Mistral-7B fine-tuned on Alpaca GPT-4 using LoRA. SPD, and two variants of CWD. Both SPD and CWD improve over the LoRA bascline, and the proposed elementwise CWD matches or outperforms SPD and the inner- product variant most benchmarks.
8

---

Published as a conference paper at ICLR 2026
AdamW (wd=0.0)
Muoon (wdw0
Perplexity (log scale)
AdamW (wd-0.1)
Perplexity (log scale)
10
Muon (wd-0.1)
Ours (wd=0.1)
Ours (wd=0.1)
6-
w
50K
100K
200K
250K
300K
100K 150K
250K
Training Steps
Training Steps
Figure Training loss of OLMo 1B on 100B tokens. Left: ADAMW. Right: MUON.
Scaling with model size. We measured final validation loss for ADAMW and CWD ("Ours") at 111M, 338M, 986M, and 2B parameters the same dataset (Figure 6). At every scale, CWD attains a lower final validation loss than ADAMW, and the gap remains stable or even widens with model size, indicating that the advantages of cautious weight decay persist into the large-model regime.
Figure 6: Final validation loss versus model size for ADAMW and ADAM + CWD ("Ours") on the same pretraining dataset, 11 338M, 986M, and 2B parameters. Across all scales, ADAM CWD achieves consistently lower final validation loss than ADAMW, and the gap remains stable or slightly increases with model size, suggesting that the benefits of cautious weight decay persist in the large- model regime.
Final validation loss
3.2
Baseline
2.6
11.1 11.1M.38M
986M
2B
Model size (M parameters)
Model Size
Optimizer
Final Validation Loss
338M
ADAMW (baseline)
3.0136
338M
ADAM + CWD
3.0059
338M
ADAMC
3.0087
338M
ADAMC CWD
2.9915
Table 6: Final validation loss for Hammaabasse models with different optimizers 338M param- eters. Adding cautious weight decay improves both ADAMW and ADAMC, with ADAMC CWD achieving the lowest loss.
Model Size Optimizer
Scheduler
Final Validation Loss
338M
ADAMW (baseline)
Cosine
3,0136
338M
ADAMW (baseline)
WSD
3,0101
338M
ADAM + CWD
Cosine
3.0059
338M
ADAM + CRD
WSD
3.0014
Table 7: Final validation loss for 338M model under different optimizer-scheduler combinations. ADAM + CWD improves over the ADAMW baseline for both cosine and WSD schedules, with the best performance obtained by ADAM + CWD with WSD.
Compound improvement with other techniques. We also observe compounding gains when com- bining cautious weight decay with other optimizer techniques Gemma-base model architectures. In particular, we compare ADAM W, ADAMC (Defazio 2025). and their variants with CWD at 338M parameters (Table

---

Published as a conference paper at ICLR 2026
Robustness to learning-rate schedules. We observe that CWD improves performance for both co- sine scheduling and the Sarmup-Stablle (WSD) schedule, which uses a 10% warmup, a long stable phase, and a 20% final decay 7
RELATED WORK
Weight decay. Weight decay originated as lS penalty for ill-posed problems and ridge regression
[Tikhonov] 1963; Hoerl & Kennard, 1970) and was introduced to neural networks as a generaliza-
tion tool to mitigate overfitting (Hanson & Pratt 1988, WWigend al. 1990) Krogh & Hertz 1991), ) Loshchilov & Hutter (219)) showed that, for adaptive methods, weight decay and E2 regularization
not equivalent, motivating the decoupled formulation in ADAMW; subsequent work established
decoupled decay as a standard feature of modern optimizers (Chen et al. 2023 2024 Liu et al.)
2025). Recent analyses suggest that in contemporary networks, weight decay functions more as a trainnn accelerator and stabilizer than as explicit regularization (Krizhevsky et al., 2017 [oof-] mann et al., 2022, Pan & Cao 2023 Angelo et al.) 2024). Interactions with early-stage training (Bjorck et al. (021), normalization layers, and learning rate schedules (Defazio) 2025) have also
been clarified, and architectural designs can obviate explicit decay (Loshchilov et al. 2025).
Weight decay variants. Various efforts have been made to develop different adaptive variants of weight decay. For example, Xie et al. (2023) found that weight decay can lead to large gradient norms the final phase of training and proposed Scheduled Weight Decay (SWD) dynamically adjust weight decay strength based on gradient norms. Kosson et al. (2024) investigate how weight decay affects individual neuron updates, revealing rotational equilibrium states that balance learning across layers and neurons. Ghiasi et al. (2023) introduce adaptive weight decay that automatically tunes the hyperparameter during training based classification and regularization loss gradients, achieving significant improvements in adversarial robustness, Tian et al. (2024) introduce Selective Projection Decay (SPD) for robust fine-tuning, featuring selective weight decay via a mask that is somewhat similar to CWD. However, SPD and CWD differ in significant ways, including the intended setting, mechanism, theoretical properties, and empirical performance.
Masked or conditional updates. Several works have explored the bign-based conditioning of op- timizer updates. Riedmiller & Braun (1993) introduced RPROP, which adjusted step sizes based on current gradient and past gradient sign agreement. [Liang et al.) (2024a) propose the cautious opti- mizer, which restricts updates to dimensions where the proposed update and current gradient share the same sign. Wang et al. (2024) apply similar mask to ADAM to improve robustness in online learning.
Constrained and bilevel optimization. Decoupled weight decay can be interpreted through the lens of Frank-Wolff algorithms for constrained optimization (Frank & Wolfe 1956, Jaggi 2013 Sfyraki Pethick et al. 2025 This connection suggests that optimizers with decoupled
& Wang 2025
weight decay implicitly solve constrained optimization problems, which was shown to be the case
for LION (Chen et al. [2024] Sfyraki & Wang 2025 Pethick et al. 2025). ADAMW (Xie & Li
2024
Bernstein & Newhouse 2024). and MUON (Chen et al.) 2025 Sfyraki & Wang 2025 Lau
al. 2025). In contrast, optimizers with cautious weight decay perform bilevel optimization,
framework from classical optimization (Solodov) 2007a Sabach & Shtemn 2017) that has been
recently explored in machine learning (Gong al. 2021 Liu
2022
Petrulionyte al. 2024).
CONCLUSION
We introduce cautious weight decay and formalize it as a simple, optimizer-agnostic modification of decoupled weight decay that preserves the optimization objective while retaining the practi- cal benefits of weight decay. For standard optimizers (SGD, ADAM, and LION-K), we show the bilevel optimization structure of cautious weight decay and establish convergence guarantees in both continuous- and discrete-time regimes. Across diverse tasks and benchmarks, cautious weight decay consistently improves training dynamics compared to no decay and traditional decoupled de- cay, yielding faster loss reduction and more stable trajectories without changes to hyperparameters or model architectures. Our results indicate that cautious weight decay is theoretically principled and empirically effective technique that retains the benefits of weight decay while addressing its fundamental limitations.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported in part by the Institute for Foundations of Machine Learning (IFML) and the Office of Naval Research (ONR) under Grant No. N00014-25-1-2354.
REFERENCES
Yossi Arjevani, Yair Carmon, John C. Duchi, Dylan J. Foster, Nathan Srebro, and Blake E. Wood- worth. Lower bounds for non-convex stochastic optimization. Math. Program., 199(1):165-214, 2023.
Andrea Bacciotti and Francesca Ceragioli. Stability and stabilization of discontinuous systems and nonsmooth Iyapunov functions. ESAIM: Control, Optimisation and Calculus of Variations, 4: 361-376, 1999.
Anas Barakat and Pascal Bianchi. Convergence and dynamical behavior of the ADAM algorithm for nonconvex stochastic optimization. SIAM J. Optim., 31(1):244 274, 2021.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. CoRR, abs/2409.203255, 2024.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. SIGNSGI: compressed optimisation for non-convex problems. In Proceedings of the 35th In- ternational Conference on Machine Learning, ICML 2018, volume 80 of Proceedings of Machine Learning Research, pp. 559-568, 2018.
Johan Bjorck, Kilian Q. Weinberger, and Carla P. Gomes. Understanding decoupled and early weight decay. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty- Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, 1, pp. 6777-6785, 2021.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler. Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. Advances in Neu- ral Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, 2020.
Congliang Chen, Li Shen, Fangyu Zou, and Wei Liu. Towards practical adam: Non-convexity, convergence theory, and mini-batch acceleration. J. Mach. Learn. Res., 23(229):1-47, 2022.
Lizhang Chen, Bo Liu, Kaizhao Liang, and Qiang Liu. Lion secretly solves a constrained optimiza- tion: As аyаpunov predicts. In The Twelfth International Conference on Learning Representattions, ICLR 2024, 2024.
Lizhang Chen, Jonathan Li, and Qiang Liu. Muon optimizes under spectral norm constraints. CoRR, 150544 2025.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang. Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe Le. Symbolic discovery of optimization algorithms. In Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NeurlPS 2023, 2023.
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of A class of adam-type algorithms for non-convex optimization. In 7th International Conference Learning Represen- tations, ICLR 20/9, 2019.
Francesco D' Angelo, Maksym Andriushchenko, Aditya Vardhan Varre, and Nicolas Flammarion. Why do need weight decay in modern deep learning? In Advances in Neural Information Processing Systems 38: Anmual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, 2024.
11

---

Published as a conference paper at ICLR 2026
Aaron Defazio. Why gradients rapidly increase near the end of training. CoRR, abs/2506.02285, 2025.
Alexandre Défossez, Léon Bottou, Francis R. Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Trans. Mach. Learn, Res., 2022, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference Computer Viston and Puttern Recognition (CVPR 2009), pp. 248-255, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, 2021.
John C. Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. J. Mach. Learn. Res., 12:2121 2159, 2011.
Aleksej F. Filippov. Differential Equations with Discontinuous Righthand Sides, volume 18 of Math- ematics and Its Applications. Springer, 1988.
Marguerite Frank and Philip Wolfe. algorithm for quadratic programming. Nav. Res. Logist. Q., 3(1 3(1-2):95-110, 1956.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. CoRR, abs/2101.00027, 2021.
Saeed Ghadimi and Guanghui Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic programming. SIAM J. Optim., 23(4):2341-2368 2013.
Amin Ghiasi, Ali Shafahi, and Reza Ardekani. Improving robustness with adaptive weight decay. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurlPS 2023, 2023.
Chengyue Gong, Xingchao Liu, and Qiang Liu. Automatic and harmless regularization with con- strained and lexicographic optimization: A dynamic barrier approach. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Sys- tems 2021, NeurlPS 2021, pp. 19630-29642 2021.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In Proceedings of the 35th International Conference Machine Learning, ICML 2018, volume 80 Proceedings of Machine Learning Research, pp. 1837-1845, 2018.
Stephen Jose Hanson and Lorien Y. Pratt. Comparing biases for minimal network construction with prok-propagttion. In Advances in Neural Information Processing Systems 1, NIPS Conference, pp. 177-185, 1988.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, pp- 770-778, 2016.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In 9th International Confer- ence on Learning Representations, ICLR 2021, 2021.
Arthur E. Hoerl and Robert W. Kennard. Ridge regression: Biased estimation for nonorthogonal problems. Technometrics, 12(1):55-67 1970.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre.
12

---

Published as a conference paper at ICLR 2026
An empirical analysis of compute-optimat large language model training. In Advances in Neu- ral Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, 2022.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth Interna- tional Conference on Learning Representations, ICLR 2022, 2022.
Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, and Shiwei Liu. SPAM: spike- aware adam with momentum reset for stable LLM training. In The Thirteenth International Con- ference on Learning Representations, ICLR 2025, 2025.
Martin Jaggi. Revisiting frank-wolfe: Projection- free sparse convex optimization. In Proceedings of the 30th International Conference on Machine Learning, ICML 2013, volume 28 of JMLR Workshop and Conference Proceedings, pp. 427-435, 2013.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Ann Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William Sayed. Mistral 7b. CoRR, 10./2210 2023.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.
Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej. Sarah Perrin, Ta- tiana Matejovicova, Alexandre Ramé, Morgane Rivière, Louis Rouillard, Thomas Mesnard, Geof- frey Cideron, Jean-Bastien Grill, Sabela Ramos, Edouard Yvinec, Michelle Casbon, Etienne Pot, Ivo Penchev, Gaël Liu, Francesco Visin, Kathleen Kenealy, Lucas Beyer, Xiaohai Zhai, Anton Tsitsulin, Róbert Busa-Fekete, Alex Feng, Noveen Sachdeva, Benjamin Coleman, Yi Gao, Basil Mustafa, lain Barr, Emilio Parisotto, David Tian, Matan Eyal, Colin Cherry, Jan-Thorsten Peter, Danila Sinopalnikov, Surya Bhupatiraju, Rishabh Agarwal, Mehran Kazemi, Dan Malkin, Ravin Kumar, David Vilar, Idan Brusilovsky, Jiaming Luo, Andreas Steiner, Abe Friesen, Abhanshu Sharma, Abheesht Sharma, Adi Mayrav Gilady, Adrian Goedeckemeyer, Alaa Saade, Alexan- der Kolesnikov, Alexei Bendebury, Alvin Abdagie, Amit Vadi, András György, André Susano Pinto, Anil Das, Ankur Bapna, Antoine Miech, Antoine Yang, Antonia Paterson, Ashish Shenoy, Ayan Chakrabarti, Bilal Piot, Bo Wu, Bobak Shahriari, Bryce Petrini, Charlie Chen, Charline Le Lan, Christopher A. Choquette-Choo, CJ Carey, Cormac Brick, Daniel Deutsch, Danielle Eisen- bud, Dee Cattle, Derek Cheng, Dimitris Paparas, Divyashree Shivakumar Sreepathihalli, Doug Reid, Dustin Tran, Dustin Zelle, Eric Noland, Erwin Huizenga, Eugene Kharitonov, Frederick Liu, Gagik Amirkhanyan, Glenn Cameron, Hadi Hashemi, Hanna Sllmezak--lucinskke Harman Singh, Harsh Mehta, Harshal Tushar Lehri, Hussein Hazimeh, Ian Ballantyne, Idan Szpektor, Ivan Nardini, Jean Pouget-Abadie, Jetha Chan, Joe Stanton, John Wieting, Jonathan Lai, Jordi Orbay, Joseph Fernande7, Josh Newlan, Ju-yeong Ji, Jyotinder Singh, Kat Black, Kathy Yu, Kevin Hui, Kiran Vodrahalli, Klaus Greff, Linhai Qiu, Marcella Valentine, Marina Coelho, Marvin Ritter, Matt Hoffmman, Matthew Watson, Mayank Chaturvedi, Michael Moynihan, Min Ma, Nabila Babar, Natasha Noy, Nathan Byrd, Nick Roy, Nikola Momchev, Nilay Chauhan, Oskar Bunyan, Pankil Botarda, Paul Caron, Paul Kishan Rubenstein, Phil Culliton, Philipp Schmid, Pier Giuseppe Sessa, Pingmei Xu, Piotr Stanczyk, Pouya Tafti, Rakesh Shivanna, Renjie Wu, Renke Pan, Reza Rokni, Rob Willoughby, Rohith Vallu, Ryan Mullins, Sammy Jerome, Sara Smoot, Sertan Girgin, Shariq lqbal, Shashir Reddy, Shruti Sheth, Siim Pöder, Sijal Bhatnagar, Sindhu Raghuram Pa- nyam, Sivan Eiger, Susan Zhang, Tianqi Liu, Trevor Yacovone, Tyler Liechty, Uday Kalra, Utku Evci, Vedant Misra, Vincent Roseberry. Vlad Feinberg, Vlad Kolesnikov, Woohyun Han, Woosuk Kwon, Xi Chen, Yinlam Chow, Yuvein Zhu, Zichuan Wei, Zoltan Egyed, Victor Cotruta, Minh Giang, Phoebe Kirk, Anand Rao, Jessica Lo, Erica Moreira, Luiz Gustavo Martins, Omar Sanse- viero, Lucas Gonzalez, Zach Gleicher, Tris Warkentin, Vahab Mirrokni, Evan Senter, Eli Collins, Joelle K. Barral, Zoubin Ghahramani, Raia Hadsell, Yossi Matias, D. Sculley, Slav Petrov, Noah Fiedel, Noam Shazeer, Oriol Vinyals, Jeff Dean, Demis Hassabis, Koray Kavukcuoglu, Clément Farabet, Elena Buchatskaya, Jean-Baptiee Alayrac. Rohan Anil, Dmitry (Dima) Lepikhin, Sebas- tian Borgeaud, Olivier Bachem, Armand Joulin, Alek Andreev, Cassidy Hardin, Robert Dadashi, and Léonard Hussenot. Gemma technical report. CoRR, abs/2503. 19786, 2025.
13

---

Published as a conference paper at ICLR 2026
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd Interna- tional Conference on Learning Representations, ICLR 2015, 2015.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay bal- ances learning across neural networks. In Forty-first International Conference on Machine Learn- ing, ICML 2024, 2024.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convo- lutional neural networks. Commun. ACM, 60(6):84-90, 2017.
Anders Krogh and John A. Hertz. A simple weight decay can improve generalization. In Advances in Neural Information Processing Systems 4, NIPS Conference, pp. 950-957, 1991.
Joseph P. LaSalle. Some extensions of liapunov's second method. IRE Transactions on Circuit Theory, 7(4):520-527, 1960.
Tim Kiz Lau, Qi Long, and Weijie Su. Polargrad: A class matrix graddent optimizers from a unifying preconditioning perspective. CoRR, abs/2505.21799, 2025.
Chen Liang, Da Huang, Chengrun Yang, Xiaomeng Yang. Andrew Li, Xinchen Yan, and Simply Contributors. Simply: an experiment to accelerate and automate AI research. GitHub repository, 2025. URL /
wwt.//o////o/o. deweee-eededddddeeeeee simply.
Kaizhao Liang, Lizhang Chen, Bo Liu, and Qiang Liu. Cautious optimizers: Improving training with line of code. CoRR, abs/241 16008,, 2024a.
Kaizhao Liang, Bo Liu, Lizhang Chen, and Qiang Liu. Memory-efficient LLM training with online subspace descent. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, 2024b.
Bo Liu, Mao Ye, Stephen Wright, Peter Stone, and Qiang Liu. Bome! bilevel optimization made easy: simple first-order approach. In Advances in Neural Information Processing Systems 35; Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022.
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Conference on Learning Representations, ICLR 2024, 2024.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang. Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for LLM training. CoRR, abs/2502. 2025.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
Ilya Loshchilov, Cheng-Ping Hsieh, Simeng Sun, and Boris Ginsburg. ngpt: Normalized transformer with representation learning the hypersphere. In The Thirteenth International Conference on Learning Representations, ICLR 2025, 2025.
Qijun Luo, Hengxu Yu, and Xiao Li. Badam: A memory efficient full parameter optimization method for large language models. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024. 2024.
James Martens and Roger B. Grosse. Optimizing neural networks with coneccer-ffcttoedd approx- imate curvature. In Proceedings of the 32nd International Conference on Machine Learning, ICMI. 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 2408-2417, 2015.
Son Nguyen, Lizhang Chen, Bo Liu, and Qiang Liu. H-fac: Memory-efficient optimization with factorized hamiltonian descent. CoRR, 0bs/2406.09955. 2024.
14

---

Published as a conference paper at ICLR 2026
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo. Shane Arora, Akshita Bha- gia, Yuling Gu, Shengyi Huang. Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi. Nouha Dziri, Michal Guerquin, Hamish Ivison, Pang Wei Koh, Jiacheng Liu, Saumya Malik, William Merrill, Lester James V. Miranda, Jacob Morrison, Tyler Murray, Crystal Nam, Valentina Py- atkin, Aman Rangapur, Michael Schmitz, Sam Skjonsberg, David Wadden, Christopher Wilhelm, Michael Wilson, Luke Zeu, Ali Farhadi, Noah A. Smith, and Hannanch Hajishirzi. olmo 2 furious. CoRR, 0bs/2501.006666 2025.
OpenAl. GPT-4 technical report. CoRR, abs/2303.08 087744 2023.
Leyan Pan and Xinyuan Cao. Towards understanding neural collapse: The effects of batch normal- ization and weight decay. CoRR, abs/2309.04644 2023.
Baolin Peng, Chunyuan Li, Pengcheng He, Michel Galley, and Jianfeng Gao. Instruction tuning with GPT-4, CoRR, abs/2304.032777, 2023.
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained Imos. In Forty-second International Conference on Machine Learning, ICML 2025, 2025.
leva Petrulionyte, Julien Mairal, and Michael Arbel. Functional bilevel optimization for machine learning. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., (11440)11-61 2020.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In 6th International Conference on Learning Representations, ICLR 2018, 2018.
Martin A. Riedmiller and Heinrich Braun. A direct adaptive method for faster backpropagation learning: the RPROP algorithm. In Proceedings of International Conference on Neural Networks (ICNN'88) San Francisco, CA, USA, March 28- April 1, 1993, pp. 586-59666 591. IEEE, 1993.
Shoham Sabach and Shimrit Shtemn. A first order method for solving convex bilevel optimization problems. SIAM J. Optim,, 27(2):640 2017.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: an adver- sarial winograd schema challenge at scale. Commun. ACM, 64(9):99-106 106, 2021.
Kevin Scaman and Cédric Malherbe. Robustness analysis of non-convex stochastic gradient descent using biased expectations. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, 2020.
Maria-Elen Sfyraki and Kun Wang. Lions and muons: Optimization via stochastic frank-wolfe. CoRR, 0aes//506.041999 2025.
Daniel W. Shevitz and Brad Paden. Lyapunov stability theory of nonsmooth systems. IEEE Trans. Autom. Control., 1910 1994.
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Au- thur, Ben Bogin, Khyathi Raghavi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Harsh Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Ian Magnusson, Ja- cob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Abhi- lasha Ravichander, Kyle Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Pete Walsh, Luke Zettlemoyer, Noah A. Smith, Hannanch Hajishirzi. Iz Beltagy, Dirk Groeneveld, Jesse Dodge, and Kyle Lo. Dolma: an open corpus of three trillion tokens for lan- guage model pretraining research. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), ACL 2024, pp. 15725-15788, 2024.
Mikhail V. Solodov. A bundle method for a class of bilevel nonsmooth convex minimization prob- lems. SIAM J. Optim., 18(1):242225, 2007a.
15

---

Published as a conference paper at ICLR 2026
Mikhail V. Solodov. An explicit descent method for bilevel convex optimization. J. Convex Anal., 14(2):227-233 238, 2007b.
Ilya Sutskever, James Martens, George E. Dahl, and Geoffrey E. Hinton. On the importance of ini- tialization and momentum deep learning. In Proceedings of the 30th International Conference Machine Learning, ICMI. 2013, volume 28 of JMLR Workshop and Conference Proceedings, pp. 1139-1147, 2013.
Junjiao Tian, Chengyue Huang, and Zsolt Kira. Rethinking weight decay for robust fine-tuning of foundation models. In Advances in Neural Information Processing Systems 38: Annual Confer- ence on Neural Information Processing Systems 2024, NeurIPS 2024, 2024.
Andrey Tikhonov. On the solution of ill-posed problems and the method of regularization. Dokl. Akad. Nauk SSSR, 151(3):001-504, 1963.
Hugo Touvron, Louis Martin, Kevin Stone. Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao. Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuming Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CoRR, 0tss/2307.09288 2023.
Shaowen Wang, Anan Liu, Jian Xiao, Huan Liu, Yuekui Yang, Cong Xu, Qianqian Pu, Suncong Zheng, Wei Zhang, and Jian Li. Cadam: Confidence-basse optimization for online learning. CoRR, abs/2411.196477 2024.
Andreas S. Weigend, David E. Rumelhart, and Bernardo A. Huberman. Generalization by weight- elimination with application forecasting, In Advances in Neural Information Processing Sys- tems 3, NIPS Conference, pp. 875-882, 1990.
Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. CoRR, abs/2509.0204g 2025.
Shuo Xie and Zhiyuan Li. Implicit bias of adamw: l -nomm constrained optimization. In Forty-first International Conference on Machine Learning, ICML 2024, 2024.
Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Trans. Pattern Anal. Mach. Intell., 46(12):9508-955 9520,,9500000 2024.
Zeke Xie, Zhiqiang Xu, Jingzhao Zhang, Issei Sato, and Masashi Sugiyama. On the overlooked pitfalls of weight decay and how to mitigate them: A gradient-norm perspective. Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Process- ing Systems 2023, NeurIPS 2023, 2023.
An Yang. Anfeng Li, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng. Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jian Yang, Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang. Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang, Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren. Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yuqiong Liu, Zekun Wang. Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report. CORR, abs/ 250.0/90509888 2025.
16

---

Published as a conference paper at ICLR 2026
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael W. Mahoney. ADAHESSIAN: an adaptive second order optimizer for machine learning. In Thirty-Fifih AAAI Conference on Artificial Intelligence, AAAI 2021. Thirty-Thire Conference Innovative Appli- cations of Artificial Intelligence, IAAI 2021, I, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, pp. 10665-10673, 2021.
Manzil Zaheer, Sashank J. Reddi, Devendra Singh Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NewrIPS 2018, pp. 9815- 9825, 2018.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. CoRR, ftbs/2001.02385, 2024.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P. Kingma, Yinyu Ye. Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates gain more. In The Thirteenth International Conference on Learning Representations, ICLR 2025, 2025.
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham M. Kakade. Decon- structing what makes good optimizer for autoregressive language models. The Thirteenth International Conference Learning Representations, ICLR 2025, 2025.
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin Wang. Amin Saied, Weizhu Chen, and Nan Duan. Agieval: A human-centric benchmark for evaluating foundation models. In Findings of the Association for Computational Linguistics: NAACL 2024, pp. 2299- 2314, 2024.
17

---

Published as a conference paper at ICLR 2026
A NOTATION AND DEFINITIONS
= (1,2,3,... denotes the natural numbers. For n E N, [n] denotes the set 2.. n). Vectors are denoted in lowercase boldface, and matrices denoted in capital boldface. 0 and de- note the all-zeros and all-ones tensors of appropriate dimension, respectively. Scalar operations and functions, e.g. multiplication, division, and square roots, understood to be performed entrywise when applied to vectors. We also use @ to explicitly denote the entrywise product. x* denotes the positive part of 2, i.e.
:
max(0, = f0r {{o if 0
otherwise
I|-llp denotes the ep norm for p E [1, so]. (: denotes the standard inner product on R. [x]; denotes the jth entry of a vector x. diag (x) denotes the diagonal matrix with diagonal entries given by x. I[x 2 0) denotes the indicator tensor that in a coordinate if nonnegative in that coordinate and 0 otherwise. If K IRd - IR is convex, we let 0K((x) denote the set of subgradients of K. at x and overload Vx((x) to denote an element of [x()
Definition (L-smoothness). A function f IR + R is L-smooth if it is differentiable and [VV V f(x)ll2 - L Iyy x||2 for all x, E g
If f is L-moooth, then
f(y) S f(x) fx()), x) + bly x12 for all x, E R.. Definition (Coerciveness). function f :Rd + coercive if f(x) xc as |x|| oo.
B PSEUDOCODE OF OPTIMIZERS WITH CWD
B. SGD WITH MOMENTUM
Algorithm 2 SGD with momentum and cautious weight decay
I: given learning rates {m JeEN C Rso, momentum coefficient /3 E [0, 1), weight decay coefficient À > 0
2: initialize time step 1, parameters E Rdd first moment mo 0
repeat
4: g: f StochasticGradien(xx)
5: m. + 3m;-1 + 1 8)g:
6: X
m, +mm,,,,, 0)x,
entrywise multiplication
t+t+1
until stopping criterion is met
9: return optimized parameters X
B.2 LION-IC
Algorithm LION-K with cautious weight decay
I: given learning rates heN C 1x,. momentum coefficients 31, 82 E [0, 1). convex KC : led - 18 with subgradient VK. weight decay coefficient À > 0
initialize time step I, parameters X: E Idd, first moment m. <-
3: repeat
4: ge <- StochasticGradient
5: mi+1 + 32m, (1 :2)8:
6: m + 8.m, (1 f1)g.
7:
Xt+1 f- x: + np. (V((m)++1) (I(()((m))) 5 0)x:
b. entrywise multiplication
8: 1+1+1
9: until stopping criterion is met
10:
return optimized parameters x;
18

---

Published as a conference paper at ICLR 2026
B.3 LION
Algorithm LION with cautious weight decay
given learning rates (m JeN Rzo. momentum coefficients 31, 32 [0, 1),
weight decay coefficient >
2: initialize time step 1, parameters é I2 first moment mo <- 0
3: repeat
gi StochasticGrradenn(t)
m. + m + 31)g.
6: f- 1/: (gg((m)) + 2 0)x,
entrywise sgn and multiplication
7: m. - -2 + 82)g.
++++1
9:
until stopping criterion is met
10: return optimized parameters X,
B.4 MUON
Algorithm MUON with cautious weight decay
I: given learning rates {Nc)ee Rso. momentum coefficient & E [o, 1). weight decay coefficient
2: initialize time step + 1, parameters X1 E RR**" first moment Mo + 0
3: repeat
4: G, StochasticCradient
M. 8M G.
6: O, + NewtonSchulz(Me.
approximation of matrix sign
7:
X,+1 X, (O, +AT(O,X, 2 O)X,
entrywise matrix multiplication
8: 16 1+1
until stopping criterion is met
10: return optimized parameters X,
B.5 ADAM
Algorithm ADAM with cautious weight decay
I: given learning rates (ne C R20, momentum coefficients 0 81 1, numerical stability constant 2 0, weight decay cocfficient
2: initialize time step +- 1, parameters x E IRd first moment mo <- 0, second moment Vo <- 0
3: repeat
4: gi StochasticGrradienff
m. f- B. +
31 )g
6: V: + B2V + (1- 32)gi
> entrywise multiplication
7: ii + (1 - {9 m
8:
v. + (8) -ive
9:
X++1
1/ teltel
Am 2 0)x,
entrywise operations
1+
until stopping criterion is met
12: return optimized parameters X;
C FIXED-POINT ANALYSIS
Revisiting the fixed-point analysis in Section 2.2 for ADAMW, suppose the trajectory of ADAMW
converges fixed point (x*, m V*), that m V (x(x*) and v (x(x*)² Passing the
19

---

Published as a conference paper at ICLR 2026
limit > O, the fixed-poing condition gives
V fx(x)
+ Ax + sgn + Ax" 0.
fx))* +e1
Taking inner products with V fx((x) yields [| (Xxx Vf((')) 0,
0, which shows that x" a Karush-Kuhn Kuhn-Tucker (KKT) point of the constrained optimization problem
min f(x xER
s.t. ||x|
-
(3)
by Lemma 3.8 of Xie Li (2024). Intuitively, ADAMW normalizes the gradient its coordinate- wise sign at stationarity and then balances it against the linear pull of the decoupled weight decay, which enforces a box constraint on the parameters. Xie & Li (2024) formalize this intuition and show that whenever the iterates of ADAMW converge, the limit point is KKT point of the box- constrained problem (3.. However, this guarantee holds only under the assumption of convergence,
and ADAMW is not known to converge in general.
We remark that we can adapt this argument for another, more heuristic insight into why optimizers with cautious weight decay perform unbiased optimization. Suppose ADAM with cautious weight decay reaches a fixed point, so that
ff(x*) f((x)) ) +el XIX''((x'x'' 2 0)x"
For fixed point of LION-K with cautious weight decay, we have
-Vx(=V/(x*) = X[(x<(-V/(x')> S 50)x"
In either situation, casework on the signs of the update and x shows that both sides must be 0. follows that Vf(x*) = 0 for ADAM and VK(-V/(x*))
0 for LION-K, and if K. is convex function that achieves a unique minimum at 0 (e.g. a norm), then this condition becomes V fx(")) = 0 well. Hence, the fixed-poing analysis suggests that ADAM and LION-K with cautious weight decay find a stationary point of the original objective f.
D LYAPUNOV FUNCTIONS
Throughout this section, vector variables implicitly dependent on / when clear from context, and we drop the subscript for notational simplicity.
D.1
SGD
SGD with cautious weight decay admits the continuous-time dynamics
x V (x) XI(V/(()x 0)x,
which has a Lyapunov function 7xxx) = f(x), since
dHf dt -V(x)) AI(V
0)x) = = -VV [VV](())) s 0.
D.2 SGD WITH MOMENTUM
When SGD is equipped with momentum (Sutskever al. 2013) and cautious weight decay, the continuous-time dynamics becomes
-mm-- AI(mx 2 0)x m = 3V m),
which has a Lyapunov function
H(x, m) Bf(x) + |m| + 1 [mxx)
since dHd
dt (1VV((x) + AlImmx 2 0)m, -m - AI(mx 2 0)x) + (m + AI(mx 2 0)x, B(V(() - -m))
-AT(mxx 2 0) 31, m²) X(3 2) ||mm+)) 5.
20

---

Published as a conference paper at ICLR 2026
D.3 LION-K
We assume that is convex and satisfies sgn( (X((m)) = sgn m for all m 12d, This assumption mild and that holds for every example of K. given by Chen al. (2024).
The continuous- time dynamics LION-K without gradient enhancement is given by
x VK(m) λx
(4)
I -aVf(x) - ym.
Chen et al. 2024) showed that this system has a Lyapunov function
H((x m) = ax((x + 1x (Ax) + (Cx) + K(m) - (m, Xx);
thereby elucidating the origin of the K (Ax) regularization term. However, when equipped with cautious weight decay, becomes
x = VK(m) - Al(mmx < <0xx
(5)
-VVf(x) - ym
and admits Lyapunov function
H(x, m) ax((x) + Kmm) + + [l[(=x))+
(6)
which corresponds to optimizing the original objective f. To see that (6) is a Lyapunov function for 5), note that
HH dt
= (xV AI(mx & 0)m, VK(m) AI(mmx & 0)x)
+ (VK(m)) - AT(mxx & 0)x, -aVf(x)) -ym)
(VK(()) AI(mx & 0)x, (AI(mmx 0) ++)1m)
(AI(mx s 0) +1, VK(mm)) A(A 7) (mx) + ||| S. 0.
D.4 ADAM
The continuous-time time limit of ADAM with cautious weight decay yields the system of ordinary dif- ferential equations (cf. Barakat Bianchi (201))
(1 - exp(- 7 m
V(1 exp( l1 AI(mx 2 0)x
m a (V)(x) m)
xx)(xx2 v).
We assume that 0 < < 4a, which is satisfied by standard implementations of ADAM in practice.
This system admits the Lyapunov function
H(x, v,t af(x)) + 2( + e1)
Og m
+ (mx) 111
(8)
where
a := exp(-ol)) and
7t = (1 exp(-2())-1
21

---

Published as a conference paper at ICLR 2026
To see that ( is Lyapunov function for 77. note that H. is lower bounded by af and dtd
OH
di (V*H, + V m) + HV, v) +
[xV((x) + XIx 2 0)m,
a
Al(mx 2 0)x
v + +e
at /
a.m
Allmmx 2 aVx(((( - m)
4/v 70V + EE) (V)))>))
TVV +el
m
2a exx e1) - a 'iri""erse "exp((>()))))) 1
2
2 (a, (V +c))
a.m2
+ M(xx + () + O.7 CNTY m m
Vx
(a1 + AII(mxx 2 0))
4/v (VTVV + +e) 2,
V/iv d
017m² VTVV
m²
20exp(- +V + - a expl (xxp(-xxppppp))ppp -74)
(V E1) 2:
1
2
( 10.'"(7x)) +
a.m
az (22,000 -at) 77 eyepp( (app(--t))m?
( a AI(mx 0), nv + el
.1
4( nV + +1)
2 a 2(exp(at) - 1) 4
a
a.m²
1 --A(mmx 2 0),
4(exp(7t)-) -
v V +11
00,
where the first inequality drops some nonpositive terms and uses V/VV S VAv+ +el and the second inequality uses
? 2(exp(at) - + 4(exp(rt) - 0
for 0 < S 4a and
Remark 1. Cautious weight decay can be seen as an attempt to fix the asymptotic instability of ADAMW via Lyapunov function. Consider the simplified continuous ADAMW dynamics
m
Ax
Vv
m vf(x) m v)(xx)
(9)
and the function
m²
H((,, m, v) = f(x) +
+ (m,
2/V 1
By straightforward computation,
dH dt
m²
f + Am,
Ax
Ax, Vf(x) m
+
vf(x)²
V
4v
( X
m²
fx)
A(A + 1)mx
4v
3)
m²
1 m (
A + 4
v - 1)
4
Note that H is not guaranteed to be lower bounded and since (m, has unknown sign. This motivates the introduction of a mask I(mx 2 0) to the weight decay term and a slight adjusiment H that the result Lyapunov function for (
dHl is not guaranteed to be nonnegative,
de
For expositional clarity, we treat the ODEs and Lyapunov candidates in this section as smooth, even though the dynamics include the discontinuous indicator function II(ux 2 0). A fully rigorous analysis can be developed by interpreting the systems in the sense of differential inclusions, specifically, using Filippov's 's framework (Filippov) 1988). and by applying specialized tools from nonsmooth Lyapunov stability theory to obtain convergence guarantees (Shevitz Paden 1994 Bacciott & Ceragioli 1999..
Remark 2.
22

---

Published as a conference paper at ICLR 2026
E CONVERGENCE RATE OF ADAM WITH CAUTIOUS WEIGHT DECAY
In this section, we show that under the following assumptions, ADAM with cautious weight decay (Algorithm [6) achieves a convergence rate on the squared gradient norm and an additional station- arity measure.
Assumption (Smoothness). f is L-smooth and lower bounded by finite constant f".
Assumption (Bounded variance). The stochastic gradient g: satisfies
o²
[g xx] (xx) and Var(g:) = E || V. X
1bbatch
where a is constant and Abatch denotes the batch size.
Assumption 3 (Bounded iterates and bounded gradients). There exist constants R and G such that | llx I G a.s. for all t E N.
||xe|l> & R and
Assumptions 1 and are standard and often used in the analysis of stochastic gradient algorithms
(Ghadimi & Lan 2013 Barakat & Bianchi 2021 Défossez et al.) 2022 Arjevani et al., [202)) Assumption 3 can be justified using the Lyapunov function (8) if f is additionally assumed to be
coercive, since Goobins-Sieegaat
argument with sufficiently small ry shows that the optimizer
states remain in a compact sublevel set of H. a.s. For the sake of clarity, here we take it as an explicit
assumption. Similar assumptions have often been used for the analysis of ADAM-style algorithms (Kingma & Ba
2015; Reddi et al. 2018 Zaheer et al. 2018 Chen et al.
2019
Défossez et al,
2022
Chen al. [2022).
Theorem 1. Under Assumptions 7 2 and
3
let B1 S 82 <1, X2 A 0, 0, and n 17 0. Then for all T N.
and suppose x4 is updated using Algorithm
6
Ki
K₂ T
K40
M KE
IV + (V 11]
S
+
+ Kan+ +
T
n/
Nbatch
where K1, K2, K3, and KA are constants depending only on L, R, G. B2 and f(xi)-f'.
Remark 3. The first term on the left-hand side, [x reflects how much is optimized, while the second term, (x))(x))x,)) II reflects the degree of conftict between the objective f and the parumeter magnitudes. If V f(x:)xx Do O, then there is room to jointly decrease both I and the magnitudes. Thus, small value of (xx)(x reached a state where it is difficult to further decrease f and shrink the magnitudes simultaneously. This corresponds convergence toward a Pareto front, where trade-offs between the objectives become unavoidable.
indicates that the optimizer has
Remark 4. In the setting of Theorem let T E N and 17 = (v+) Then
T Σ E + CE|T]
+A](V =
o VT /Tbatch
An OTT 2) bound can be obtained by making the unrealistic assumption Abatch = OT)) However, even without this assumption, the stated bound is of theoretical interest. For additional discussion, see Bernstein et al. (2018) Zaheer et al. (2018): Chen et al. (2024).
Lemma 1. For all t E N.
m
31
=: C. B2
Proof. It suffices to work an arbitrary coordinate i. Let m = [mr]ss v = [eels, and ge := [g]]- By expanding the update rules for m and v, we obtain
1
3
3 B gs
and v
1- :2
Σ
3 gk.
1 g
1 135
kE
kE((]
23

---

Published as a conference paper at ICLR 2026
Now by Cauchy-Schwarz.
m
(1 - 31)²
1 1 $2
P
t-
(1- (1-81)
1 1 $1
Σ 39
(1 81)²
1 B2
B2
(1
B2
kElt]
kE(d
(1 - 3,)²
1 1-
82
1 1 1
1 31
1 - 31
1 B1
81)2
1 B2
1
31
32
3
B2
The conclusion follows from
m
v m
1 31
V +
V
1 - 32
Fact 1 (Lemma F.1. Bernstein et al. (2018). For all E N, i [d], and on 02,...... Or E R,
E
an ([gk] - [V((x)))))
ak.
Thatch kE[(]
AE
Lemma 2. For all / E N,
E [(l()(() mellil G/G++ 3inLd(C + +AR)
ad
1-
Nbatch (1 + 31)
Proof. Note that
m -8iVf(()
Σ
Σ kE[t]
0 (g: (f(x))
RE[t-1]
(10)
By smoothness, Lemma] D and Assumption B
we have
[[v fx(+++1) s vdllVf(xx) - vf(x++1)l22 1 LVdd|xxHr - -xxll <LL((CARR)
(11)
By Jensen's inequality and Fact I
E
Σ 3 k ([Ex): (x((x)))))
IE
Σ kE(C]
B ((gx): [x(xx)]))
[E(()
(12)
(8f))-b <
σ
Nbatch
v Mbatch (1
137 33)
kE((t]
Taking [|]]-|1]] of (10) and applying I1 and 12
B1 B. AR)
B|ll fxx)) milldl 8 [V(())]
(1 B1)
Σ 3 kgr vf((x)) kE/[]
1-
BinLd + AR)
od
Gd
1 1-
Mbatch (1 (1+
as desired.
Lemma 3. For all ( E N,
E | (x) 11
BGGd BInGLd(C+ + AR)
oGd
E
V vx(x))
<
VV. +
G+c
(1-Bi)e
ev Tbatch (1 + +3)
24

---

Published as a conference paper at ICLR 2026
Proof. We have
v
(xxt)
m.
Vx(x))
Vf(x.))
m, - vs(xi)
VV. +
VV. +
1
(x(( Vs(xx) m
[V
G+c
Vv, +1
1 |V
Vf(x))
[Vf(x,) melli
G+E
VV, el
The result follows by Lemma 4, For all m, E.R
vf(x))
S and Lemma
Wi+e
Kl[mm 2 0) Mgr 2 0))ar S I1(gg S [)]z]
Proof. If x 0, then the inequality is trivially valid, so suppose + 0. We proceed by casework on the sign of mgg
If mg > 0, then m and g have the same sign, and the conditions ma and ga 2 0 equivalent. Thus I(mrr 2 0) - [gg 2 0) = 0, and the inequality holds.
If my 0, then T(mg S 0) 1. It remains show ([mm 2 0) Mggar 2 0))xl S [z]] which
follows upon realizing L(max 2 0) - 1(gz 0) E 1,.
Lemma 5. For all t E N.
3n + AR)
RRd
El- (V)(())
- --
1-31
Mbatch (1 + 31)
Proof. We have
- (Vf(x)) 2 0)x;)
(V 0)x + (I(mxx 0) V((x,) 2 0))x;)
[VX( 2 0) T(m,x; 0) x))) [xxx I
< (I] 2 0) [(m,x 2 0))x - (x, 11
( 5 [0) xx/) (V [x,
(13)
where the fourth line uses Lemma the inner product,
Taking the expectation
conditioned on x, and expanding
(xx) V (xx,) - 0)(x)))
| xx] ((V xx) S 0) x 1x,)
[V
X B[[ [m. [V/(x))] 5 xx]
[xxxx))
[x.]a] | [xx))] 0 x.)
- [mi]]] | 2 [x)(())) |x.)
[E[d]
> [xx]]al [[ fxx
[m.] xx]
*E[d]
< R
[[] - m.ll xx),
(14)
where the fifth line uses Markov's inequality. Taking the expectation of 14 and applying Lemma? ?
Ln + AR)
RRd
E| - (V,( 2 0)x,) < - -
1 1- 31
Nbatch (1 B1)
as desired.
n
25

---

Published as a conference paper at ICLR 2026
We are now ready to prove Theorem |
Proof of Theorem
! Let
m
A = [xx++1) (xx)) and d,
+ Xm 2 0)x
vv, e
By smoothness,
A S (Vf( X -xi) + 2
|x++1 - -x
7 (V 61) + L ||ōe112
2
(15)
-- (V((x))
m
(V ((mx 0)x,) +
118:1/2
Vv. + c1
2
fx(x),
m
77
(V)(() 0)x,) 1/
1
1- 39
Vv2
Taking the expectation of 15 and applying Lemmas
and
I (xx))
B
BnGGLd(C + AR)
oGd
E[Ad] <
1 B
G+e
B1)e
Mbathh (1 + 31)
SILLd(CC + XR)
o RRd
+ yd [B[]](V/(xx,)))
G|Rd++
1 1-31
Abatch (1 + B1)
(16) and
We can assume G 2 T- without loss of generality. Rearranging [6.. using 1 1-3 3i
(1 - (G + c) 2 1, summing over T iterations, and dividing both sides by T gives
1 Σ E[S(x)) G+E -() T
ft) G+e Σ
B&Gdd +
nT
(1-
EE[T]
GdiG +6)
A(G cc)
GRdd Ao +c)
EV Pbacch (1+
T
Vmbatch (1 + B1)
TE(T)
SInALRd(C + (G +6)
+
+ ( ()
1-
K1
K2
K40
<
+
K2n
nT
TDbttch
where the fourth line uses E 8 < T-
B and
S(x,) = [V + (V
K, = (G (G++) ()) - 1r)
Gd(
( + AR
K₂ = 1- B
K3 = Ld + (R)(G++))
G
+ + Ld +000=++21100 (G+
1- Bi
K4 = d(G+c) G
+ AR) +
VI+ Bi
26

---

2.3B Model 986M Model 338M Model 111M Model Published as a conference paper at ICLR 2026
Table Hyperparameter configurations for the different model sizes. All models use an expansion factor of 8 and a vocabulary size of 100,864.
Hyperparameter
Model Architecture
Total Parameters
2,321.38M
985.89M
338.44M
110.55M
Model Dimension
2048
1536
1024
512
Number of Layers
18
12
Number of Heads
8
8
8
8
Per Head Dimension
256
256
128
64
Sequence Length
2048
2048
2048
2048
Validation Setup
Evaluation Batch Size
1024
512
128
256
Number of Eval Steps
4
Evaluation Interval
1000 steps
1000 steps
500 steps
500 steps
F MODEL. & EXPERIMENT CONFIGURATIONS
We evaluate cautious weight decay (CWD) across two experimental setups: (1) transformer models ranging from 111M to 2.3B parameters, and (2) the OLMo-1B architecture. All models employ SwiGLU activations and rotary position embeddings (RoPE). To ensure fair comparison, duct extensive grid searches to optimize hyperparameters for each baseline optimizer (ADAMW, LION, and MUON) before applying CWD with identical settings. Table 8 details the scaled model configurations, Table / presents the OLMo-IB IB architecture, and the following subsection describes our hyperparameter search methodology.
We conducted an extensive grid search to determine optimal hyperparameters for ADAMW, LION, and MUON optimizers. Our learning rate search employed a quasi-logarithmic grid spanning four orders of magnitude from 1 x 10-5 to x 10-1, with denser sampling in the critical 10-1 to 10-2 range where transformer models typically achieve optimal performance. The grid included standard decade values (e.g.. 0.001, 0.01) well as intermediate points within each logarithmic interval (e.g., 1g.,0 0.2, 0.3, 0.5, 0.8 scaled to each decade) to capture potential performance peaks between order- of of-magnitude boundaries, totaling 24 distinct learning rate values. For the learning rate schedule, we systematically evaluated warmup ratios of {0, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5), corresponding to 0% to 50% of total training steps dedicated to linear warmup, followed by cosine annealing decay. For ADAMW, we additionally performed a grid search over the momentum parameters Bi and 82. evaluating combinations of 31 E (0.85, 0.9, 0.95) and B2 E (0.95, 0.98, 0.99, 0.995, 0.999). Our experiments identified 31 = 0.9 and 3g = 0.95 as the optimal configuration. For LION, we swept 31 {0.85, 0.9, 0.95) and B2 e (0.95, 0.98, 0.99). finding
0.9 and B3 0.95 be optimal. For MUON, we similarly swept momentum coefficients and confirmed 0.95 as optimal.
G ADDITIONAL EXPERIMENT RESULTS
This section provides supplementary experimental analyses that further characterize the behavior of cautious weight decay (CWD) across different optimizers and training dynamics. We present detailed visualizations of the mask activation patterns (Figure 7 showing how the fraction of parameters re- ceiving weight decay evolves during training for both ADAMW and MUON optimizers. Additionally, we include comprehensive loss and accuracy curves for all three optimizers (ADAMW, LION, and MUON) across model scales from 11 to 2.3B parameters (Figures IO demonstrating consis- tent improvements with CWD. Finally, Figure 13 tracks the evolution of parameter norms throughout training, revealing that CWD maintains stable regularization comparable to standard weight decay while achieving superior performance.
27

---

Published as a conference paper at ICLR 2026
Table 9: Model Architecture Configuration for OLMo-
Hyperparameter
Value
Architecture
Hidden dimension (dmodal)
2048
Number of attention heads
16
Number of layers
MLP ratio
8
Vocabulary size
50,280
Embedding size
50,304
Max sequence length
2048
Attention Mechanism
Positional encoding
RoPE
Flash attention
Multi-query attention
ALiBi
Attention dropout
0.0
Attention layer norm
x
Model Components
Activation function
SwiGLU Sequential
Block type
Weight tying
Include bias
Layer norm type
Default
Layer norm with affine
Residual dropout
0.0
Embedding dropout
0.0
Initialization
Initialization method
Mitchell
Initialization device
CUDA
Training Trajectory
Steps)
Pull Training Trajectory
Steps)
0.22
Ratio
Ratto
2:00
Ratio
Ratio
Ranto WAdamM
3.20
Rat (Adami))
Ratio (Muoni
Ratio MMnn
Training Steps
Training Steps
Training Steps
Training Steps
Figure 7: Masked weight-decay activation ratio : [f(,, i.e., the fraction of parameters for which the sign-selective mask active at step (d number of parameters). Left: ADAMW: right: MUON. Insets zoom into the first 2.5k steps to highlight carly-training behavior. Model: Qwen-0.6B Yang et al. 22025) trained on The Pile (Gao et al. 2021).
28

---

Published as a conference paper at ICLR 2026
Accuracy (higher is better)
GPT
ADAMW
LION
MUON
Model Size
Ours
Base
Ours
Base
Ours
Base
338M
0.4232
0,4221
0.4230
0,4211
0.4256
0,4252
986M
0.4566
0,4556 0.4831
0.4552 0.4839
0.4545
0.4589
0.4575
2
0.4847
0,4830
0.4873
0,4858
Loss (lower is better)
GPT
ADAMW
LION
MUON
Model Size
Ours
Base
Ours
Base
Ours
Base
338M
3.0059
3.0136
3.0012
3.0121
2.9851
2.9896
986M
2.7053
2.7142
2.7171
2.7231
2.6873
2.6968
2B
2.4881
2.4973
2.4961
2.5012
2.4703
2,4803
Table 10: Final evaluation accuracy (higher is better) and loss (lower is better) comparisons across different model sizes, expanded to the full text width. Our proposed method benchmarked against three baseline optimizers: ADAMW, LION, and MUON. The best result in each pair is bolded.
Eval Loss Comparison
Eval Loss Comparison
Eval Loss Comparison
3.0
3.5
Ours
3.5
Ours
Ours
Muon
3.4
Muon
2.9
Muon
3.4
2.8
Loss
Loss
3.0
2.9
3.1
2.5
3.0
5000
10000 15000
20000
5000
10000
15000
5000
10000
15000 20000
Step
Step
Step
(a) 338M parameters
(b) 986M parameters
(c) 2B parameters
Figure Training dynamics across model scales with MUON optimizer. Baseline MUON (dashed) vs. MUON with CWD (solid).
Eval Loss Comparison
Eval Loss Comparison
Eval Loss Comparison
Ours
Ours
Ours
3.5
AdamW
AdamW
3.0
AdamW
3.1
Loss
Loss
3.3
2.8
2.9
2.7
3.3
2.7
25
5000
10000 15000
20000
50000
10000
15000
5000
10000
15000
20000
Step
Step
Step
(a) 338M parameters
(b) 986M parameters
(c) 2B parameters
Figure 9: Training dynamics across model scales with ADAMW optimizer. We compare baseline ADAMW (dashed) against ADAMW with CWD (solid) on models ranging from 338M to 2B param- eters.
29

---

Published as a conference paper at ICLR 2026
Eval Loss Comparison
Eval Loss Comparison
Eval Loss Comparison
Ours
Ours
Ours
Lion
29
Lion
3.6
3.5
LO55
Loss
3.4
2.9
2.6
2.8
3.0
15000 Step
20000
Step
Step
(a) 338M parameters
(b) 986M paramcters
(c) 2B parameters
Figure 10: Training dynamics across model scales with LION optimizer. Baseline LION (dashed) vs. LION with CWD (solid).
Eval Loss
Training Loss
RMS Norm of Params
Ours
RMS Norm of Params
AdamW
31
Adam
3.3
Lo50
3.0
Loss
3.0
2.9
299
0.04
2.8
2.8
2.7
0.02
5300
10000
15000
10000
5000
10000
15000
Step
Step
Step
Figure 11: Training dynamics for the 986M parammeeee Gemma model.
30

---

Published as a conference paper at ICLR 2026
111
338M
986M
2B Global Grad RMS Norm
Global Grad RMS Norm
Global Grad RMS Norm
Global Grad RMS Norm
Adamm
AdamW
AdamW
AdamW
AdamW
10-d
20
10009
Stop
Stop
Stop
Global Grad RMS Norm
Global Grad RMS Norm
Global Grad RMS Norm
Global Grad RMS Norm
Ours
Ours Lion
10
Ours Lion
1p
Lion
10
10a4
Lonk
Reddec
10*h
0
1003
2000 Step
3300
20020 Stop
20002
5003 10200 25005 23000
5000 20020
23003
Stop
Stop
Global Grad RMS Norm
Global Grad RMS Norm
Global Grad RMS Norm
obal Grad RMS Norm
Muon
10-4
Muon
20"h
Muon
Muan
Muon
10-
10-1
Stop
Step
Step
Figure 12: Comparison of gradient norms using RMS normalization across four model sizes: 111M, 338M, 986M, and 2B. All models are trained under Chinchilla settings. CWD achieves lower gradient norms across all configurations.
31

---

Published as a conference paper at ICLR 2026
block 0/attn/a
black Ofattn///e dim scalefscale
block 0/attn/aky
block 0/ffn
RUPS
duta
adammw
adamm adami
acamt
adamm
adamn
value
adam
block 0/ffn gate/w
block O/ffn 1/w
black 0/post 0/scale
blockk 0/post In. 1/scale
644
addrmw
ouPS
ourz
scama
adtem
adam
363m
block Sookk_0/peerr____ 0/scale
block 0/pre 1/scale
block 1/attn// _rooj
block 1/attn/per dim scale/sci outs
ours
adammw
adamw
adami
valus
adam
36am
thfe
1999 the
1/attt/gkv_pptttteeeeeeaec proj
block_I/fff
block O_gate/w
1/W
OurS
OLTS
ours
OLTS
adamm
acarpw
adamow
20omn
543
1130
8550
1700
700 1000
Figure 13: Evolution of parameter norm (RMS) during training for 986M parameter model. We compare three optimization strategies: ADAMW with weight decay 0.1 (orange), our proposed method (blue), and ADAM without weight decay (green). Our method maintains stable parame- ter norms comparable to ADAMW while achieving improved performance.
32


Reviews:
[{"review_id": "4jKpS3To33", "reviewer": "Reviewer_XuM6", "summary": "This paper proposes Cautious Weight Decay (CWD) — a simple, optimizer-agnostic modification that applies weight decay only when the update direction matches the parameter sign. Unlike standard weight decay, CWD preserves the original objective and admits a bilevel/sliding-mode interpretation. It requires no extra hyperparameters and works as a drop-in replacement for AdamW, Lion, and Muon. Experiments on large-scale language model pretraining and ImageNet classification show consistent improvements in loss and accuracy.", "questions": "Please refer to the weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. Well written, easy to follow, and clearly presented.\n2. Theoretical foundations are solid and supported by extensive experiments across both vision and language tasks.", "weaknesses": "While the proposed method is elegant and broadly applicable, its novelty may be somewhat limited. Conceptually, it appears closely related to Selective Projection Decay (SPD) [1], although applied in a different setting. SPD selectively applies regularization to layers whose updates are inconsistent, determined by the sign of the inner product between the negative gradient direction $(-g_t)$ and the accumulated parameter change $(\\theta_{t-1} - \\theta_0)$. A positive inner product indicates progress toward lower loss, whereas a negative value implies potential movement toward higher loss, prompting SPD to impose stronger penalties on such layers. The main difference is that SPD is designed for robust fine-tuning and therefore regularizes the $L_2$ distance between pretrained and finetuned weights $(\\theta_{t-1} - \\theta_0)$, while CWD regularizes the $L_2$ norm of the weights themselves $(\\theta_{t-1})$. Under this substitution, the core mechanism becomes mathematically similar, making the originality of CWD somewhat questionable, despite the fact that this paper provides stronger theoretical justification and demonstrates wider applicability across optimizers and domains.\n\n[1] Tian, Junjiao, Chengyue Huang, and Zsolt Kira. \"Rethinking weight decay for robust fine-tuning of foundation models.\" Advances in Neural Information Processing Systems 37 (2024): 22418-22440.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lll59v7uYB", "reviewer": "Reviewer_Jtg5", "summary": "Cautious weight decay describes a one-line modification of weight decay such that it is applied only on parameters that are aligned with the update direction.\nThe paper shows that this is related to minimizing the unregularized loss, with limiting point such that all weights have minimal magnitude locally.\nLarge-scale experiments show the effectiveness of the method.\n\n**Overall,** the contributions and evaluations in this paper are of highest quality and the presented technique is likely to be adapted in practice quickly.", "questions": "* Did you test the effect of CWD for different learning-rate schedules (e.g. cosine and WSD)? Did you test whether CWD is compatible with AdamC by Defazio, 2025, and if there is a compound improvement of using both?\n* Does the floating-point format have an effect on CWD, in the sense that using low-precision formats might introduce a harmful bias in the direction/weights?\n* Figure 3 left: why does the CWD trajectory not stop once it reaches the Pareto set? \n* Assumption 1 states that the iterates are bounded due to coercivity. Can you provide a short proof for this, and also clarify in which probabilistic sense (as the iterates are random variables)?\n\nMinor:\n\n* line 220: why does $m$ necessarily decay to zero? In the stochastic case, with non-vanishing noise, this is not clear to me.\n* The entire paper uses the notation $(\\nabla f(x_t) x_t)^+$, but this is missing the transposition operator on the gradient.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper presents a seemingly simple idea that can be implemented with changing one line of code change. The effects of this modification are extremely well explained, and justified through experiments on industry-scale problems for language model training; the method also performs well on standard classification models for Imagenet. Several ablations are provided that justify the exact construction of cautious weight decay.\n\nBesides the clarity and convincing experiments, the paper also contains convergence analysis of the method, and extensive motivation and background through Lyapunov analysis.", "weaknesses": "The paper has no major weaknesses. The only minor point/question is on the motivation of CWD through Lyapunov analysis: the design of CWD is not explicitly reflected in the choice of its Lyapunov function. This leaves open the question whether other modifications of weight decay would allow for the same Lyapunov function, and if these alternatives would be interesting/competitive methods. From the current presentation, I could not find an argument why this modification of weight decay would be the singular one to yield improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VeBETOeQbq", "reviewer": "Reviewer_hDqM", "summary": "The authors propose a simple modification to weight decay, only applying the decay if the optimizer updates and the weights are aligned. The show mathematically that this leads to optimizing the original loss rather than a regularized loss. Through extensive empirical work, the authors demonstrate: 1) Adding the fix to Adam/muon/lion all improves convergence speed for LLM training without retuning hyperparameters. (fig4) 2) The performance is consistently better across levels of weight decay (fig1), 3) ImageNet accuracy is also improved.", "questions": "1. Are you able to provide scaling plots of final-loss-vs-model size for the baseline and your method? It would be interesting to see if the \"boost\" in performance is retained as the model is scaled.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The method is very simple\n- The empirical results are very convincing, especially since they use multiple codebases\n- The paper is well written and easy to follow", "weaknesses": "- Personally, I think the theoretical results are distracting. I'd prefer that they be moved to the appendix and that the main paper is focused on the empirical results. That is what most people will care about.\n- There is some missing related work (e.g., \"Understanding decoupled and early weight decay\")", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TQixSwrO5H", "reviewer": "Reviewer_yQsL", "summary": "This paper proposed a novel way of weight decay, which only implements weight decay on the coordinates that the direction aligns with the gradient. The new method is flexible to work with any optimizer and can be implemented easily and efficiently. It can converge to the minimizer manifold, unlike decoupled weight decay will add some additional constraint and can’t reach the minimizer.  The paper theoretically shows the converged point is a Pareto optimal point in the manifold. The experiments demonstrate it empirically works better than decoupled weight decay or no weight decay.", "questions": "1. Can you address the issues about convergence rate as I mentioned in weaknesses?\n2. I don’t understand the interpretation of the second term in remark 1. Are you suggesting that it will converge to a Pareto front? This claim can only be rigorously shown by a proof similar with Xie and Li 2024. \n3. This kind of cautious weight decay makes SGD not rotation-invariant. Have you thought about global cautious weight decay, which implement the weight decay when its inner product with gradient is positive?\n\nThere are some possible typos or things that are not clear enough to me:\n1. Page 4 line 181: the set of stationary points **is** a union of connected\n2. Page 5 line 263: the order comparison between two vectors is a bit ambiguous. I suggest write $y_i$ and $x_i$ explicitly in the scalar order\n3. Figure 3: left function should be $f(x,y)=((y-3)^2+(x-3)^2-1)^2$ in the caption. It’s also better to indicate that purpler region is global minimizers. \n4. Page 22 line 1153: the last line seems wrong. It seems you are using $||a+b||_2^2 \\leq ||a||_2^2 +||b||_2^2$ here. It can be correct if multiplied by 2.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The theoretical analysis on the continuous part shows the advantage of cautious weight decay, which can be an improvement over decoupled weight decay. The converged point of cautious weight decay has some good implicit bias. \n2. The extensive experiment results show the advantage of weight decay is promising. \n3. The comparison between different weight decay is very clear with the visualization of trajectories.", "weaknesses": "The statement of theorem 1 and the noise assumption are problematic. The four constants $K_1, K_2, K_3, K_4$ should depend on $L, \\sigma$ and initial loss gap, which needs to be shown explicitly to check whether the convergence rate is valid. When I read the proof, more problems arise. \n1. The noise assumption depends on the batch size, which is unconventional. In other papers that show convergence rate, the variance of the stochastic gradient is fixed $\\sigma^2$. Then they can still show $T^{-1/2}$ convergence rate. If this paper’s result is presented with the standard noise assumption, the convergence rate will have a constant term $\\sigma$ and therefore is even $\\Theta(1)$ for the dependence on $T$. \n2. I don’t understand why the iterates are bounded in Lemma 2. Also I don’t know the exact meaning of $G$. If it is a constant independent of $T$ and other problem parameters such as $L$ and $\\sigma$, then you are basically cheating with this lemma because you essentially assume the loss function is $G$-Lipschitz rather than $L$-smooth, which is a much stronger assumption. And you also assume the region is bounded, i.e., $||x_t||_\\infty \\leq R$, which is used in the proof of lemma 6. These should be mentioned explicitly in the main text because they are stronger assumptions than smoothness and noise variance. If $G$ can have some dependence on $T$, then $K_1, K_2, K_3, K_4$ in theorem 1 are not constant and the final convergence rate will have a much worse dependence on $T$. \n3. The constants $K_2, K_3, K_4$ are $O(d)$, which means the convergence rate has a very bad dependence on $d$. This is not ideal when $d$ can be much larger than $T$ in reality. Also they are $O(1/\\epsilon)$ so the dependence on $\\epsilon$ is also not ideal. But this issue is much less serious than the two issues above. It’s good to improve the dependence on $d$ and $\\epsilon$ but I can accept the current dependence as long as the proof fixes issue above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)