Two things matter most: StoSignSGD stays numerically stable in our radical low-precision setting, and its stochastic sign fixes a basic non-convergence issue of deterministic SignSGD.
We store gradients and persistent optimizer states directly in low precision. StoSignSGD remains stable in these experiments without optimizer-specific stabilization.
See low-precision results →Structural stochasticity removes the bias of deterministic sign updates. StoSignSGD converges on the simple failure case below and achieves sharp convex and non-convex guarantees.
See theory →| Optimizer | Train loss | Val loss |
|---|---|---|
| AdamW | NaN | NaN |
| Muon | NaN | NaN |
| SignSGD | 3.539 | 3.546 |
| Lion | 3.452 | 3.460 |
| StoSignSGD | 3.394 | 3.405 |
AdamW and Muon diverge; StoSignSGD stays stable and reaches the best final loss among the stable baselines.
The basic idea is to replace the deterministic sign with a stochastic sign operator. For a coordinate-wise scale \(\mathbf G\succeq |\mathbf x|\), sample \(\mathbf n\sim\Unif([-1,1]^d)\) and define
The practical optimizer keeps momentum and tracks a damped coordinate-wise maximum. The update used throughout this page is:
We pretrain GPT-2 on OpenWebText in an intentionally aggressive FP8 setting: not only low-precision linear operations, but also gradients and persistent optimizer states are stored in FP8. AdamW and Muon diverge early. StoSignSGD stays stable and reaches the same validation-loss target with 53% fewer tokens than SignSGD and 30% fewer than Lion, corresponding to 2.14× and 1.44× token-efficiency speedups.
In the diagnosed run, the squared gradients fall below the minimum representable FP8 value, so the second-moment buffer \(\mathbf v_t\) is flushed to zero and the normalized update becomes unstable. StoSignSGD avoids a squared-gradient variance buffer: it only uses momentum, a max-buffer, and a sign update.
We then push the same idea to NVFP4 pretraining with the OLMo2 series. Gradients and optimizer states are again stored in FP4. AdamW and Muon diverge to NaN at every scale, while StoSignSGD remains stable and consistently outperforms Lion. The gain in average downstream accuracy is +0.45, +0.44, and +1.13 points from OLMo2-70M to OLMo2-370M.
OLMo2 model names count non-embedding parameters. Including embeddings, the three models contain 136M / 267M / 474M parameters and are trained on 4.19B / 9.23B / 20.97B tokens, respectively.
| Model | Method | ARC-E | ARC-C | OBQA | HellaSwag | PIQA | SIQA | BoolQA | Avg. |
|---|---|---|---|---|---|---|---|---|---|
| OLMo2-70M | AdamW | -- | -- | -- | -- | -- | -- | -- | -- |
| Muon | -- | -- | -- | -- | -- | -- | -- | -- | |
| Lion | 33.86 | 21.40 | 25.40 | 26.09 | 55.28 | 38.79 | 52.26 | 36.15 | |
| StoSignSGD | 33.33 | 20.74 | 26.40 | 26.08 | 55.60 | 39.25 | 54.83 | 36.60 | |
| OLMo2-150M | AdamW | -- | -- | -- | -- | -- | -- | -- | -- |
| Muon | -- | -- | -- | -- | -- | -- | -- | -- | |
| Lion | 39.47 | 19.73 | 27.60 | 26.85 | 59.41 | 39.66 | 61.35 | 39.15 | |
| StoSignSGD | 41.93 | 21.40 | 26.60 | 27.06 | 58.81 | 40.69 | 60.67 | 39.59 | |
| OLMo2-370M | AdamW | -- | -- | -- | -- | -- | -- | -- | -- |
| Muon | -- | -- | -- | -- | -- | -- | -- | -- | |
| Lion | 51.58 | 26.76 | 29.00 | 33.54 | 63.66 | 42.84 | 55.78 | 43.31 | |
| StoSignSGD | 52.11 | 25.42 | 32.20 | 33.80 | 64.74 | 41.81 | 61.01 | 44.44 |
The effect is not limited to pretraining. On mathematical reasoning SFT, StoSignSGD is particularly strong on foundational math tasks. For Qwen2.5-7B trained on NuminaMath-CoT, it improves GSM8K by about 3 points over AdamW while matching AdamW on MATH.
| Optimizer | GSM8K | MATH |
|---|---|---|
| AdamW | 74.37 ± 1.20 | 48.86 ± 0.66 |
| SignSGD | 71.95 ± 1.24 | 47.66 ± 0.66 |
| StoSignSGD | 77.33 ± 1.15 | 48.88 ± 0.65 |
We use an unbiased sign conversion framework to turn a general optimizer update into a stochastic sign update. This lets us isolate three ingredients: structural noise, coupling the scale to momentum, and using an infinity-norm / max-based scale.
| Optimizer | Structural noise | \(\boldsymbol\sigma_t\) depends on \(\mathbf m_t\) | Inf-norm scale | RMS norm of \(\mathbf m_t/\boldsymbol\sigma_t\) |
|---|---|---|---|---|
| SignSGD | ✗ | ✗ | ✗ | 1 |
| AdamW | ✗ | ✗ | ✗ | ≈ 0.2 |
| AdaMax | ✗ | ✗ | ✓ | ≈ 0.1 |
| IE-StoSignSGD | ✗ | ✓ | ✓ | ≈ 0.4 |
| SignAdamW | ✓ | ✗ | ✗ | 1 |
| SignAdaMax | ✓ | ✗ | ✓ | 1 |
| StoSignSGD | ✓ | ✓ | ✓ | 1 |
Controlled ablations show that the full combination is what consistently gives the strongest results.
This project started from a theoretical issue: deterministic SignSGD can fail even on simple convex non-smooth objectives. StoSignSGD removes that bias and gives sharp guarantees in both convex and non-convex settings.
Under convexity and coordinate-wise Lipschitzness, set \(\eta_t=D_\infty/\sqrt{2t}\). Online StoSignSGD ensures
\[ \E\!\left[\max_{\mathbf x\in\X}\Regret_T(\mathbf x)\right] \le (2+\sqrt2)D_\infty\lVert\mathbf L\rVert_1\sqrt T. \]The paper also proves the matching lower bound \(\E[\max_{\mathbf x\in\X}\Regret_T(\mathbf x)]\ge D_\infty\lVert\mathbf L\rVert_1\sqrt{T/8}\), so the dependence on \(T\), \(D_\infty\), and \(\lVert\mathbf L\rVert_1\) is tight up to constants.
Under coordinate-wise Lipschitzness, non-convexity, and well-behavedness, and assuming \(f(\mathbf x_0)-f^*\le\Delta_f\), set
\[ \begin{aligned} K&=\left\lceil\frac{7\sqrt{14}\,\Delta_f\delta^{1/2}}{2\epsilon^{3/2}}\right\rceil,\qquad N=\left\lceil751\lVert\mathbf L\rVert_1^2\epsilon^{-2}\right\rceil,\\ D_\infty&=\frac{\sqrt\epsilon}{\sqrt{14\delta}\,N},\qquad \eta_t=\frac{\sqrt2D_\infty}{\sqrt{((t-1)\bmod N)+1}}. \end{aligned} \]Then StoSignSGD finds a \((\delta,\epsilon)\)-\(\ell_{1,\infty}\)-stationary point within
\[ T=KN\le39326\,\Delta_f\lVert\mathbf L\rVert_1^2\delta^{1/2}\epsilon^{-7/2} \]stochastic gradient evaluations. This removes the extra explicit \(\sqrt d\) factor in the cited prior \(\ell_{1,2}\) Adam upper bound, while matching the known \(\ell_{2,2}\) lower-bound scaling in \(\Delta_f,\delta,\epsilon\).
@article{yu2026stosignsgd,
title={StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models},
author={Dingzhi Yu and Rui Pan and Yuxing Liu and Difan Zou and Tong Zhang},
journal={arXiv preprint arXiv:2604.15416},
year={2026}
}