Gaussian mixture models improve predictions for constant-step learning algorithms
Moment-Accurate Gaussian Mixtures for Constant-Step Stochastic Approximation
Machine Learning
Summary
Many machine learning methods update their guesses step-by-step using a fixed step size, but it's hard to accurately predict how these guesses vary. The authors show that by using a mix of specific Gaussian models, one can make much better predictions about the expected errors and variability in these guesses. Their approach deals with technical challenges others have missed and can handle complex scenarios where the typical assumptions don’t hold. They confirmed their ideas with numerical experiments, helping better understand and control the learning process in algorithms like stochastic gradient descent.
What this means in practice
- •For machine learning engineers: Predict variability and expected errors in optimization steps of learning algorithms across different step sizes and batch sizes.
- •For statistical software developers: Implement improved Gaussian mixture models to simulate and analyze constant-step stochastic approximation behavior in training systems.
Authors
Xiaoli Li, Wei Biao Wu
Abstract
Local Gaussian models of constant-step learning predict output variability and expected losses, but weak convergence alone does not justify these moment predictions. We establish moment-accurate Gaussian mixtures by matching stationary energy with local Ornstein--Uhlenbeck limits, ruling out quadratic tail mass invisible to weak convergence. For step size $a$, the second-order Wasserstein error is $o(\sqrt a)$, uniformly over invariant laws, using each law's actual root weights. The assumptions combine confinement, descent, finitely many hyperbolic equilibria and root continuity with finite-variance innovations. The result yields observable covariances, expected objective gaps and first-order mean shifts, while allowing singular covariances, compatible saddles and weights without a limit. For additive noise given by a fixed invertible transform of independent standardized Student $t_3$ coordinates, symmetry gives an order-sharp $\sqrt a$ smooth-test bound. Numerical transport calculations demonstrate the value of root-specific covariances; controlled SGD studies assess observable predictions across step sizes, batch sizes and model geometries.