Direction Matters: On the Implicit Bias of Stochastic Gradient Descent with Moderate Learning Rate

TL;DR

This paper investigates the implicit directional bias of stochastic gradient descent (SGD) with moderate learning rates, revealing how it differs from gradient descent (GD) and impacts early stopping and hyperparameter tuning in overparameterized linear regression.

Contribution

It characterizes the directional bias of SGD in the moderate learning rate regime and explains its practical implications for early stopping and hyperparameter choices.

Findings

SGD converges along large eigenvalue directions, unlike GD.

Directional bias affects the optimality of early stopping.

The theory explains common SGD hyperparameter tuning practices.

Abstract

Understanding the algorithmic bias of \emph{stochastic gradient descent} (SGD) is one of the key challenges in modern machine learning and deep learning theory. Most of the existing works, however, focus on \emph{very small or even infinitesimal} learning rate regime, and fail to cover practical scenarios where the learning rate is \emph{moderate and annealing}. In this paper, we make an initial attempt to characterize the particular regularization effect of SGD in the moderate learning rate regime by studying its behavior for optimizing an overparameterized linear regression problem. In this case, SGD and GD are known to converge to the unique minimum-norm solution; however, with the moderate and annealing learning rate, we show that they exhibit different \emph{directional bias}: SGD converges along the large eigenvalue directions of the data matrix, while GD goes after the small eigenvalue directions. Furthermore, we show that such directional bias does matter when early stopping is adopted, where the SGD output is nearly optimal but the GD output is suboptimal. Finally, our theory explains several folk arts in practice used for SGD hyperparameter tuning, such as (1) linearly scaling the initial learning rate with batch size; and (2) overrunning SGD with high learning rate even when the loss stops decreasing.