PaperlayerView exact revision on arXiv

Source-linked AI summary

Cautious Weight Decay

Lizhang Chen, Jonathan Li, Kaizhao Liang, Baiyu Su, Cong Xie, Nuo Wang Pierse, Chen Liang, Ni Lao, Qiang Liu

arXiv:2510.12402v2cs.LGmath.OCstat.ML

TL;DR

Standard decoupled weight decay can implicitly impose regularized or constrained objectives, while CWD applies decay only where parameter and update signs align and preserves the original loss. The paper frames its regularization effect as bilevel and reports improved training outcomes across language modeling and ImageNet classification.

  • Problem

    Decoupled weight decay can implicitly impose regularized or constrained objectives, and decay may resist beneficial movement when parameter and update directions differ.

    Source: §1 · p. 2

  • Method

    CWD applies weight decay only where parameter and update signs align, preserving the original objective while giving its regularization effect a bilevel interpretation.

    Source: §1 · p. 2

  • Results

    Across language modeling and ImageNet classification, CWD generally accelerates convergence and lowers final validation loss for AdamW, Lion, and Muon, with higher zero-shot accuracy from 338M to 2B parameters.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The reported improvements span benchmarks and architectures without retuning baseline settings, supporting CWD as a drop-in optimizer modification.

    Source: §1 · pp. 2–3

    Convergence cannot be established for the AdamW characterization because a Lyapunov function is lacking.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Decoupled weight decay can oppose optimizer movement in some coordinates and implicitly impose regularized or constrained objectives. Cautious weight decay (CWD) selectively applies decay, preserving the original loss while producing theoretical and empirical benefits.

1 Introduction

  • Decoupled weight decay is presented as a way to improve training stability and generalization by preventing adaptive learning rates from interfering with regularization.The introduction points to AdamW’s success in large-model training and later optimizers including Lion and Muon.Source: §1 · p. 2
  • Decoupled weight decay can resist beneficial movement when parameter and update directions differ, and it implicitly imposes regularization or parameter-norm constraints.The cited discussion names AdamW, Lion, and Muon among the optimizers associated with such constraints.Source: §1 · p. 2
  • CWD applies decay only in coordinates where parameter and update signs align, as a one-line modification without additional hyperparameters.The introduction presents it as a refinement of decoupled weight decay.Source: §1 · p. 2
  • CWD is analyzed as optimizing the original loss rather than a regularized surrogate, with its regularization effect framed as a bilevel search for locally Pareto-optimal stationary points.The convergence rate for discrete-time Adam is established in the smooth nonconvex setting under additional assumptions.Source: §1 · pp. 2–3
  • On the toy example, Adam + CWD continues with sliding-mode dynamics within the minimizer manifold, unlike Adam, which halts at a minimizer.The figure contrasts these trajectories with AdamW, which minimizes within a constrained region.Source: Figure 2 · p. 2
  • CWD generally accelerates convergence and lowers final validation loss for AdamW, Lion, and Muon, with higher zero-shot accuracy across benchmarks from 338M to 2B parameters.The reported improvements cover language modeling and ImageNet classification without retuning baseline settings.Source: §1 · p. 3

2 Background and Motivation

The section describes decoupled weight decay as an optimizer update augmented by direct parameter decay, alongside an adaptive update vector and its learning-rate and decay-coefficient terms.

2.1 Decoupled weight decay

  • The decoupled weight-decay framework augments a gradient-based optimizer update with a direct parameter-decay term.The framework is presented as applicable to a range of standard optimizers.Source: §2.1 · p. 3
  • The update vector can be adaptive and sign-normalized, built from first- and second-moment estimates such as momentum buffers and diagonal preconditioners.The learning rate η_t is positive, and λ is the nonnegative decoupled weight-decay coefficient.Source: §2.1 · p. 3
  • For Adam, the update description uses a diagonal preconditioner and a bias-corrected first-moment estimate, with β1, β2, and ϵ identified as momentum coefficients and a numerical-stability constant.The passage also specifies that the moment estimates are bias-corrected.Source: §2.1 · p. 3

Lion-K.

Lion-K forms a momentum-filtered update using a convex-function subgradient. The section relates decoupled weight decay to regularized or constrained objectives, motivating a mechanism that retains its practical benefits while optimizing the original objective.

