Cosine relations reshape momentum for better training outcomes
COREM: Cosine-Relation Momentum Reshaping with Stateful Writeback
Machine Learning
Summary
Training AI models involves adjusting many settings, called optimizer states, which can be thought of as big matrices. The authors looked at whether the relationships inside these matrices can be better used rather than treating each piece separately. They created a new method called COREM that reshapes the momentum during training by examining how parts of these matrices relate via cosine similarity, then updates the training process based on that. Their tests showed COREM improves model accuracy later in training and uses less compute compared to a similar method called Muon.
What this means in practice
- •For machine learning engineers: Enhance training efficiency and final accuracy of neural networks by integrating COREM for better momentum updates.
- •For deep learning framework developers: Implement COREM into optimization libraries to provide improved momentum handling with lower computational costs than existing methods.
Authors
Yan Wang, Xiaochuan Wang, Yuxiang Sun
Abstract
Matrix-valued optimizer states may contain relational structure that is not captured by treating their entries independently. We study whether relations within matrix-valued optimizer states can be exploited to improve optimization. To this end, we introduce a unit-relation-transform abstraction and instantiate it as COREM, a Cosine-Relation Momentum Reshaping method with stateful writeback. COREM partitions the momentum state into update units, computes cosine relations among them, and uses these relations to reshape the momentum before writing the transformed state back to the optimizer. This stateful mechanism allows the reshaped momentum to affect not only the current update but also future optimization dynamics. We evaluate COREM on CIFAR-10 with an MLP and on enwik8 with a Transformer. Compared with Muon, COREM shows lower early-stage step efficiency but stronger improvement in the mid-to-late stages of training, achieving better final validation performance on CIFAR-10 and comparable final performance on enwik8. Spectral diagnostics on enwik8 show that COREM consistently increases entropy effective rank and reduces the concentration of singular energy in dominant modes, while preserving an anisotropic spectrum. For square matrix updates, COREM requires approximately 13.3% of the transformation FLOPs of Muon with five Newton-Schulz iterations.