Muon-RT Reuses Teacher Trajectories to Eliminate Momentum Buffers in LLM Self-Distillation

Summary:Muon-RT targets a persistent-memory bottleneck in LLM self-distillation: conventional Muon maintains a separate momentum buffer even when an exponential-moving-average teacher already preserves a history of student updates. The method reuses the teacher-student weight displacement, removes its radial component, applies spectral normalization, and mixes the resulting trajectory with the current gradient after Newton-Schulz orthogonalization. On SDFT with Qwen2.5-7B-Instruct, the paper reports a

---

Core Idea: Eliminate a Redundant Momentum History

Muon-RT targets persistent optimizer-state overhead in LLM self-distillation. An exponential-moving-average teacher already evolves with the student and therefore preserves information about historical updates. Conventional Muon nevertheless maintains a separate momentum buffer, duplicating part of that temporal signal in a different form.

Muon-RT reuses the teacher-student displacement

[ \Delta_t=\theta_{teacher,t}-\theta_{student,t}, ]

as an alternative to the momentum buffer. Under a local model, the paper argues that this displacement encodes an exponentially weighted history of applied updates and smooths temporally uncorrelated noise.

The direct benefit is reduced persistent training memory. For a parameter matrix (W\in\mathbb{R}^{m\times n}), a first-order momentum buffer normally requires another tensor with the same shape and precision as the parameter. As hidden weight matrices accumulate, this buffer consumes substantial memory and may increase optimizer-state sharding, checkpoint size, and communication overhead. Reusing teacher displacement removes that redundant state, although it introduces dependence on the EMA coefficient, teacher update cadence, and teacher-student synchronization.

Update Mechanism: Remove, Normalize, Then Mix

Muon-RT does not use the raw teacher-student displacement as the update direction.

1. Remove the Radial Component

If the displacement is strongly aligned with the current parameter or gradient direction, applying it directly may mostly rescale the parameter rather than improve the optimization geometry. The method removes this radial component to reduce conflict between the historical trajectory and the current descent direction.

2. Apply Spectral Normalization

Large differences in singular-value scales can cause matrix updates to be dominated by a small number of directions. Spectral normalization reduces this imbalance and makes the update less dependent on raw displacement magnitude.

The tradeoff is additional computation. Depending on the implementation, spectral processing may require matrix decompositions, iterative approximation, or related numerical routines. Reusing a state signal is therefore not automatically faster than maintaining momentum, even if it is cheaper in memory.

3. Mix with the Current Gradient

The normalized historical trajectory is combined with the current gradient after Newton-Schulz orthogonalization. The displacement contributes temporal smoothing, while the gradient retains responsiveness to the current batch.

This design implicitly introduces several hyperparameters: EMA decay, teacher update frequency, trajectory mixing weight, spectral normalization settings, and the Newton-Schulz iteration count. Poor alignment between these settings could cause either stale updates or excessive noise amplification.

Theoretical Basis: Linear Contraction Under Explicit Assumptions

Under stated regularity and alignment assumptions, the paper establishes linear contraction of an expected objective-gap bound toward a residual error level. This supports the claim that the teacher displacement can act as a meaningful smoothing signal under specific local conditions rather than being an arbitrary heuristic substitute for momentum.

The guarantee should not be read as unconditional global convergence. It depends on local-model approximations, teacher-student alignment, gradient properties, and step-size conditions. A teacher that updates too quickly may retain mostly recent noise, while one that updates too slowly may provide stale information. The residual error level also depends on the regularity terms used in the analysis.

Experimental Results: Accuracy and Persistent Memory Improve Together

SDFT on Qwen2.5-7B-Instruct

In Self-Distillation Fine-Tuning, Muon-RT improves average accuracy over AdamW by 2.6 points while reducing persistent training memory by 42%. Against momentum Muon, it improves average accuracy by 5.4 points and reduces persistent memory by 27%.

This is the paper's strongest result because the proposed method improves quality and memory efficiency within the same self-distillation framework. However, the reported comparisons focus on AdamW and standard Muon. The summary does not establish whether the gains persist across additional modern matrix optimizers, multiple random seeds, or equal hyperparameter-search budgets.

SDPO on Olmo-3-7B-Instruct

In Self-Distillation Policy Optimization, Muon-RT improves average accuracy over AdamW by 2.0 points. This indicates that the approach is not restricted to supervised fine-tuning and can also be integrated into a policy-optimization workflow.

For RL-based training, average downstream accuracy alone is not enough to judge the trade-off. Reward variance, rollout cost, training stability, convergence speed, and the comparison against momentum Muon are also needed. Those details are not available in the supplied summary.

Engineering Trade-offs

Muon-RT's clearest benefit is lower persistent optimizer memory, but the total training cost depends on several factors:

• Memory savings: Hidden weight matrices no longer require a separate first-order momentum buffer. The benefit should generally grow with the number and size of trainable matrices.
• Compute overhead: Each step must compute the teacher-student displacement, remove its radial component, normalize it spectrally, and mix it with an orthogonalized gradient.
• Distributed consistency: All parameter replicas need synchronized teacher weights and deterministic update behavior. Asynchronous EMA or delayed teacher synchronization changes the semantics of the reused trajectory.
• Checkpointing: Smaller optimizer state can reduce checkpoint and restore costs, although the EMA teacher itself still needs to be saved or reconstructed.
• Hyperparameter sensitivity: EMA decay, update frequency, mixing ratio, and spectral processing settings can materially affect stability.

The method is therefore better positioned as a memory-oriented optimization for large-scale post-training than as a universally faster optimizer.

Deployment Recommendations

Teams reproducing the work should measure:

1. Persistent and peak memory separately, including parameters, gradients, EMA teacher weights, optimizer state, and activations.
2. Per-step wall-clock time, isolating teacher updates, displacement computation, and spectral processing.
3. Variance across random seeds, especially for the reported 2.0-5.4 point accuracy gains.
4. Equal-budget optimizer comparisons, with comparable learning-rate search, step counts, and evaluation protocols.
5. Checkpoint recovery, ensuring EMA teacher state is restored correctly.
6. Larger-scale regimes, including full-parameter tuning, continued pretraining, and RL post-training.

Assessment

Muon-RT's main contribution is to reuse a state signal that self-distillation already maintains rather than independently accumulating a momentum buffer. By treating the teacher-student displacement as a spectral, history-aware direction and mixing it with the current gradient, the method reports better accuracy and lower persistent memory on two self-distillation settings.

The important engineering question is whether the extra displacement and spectral operations cost less than the savings from removing momentum state. The current evidence is promising but limited to roughly 7B-scale models and a small comparison set. Larger-scale replication, end-to-end throughput measurements, and distributed-training validation are needed before treating Muon-RT as a drop-in optimizer replacement.

---

Source: Hugging Face Daily Papers (2026-10-30)

— END —

Tags: #Muon-RT Follow #Muon Optimizer Follow #LLM Self-Distillation Follow #EMA Teacher Follow #Training Memory Optimization Follow