Provable Benefit of SignGD: A Minimal Model Under Heavy-Tailed Class Imbalance

Robin Yadav; Shuo Xie; Tianhao Wang; Zhiyuan Li

cs.LGNov 30, 2025

Provable Benefit of SignGD: A Minimal Model Under Heavy-Tailed Class Imbalance

Authors: Robin Yadav, Shuo Xie, Tianhao Wang, Zhiyuan Li

Organizations: Toyota Technological Institute at Chicago · University of British Columbia · University of California, San Diego

Abstract

Adaptive and non-Euclidean optimizers often outperform Euclidean methods such as stochastic gradient descent (SGD) in language modeling by a large margin. Existing theory usually explains this gap by assuming favorable smoothness geometry or noise structure tailored to the specific optimizer. We instead ask whether such geometry can be induced from a concrete learning setting. Starting from an optimizer gap that persists across realistic language-modeling experiments, we progressively remove sequence dependence, architectural complexity, and stochasticity. We find that the gap exists in a minimal setting: the softmax unigram model with heavy-tailed data. This model exposes a simple deterministic mechanism under heavy-tailed class imbalance. We prove that GD learns rare tokens slowly because the corresponding logits receive only tiny updates, while SignGD removes this magnitude dependence and moves rare and common coordinates on a more comparable scale. We make this precise with upper and lower bounds for the convergence rate of GD and upper bounds for the convergence of SignGD. Our stochastic bounds contain additional noise-dependent terms that can obscure this advantage in the convergence guarantees and can be reduced by increasing the batch size

Figures & tables

Panel of Figure 1 — Figure 1: GD struggles to learn the tail token probabilities on a simple deterministic softmax unigram model with heavy-tailed class imbalance . The softmax unigram model objective is defined as $f_{\mathrm{uni}}(\theta)=\KL(p\|\operatorname{softmax}(\theta))$ for logits $\theta$ . (a) Adam and SignGD outperform GD and normalized GD by a significant margin. (b) True token distribution for the softmax unigram model; black indicates head tokens, and yellow indicates tail tokens. (c) and (d) KL contributions grouped by token rank for GD and SignGD, respectively, where each group has mass 0.1. The contribution of token $i$ is $p_{i}\log(p_{i}/q_{i})-p_{i}+q_{i}$ for $q=\operatorname{softmax}(\theta)$ . These non-negative coordinate-wise terms, which sum to the total KL, show that GD fails to reduce loss on tail tokens, while SignGD makes more uniform progress across ranks. Although tail contributions are small, the large number of tail tokens renders them non-negligible.

Panel of Figure 2 — Figure 2: SignSGD makes more uniform progress across token frequencies than SGD when training a transformer model. (a) Training loss for SGD, Adam, normalized SGD, and SignSGD when training all parameters of a one-layer transformer on the PTB dataset. (b,c) Training loss per token-frequency group for SGD and SignSGD, respectively. Tokens are sorted by frequency and partitioned into ten groups, each accounting for roughly $10\%$ of the training samples. Dark curves represent the most frequent tokens and light yellow curves the least frequent tokens. SGD leaves the low-frequency groups at high loss, whereas SignSGD makes progress across all frequency groups. The overall SGD and SignSGD losses in (a) correspond to the averages of the groups in (b) and (c) , respectively.

Panel of Figure 3 — Figure 3: The SignSGD and SGD gap persists as the language-modeling setup is progressively simplified, starting from a one-layer transformer. (a) We freeze the embedding and attention layers and only train the classification head. (b) We remove the attention layer from (a) and train only the output classifier, essentially yielding a bigram model. (c) We collapse all token embeddings to a shared vector, reducing the task to unigram softmax learning. The red dashed line shows the minimum achievable cross-entropy loss for the unigram model, i.e., the entropy of the token distribution. SGD initially descends faster but plateaus above this minimum, while SignSGD approaches it. (d) The same experiment as in (c) , but we plot the function sub-optimality, i.e., the difference between the cross-entropy loss and the entropy of the token distribution, to clearly visualize the gap.

Figure 4

Figure 5

Panel of Figure 5 — Figure 5: Batch-size comparison on the stochastic softmax unigram model with $p$ following Definition 3.1 with $d=10^{4}$ . The panels show the population KL loss for SGD, Adam, normalized SGD, and SignSGD using different mini-batch sizes.

Panel of Figure 6 — Figure 6: Batch-size comparison for a two-layer transformer trained on TinyPTB. The panels show training cross-entropy for SGD, Adam, normalized SGD, and SignSGD; curves are medians over three seeds and shaded regions show the minimum-to-maximum range.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Grid search with vanilla GD in the intermediate- and large-step-size regimes for the softmax unigram model with $d=10^{4}$ . Each point reports the final KL divergence for a learning rate $\eta=c\log d$ . Large learning-rate multipliers lead to instability, consistent with Theorem D.4 .

Experiment	Model	Mini-batch	Horizon and seeds
Figure 2	One-layer causal transformer with embedding dimension $512$ , $8$ attention heads, feed-forward width $256$ , and dropout $0.2$ . The embedding, transformer layer, and output classifier are all trained. The model has approximately $11.49$ M trainable parameters.	$1{,}024\times 35$	$100$ epochs; seeds $0,1,2$ for the optimizer comparison and seed $0$ for the per-frequency-group runs.
Figure 3 (a)	One-layer transformer with embedding dimension $256$ , one attention head, feed-forward width $128$ , and dropout $0.2$ . The token embedding and complete transformer encoder are frozen; only the softmax classification head is trained. The model has approximately $5.42$ M parameters, of which $2.55$ M are trainable.	$1{,}024\times 35$	$100$ epochs; seeds $0,1,2$ .
Figure 3 (b)	Bigram model obtained by removing attention. It consists of a fixed $9{,}922\times 256$ token-embedding matrix followed by a trainable, bias-free $256\times 9{,}922$ output classifier. It has approximately $5.08$ M parameters, of which $2.54$ M are trainable.	$1{,}024\times 35$	$100$ epochs; seeds $0,1,2$ .
Figure 3 (c,d)	Collapsed softmax unigram model. Every input token is replaced by the same index, leaving one trainable vector of $9{,}922$ logits. The logits are initialized to zero.	$1{,}024\times 32$	$200$ epochs; seeds $0,1,2$ .
Figure 1	Deterministic softmax unigram model with $d=10{,}000$ trainable logits, initialized to zero, and target probabilities $p_{i}\propto i^{-1}$ .	Full batch	$600$ gradient steps; seed $0$ .

