The Implicit Bias of Steepest Descent with Mini-batch Stochastic Gradient

A variety of widely used optimization methods like SignSGD and Muon can be interpreted as instances of steepest descent under different norm-induced geometries. In this work, we study the implicit bias of mini-batch stochastic steepest descent in multi-class classification, characterizing how batch size, momentum, and variance reduction shape the limiting max-margin behavior and convergence rates under general entry-wise and Schatten-$p$ norms. We show that, without momentum, worst-case convergence and successful classification can only be guaranteed with full-batch gradient. In contrast, momentum enables small-batch convergence to an approximate max-margin solution through a batch-momentum trade-off, though it slows convergence. This approach provides fully explicit, dimension-free rates that improve upon prior results. Moreover, we prove that variance reduction can recover the exact full-batch implicit bias for any batch size, albeit at a slower convergence rate. Finally, we further investigate the batch-size-one steepest descent without momentum, and reveal its convergence to a fundamentally different bias via a concrete data example, which reveals a key limitation of purely stochastic updates. Overall, our unified analysis clarifies when stochastic optimization aligns with full‑batch behavior, and paves the way for perform deeper explorations of the training behavior of stochastic gradient steepest descent algorithms.

A variety of widely used optimization methods like SignSGD and Muon can be interpreted as instances of steepest descent under different geometries. In this work, we study how mini-batch stochastic steepest descent behaves in multi-class classification, focusing on the effects of batch size, momentum, and variance reduction on the solutions learned during training. We show that without momentum, worst-case convergence and successful classification can only be guaranteed when the algorithm uses the full-batch gradient, while momentum enables convergence to an approximate max-margin solution even with small batches through a batch–momentum trade-off. We further show that variance reduction can recover the same max-margin behavior as full-batch training for any batch size. Finally, we show that purely stochastic updates with batch size one can converge to fundamentally different solutions, revealing a key limitation of highly stochastic training. Overall, our unified analysis clarifies when stochastic optimization aligns with full‑batch behavior, and paves the way for perform deeper explorations of the training behavior of stochastic gradient steepest descent algorithms.