StoSignSGD
Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models
Can an optimizer still work when gradients and optimizer states themselves are stored in FP8 or FP4? StoSignSGD is a sign-based optimizer built for this regime. In our aggressive low-precision pretraining experiments, it remains stable without optimizer-specific stabilization, while AdamW and Muon encounter NaN; it also improves over stable sign-based baselines in both token efficiency and downstream accuracy.
StoSignSGD in one equation
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:
FP8: AdamW underflows; sign-based methods stay stable
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.
Why does AdamW fail here?
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.
| Optimizer | Train loss | Val loss | Train ppl. | Val ppl. |
|---|---|---|---|---|
| AdamW | NaN | NaN | NaN | NaN |
| Muon | NaN | NaN | NaN | NaN |
| SignSGD | 3.539 | 3.546 | 34.43 | 34.67 |
| Lion | 3.452 | 3.460 | 31.55 | 31.82 |
| StoSignSGD | 3.394 | 3.405 | 29.78 | 30.12 |
FP4: the gap grows at the largest scale
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 |
7B+ LLM post-training
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 |
What makes StoSignSGD work?
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.
Theory: fixing non-convergence on non-smooth objectives
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\).
BibTeX
@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}
}