Tamed subgradient Langevin optimizer
Implementation & benchmark of arXiv:2608.06283 — The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity
Source paper: The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity arXiv:2608.06283 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace the raw subgradient step by a state-dependent tamed step that is approximately linear for small subgradients but saturates for superlinear ones, and optionally add Langevin noise. Unlike ordinary fixed gradient clipping, the taming threshold is coupled to the step size, so the modification becomes small in the small-step regime while preventing a single nonsmooth or exploding coordinate from destabilizing training.
Formulas
Mathematical statement
The paper defines the optimization target as \(\theta^\star\in\operatorname*{arg\,min}_{\theta\in\mathbb{R}^d}u(\theta)\), where \(u:\mathbb{R}^d\to\mathbb{R}\) is the potential and \(d\) is the parameter dimension. It associates the Gibbs distribution \(\pi_\beta(\mathrm{d}x)\propto e^{-\beta u(x)}\,\mathrm{d}x\), where \(\beta>0\) is inverse temperature; large \(\beta\) concentrates samples near minimizers. At a nonsmooth point, \(g(\theta)\in\partial u(\theta)\) is a subgradient. The implementable tamed drift is \(T_{\lambda,\alpha}(g)_i=g_i/(1+\lambda^\alpha|g_i|)\), where \(\lambda>0\) is the step size, \(\alpha\ge0\) is the taming exponent, and \(i\) indexes coordinates. It satisfies \(|T_{\lambda,\alpha}(g)_i|\le\lambda^{-\alpha}\) and agrees approximately with \(g_i\) when \(\lambda^\alpha|g_i|\ll1\). The stochastic update is \(\theta_{n+1}=\theta_n-\lambda T_{\lambda,\alpha}(g_n)+\sqrt{2\lambda/\beta}\,\xi_n\), with \(\xi_n\sim\mathcal N(0,I_d)\). The paper's stated contribution is a non-asymptotic Wasserstein-2 and excess-risk analysis for this combination of nonsmoothness, nonconvexity, and superlinear growth; the supplied extraction does not contain the detailed theorem constants.
Implementation notes
Integrate this optimizer immediately after computing a minibatch loss and its gradients, replacing the AdamW or SGD parameter update. For each parameter tensor, obtain the autodiff subgradient \(g_t\); at ReLU or absolute-value kinks, the framework's chosen derivative is used as the subgradient. Use this pseudocode: sample minibatch \(B_t\); compute \(u_t\) and \(g_t=\partial u_t(\theta_t)\); choose step size \(\lambda_t\), inverse temperature \(\beta_t\), and exponent \(\alpha\); compute \(h_{t,i}=g_{t,i}/(1+\lambda_t^\alpha|g_{t,i}|)\); draw independent standard-normal noise \(\xi_t\); update \(\theta_{t+1}=\theta_t-\lambda_t h_t+\sqrt{2\lambda_t/\beta_t}\xi_t\). Apply decoupled weight decay separately if required. Start with \(\alpha=1/2\). Choose constant \(\beta\) so the RMS Langevin noise is 1--5% of the RMS deterministic step, then test an annealed schedule increasing \(\beta_t\) over training. The exact mathematical computations are the tamed drift and Gaussian diffusion. Estimate gradient scale with an exponential moving average if raw parameter units cause poor conditioning; use this only as an ablation, because the core method is the unnormalized formula above. First train a 6-layer Transformer on WikiText-2, or a small CNN on CIFAR-10, comparing SGD, AdamW, and SG-TULA at matched update FLOPs. Record training loss versus steps and wall-clock time, validation perplexity or accuracy, gradient-spike frequency, update norms, and NaN/Inf events. Success means fewer catastrophic loss spikes and equal-or-better validation performance, especially at learning rates where ordinary SGD or AdamW becomes unstable; a further signal is improved validation performance from the annealed-noise schedule at equal compute.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Significantly worse ✗
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built the coordinatewise TamedSubgradientLangevin PyTorch optimizer and deterministic toy verification. The taming bound was reached within 0.01%, the half-response transition matched lambda^(-alpha) within 0.6%, and saturated displacement scaling matched the predicted lambda^(1-alpha) ratio to numerical precision. In the aggressive toy sweep, raw SGD exploded for learning rates 0.5 through 5, while the tamed method remained finite for 100 steps with final losses below 0.052; this is a clear stability signal, not evidence of a generalization win.
- Agent confidence
- 9/10
- Baseline
- Raw SGD: stable at lr=0.2 with final loss 0.00268, but exploded rapidly for lr>=0.5, reaching losses from 1.22e49 to 1.02e87.
- Idea
- Tamed alpha=0.5: |T(g)|<=lambda^-alpha verified with max/bound 0.99990--0.999999; half-response matched lambda^-alpha within 0.6%; saturated displacement ratios matched lambda^0.5 to about 2e-10 relative error. It remained finite for 100 steps at all tested learning rates, with final losses 0.00193, 0.00532, 0.00689, 0.01884, 0.02796, and 0.05195 for lr 0.2, 0.5, 1, 2, 3, and 5.
Limitations: The experiment used a one-dimensional deterministic quartic-plus-absolute-value toy objective, not CIFAR-10, WikiText-2, or a neural network. Langevin noise, annealing, minibatch gradients, AdamW comparison, wall-clock/FLOP matching, validation metrics, and high-dimensional conditioning were not tested.
How to run: python3 experiment.py && /home/maxwelhelp/main/bin/python3 aggressive_sweep.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.06283)
Significantly WORSE than the tuned baseline (+7.4%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Tabular regression (Friedman)
- Model
- mlp_tiny
- Paired seeds
- 8
- Baseline mean
- 11.0541
- Idea mean
- 11.8734
- Effect (Δ)
- 0.8194 (+7.4%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±0.8%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
Only the registered built-in tabular optimizer track was tested. Langevin noise and annealing, AdamW, weight decay, wall-clock/FLOP matching, larger models, and the vision/sequence/dynamics tracks were not tested. The implementation used deterministic tamed SGD (beta=infinity), so conclusions apply to the taming component rather than the full noisy SG-TULA variant.
How to run: python3 bench_run.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 10/10
Implemented tamed coordinatewise SGD on the registered tabular Friedman#1 track with the shared mlp_tiny architecture and matched training settings. Across 8 paired seeds, taming significantly worsened test MSE: 11.8734455109 versus 11.0540536642 for SGD, paired delta +0.8193918467 with permutation p=0.0081. The trained-model update-cap signature was confirmed, but the benchmark task metric showed no win.
- Baseline
- SGD, mlp_tiny, lr=0.003, 18 epochs: mean test MSE 11.054053664207458.
- Idea
- Tamed alpha=0.5, mlp_tiny, lr=0.003, 18 epochs: mean test MSE 11.873445510864258; paired delta +0.8193918466567993, p=0.0081, 0/8 wins.
Benchmark result
Significantly WORSE than the tuned baseline (+7.4%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Tabular regression (Friedman)
- Model
- mlp_tiny
- Paired seeds
- 8
- Baseline mean
- 11.0541
- Idea mean
- 11.8734
- Effect (Δ)
- 0.8194 (+7.4%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±0.8%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
Only the registered built-in tabular optimizer track was tested. Langevin noise and annealing, AdamW, weight decay, wall-clock/FLOP matching, larger models, and the vision/sequence/dynamics tracks were not tested. The implementation used deterministic tamed SGD (beta=infinity), so conclusions apply to the taming component rather than the full noisy SG-TULA variant.
How to run: python3 bench_run.py
Artifacts
- aggressive_results.json 2.4 KB View Raw
- aggressive_sweep.py 0.9 KB View Raw
- bench_report.json 5.9 KB View Raw
- bench_run.py 4.4 KB View Raw
- experiment.py 4.0 KB View Raw
- report.md 1.8 KB View
- report_bench_2026-09-01T010934.md 4.0 KB View
- results.json 6.0 KB View Raw
- tamed_optimizer.py 1.2 KB View Raw