Lion-K.

  • Lion-K forms a momentum-filtered step preconditioned using a subgradient of a convex function.Lion is given as the case K = ||·||1, and Muon as K = ||·||tr for matrix-valued parameters.Source: Lion-K · p. 3

2.2 Implicit regularization effects of weight decay

  • For SGD, decoupled weight decay is exactly a step on an ℓ2-regularized objective; for Lion-K, a fixed point suggests optimization of a regularized objective determined by K.The Lion-K fixed-point argument uses the convex conjugate K* to express the objective’s gradient condition.Source: §2.2 · pp. 3–4
  • For Lion and Muon, the implicit regularization corresponds to constrained optimization involving dual-norm balls, while AdamW is characterized as box-constrained.The section notes that convergence for the AdamW characterization cannot be established because a Lyapunov function is lacking.Source: §2.2 · p. 4
  • Because AdamW and Lion-K implicitly optimize λ-dependent regularized surrogates, the section motivates retaining decoupled decay’s practical benefits while optimizing the original objective.The named benefits include regularization and training acceleration.Source: §2.2 · p. 4

3 Cautious Weight Decay

Cautious weight decay is a one-line modification compatible with gradient-based optimizers, and its trajectories remain unbiased with respect to the original loss. On the stationary manifold, it induces sliding-mode dynamics that reduce parameter magnitudes toward locally Pareto-optimal stationary points; the paper also extends its analysis to discrete time.

3 Cautious Weight Decay

  • CWD is a one-line modification using entrywise multiplication and is described as implementation-trivial and universally compatible with gradient-based optimizers.The update-rule modification is introduced in Section 3; the supplied passages do not give its full formula.Source: §3 · p. 4

3.1 Convergence to the stationary manifold

  • Lyapunov analysis shows convergence to stationary points of the original loss for SGD, SGDM, Lion-K, and Adam, with momentum vanishing for momentum-based methods.For SGD, the decay is deactivated whenever it conflicts with the main objective; the convergence claim holds under the base optimizer's stated convergence conditions.Source: §3 · p. 4, §3.1 · p. 5, §3.2 · p. 6

3.2 Sliding mode dynamics

  • Within the stationary manifold, CWD induces sliding-mode motion that reduces parameter magnitudes until no simultaneous coordinatewise reduction remains, yielding a locally Pareto-optimal stationary point.The local Pareto ordering prioritizes smaller parameter magnitudes.Source: §3 · p. 4, §3.2 · pp. 5, 7
  • On two toy objectives, Adam reaches a generic minimizer-manifold point, AdamW solves a box-constrained problem, and Adam+CWD reaches the manifold's Pareto front.The exact CWD limit may depend on initialization and discretization because the Pareto set need not be a singleton.Source: Figure 3 · p. 7, §3.2 · p. 7

3.3 Discrete-time analysis

  • Discrete-time analysis extends the convergence guarantees to various optimizers and gives an explicit convergence rate for Adam; Figure 4 plots C4 evaluation loss across model scales and optimizers.Its curves compare CWD with tuned standard-decay baselines while reusing baseline hyperparameters for CWD.Source: §3.3 · p. 7, Figure 4 · p. 8

4 Experiments

