Courtney Paquette

2026-07-12

Mathematical Programming (published)

Begoña García Malaxechebarría

High-dimensional Limit of SGD for Diagonal Linear Networks

Maryam Fazel

Dmitriy Drusvyatskiy

Understanding the behavior of stochastic gradient methods is a central problem in modern machine learning. Recent work has highlighted diago… (see more)nal linear networks as a simplified yet expressive setting for analyzing the optimization and generalization properties of neural models. In this work, we show that in the high-dimensional regime, stochastic gradient descent on diagonal linear networks is well-approximated by continuous dynamics governed by a stochastic differential equation (SDE), which explicitly decouples the drift from the gradient noise. We further derive a deterministic partial differential equation whose solution propagates the relevant state of the iterates and characterizes the time evolution of a broad class of observable statistics, including the risk, curvature, and other metrics for optimality. Finally, we show that, under a suitable parametrization, the stochastic dynamics are globally well posed and converge exponentially fast to zero risk with high probability, yielding a fully explicit non-asymptotic description of their long-time behavior. Numerical simulations corroborate our theoretical findings.

2026-05-15

arXiv (preprint)

Phases of Muon: When Muon Eclipses SignSGD

Lucas Benigni

Atish Agarwala

Recently, Muon and related spectral optimizers have demonstrated strong empirical performance as scalable stochastic methods, often outperfo… (see more)rming Adam. Yet their behaviour remains poorly understood. We analyze stochastic spectral optimizers, including Muon, on a high-dimensional matrix-valued least squares problem. We derive explicit deterministic dynamics that provide a tractable framework for studying learning behaviour with a focus on (stochastic) SignSVD, which Muon approximates, and (stochastic) SignSGD, the latter serving as a proxy for Adam. Our analysis shows that for large batch size, SignSVD performs a square-root preconditioning with respect to the data covariance spectrum, while for small batch size smaller eigenmodes behave like SGD, slowing down convergence. We contrast with SignSGD which for generic covariance performs no preconditioning and has no transition, leading to different optimal learning rates and convergence characteristics. The two methods match up to a constant factor with isotropic data, but behave differently with anisotropic data. An analysis of a power law covariance model with data exponent

2026-05-09

arXiv (preprint)

Mirror Descent Algorithms with Nearly Dimension-Independent Rates for Differentially-Private Stochastic Saddle-Point Problems

Tomas Gonzalez

Cristobal Guzman

2026-02-19

SIAM Journal on Optimization (published)

Logarithmic-time Schedules for Scaling Language Models with Momentum

In practice, the hyperparameters …

2026-02-04

arXiv (preprint)

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-09-17

NeurIPS.cc/2025/Conference (spotlight)

Two-point deterministic equivalence for SGD in random feature models

Alexander Atanasov

Blake Bordelon

Jacob A Zavatone-Veth

Cengiz Pehlevan

2025-06-08

ICML.cc/2025/Workshop/HiLD (poster)

Implicit Diffusion: Efficient Optimization through Stochastic Sampling

Pierre Marion

Anna Korba

Peter Bartlett

Mathieu Blondel

Valentin De Bortoli

Arnaud Doucet

Felipe Llinares-López

Quentin Berthet

2024-12-31

AISTATS (published)

proceedings.mlr.press

High Dimensional First Order Mini-Batch Algorithms on Quadratic Problems

Andrew Nicholas Cheng

Kiwon Lee

We analyze the dynamics of general mini-batch first order algorithms on the …

2024-10-09

NeurIPS.cc/2024/Workshop/OPT (published)

4+3 Phases of Compute-Optimal Neural Scaling Laws

Elliot Paquette

Lechao Xiao

Jeffrey Pennington

We consider the solvable neural scaling model with three parameters: data complexity, target complexity, and model-parameter-count. We use t… (see more)his neural scaling model to derive new predictions about the compute-limited, infinite-data scaling law regime. To train the neural scaling model, we run one-pass stochastic gradient descent on a mean-squared loss. We derive a representation of the loss curves which holds over all iteration counts and improves in accuracy as the model parameter count grows. We then analyze the compute-optimal model-parameter-count, and identify 4 phases (+3 subphases) in the data-complexity/target-complexity phase-plane. The phase boundaries are determined by the relative importance of model capacity, optimizer noise, and embedding of the features. We furthermore derive, with mathematical proof and extensive numerical evidence, the scaling-law exponents in all of these phases, in particular computing the optimal model-parameter-count as a function of floating point operation budget.

2024-09-24

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 W. Mackenzie

Elliot Paquette

We develop a framework for analyzing the training and learning rate dynamics on a large class of high-dimensional optimization problems, whi… (see more)ch we call the high line, trained using one-pass stochastic gradient descent (SGD) with adaptive learning rates. We give exact expressions for the risk and learning rate curves in terms of a deterministic solution to a system of ODEs. We then investigate in detail two adaptive learning rates -- an idealized exact line search and AdaGrad-Norm -- on the least squares problem. When the data covariance matrix has strictly positive eigenvalues, this idealized exact line search strategy can exhibit arbitrarily slower convergence when compared to the optimal fixed learning rate with SGD. Moreover we exactly characterize the limiting learning rate (as time goes to infinity) for line search in the setting where the data covariance has only two distinct eigenvalues. For noiseless targets, we further demonstrate that the AdaGrad-Norm learning rate converges to a deterministic constant inversely proportional to the average eigenvalue of the data covariance matrix, and identify a phase transition when the covariance density of eigenvalues follows a power law distribution. We provide our code for evaluation at https://github.com/amackenzie1/highline2024.

2024-09-24

NeurIPS.cc/2024/Conference (poster)

Mirror Descent Algorithms with Nearly Dimension-Independent Rates for Differentially-Private Stochastic Saddle-Point Problems extended abstract

Tomas Gonzalez

Cristobal Guzman