aiqbench analysis
COREM introduces a novel optimization technique that leverages relational structures within matrix-valued optimizer states to reshape momentum before state writeback. While it incurs higher computational costs in early training stages compared to Muon, the method demonstrates superior convergence and validation performance in mid-to-late training phases on both MLP and Transformer architectures. This approach offers a promising avenue for improving optimization dynamics by explicitly modeling correlations that standard independent-entry methods overlook.
COREM: Cosine-Relation Momentum Reshaping with Stateful Writeback
arXiv:2609.22487v1 Announce Type: new 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.
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.
Machine Learning (cs.LG)
arXiv:2609.22487 [cs.LG]
(or arXiv:2609.22487v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2609.22487 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 284 words
Related coverage
Game Design Creativity evals
LessWrong · Oct 12, 2026
Futures Drop, Yields And Oil Jump Ahead Of Iran UN Address As Trump - Xi Summit Looms
zerohedge.com · Sep 23, 2026
Digital States 2026: AI Moves From Experiment to Infrastructure
govtech.com · Sep 24, 2026
Joint Affine Spectral Shaping: Coupling Weight and Bias Updates Beyond Weight-Only Muon
arXiv · Aug 5, 2026