Across the evaluated language-model settings, cautious weight decay (CWD) improves loss and downstream accuracy over standard decay, including under matched baseline settings. Ablations and training dynamics indicate that its gains depend on selective decay, and results extend to ImageNet, schedule variants, and larger models.

  • The language-model setup uses 20 tokens per parameter on C4 for 338M, 986M, and 2B models, reusing tuned baseline settings for CWD without retuning.The experiments also include OLMo-1B trained on 100B Dolma tokens, or 100 tokens per parameter.Source: §4 · p. 8
  • ImageNet experiments report similar gains across ViT and ResNet architectures.Table 3 reports validation accuracy across architectures and optimizers, comparing tuned-decay baselines with CWD at the same coefficient.Source: §4 · p. 9, Table 3 · p. 9
  • At fixed weight-decay coefficients, CWD lowers final loss while leaving the minimizing coefficient essentially unchanged, enabling reuse of tuned values without additional sweeps.The sweep covers λ∈[0, 0.4] for Muon and AdamW and λ∈[0, 3.0] for Lion.Source: §4 · p. 9
  • Random masks matched to CWD’s sparsity, gradient-based masks, and no decay all underperform the update-based mask.Random masking changes AdamW loss from 2.56 to 2.82 and Muon loss from 2.42 to 2.73; the table compares OLMo-1B trained on 100B tokens.Source: §4 · p. 9, Table 2 · p. 9
  • On 1B models trained for 100B tokens, CWD improves loss trajectories versus tuned AdamW and Muon; at 986M it ends with lower loss and an intermediate parameter norm.Without decay, loss falls faster mid-training but plateaus earlier at a higher value and largest norm; tuned AdamW ends with the smallest norm.Source: §4 · pp. 9–10, Figure 5 · p. 10
  • CWD consistently lowers evaluation loss for AdamW, Lion, and Muon across 338M–2B models while increasing downstream accuracy.The OLMo-1B, 100B-token setting also evaluates downstream reasoning benchmarks.Source: §4 · p. 10, Table 4 · p. 10
  • At 111M, 338M, 986M, and 2B parameters, Adam + CWD has lower final validation loss than AdamW, with the gap stable or slightly increasing.Figure 6 compares the methods on the same pretraining dataset.Source: §4 · p. 11, Figure 6 · p. 11
  • Adding CWD improves AdamW and AdamC, with AdamC + CWD achieving the lowest loss; Adam + CWD also improves over AdamW under cosine and WSD schedules.The best schedule result is Adam + CWD with WSD, which uses 10% warmup, a long stable phase, and 20% final decay.Source: Table 6 · p. 11, §4 · p. 12, Table 7 · p. 12

5 Related Work

Related work situates weight decay in modern optimization, alongside adaptive decay variants and sign-conditioned updates. The paper contrasts CWD’s bilevel interpretation with constrained-optimization interpretations of standard decoupled decay.

  • Decoupled weight decay became standard after AdamW separated decay from adaptive updates, with later optimizers including Lion [CLH+23], Lion-K, and Muon [LSY+25].Recent analyses describe weight decay in contemporary networks as a training accelerator and stabilizer, alongside work on schedules and normalization [PC23].Source: §5 · p. 12
  • Prior variants include scheduled decay, neuron-update analyses [KMJ24], adaptive decay [GSA23], and Selective Projection Decay for robust fine-tuning [THK24].The paper distinguishes SPD from CWD by intended setting, mechanism, theoretical properties, and empirical performance.Source: §5 · p. 12
  • Earlier sign-conditioned methods include Rprop’s gradient-sign step sizing, cautious optimizer updates [LCLL24], and Adam masking for online-learning robustness [WLX+24].Source: §5 · pp. 12–13
  • Decoupled decay has been interpreted as constrained optimization for Lion [PXA+25], AdamW [BN24], and Muon [CLL25] [LLS25], whereas CWD is framed as bilevel optimization.Source: §5 · p. 13

6 Conclusion

The conclusion presents CWD as an optimizer-agnostic modification that preserves the objective while retaining practical weight-decay benefits. The appendix provides notation and algorithm listings for several optimizer variants.

6 Conclusion

  • CWD preserves the optimization objective while retaining practical weight-decay benefits, with bilevel structure and convergence guarantees in continuous and discrete time.The conclusion reports faster loss reduction and more stable training trajectories without changing hyperparameters or model architectures.Source: §6 · p. 13

A Notation and Definitions

  • The notation defines entrywise operations, vector and matrix conventions, norms, inner products, indicators, smoothness, and coerciveness.Source: A Notation and Definitions · p. 22

B.1 SGD with momentum

  • The appendix lists an SGD-with-momentum variant with cautious weight decay, specifying learning rates, momentum, and a decay coefficient.Source: §B.1 · p. 23

B.2 Lion-K

  • The Lion-K listing specifies learning rates, two momentum coefficients, a convex function with subgradient, and a decay coefficient.Source: §B.2 · p. 23

B.3 Lion

  • The Lion listing gives its learning rates, momentum coefficients, decay coefficient, stochastic-gradient step, and parameter update.Source: §B.3 · p. 23

