MUON-RT Reusing the Teacher's Trajectory for Momentum-Free Muon for LLM Self-Distillation

Muon-RT reuses the EMA teacher's displacement to remove Muon momentum buffers in self-distillation, reducing persistent memory and improving accuracy over AdamW and momentum Muon.

Optimizing EMA-based self-distillation requires student and teacher parameters plus optimizer state. Muon stores one momentum buffer per hidden weight matrix, so with an EMA teacher it needs three persistent states for each hidden matrix, while AdamW needs four including its two moment buffers. The paper asks whether the teacher's existing parameter trajectory can replace Muon's momentum buffer to reduce persistent memory while improving task accuracy.

Muon-RT replaces Muon's momentum buffer for each hidden weight matrix with the teacher–student displacement. It removes the component of the displacement parallel to the student weights, estimates its largest singular value with power iteration, and spectrally normalizes the remainder. The current gradient is orthogonalized by Newton–Schulz without momentum, and the two directions are mixed after orthogonalization using a small fixed coefficient. Only one power-iteration vector per matrix is retained for warm-starting, not additional matrix-sized persistent state. The method applies where a parameter-aligned EMA teacher already exists, including SDFT and SDPO.

• With SDFT on Qwen2.5-7B-Instruct, Muon-RT improves average accuracy over AdamW by 2.6 points with 42% less persistent memory, and over momentum Muon by 5.4 points with 27% less persistent memory. §Abstract
• With SDPO on Olmo-3-7B-Instruct, Muon-RT improves average accuracy over AdamW by 2.0 points. §Abstract
• Muon-RT achieves 74.2% on tool use and 71.8% on science, compared with 70.6% and 70.2% for SDFT with AdamW, and improves on momentum-free Muon by 5.1 and 2.8 points. §4.2
• Measured persistent memory per GPU is 37.7 GB for Muon-RT, 13.9 GB less than Muon with standard momentum and 26.8 GB less than AdamW. §4.4
• With SDPO on Olmo-3-7B-Instruct, Muon-RT is 2.0 points above AdamW on average; its margin is 0.7–1.1 points after 1 hour and 2.4–3.7 points after 5 hours, and it reaches 82.8% on Materials after 5 hours. §4.3
• At 10T parameters, Muon-RT needs 121 TB, 25% less than Muon or MuonClip with momentum and 40% less than AdamW. §4.4

• Experiments cover 7B-parameter models from 2 families and 4 tasks; larger models and further domains are left for future work.
• The convergence analysis relies on smoothness, PL and alignment assumptions whose constants are not estimated for LLMs.
• Muon-RT applies when a parameter-aligned EMA-teacher already exists; a teacher kept only to replace momentum would not save memory.
• Lemma 2 is exact for α=0 and a small-α approximation otherwise; the variance factor is not an exact reduction for the final normalized direction because the operations are nonlinear.

• Unrolling the EMA update shows the teacher–student displacement is a discounted history of applied update directions, not raw gradients. §3.1
• Under a local fixed direction and temporally uncorrelated noise, the teacher displacement has mean scaled by 1−τ^s and variance reduced by a factor up to (1−τ)/(1+τ). §3.4
• Muon-RT orthogonalizes the current gradient without momentum, removes the displacement's radial component, spectrally normalizes the remainder, and mixes the two directions after Newton–Schulz. §1
• Muon-RT reduces matrix-sized persistent tensors from 4 to 3 and keeps only one power-iteration vector of length B per matrix for warm-starting. §3.3
• Proposition 1 shows the Newton–Schulz polar direction is optimal for the current gradient, so the teacher branch can only add temporal information. §A.1
• The teacher weight α=0.1 is best on both SDFT tasks and is used unchanged for SDPO. §A.11
• QK-Clip is not needed in this setting and is harmful with Muon-RT; logits stay within ±1% of the base value for AdamW, Muon and Muon-RT. §A.8