Matrix Shampoo optimizer achieves best tracking in changing problems
Minimax Optimality of Shampoo for Nonstationary Stochastic Optimization
Machine Learning
Summary
Tracking constantly changing targets in optimization is tricky because noise and changes can interfere. This paper studies the Matrix Shampoo optimizer and shows how it balances noise smoothing and quick adaptation to changes. The authors prove mathematically that Shampoo is as good as any optimizer can be in this setting. They also explain how different settings like momentum and stabilization affect its behavior. Their results apply broadly to nonstationary problems and are tested on both synthetic and real data.
What this means in practice
- •For machine learning engineers: Optimize training on time-varying data or models with smoother noise reduction and faster response to changes using Shampoo's optimized momentum and stabilization.
- •For control system developers: Improve adaptive algorithms that track changing system dynamics by applying Shampoo’s proven tracking guarantees under varying noise and drift conditions.
Authors
Abir Sarkar, Nirmalya Mandal, Martin T. Wells
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 \(β_1\) and stabilization parameter \(\varepsilon\). These bounds reveal an explicit drift--noise and stability--geometry tradeoffs: noise-dominated regimes favor larger $β_1$, which provides stronger first-moment averaging and greater suppression of gradient noise, whereas drift-dominated regimes favor smaller $β_1$, allowing faster adaptation to moving optimizer. Larger $\varepsilon$ 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.