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

Adaptive and non-Euclidean optimizers often outperform 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 problem. Starting from an optimizer gap on 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 class imbalance. We prove that GD learns rare tokens slowly because the corresponding logits receive small 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. In the stochastic setting, we prove that mini-batch noise masks rather than creates this advantage, with sufficiently large batches yielding a separation between SignSGD and GD.