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

TL;DR

This paper analyzes how mini-batch stochastic steepest descent methods, including momentum and variance reduction, influence the implicit bias and convergence behavior in multi-class classification, revealing conditions for full-batch-like behavior and limitations of stochastic updates.

Contribution

It provides a unified, explicit analysis of the implicit bias and convergence rates of mini-batch stochastic steepest descent, including effects of batch size, momentum, and variance reduction.

Findings

Large batches are needed for convergence without momentum.

Momentum enables small-batch convergence but slows it down.

Variance reduction can recover full-batch implicit bias for any batch size.

Abstract

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, convergence only occurs with large batches, yielding a batch-dependent margin gap but the full-batch convergence rate. In contrast, momentum enables small-batch convergence 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.