[Paper Review] The Implicit Regularization of Stochastic Gradient Flow for Least Squares

Alnur Ali, Edgar Dobriban+1 more|arXiv (Cornell University)|Mar 17, 2020

Stochastic Gradient Optimization Techniques90 references19 citations

TL;DR

This paper establishes a tight theoretical connection between mini-batch stochastic gradient flow (SGF) and ridge regression in least squares problems. Using a continuous-time stochastic differential equation model, it proves that the excess risk of SGF at time $ t $ is bounded above by that of ridge regression with $ \lambda = 1/t $, with explicit dependence on step size, mini-batch size, and variance, showing SGF implicitly regularizes like ridge with minimal computational cost.

ABSTRACT

We study the implicit regularization of mini-batch stochastic gradient descent, when applied to the fundamental problem of least squares regression. We leverage a continuous-time stochastic differential equation having the same moments as stochastic gradient descent, which we call stochastic gradient flow. We give a bound on the excess risk of stochastic gradient flow at time $t$, over ridge regression with tuning parameter $λ= 1/t$. The bound may be computed from explicit constants (e.g., the mini-batch size, step size, number of iterations), revealing precisely how these quantities drive the excess risk. Numerical examples show the bound can be small, indicating a tight relationship between the two estimators. We give a similar result relating the coefficients of stochastic gradient flow and ridge. These results hold under no conditions on the data matrix $X$, and across the entire optimization path (not just at convergence).

Motivation & Objective

• To understand the implicit regularization effect of mini-batch stochastic gradient descent (SGD) in least squares regression.
• To establish a precise, time-dependent relationship between the risk of stochastic gradient flow (SGF) and ridge regression with $ \lambda = 1/t $.
• To quantify how algorithmic hyperparameters—step size, mini-batch size, and time—drive the excess risk of SGF relative to ridge.
• To demonstrate that SGF achieves near-optimal statistical performance with significantly reduced computation time compared to ridge.
• To provide bounds valid across the entire optimization path, not just at convergence, under no assumptions on the data matrix $ X $.

Proposed method

• Models SGD via a continuous-time stochastic differential equation called stochastic gradient flow (SGF), which matches the first and second moments of mini-batch SGD.
• Derives an exact risk decomposition for SGF, showing excess risk over ridge regression with $ \lambda = 1/t $ as a sum of three terms: variance of ridge, variance from mini-batching, and a fluctuation term from constant step size.
• Uses spectral decomposition of the data matrix $ X $ and singular values to bound the difference between SGF and ridge coefficients.
• Introduces a function $ f(x) = (1 - e^{-x})(1+x)/x $ to analyze the coefficient distance, leveraging its unimodal and bounded nature to derive a uniform bound.
• Employs matrix inequalities and concentration arguments to show that the coefficient difference is bounded by $ (g(t) - 1) \| \hat{\beta}^{\text{ridge}}(1/t) \|_2 $, where $ g(t) $ is a time-dependent factor.
• Validates theoretical bounds with numerical experiments on Gaussian, Student-t, and Bernoulli data across varying $ n, p, m $, and $ \epsilon $.

Experimental results

Research questions

• RQ1 How does the excess risk of stochastic gradient flow at time $ t $ compare to that of ridge regression with tuning parameter $ \lambda = 1/t $?
• RQ2 What is the explicit dependence of the excess risk on algorithmic hyperparameters such as step size, mini-batch size, and time?
• RQ3 Can the implicit regularization of SGF be tightly bounded and quantified across the entire optimization path, not just at convergence?
• RQ4 How does the performance of SGF compare to ridge regression in terms of risk and computation time, especially when both are optimally stopped?
• RQ5 What is the relationship between the coefficient paths of SGF and ridge regression with $ \lambda = 1/t $, and how tight is this bound?

Key findings

• The excess risk of stochastic gradient flow at time $ t $ is bounded above by that of ridge regression with $ \lambda = 1/t $, with the bound decomposed into three interpretable terms: ridge variance, stochasticity variance, and a fluctuation term.
• The bound is tight in practice—numerical examples show the excess risk is as low as 1.0032 times the optimal ridge risk, indicating a strong implicit regularization effect.
• The coefficient difference between SGF and ridge regression with $ \lambda = 1/t $ is bounded by $ (g(t) - 1) \| \hat{\beta}^{\text{ridge}}(1/t) \|$ where $ g(t) \leq 1.2985 $, showing the two estimators are close across all $ t $.
• The fluctuation term in the risk bound vanishes as $ t \to \infty $, and is zero in the overparametrized regime, reflecting convergence to the least squares solution.
• SGF achieves near-optimal risk with significantly less computation time than ridge regression when both are stopped at their optimal time, demonstrating a favorable computational-statistical trade-off.
• The results hold without any assumptions on the data matrix $ X $, and are valid across the entire optimization path, not just at convergence.