Appendix

Table 1: Model and training details for the experiments in Section 3.

Experiment	Optimizer	Search grid	Selected learning rate
Figure 2	SGD	$\mathcal{G}(-3,0,1/4)$	$10^{-1/2}\approx 0.316$
	Adam	$\mathcal{G}(-4,-2,1/4)$	$10^{-11/4}\approx 0.00178$
	normalized SGD	$\mathcal{G}(-2,1,1/4)$	$10^{1/4}\approx 1.78$
	SignSGD	$\mathcal{G}(-5,-2,1/4)$	$10^{-5/2}\approx 0.00316$
Figure 3 (a)	SGD	$\mathcal{G}(-3,0,1/2)$	$10^{-1/2}\approx 0.316$
	SignSGD	$\mathcal{G}(-5,-2,1/2)$	$10^{-5/2}\approx 0.00316$

Appendix

Table 2: Per-optimizer learning-rate tuning for the Section 3 experiments. For Adam, we use

\beta_{1}=0

\beta_{2}=0.999

, and

\epsilon=10^{-8}

Batch size	SGD	Adam	Normalized SGD	SignSGD
$64$	$10$ $(2)$	$10^{-3/2}$ $(2)$	$10^{1/2}$ $(2)$	$10^{-3/2}$ $(2)$
$1024$	$10^{3/2}$ $(0)$	$10^{-3/2}$ $(0)$	$10$ $(2)$	$10^{-2}$ $(0)$
$32768$	$10^{3/2}$ $(0)$	$10^{-3/2}$ $(0)$	$1$ $(0)$	$10^{-3/2}$ $(1)$

Appendix

Table 3: Selected learning rates for the stochastic softmax unigram experiment. Selected seeds are shown in parentheses.

Batch size	SGD	Adam	Normalized SGD	SignSGD
$32$	$10^{-1}$	$10^{-4}$	$10^{-1/2}$	$10^{-7/2}$
$256$	$10^{-1}$	$10^{-7/2}$	$1$	$10^{-3}$
$1024$	$10^{-1}$	$10^{-3}$	$10^{-1}$	$10^{-3}$

Appendix

Table 4: Selected learning rates for the TinyPTB transformer experiment.

Explore similar work

StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models
Apr 16, 2026Dingzhi Yu, Rui Pan, Yuxing Liu +1Stochastic Gradient Descent Stochastic Convex Optimization
University of Illinois Urbana-Champaign
Open Problem: Is AdamW Effective Under Heavy-Tailed Noise?
Jun 22, 2026Dingzhi Yu, Hongyi Tao, Yuanyu Wan +2Adam Long-Tailed Distribution
1Nanjing University · 2Zhejiang University · 3Fudan University
When and Why SignSGD Outperforms SGD: A Theoretical Study Based on $\ell_1$ -norm Lower Bounds
May 7, 2026Hongyi Tao, Dingzhi Yu, Lijun ZhangStochastic Gradient Descent Signsgd
State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing 210023, China · School of Artificial Intelligence, Nanjing University, Nanjing 210023, China

Apr 16, 2026, cs.LGJ/K move · Enter open · S save

StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models

Dingzhi Yu, Rui Pan, Yuxing Liu, Tong Zhang

University of Illinois Urbana-Champaign

Stochastic Gradient Descent Stochastic Convex Optimization Signsgd Large Language Model Training Sign Stochastic

Open paper AskPDF arXiv

Sign-based optimization algorithms, such as SignSGD, have garnered significant attention for their remarkable 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 in modern machine learning due to ReLUs, max-pools, and mixture-of-experts. To overcome this fundamental limitation, we propose \textbf{StoSignSGD}, an algorithm that injects structural stochasticity into the sign operator while maintaining an unbiased update step. In the regime of (online) convex optimization, our theoretical analysis shows that 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 exhibits robust stability and superior efficiency across diverse large language model (LLM) training regimes. Notably, in low-precision FP8 pretraining -- a setting where AdamW fails catastrophically -- StoSignSGD remains highly stable and yields a remarkable 1.44

\times

to 2.14

\times

speedup relative to established baselines. Furthermore, when fine-tuning 7B LLMs on mathematical reasoning tasks, StoSignSGD delivers substantial performance gains over both AdamW and SignSGD. Finally, to dissect the mechanisms driving its success, we develop a sign conversion framework capable of transforming any general optimizer into its unbiased, sign-based counterpart. Utilizing this framework, we deconstruct the core components of StoSignSGD and present a comprehensive ablation study to empirically validate our algorithmic design choices.

Abstract

Explore similar work

StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models

Open Problem: Is AdamW Effective Under Heavy-Tailed Noise?

When and Why SignSGD Outperforms SGD: A Theoretical Study Based on ℓ1\ell_1ℓ1​-norm Lower Bounds

When and Why SignSGD Outperforms SGD: A Theoretical Study Based on $\ell_1$ -norm Lower Bounds