# StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models arXiv: https://arxiv.org/abs/2604.15416 · HTML: https://mlsanity.com/paper/2604.15416/ ## Key facts - Cited by 5 papers in the map (data through 14 Oct 2026). - Builds on 123 of its 170 references that are in the map. - First posted on arXiv 16 Apr 2026. ## Builds on (most cited first, 30 of 123) - [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015) - [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020) - [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019) - [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023) - [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024) - [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023) - [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023) - [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022) - [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021) - [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025) - [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021) - [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016) - [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024) - [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018) - [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019) - [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022) - [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017) - [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016) - [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023) - [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021) - [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022) - [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019) - [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023) - [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019) - [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018) - [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019) - [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025) - [LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models](https://mlsanity.com/paper/2403.13372/) (2024) - [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019) - [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017) ## Built on by (most cited first, 5 of 5) - [Optimizer-Model Consistency: Full Finetuning with the Same Optimizer as Pretraining Forgets Less](https://mlsanity.com/paper/2605.06654/) (2026) - [When and Why SignSGD Outperforms SGD: A Theoretical Study Based on ℓ1-norm Lower Bounds](https://mlsanity.com/paper/2605.06615/) (2026) - [Open Problem: Is AdamW Effective Under Heavy-Tailed Noise?](https://mlsanity.com/paper/2606.23676/) (2026) - [Adam Converges in Nonsmooth Nonconvex Optimization](https://mlsanity.com/paper/2606.22326/) (2026) - [Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon](https://mlsanity.com/paper/2610.19844/) (2026) ## Abstract (arXiv) Sign-based optimization algorithms, such as SignSGD, have garnered attention for their performance in distributed learning and training large foundation models. Despite their empirical superiority, SignSGD is known to diverge on non-smooth objectives, which are ubiquitous due to ReLUs, max-pools, and mixture-of-experts. To overcome this limitation, we propose StoSignSGD, an algorithm that injects structural stochasticity into the sign operator while maintaining an unbiased update step. In the regime of (online) convex optimization, StoSignSGD rigorously resolves the non-convergence issues of SignSGD, achieving a sharp convergence rate matching the lower bound. For the more challenging non-convex non-smooth optimization, we introduce generalized stationary measures that encompass prior definitions, proving that StoSignSGD improves upon the best-known complexity bounds by dimensional factors. Empirically, StoSignSGD is stable and efficient across diverse large language model (LLM) training regimes. In aggressive low-precision pretraining, which spans both FP8 and the far more demanding FP4 regime where AdamW fails catastrophically, StoSignSGD stays stable and consistently performs the best. It attains a 1.44x to 2.14x speedup over established baselines under FP8. Under 4-bit precision, it improves downstream accuracy on the largest OLMo2-370M model by 1.13 points over the strongest stable baseline, and this advantage grows as both the model size and the data scale up. When fine-tuning 7B LLMs on mathematical reasoning tasks, StoSignSGD also delivers clear gains over both AdamW and SignSGD. Finally, to explain why it works, we develop a sign conversion framework that turns any general optimizer into its unbiased, sign-based counterpart. Using this framework, we decompose the core components of StoSignSGD and run a comprehensive ablation study to validate our design choices.