Initial-Only Weight Decay with Tail Averaging
Implementation & benchmark of arXiv:2608.22953 — Stochastic gradient descent with initial regularization
Source paper: Stochastic gradient descent with initial regularization arXiv:2608.22953 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace constant weight decay by the paper's initial-regularization schedule: apply isotropic shrinkage only for the first m optimizer steps, then set the regularization coefficient exactly to zero. Collect checkpoints only in a later interval, such as steps 2m through 3m, and average them to reduce stochastic variance. This creates an explicitly separated representation-shaping phase and fitting phase rather than forcing one regularization strength throughout training.
Formulas
Mathematical statement
The core SGDIR recursion is equation (1): \(\theta_{t+1}=(1-\gamma\lambda_t)\theta_t-\gamma(\langle x_t,\theta_t\rangle-y_t)x_t\), where \(\theta_t\) is the parameter vector, \(x_t\) is the sampled feature/input vector, \(y_t\) is its target, \(\gamma>0\) is the constant step size, and \(\lambda_t=\Lambda\mathbf{1}(t<m)\) applies regularization of strength \(\Lambda>0\) only before step \(m\). The paper analyzes the tail average \(\bar{\theta}_{2m,3m}=m^{-1}\sum_{t=2m}^{3m-1}\theta_t\). For the error relative to the population minimizer \(\theta^*\), equation (37) gives \(w_{t+1}-\theta^*=P_t(w_t-\theta^*)-\gamma\lambda_t\theta^*\), with \(P_t=I-\gamma(x_t\otimes x_t+\lambda_t I)\). Here \(x_t\otimes x_t\) is the rank-one covariance operator and \(w_t\) denotes the regularized iterate. If \(\Sigma=\mathbb{E}[x_t\otimes x_t]\), then \(A_t=\mathbb{E}[P_t]=I-\gamma(\Sigma+\lambda_t I)\). Lemma 9 states \(A_t\succeq0\) and \(\mathbb{E}[P_t^2]\preceq(1-\gamma\lambda_t)A_t\preceq(1-\gamma\lambda_t)^2I\). The adaptation applies the same multiplicative shrinkage to neural-network weights, treating the local stochastic update as the nonlinear analogue of the linear recursion; the inequality motivates choosing \(\gamma\) so that the initial phase has \(0\leq\gamma\Lambda\leq1\), while ordinary gradient stability is controlled empirically or by a short warmup.
Implementation notes
1. Integration point: modify the SGD or momentum-SGD parameter update, immediately after computing the minibatch gradient and before the parameter update. Exclude biases, normalization affine parameters, and optionally embeddings from shrinkage, matching common decoupled weight-decay practice. Let \(m\) be the initial-regularization horizon and \(\Lambda\) the decay coefficient. 2. Pseudocode: initialize parameters \(\theta\) with the chosen network initialization; for optimizer step \(t=0,\ldots,T-1\), compute minibatch gradient \(g_t\), set \(\lambda_t=\Lambda\) if \(t<m\), otherwise set \(\lambda_t=0\), and perform \(\theta\leftarrow(1-\gamma\lambda_t)\theta-\gamma g_t\). Store parameter checkpoints for \(t\in[2m,3m)\), and at the end set \(\theta_{out}=m^{-1}\sum_{t=2m}^{3m-1}\theta_t\). In a deep network, maintain the running average in FP32 and evaluate the averaged copy after updating batch-normalization statistics. 3. Mathematical quantities: \(\Lambda,m,\gamma\), the exact shrinkage factor \(1-\gamma\Lambda\), and the tail-average window come directly from the construction. The covariance operator \(\Sigma\), eigenvalues, and theoretical excess-risk constants need not be estimated for the MVP. Estimate stability empirically using minibatch gradient norms and the largest observed parameter-update ratio; require \(0\leq\gamma\Lambda\leq1\) and clip if violated. 4. First experiment: train a 2-layer MLP and a small ResNet-18 on CIFAR-10, plus a linear or kernelized regression model where the assumptions are closest to the paper. Compare constant weight decay, cosine-decayed decay, no decay, and initial-only decay at matched learning rate, total steps, and batch size. Sweep \(m/T\in\{0.05,0.1,0.25,0.5\}\) and \(\gamma\Lambda\in\{0.01,0.05,0.1,0.5\}\). Measure training loss, validation accuracy, sharpness proxy, gradient norm spikes, and final test accuracy. Success means lower validation error or better late-stage loss than constant decay, with fewer early instability spikes and no extra forward/backward FLOPs; additionally test whether tail averaging improves accuracy at fixed compute.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 7/10
Built a CPU-only SGDIR-style linear regression MVP with exact initial-only shrinkage, delayed tail averaging, and a numerical diagonal-operator check. The core inequality passed, with maximum estimated violation -0.000919, and initial-only decay strongly beat constant decay on test MSE (0.06390 vs 0.32036). However, no decay was marginally better (0.06373), so the evidence supports avoiding persistent decay more clearly than proving an advantage over no regularization.
- Agent confidence
- 7/10
- Baseline
- Constant decay: final test MSE 0.3204 +/- 0.0714; tail MSE 0.3126 +/- 0.0714 across 8 seeds.
- Idea
- Initial-only decay: final test MSE 0.06390 +/- 0.00144; tail-averaged MSE 0.06370 +/- 0.00151. No-decay reference: final 0.06373; tail 0.06328.
Limitations: Only synthetic linear squared-loss regression was tested; no MLP, CIFAR-10, momentum, hyperparameter tuning, FLOP benchmarking, sharpness measurement, or tuned constant-decay comparison was included. The operator check uses finite Monte Carlo sampling.
How to run: python3 experiment.py --runs 8 --out results.json
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Super-audit (legacy) latest Failed ✗
Agent confidence: 9/10
Built and ran the specified 784-128-128-10 ReLU MLP with AdamW, cosine learning-rate annealing, 3,000 steps, and five seeds on real MNIST. The operator inequality sanity check passed, but initial-only decay did not produce a distinct improvement: initwd_tail reached 0.97744 +/- 0.00110 accuracy versus 0.97720 +/- 0.00106 for constant WD, while constant-WD tail averaging also reached 0.97744 +/- 0.00109. Thus the observed small gain is attributable to tail averaging, not the proposed initial-only schedule; no faster best step was observed.
- Baseline
- MNIST, 5 seeds: baseline constant WD test accuracy 0.97720 +/- 0.00106, best accuracy 0.97764 +/- 0.00058, steps-to-best 2700 +/- 300, weight norm 15.5723 +/- 0.1096, mid-gradient norm 0.5346 +/- 0.0108, runtime 23.06 +/- 1.37 sec/mode-seed. wd_tail: 0.97744 +/- 0.00109; paired gain over baseline +0.00024 +/- 0.00009.
- Idea
- initwd_tail: test accuracy 0.97744 +/- 0.00110, best accuracy 0.97764 +/- 0.00059, steps-to-best 2700 +/- 300, weight norm 15.7558 +/- 0.1114, mid-gradient norm 0.5352 +/- 0.0105, runtime 23.62 +/- 1.25 sec/mode-seed. Paired gain over baseline +0.00024 +/- 0.00025, over initwd +0.00020 +/- 0.00032, and versus wd_tail approximately 0.00000 +/- 0.00019 (2/5 wins). Operator check passed with gaps -0.000371 and -0.000376.
Tested one real dataset, one MLP width, one prescribed weight-decay/horizon/window configuration, and five seeds; did not test CIFAR-10, ResNet, momentum SGD, hyperparameter sweeps, sharpness, or matched-FLOP/steps-to-target analyses. MNIST evaluation sampled every 300 steps, and the operator inequality uses finite Monte Carlo sampling. The reported final weight norm for tail modes is the non-averaged training norm.
How to run: python3 exp85_initial_only_weight_decay_tail_averaging.py --runs 5 --T 3000 --width 128 --out super_results.json
Artifacts
- exp85_initial_only_weight_decay_tail_averaging.py 6.4 KB View Raw
- experiment.py 4.5 KB View Raw
- report.md 3.9 KB View
- report_super_2026-08-30T083749.md 5.7 KB View
- results.json 13.6 KB View Raw
- super_results.json 8.9 KB View Raw