Elliot Paquette

Dimension-adapted Momentum Outscales SGD

We investigate scaling laws for stochastic momentum algorithms with small batch on the power law random features model, parameterized by dat… (see more)a complexity, target complexity, and model size. When trained with a stochastic momentum algorithm, our analysis reveals four distinct loss curve shapes determined by varying data-target complexities. While traditional stochastic gradient descent with momentum (SGD-M) yields identical scaling law exponents to SGD, dimension-adapted Nesterov acceleration (DANA) improves these exponents by scaling momentum hyperparameters based on model size and data complexity. This outscaling phenomenon, which also improves compute-optimal scaling behavior, is achieved by DANA across a broad range of data and target complexities, while traditional methods fall short. Extensive experiments on high-dimensional synthetic quadratics validate our theoretical predictions and large-scale text experiments with LSTMs show DANA's improved loss exponents over SGD hold in a practical setting.

2025-05-01

arXiv (published)

doi.org

4+3 Phases of Compute-Optimal Neural Scaling Laws

Lechao Xiao

Jeffrey Pennington

2024-09-25

NeurIPS.cc/2024/Conference (spotlight)

Elizabeth Collins-Woodfin

The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms

Inbar Seroussi

Begoña García Malaxechebarría

Andrew Mackenzie

2024-09-25

NeurIPS.cc/2024/Conference (poster)

Elizabeth Collins-Woodfin

Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models

Inbar Seroussi

2023-08-17

ArXiv (preprint)

doi.org

Homogenization of SGD in high-dimensions: Exact dynamics and generalization properties

Ben Adlam

Jeffrey Pennington

2022-05-14

ArXiv (preprint)

Homogenization of SGD in high-dimensions: Exact dynamics and generalization properties

Ben Adlam

Jeffrey Pennington

We develop a stochastic differential equation, called homogenized SGD, for analyzing the dynamics of stochastic gradient descent (SGD) on a … (see more)high-dimensional random least squares problem with

2022-05-14

ArXiv (preprint)

Implicit Regularization or Implicit Conditioning? Exact Risk Trajectories of SGD in High Dimensions

Ben Adlam

Jeffrey Pennington

Stochastic gradient descent (SGD) is a pillar of modern machine learning, serving as the go-to optimization algorithm for a diverse array of… (see more) problems. While the empirical success of SGD is often attributed to its computational efficiency and favorable generalization behavior, neither effect is well understood and disentangling them remains an open problem. Even in the simple setting of convex quadratic problems, worst-case analyses give an asymptotic convergence rate for SGD that is no better than full-batch gradient descent (GD), and the purported implicit regularization effects of SGD lack a precise explanation. In this work, we study the dynamics of multi-pass SGD on high-dimensional convex quadratics and establish an asymptotic equivalence to a stochastic differential equation, which we call homogenized stochastic gradient descent (HSGD), whose solutions we characterize explicitly in terms of a Volterra integral equation. These results yield precise formulas for the learning and risk trajectories, which reveal a mechanism of implicit conditioning that explains the efficiency of SGD relative to GD. We also prove that the noise from SGD negatively impacts generalization performance, ruling out the possibility of any type of implicit regularization in this context. Finally, we show how to adapt the HSGD formalism to include streaming SGD, which allows us to produce an exact prediction for the excess risk of multi-pass SGD relative to that of streaming SGD (bootstrap risk).

Trajectory of Mini-Batch Momentum: Batch Size Saturation and Convergence in High Dimensions

Kiwon Lee

Andrew Nicholas Cheng