A Simplified Analysis of SGD for Linear Regression with Weight Averaging

TL;DR

This paper offers a simplified, accessible analysis of stochastic gradient descent (SGD) for linear regression, recovering sharp bias-variance bounds without complex operator manipulations, aiding future optimization research.

Contribution

It provides a straightforward linear algebra-based analysis of SGD for linear regression that reproduces existing bounds, simplifying understanding and further research.

Findings

Recovered bias and variance bounds using simple linear algebra

Simplified analysis bypasses complex PSD operator manipulations

Facilitates analysis of mini-batching and learning rate schedules

Abstract

Theoretically understanding stochastic gradient descent (SGD) in overparameterized models has led to the development of several optimization algorithms that are widely used in practice today. Recent work by~\citet{zou2021benign} provides sharp rates for SGD optimization in linear regression using constant learning rate, both with and without tail iterate averaging, based on a bias-variance decomposition of the risk. In our work, we provide a simplified analysis recovering the same bias and variance bounds provided in~\citep{zou2021benign} based on simple linear algebra tools, bypassing the requirement to manipulate operators on positive semi-definite (PSD) matrices. We believe our work makes the analysis of SGD on linear regression very accessible and will be helpful in further analyzing mini-batching and learning rate scheduling, leading to improvements in the training of realistic models.