B.4 Muon

  • The Muon listing uses momentum and a Newton–Schulz approximation of the matrix sign in its optimizer steps.Source: §B.4 · p. 24

B.5 Adam

  • The Adam listing specifies learning rates, two momentum coefficients, a numerical-stability constant, and a decay coefficient.Source: §B.5 · p. 24

C Fixed-Point Analysis

The fixed-point analysis distinguishes AdamW’s conditional convergence to a constrained KKT point from cautious-decay fixed points that are stationary for the original objective. The section also formulates continuous-time dynamics for SGD and momentum SGD.

C Fixed-Point Analysis

  • AdamW’s fixed point is a KKT point of a box-constrained problem, but this conclusion is conditional on convergence, which is not known in general.At the fixed point, the analysis takes the limit as ϵ approaches zero and uses the resulting condition to establish the KKT property.Source: C Fixed-Point Analysis · pp. 24–25
  • For Adam and Lion-K with cautious weight decay, fixed-point sign casework implies stationarity of the original objective when K has a unique minimum at 0.For Lion-K, the intermediate condition is ∇K(−∇f(x⋆)) = 0; the stated condition on K makes this ∇f(x⋆) = 0.Source: C Fixed-Point Analysis · p. 25

D.1 SGD

  • The section formulates continuous-time dynamics for SGD with cautious weight decay.Source: §D.1 · p. 25

D.2 SGD with momentum

  • For momentum SGD with cautious weight decay, the continuous-time dynamics are paired with a Lyapunov function.The passage also gives an expression for the derivative along these dynamics.Source: §D.2 · pp. 25–26

D.3 Lion-K

The continuous-time Lion-K analysis gives a Lyapunov function corresponding to the original objective under cautious weight decay. The Adam analysis motivates a Lyapunov adjustment, while noting assumptions and limitations in the analysis.

D.3 Lion-K

  • For Lion-K with cautious weight decay, the Lyapunov function corresponds to optimizing the original objective.The stated assumption is that K is convex and sgn(∇K(m)) = sgn(m) for all m.Source: §D.3 · p. 26

D.4 Adam

  • The Adam analysis assumes 0 < γ ≤ 4α and gives a Lyapunov function for its continuous-time dynamics.The passage says this parameter condition is satisfied by standard Adam implementations in practice.Source: §D.4 · pp. 26–27
  • The AdamW Lyapunov candidate is not guaranteed to be lower bounded, motivating a cautious-decay adjustment to address asymptotic instability.The exposition treats the ODEs and Lyapunov candidates as smooth for clarity; a rigorous analysis can use differential inclusions and nonsmooth Lyapunov tools.Source: §D.4 · pp. 27–28

E Convergence Rate of Adam with Cautious Weight Decay

The section analyzes Adam with cautious weight decay under smoothness, lower-boundedness, bounded-variance, and bounded-iterate assumptions. Its convergence measure also tracks conflict between objective improvement and parameter magnitudes.

E Convergence Rate of Adam with Cautious Weight Decay

  • Under smoothness, lower-boundedness, bounded variance, and bounded iterates and gradients, the theorem bounds squared-gradient norm and an additional stationarity measure.The theorem also assumes 0 ≤ β1 ≤ β2 < 1, λ ≥ 0, ϵ > 0, and a constant positive step size.Source: E Convergence Rate of Adam with… · p. 28
  • The additional stationarity measure reflects conflict between optimizing f and reducing parameter magnitudes; a small value indicates a trade-off toward a Pareto front.The passage characterizes the measure as ∥(∇f(xt)xt)+∥1.Source: E Convergence Rate of Adam with… · p. 29
  • A stronger bound requires the stated unrealistic assumption that batch size grows as Θ(T); the section describes its bound as theoretically interesting without that assumption.Source: E Convergence Rate of Adam with… · p. 29

F Model & Experiment Configurations

The experiments cover scaled transformer models and OLMo-1B, with baseline hyperparameters searched before applying cautious weight decay using identical settings.

G Additional Experiment Results

Supplementary analyses track CWD’s mask behavior and training dynamics across optimizers and scales, reporting consistent improvements and stable parameter norms.

G Additional Experiment Results

Loading 2510.12402v2…