发表机构
Aalto University; Helsinki University; City University of Hong Kong（阿尔托大学; 赫尔辛基大学; 香港城市大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
COREM通过利用矩阵值优化器状态中的余弦关系重塑动量并状态回写，在CIFAR-10和enwik8上实现更优或相当的最终性能，且计算开销更低。
AI 中文摘要
矩阵值优化器状态可能包含关系结构，而独立处理其条目无法捕获这些结构。我们研究矩阵值优化器状态中的关系是否可以被利用来改进优化。为此，我们引入了一个单位关系变换抽象，并将其实例化为COREM，一种带状态回写的余弦关系动量重塑方法。COREM将动量状态划分为更新单元，计算它们之间的余弦关系，并利用这些关系在将变换后的状态写回优化器之前重塑动量。这种状态机制允许重塑后的动量不仅影响当前更新，还影响未来的优化动态。我们在CIFAR-10上使用MLP，在enwik8上使用Transformer评估了COREM。与Muon相比，COREM在早期阶段的步进效率较低，但在训练中后期表现出更强的改进，在CIFAR-10上取得了更好的最终验证性能，在enwik8上取得了相当的最终性能。enwik8上的谱诊断表明，COREM持续增加熵有效秩，并减少主导模式中奇异能量的集中度，同时保持各向异性谱。对于方阵更新，COREM所需的变换FLOPs约为使用五次Newton-Schulz迭代的Muon的13.3%。
英文摘要
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.
