✓✓ Beats tuned baseline 2026

Learning-Rate-Scaled Weight Decay

Implementation & benchmark of arXiv:2607.23777 — Scale Weight Decay and Train Better

Usefulness7/10

Difficulty2/10

Novelty6/10

Optimization Speedup Stability Generalization Optimizer Training dynamics Moe routing Optimization Dynamical systems Stochastic processes

Source paper: Scale Weight Decay and Train Better arXiv:2607.23777 ⓘ · analyzed Aug 30, 2026

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

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.

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$$

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.

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

Beats tuned baseline

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗

Stage 2 · Practical benchmark result: Significant win ✓

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

How to read this: The implementation beat the baseline, but the claimed mechanism did not reproduce at benchmark scale. The gain may come from another effect; treat the cause as unresolved.

Stage 1 — Mechanism check agent confidence 9/10

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.

Agent confidence: 9/10

Baseline: Three-seed mean final loss 0.05848 (std 0.02435), loss at step 300 0.07288, final norm 3.4731.
Idea: Three-seed mean final loss 0.02899 (std 0.00910), loss at step 300 0.04410, final norm 3.5798. Exact-decay sweep observed/predicted log-ratio errors were 0–0.075%; schedule log-shrinkage was -0.6727382872123935 observed versus -0.6727382872123923 predicted; lambda-sweep normalized slopes ranged -3.3571 to -3.3710 versus predicted -3.3564.

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.

How to run: python3 experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.23777) Sep 2, 2026

Beats the tuned baseline by 0.0% (p=0.0169, wins 7 / 8 paired seeds; metric: lower is better).

Benchmark: Tabular regression (Friedman)
Model: mlp_tiny
Paired seeds: 8
Baseline mean: 11.6623
Idea mean: 11.662
Effect (Δ): -0.0002 (−0.0%; negative = idea better)
Wins: 7 / 8 paired seeds
p-value: 0.0169 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±0.0%
Mechanism: Not confirmed ✗
Practical effect: Helps
Baseline tuning: swept over 6 configs

Limitations:

Only the registered tabular Friedman#1 track and mlp_tiny were tested. Vision, sequence, dynamics, larger networks, SGD, Muon, wall-clock/FLOP efficiency, and independent datasets were not tested. The significant effect is less than 0.001 relative percent, and the cooldown mechanism signature was not confirmed.

How to run: cd /home/maxwelhelp/all/math2nn/experiments/exp1333_learning-rate-scaled-weight-decay && /home/maxwelhelp/main/bin/python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Aug 31, 2026 latest Worked ✓

Agent confidence: 7/10

Implemented learning-rate-scaled decoupled AdamW on the registered tabular Friedman#1 track with the shared mlp_tiny architecture and parity-preserving baseline/idea sweeps. The idea achieved a statistically significant but practically tiny improvement in test MSE across 8 paired seeds: delta=-0.00022018 and permutation p=0.0169, with 7/8 wins. The trained-model mechanism signature confirmed the direct decay multiplier but did not confirm the cooldown log-ratio prediction, so the optimization win should not be generalized beyond this small benchmark.

Baseline: Best cfg lr=0.006, weight_decay=0.001; full 8-seed mean MSE=11.662253737449646, std=1.1646783408667203.
Idea: Best cfg lr=0.006, weight_decay=0.001; full 8-seed mean MSE=11.662033557891846, std=1.16472778990251; delta=-0.00022017955780029297; p=0.0169.

Benchmark result