Entrywise dynamics and universality of general first order methods

TL;DR

This paper develops a non-asymptotic, entrywise characterization of general first order methods, demonstrating their universality across diverse random matrix models and applying these results to statistical estimators and gradient descent algorithms.

Contribution

It introduces a novel non-asymptotic, entrywise analysis framework for GFOMs, extending universality results beyond Gaussian models and enabling new insights into estimators and optimization algorithms.

Findings

Entrywise universality holds for GFOMs across broad matrix models.

Provides non-asymptotic description of empirical distributions beyond Gaussian ensembles.

Achieves entrywise Gaussian approximations for gradient descent in non-convex settings.

Abstract

General first order methods (GFOMs), including various gradient descent and AMP algorithms, constitute a broad class of iterative algorithms in modern statistical learning problems. Some GFOMs also serve as constructive proof devices, iteratively characterizing the empirical distributions of statistical estimators in the large system limits for any fixed number of iterations. This paper develops a non-asymptotic, entrywise characterization for a general class of GFOMs. Our characterizations capture the precise entrywise behavior of the GFOMs, and hold universally across a broad class of heterogeneous random matrix models. As a corollary, we provide the first non-asymptotic description of the empirical distributions of the GFOMs beyond Gaussian ensembles. We demonstrate the utility of these general results in two applications. In the first application, we prove entrywise universality for regularized least squares estimators in the linear model, by controlling the entrywise error relative to a suitably constructed GFOM. This algorithmic proof method also leads to systematically improved averaged universality results for regularized regression estimators in the linear model, and resolves the universality conjecture for (regularized) MLEs in logistic regression. In the second application, we obtain entrywise Gaussian approximations for a class of gradient descent algorithms. Our approach provides non-asymptotic state evolution for the bias and variance of the algorithm along the iteration path, applicable for non-convex loss functions. The proof relies on a new recursive leave-k-out method that provides almost delocalization for the GFOMs and their derivatives. Crucially, our method ensures entrywise universality for up to poly-logarithmic many iterations, which facilitates effective $\ell_2/\ell_\infty$ control between certain GFOMs and statistical estimators in applications.