A Polyak-Ruppert Central Limit Theorem for SA-Adam with Momentum and Non-Convergent Adaptive Preconditioning

2026-06-15

AI Summary

This work investigates whether the Polyak–Ruppert averaged iterates of SA-Adam satisfy a central limit theorem (CLT) under momentum and non-convergent adaptive preconditioning. By modeling the momentum-augmented state as a time-varying linear stochastic approximation process, the paper establishes—under local stability conditions—that when the momentum gain decays sublinearly, the asymptotic marginal covariance of SA-Adam coincides with that of SGD, rendering the adaptive mechanism asymptotically invisible. This result extends to settings with L2 regularization. Leveraging a non-autonomous CLT framework, the authors prove the asymptotic normality of SA-Adam, with covariance taking the canonical form $H^{-1}SH^{-1}$, thereby providing theoretical justification for the reliability of single-pass inference in regularized optimization problems.

Abstract

Adaptive optimizers combining preconditioning, momentum, and weight decay (Adam and AdamW) are, under Polyak-Ruppert averaging, candidate engines for one-pass inference. Does the averaged iterate keep the classical Polyak-Ruppert central limit theorem (CLT), with sandwich covariance $H^{-1}SH^{-1}$ (Hessian $H$, gradient covariance $S$), under momentum and non-convergent preconditioning? The preconditioner-only analysis does not carry over: with momentum the canonical decomposition collapses to a tautology. Treating the augmented state (iterate, momentum buffer) as a time-varying linear stochastic approximation (SA), we prove (under local stabilization) positive drift stability, a non-autonomous Polyak-Ruppert CLT, and a projection identity. The upshot: the iterate-marginal covariance is exactly the plain stochastic gradient descent (SGD) sandwich $H^{-1}SH^{-1}$, so the adaptivity is asymptotically invisible. This holds for SA-Adam (sub-linearly vanishing momentum gain, $γ\in(α,1)$; the sub-linear regime is essential), not constant-$β$ deployed Adam. Coupled $L_2$ weight decay yields the ridge-penalized sandwich, extending one-pass inference to regularized problems.

Problem

• Polyak-Ruppert averaging
• central limit theorem
• adaptive preconditioning
• momentum
• stochastic approximation

Innovation

• Polyak-Ruppert CLT
• SA-Adam
• adaptive preconditioning
• momentum
• one-pass inference