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

AI summary

Overview

Research area: Machine learning systems and optimization for large language models — specifically memory-efficient optimizers for self-distillation with exponential moving average (EMA) teachers.

Technical level: Advanced. The paper assumes familiarity with the Muon optimizer, Newton–Schulz orthogonalization, singular values and spectral norms, EMA teacher–student training, and non-convex convergence analysis (Polyak–Łojasiewicz conditions). The core idea is intuitive, but the theory sections are dense.

Scope: The paper proposes Muon-RT, an optimizer that deletes Muon's momentum buffer and substitutes the already-existing teacher–student weight displacement, evaluated with SDFT on Qwen2.5-7B-Instruct and SDPO on Olmo-3-7B-Instruct.

What This Paper Is About

Self-distillation methods such as SDFT and SDPO train a student model against an EMA teacher that continuously tracks the student's own weights, so the training run already stores two full copies of the model. Adding Muon on top means storing a third state — a momentum buffer as large as each hidden weight matrix — purely to smooth updates. The authors ask whether the teacher's existing parameter trajectory can stand in for that momentum buffer, cutting persistent memory while improving task accuracy rather than trading accuracy away for memory.

Key Contributions

1. A characterization of what the EMA teacher actually stores. The authors prove that the teacher–student displacement equals a discounted (exponentially weighted) sum of the applied update directions, not of raw gradients, and that under a local model with a fixed direction and temporally uncorrelated noise it reduces noise variance by a factor of up to (1−τ)/(1+τ).

2. The Muon-RT update rule. A momentum-free Muon variant that orthogonalizes the current gradient with Newton–Schulz, separately builds a teacher direction by removing the radial (weight-parallel) component of the displacement and normalizing it by an estimated largest singular value, then mixes the two after orthogonalization with a small fixed coefficient α.

3. Convergence guarantees. Under smoothness, a PL condition, bounded teacher drift, and an alignment assumption, the authors establish linear contraction of an expected objective-gap bound toward a residual error level, plus a corollary stating when teacher reuse contracts faster than a baseline optimizer.

4. Empirical validation on two self-distillation hosts. Muon-RT reaches the top accuracy among Muon variants in all six task-and-budget settings across SDFT and SDPO, while reducing measured and projected persistent memory.

Main Findings

• SDFT accuracy gains: On Qwen2.5-7B-Instruct, Muon-RT reaches 74.2% on the tool-use task and 71.8% on science, versus 70.6% and 70.2% for SDFT with AdamW, and versus 69.1% and 69.0% for momentum-free Muon — a gain of 5.1 and 2.8 points over that Muon baseline.

• Headline memory-plus-accuracy summary: Averaged across the two SDFT datasets, Muon-RT improves over AdamW by 2.6 points with 42% less persistent memory, and over momentum Muon by 5.4 points with 27% less persistent memory.

• Measured memory: Between steps, Muon-RT needs 37.7 GB per GPU — 13.9 GB (27%) less than Muon with its standard momentum μ=0.95 and 26.8 GB (42%) less than AdamW. Momentum-free Muon and MuonClip need the same 37.7 GB but score 0.8–5.1 points lower.

• Projected memory at scale: At 10T parameters, Muon-RT needs 121 TB, 25% less than Muon or MuonClip with momentum (160 TB) and 40% less than AdamW (200 TB). For 7B persistent states the paper cites 103.2 GB (momentum Muon), 129.0 GB (AdamW), and 75.4 GB (Muon-RT).

• SDPO accuracy gains: On Olmo-3-7B-Instruct, Muon-RT is the most accurate optimizer for SDPO on both datasets and budgets, 2.0 points above AdamW on average. After 5 hours it reaches 55.2% on Biology and 82.8% on Materials, versus 52.8% and 79.1% for SDPO with AdamW. Its margin over AdamW is 0.7–1.1 points after 1 hour and widens to 2.4–3.7 points after 5 hours.

• Faster convergence: On tool use, Muon-RT reaches 74.2% after 100 of 252 steps, whereas MuonClip needs 250 steps to reach 73.2% and Muon never exceeds 69.1%.

• Retention of prior capabilities: All self-distillation runs sit within 2.4 points of one another on prior-capability averages, and Muon-RT's averages are 11.9 and 6.7 points above SFT. On tool use it retains 96% of the AdamW average (63.0 versus 65.4).

• The teacher branch must come after orthogonalization: Mixing before Newton–Schulz drops accuracy to 72.5/69.2 (tool use/science) versus 74.2/71.8 for the full method; using a Frobenius norm instead of spectral normalization gives 71.8/68.5; keeping the radial component gives 70.5/67.3; a cosine gate on α gives 69.2/66.8.

• The mixing weight is sensitive: α=0.1 is best (74.2/71.8); α=0, 0.05, 0.2, and 0.4 all fall to between 65.4 and 69.1 on at least one task.

• Gradient momentum hurts when the target moves: Muon and MuonClip with μ=0.95 score only 66.9/68.3 and 67.0/68.8, below their momentum-free counterparts and well below Muon-RT.

• Logit clipping does not compose with the method: Adding QK-Clip to Muon-RT lowers tool-use accuracy to 66.0 (student only) and 67.3 (student and teacher).

• Overhead is negligible: The extra cost is one power iteration and one extra gather per matrix; the optimizer step takes 2.3 s, about 2% of a training step dominated by generation.

Methodology in Plain English

The authors start from an observation about bookkeeping. In self-distillation, the teacher is not a fixed model — it is an exponential moving average of the student, updated as τ·θ_T + (1−τ)·θ_S. That means the gap between teacher and student weights is a running record of where the student has been: unrolling the recursion shows the gap is a decaying sum of past update directions, with older updates weighted by powers of τ.

