Decoupling Variance and Scale-Invariant Updates in Adaptive Gradient Descent for Unified Vector and Matrix Optimization
Abstract
Adaptive methods like Adam have become the de facto standard for large-scale vector and Euclidean optimization due to their coordinate-wise adaptation with a second-order nature. More recently, matrix-based spectral optimizers like Muon (Jordan et al., 2024b) show the power of treating weight matrices as matrices rather than long vectors. Linking these is hard because many natural generalizations are not feasible to implement, and we also cannot simply move the Adam adaptation to the matrix spectrum. To address this, we reformulate the AdaGrad update and decompose it into a variance adaptation term and a scale-invariant term. This decoupling produces DeVA (Decoupled Variance Adaptation), a framework that bridges between vector-based variance adaptation and matrix spectral optimization, enabling a seamless transition from Adam to adaptive spectral descent. Extensive experiments across language modeling and image classification demonstrate that DeVA consistently outperforms state-of-the-art methods such as Muon and SOAP (Vyas et al., 2024), reducing token usage by around 6.6\%. Theoretically, we show that the variance adaptation term effectively improves the blockwise smoothness, facilitating faster convergence.
Lay Summary
This paper proposes a new training method called DeVA. The main idea is to make smarter matrix updates. DeVA gives larger updates when the learning signal is clear and reliable, and smaller updates when the signal is noisy. Older methods like Adam are good at adjusting updates, but they treat model parameters mostly as vectors. Newer methods like Muon use the matrix structure inside neural networks, but they do not adapt the update size as carefully. DeVA combines both advantages in a principal way: it uses matrix structure and also adjusts updates based on reliability. In experiments, DeVA trains language models and image models more efficiently than several strong existing methods. For language modeling, it reaches the same target using about 6.6% fewer training tokens than Muon.