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.
Figures & tables
Method
Final Val. Acc.
Best Val. Acc.
Final Val. Loss
Best Val. Loss
Muon
59.74±1.01%
60.67±0.44%
1.2157±0.0322
1.1713±0.0131
COREM
61.36±0.19%
61.87±0.23%
1.1450±0.0084
1.1282±0.0040
Table 1: CIFAR-10 validation results over three random seeds. Values denote mean ± sample standard deviation.
Method
Final Val. Loss
Final Val. BPB
Final Val. Accuracy
Muon
1.095263±0.004475
1.580130±0.006457
0.678172±0.001344
COREM
1.096513±0.002470
1.581933±0.003563
0.678525±0.001018
Table 2: enwik8 validation results after 64,000 steps over three random seeds. Values denote mean ± sample standard deviation.
Figure 1: Validation trajectories for COREM and Muon on CIFAR-10 and enwik8. The first two rows use linear optimization time, while the third row shows enwik8 trajectories with logarithmic optimization steps. Shaded regions denote the three-seed uncertainty bands.
Figure 2: CIFAR-10 validation trajectories for the tuned writeback and no-writeback variants. Shaded regions denote the three-seed uncertainty bands.
Variant
Final Acc.
Best Acc.
Final Loss
Best Loss
No writeback
58.61±0.44%
59.47±0.13%
1.2130±0.0027
1.1859±0.0035
Writeback
61.36±0.19%
61.87±0.23%
1.1450±0.0084
1.1282±0.0040
Table 3: Effect of stateful momentum writeback on CIFAR-10 over three random seeds. Values denote mean ± sample standard deviation.
Figure 3: Stability map of the unnormalized COREM variant on enwik8 over a 4k-step learning-rate/relation-strength sweep. Each finite configuration is annotated with its final validation BPB, while × denotes a numerically unstable or diverged run.
Figure 4: Raw relation-scale dynamics across parameter groups, Transformer layers, and sampled optimization steps on enwik8. Each panel corresponds to one matrix-valued parameter group; rows denote Transformer layers, columns denote sampled optimization steps, and color indicates ρt=∥Ct∥∞ . Wq , Wk , Wv , and Wo denote the attention projection matrices, while mlp.0 and mlp.2 denote the first and second linear projections in the Transformer MLP block. All panels share the same color scale.
Figure 5: Spectral diagnostics across optimizer states and parameter groups on enwik8. Each panel compares Gt , Vt , and Mt over the 36 layer–checkpoint observations for each parameter group. Panels show entropy effective rank, top-10 singular-energy concentration, and the robust 10–90 spectral condition ratio. Boxes denote medians and interquartile ranges, with whiskers spanning the observed range; the condition-ratio panel uses a logarithmic y-axis.