# Learning-Rate-Scaled Weight Decay

- ID: 1333
- Canonical URL: https://synthcore.org/idea/1333/learning-rate-scaled-weight-decay
- API JSON: https://synthcore.org/api/idea/1333.json
- API Markdown: https://synthcore.org/api/idea/1333.md
- Verification status: beats_baseline
- Source: [arXiv:2607.23777](https://arxiv.org/abs/2607.23777)
- Category: optimization
- Solves: speedup, stability, generalization
- ML areas: optimizer, training-dynamics, moe-routing
- Math tags: optimization, dynamical-systems, stochastic-processes
- Ratings: usefulness 7/10; difficulty 2/10; novelty 6/10

## Idea description

Replace constant decoupled weight decay with a coefficient proportional to the current learning rate divided by the peak learning rate. The optimizer applies ordinary decay at the learning-rate peak but weakens decay during cooldown and late training, preventing unnecessary steady-state parameter-norm shrinkage while retaining early-training stabilization.

## Mathematical statement

The paper starts from the regularized objective $\mathcal{L}_{\mathrm{reg}}(W)=\mathcal{L}(W)+\lambda\|W\|^2/2$, where $W$ is the parameter tensor, $\mathcal{L}$ is the task loss, and $\lambda\geq0$ is the nominal decay coefficient. It proposes $\lambda_t=\lambda(\eta_t/\eta_{\max})$, where $\eta_t$ is the learning rate at step $t$ and $\eta_{\max}>0$ is the peak learning rate. For a decoupled optimizer with raw update $u_t$, the parameter recursion becomes $W_{t+1}=(1-\eta_t\lambda_t)W_t-\eta_tu_t=(1-\lambda\eta_t^2/\eta_{\max})W_t-\eta_tu_t$. Constant decoupled decay instead gives $W_{t+1}^{\mathrm{const}}=(1-\eta_t\lambda)W_t-\eta_tu_t$. The proposed rule therefore has identical decay at $\eta_t=\eta_{\max}$, but its total shrinkage decreases quadratically with the learning-rate fraction. The paper argues that this preserves asymptotic stationarity guarantees of the corresponding unregularized stochastic-gradient and Muon methods, whereas constant decay can introduce persistent asymptotic bias. The practical prediction is that scaled decay yields parameter norms that stabilize instead of steadily declining.

## Key formulas

- $$\mathcal{L}_{\mathrm{reg}}(W)=\mathcal{L}(W)+\frac{\lambda}{2}\|W\|^{2}$$
- $$\lambda_t=\lambda\frac{\eta_t}{\eta_{\max}}$$
- $$W_{t+1}=\left(1-\eta_t\lambda_t\right)W_t-\eta_tu_t=\left(1-\lambda\frac{\eta_t^2}{\eta_{\max}}\right)W_t-\eta_tu_t$$

## Implementation notes

1. Integration point: modify the decoupled-weight-decay operation in AdamW, momentum SGD, or Muon. Do not inject $\lambda_tW_t$ into the gradient before momentum or adaptive preconditioning. Apply decay directly to each decay-eligible parameter tensor, excluding biases, normalization scales, and any groups excluded by the baseline. Keep the original learning-rate schedule unchanged. 2. Pseudocode: inputs are parameter tensor $W$, optimizer update $u$, nominal decay $\lambda$, current learning rate $\eta$, and peak learning rate $\eta_{\max}$. Compute `r = clip(eta / eta_max, 0, 1)`, then `lambda_t = lambda * r`, followed by `W = (1 - eta * lambda_t) * W - eta * u`. In AdamW, compute moments and the bias-corrected update $u$ exactly as usual before this line. In Muon, preserve its matrix-wise orthogonalized update and replace only its decay coefficient. 3. Computed from the paper: the schedule and multiplicative factor. Estimated empirically: per-layer norms $\|W_t\|$, validation loss, gradient norms, and effective cumulative shrinkage $\sum_t\eta_t\lambda_t$. Log these for every parameter group. 4. First experiment: train a roughly 100M-parameter decoder-only Transformer on a fixed language-model corpus for three seeds, comparing standard AdamW against scaled-decay AdamW at identical peak learning rate, nominal decay, batch size, schedule, and token budget. Measure validation perplexity versus tokens and wall-clock time, final loss, norm trajectories, and instability during warmup and cooldown. A successful outcome is faster loss reduction or lower final perplexity at equal compute, with norms plateauing rather than continuously shrinking. Repeat with momentum SGD and Muon to test whether the effect is optimizer-independent.

## Verification

- Status: beats_baseline
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: helps
- Verdict: Built a decoupled AdamW variant whose decay coefficient is scaled by clip(eta/eta_max), plus exact toy verification and a matched three-seed regression experiment. The mechanism manifested quantitatively: log-shrinkage ratios matched predicted learning-rate fractions within 0.08% relative error, schedule cumulative shrinkage matched the recursion to 1.2e-15, lambda scaling was approximately linear, and lambda=0 produced identical updates. In the small regression test, scaled decay improved mean final loss from 0.0585 to 0.0290, though this is not evidence of broad generalization.

### Mechanism check

- Verdict: Built a decoupled AdamW variant whose decay coefficient is scaled by clip(eta/eta_max), plus exact toy verification and a matched three-seed regression experiment. The mechanism manifested quantitatively: log-shrinkage ratios matched predicted learning-rate fractions within 0.08% relative error, schedule cumulative shrinkage matched the recursion to 1.2e-15, lambda scaling was approximately linear, and lambda=0 produced identical updates. In the small regression test, scaled decay improved mean final loss from 0.0585 to 0.0290, though this is not evidence of broad generalization.
- Confidence: 9/10
- Limitations: Only a tiny synthetic linear-regression task was tested, not a Transformer, language modeling, Muon, SGD, multiple schedules, or wall-clock/FLOP efficiency. The optimizer implementation is a minimal AdamW reproduction and does not include parameter-group exclusions or production checkpoint/state handling.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0169
- Paired wins: 7/8
- Benchmark verdict: idea better (significant)

## Artifacts

- [bench_report.json](https://synthcore.org/code/525/bench_report.json)
- [experiment.py](https://synthcore.org/code/525/experiment.py)
- [report.md](https://synthcore.org/code/525/report.md)
- [report_bench_2026-08-31T165535.md](https://synthcore.org/code/525/report_bench_2026-08-31T165535.md)
- [results.json](https://synthcore.org/code/525/results.json)
- [stage2_bench.py](https://synthcore.org/code/525/stage2_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/525)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.