Minimax Optimality of Shampoo for Nonstationary Stochastic Optimization
Organizations: Department of Statistics and Data Science, Cornell University · Department of Statistics, NC State University
Abstract
We develop a comprehensive theoretical analysis of Matrix Shampoo for tracking time-varying optima in nonstationary stochastic optimization. We derive finite-horizon high-probability bounds for average tracking error that disentangle initialization, optimizer drift, momentum-induced drift penalty, stochastic noise, and adaptive preconditioning, with explicit dependence on momentum parameter and stabilization parameter . These bounds reveal an explicit drift--noise and stability--geometry tradeoffs: noise-dominated regimes favor larger , which provides stronger first-moment averaging and greater suppression of gradient noise, whereas drift-dominated regimes favor smaller , allowing faster adaptation to moving optimizer. Larger reduces the preconditioner-perturbation term while weakening Shampoo's ability to exploit matrix geometry, thereby reducing adaptivity. We establish a matching minimax lower bound under squared path-variation constraint, proving the information-theoretic optimality of Shampoo over the class of causal optimization procedures. Under a constant stepsize, we establish high-probability bounds for the average projected stationarity gap along general nonstationary nonconvex objectives. We validate our theoretical findings on both synthetic and multiple real-world problems.
Figures & tables
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Stochastic Optimization
Stochastic optimization studies how to find good decisions or model parameters when objectives, data, or feedback are random or noisy.
See all 209 papers