Two complications stop you from simply calling that gap "momentum." First, Muon's updates are orthogonalized before being applied, so the gap averages optimizer outputs, not raw gradients — feeding it back in directly would reinforce stale directions, which is dangerous when the distillation target itself keeps moving. Second, if you add the gap to the gradient before Newton–Schulz, the nonlinear orthogonalization can amplify weak singular directions and you lose control over how much history is mixed in.

Muon-RT therefore runs two separate branches. The gradient branch is plain momentum-free Muon: orthogonalize the current gradient with Newton–Schulz. The teacher branch takes the displacement, subtracts its component parallel to the student weights (killing scale/weight-decay effects), and divides by an estimate of its largest singular value obtained by power iteration — warm-started from the previous step's right singular vector, so only one short vector per matrix is ever stored. The two unit-scale directions are then blended as U = (1−α)O + αP and applied with the usual Muon step size and decoupled weight decay. Because the teacher and student are sharded identically under FSDP, the displacement is computed locally and only the small scalar inner products are all-reduced.

The theory formalizes why this should help. A proposition shows the Newton–Schulz direction is already the optimal instantaneous choice for the current gradient, so the teacher can only add value by contributing temporal information — which the lemmas show it does, by averaging out noise. The convergence theorem gives a linear contraction rate q_R with a residual neighborhood whose size depends on gradient noise and teacher drift; the corollary says reuse beats a baseline whenever the teacher branch raises the effective descent coefficient without inflating the update's second moment.

Experiments keep the two host objectives (SDFT and SDPO) untouched and swap only the optimizer. SDFT is run on Qwen2.5-7B-Instruct with tool use from ToolAlpaca and science questions from SciKnowEval; SDPO is run on Olmo-3-7B-Instruct with the biology and materials subsets of SciKnowEval, using each original paper's default splits and metrics. Baselines include SFT, SFT with re-invoked prior knowledge, DFT, GRPO, and SDFT/SDPO with AdamW, Muon, and MuonClip.

Why This Matters

Impact on research. The paper reframes an already-paid memory cost as a free optimization signal. Rather than compressing or quantizing optimizer state, it points out that self-distillation and mean-teacher setups already maintain a parameter-aligned slow average, and that the displacement from current weights is a usable, provably noise-reducing direction. It also supplies a specific, slightly counterintuitive empirical lesson: conventional gradient momentum hurts when the learning target is itself moving, which bears on how optimizers are chosen for any on-policy or self-targeting training loop.

Real-world applications:

• Continual adaptation on constrained hardware. Because persistent memory drops by 27–42% for hidden matrices, the same GPU fleet can adapt larger models or hold more replicas.
• On-policy post-training pipelines. SDPO-style reinforcement-style training with feedback benefits directly, with margins that grow with training time (2.4–3.7 points after 5 hours).
• Long-context / long-run fine-tuning jobs. Memory freed from optimizer state can be reallocated to activations, batch size, or sequence length.
• Any mean-teacher setup — the method applies wherever a parameter-aligned EMA teacher exists, which the authors note includes classical mean-teacher training beyond LLMs.

Industry relevance. Muon and its variants have been used to train Kimi K2 (Bai et al., 2025) and GLM-4.5 (Team et al., 2025), and the paper's projected 10T-parameter analysis (121 TB versus 160 TB and 200 TB) targets exactly the regime frontier labs operate in. The method is a drop-in optimizer swap that leaves the training objective unchanged, which lowers the barrier to adoption.

Future Directions

• Scale and breadth. The authors explicitly flag that experiments cover 7B-parameter models from two families and four tasks; larger models and further domains are left for future work.
• Tighter theory. The convergence analysis relies on smoothness, PL, and alignment assumptions whose constants are not estimated for LLMs — measuring them would turn the bound into a practical predictor.
• Recovering second-moment information. The displacement cannot reproduce AdamW's second-moment statistics; whether some auxiliary cheap statistic could close that gap is open.
• Teachers without parameter alignment. Muon-RT applies only when a parameter-aligned EMA teacher already exists; for methods lacking one, keeping a teacher solely to replace momentum would not save memory, so how to extend the idea to unaligned or independently trained teachers remains open.

Target Audience

Optimizer and training-infrastructure researchers working on LLM pre-training and post-training; practitioners running self-distillation or mean-teacher pipelines who are memory-constrained; and theoretically inclined readers interested in what weight-averaging actually encodes. Readers wanting only the practical takeaway can read the abstract, Figure 1, and Table 3; readers wanting the full argument should read Sections 3.1–3.4 and Appendix A.1–A.2.

Authors’ abstract

Self-distillation methods for large language models with an exponential moving average (EMA) teacher already retain a history of student updates. Muon nevertheless stores a separate momentum buffer, increasing persistent training memory. We propose Muon with Reused Teacher (Muon-RT), which reuses the teacher-student displacement to eliminate momentum buffers for hidden weight matrices. We show that this displacement encodes an exponentially weighted history of applied updates and smooths temporally uncorrelated noise under a local model. Muon-RT removes its radial component, normalizes it spectrally, and mixes it with the current gradient after Newton-Schulz (NS) orthogonalization. Under explicit regularity and alignment assumptions, we establish linear contraction of an expected objective-gap bound toward a residual error level. With Self-Distillation Fine-Tuning (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. With Self-Distillation Policy Optimization (SDPO) on Olmo-3-7B-Instruct, it improves average accuracy over AdamW by 2.0 points.