Stochastic algorithms improve by using first and second moment estimates

Second-Moment Stochastic Approximation Methods

Artificial Intelligence

Summary

Many algorithms that learn from noisy data adjust their steps using average errors (first moment). This paper looks at ways to improve these algorithms by also considering how those errors vary (second moment), which is what popular methods like Adam and Muon do. The authors show how these methods can be understood as solving special matrix equations and provide a detailed analysis of their convergence. They prove that practical versions of these algorithms generally get close to the correct solution and explain how the quality of the estimates affects how close they get. Finally, they offer precise mathematical bounds for how tight this convergence is for specific algorithms.

What this means in practice

  • For machine learning engineers: Improve training stability and performance of deep learning models by understanding convergence properties of optimizers like Adam and Muon.
  • For control system designers: Use enhanced stochastic approximation methods that account for second moments to better tune adaptive controllers under uncertainty.

A theory result. No direct application yet.

Authors

Tao Jiang, Lin Xiao

Abstract

Classical stochastic approximation methods rely on estimators of the first moment (mean) of a random regression function. We study methods that employ estimators of both the first and the second moments, which include modern deep-learning optimizers such as Adam and Muon as special cases. We derive second-moment stochastic approximation methods through the lens of optimal preconditioning for solving matrix equations, and develop a two-stage framework for their convergence analysis. The first stage focuses on the analysis of conceptual (impractical) methods that rely on the exact first and second moments. In the second stage, we replace the exact moments with their respective estimators, and invoke Dvoretzky's theorem to show that the resulting practical methods converge almost surely to a neighborhood of the target solution. The size of the neighborhood depends on the biases and variances of the first- and second-moment estimators. We derive concrete bounds for Muon and a spectral variant of Adam that determine the radius of their neighborhood of convergence.
View PDFOpen arXiv