[KongchangAI]
· 2 min read· 1,239 words

Muon-RT: Reusing Teacher Trajectories for Momentum-Free LLM Self-Distillation

Muon-RT: Reusing Teacher Trajectories for Momentum-Free LLM Self-Distillation

Muon-RT reuses parameter displacement between EMA teacher and student to replace Muon's independent momentum buffer, saving memory and improving accuracy.

The article introduces an optimizer improvement called Muon-RT. Its starting point: in a self-distillation framework based on an EMA teacher, the teacher parameters are already an exponentially weighted accumulation of the student's historical updates, so the independent momentum buffer maintained by the Muon optimizer is redundant. Muon-RT reuses the teacher-student displacement as a momentum substitute, applying de-radialization and spectral normalization, then mixing with the current gradient after Newton-Schulz orthogonalization, achieving momentum-like smoothing without additional persistent state. The paper proves linear contraction of the expected objective-gap bound under regularity and alignment assumptions. Experiments show that on SDFT+Qwen2.5-7B-Instruct it improves accuracy by 2.6 points over AdamW with 42% less persistent memory, and by 5.4 points over momentum Muon with 27% less memory; on SDPO+Olmo-3-7B-Instruct it improves by 2.0 points over AdamW, demonstrating cross-paradigm applicability.

Self-distillation has become an important technique for improving large language model performance, and the exponential moving average (EMA) teacher model itself already retains historical update information from the student model. However, the mainstream Muon optimizer still additionally maintains an independent momentum buffer, which incurs persistent GPU memory overhead. A recent study proposes Muon-RT, which eliminates the momentum buffer for hidden weight matrices by reusing the displacement between teacher and student, providing a more memory-efficient path for self-distillation training.

Muon-RT paper source screenshot

Core Idea: Replacing the Momentum Buffer with Teacher Trajectories

The researchers' starting point is straightforward: in a self-distillation framework based on an EMA teacher, the teacher model is essentially an exponentially weighted history of the student model's parameter updates. Given this, having the optimizer separately store a momentum buffer is redundant. The core operation of Muon-RT is to reuse the teacher-student parameter displacement as a substitute signal for momentum.

The paper points out that this displacement encodes an exponentially weighted history of applied updates and, under a local model assumption, can smooth out temporally uncorrelated noise. In other words, the teacher trajectory itself possesses the "historical smoothing" property that momentum is expected to provide—it just had not been directly exploited by the optimizer before.

Implementation: De-radialization, Spectral Normalization, and Gradient Mixing

Muon-RT's processing pipeline involves several key steps. First, it removes the radial component from the teacher-student displacement, then applies spectral normalization, and finally mixes it with the current gradient after Newton-Schulz (NS) orthogonalization. This design remains compatible with Muon's original NS orthogonalization pipeline, with changes concentrated on the source of the momentum term rather than the overall optimization framework.

This "borrow rather than build" approach allows Muon-RT to achieve momentum-like smoothing without adding extra persistent state.

Exponential Moving Average (EMA) teacher: In self-distillation, the EMA teacher is a sliding-average copy of the student model weights, updated at each step according to (\theta_t^{\text{tea}} = \beta\theta_{t-1}^{\text{tea}} + (1-\beta)\theta_t^{\text{stu}}). Since (\beta) is typically close to 1, the teacher parameters are an exponentially weighted accumulation of the student's historical updates, naturally carrying information about "the direction of parameter movement over the past several steps." This is exactly the same kind of signal encoded by first-order momentum (m_t = \mu m_{t-1} + g_t), except that momentum accumulates in gradient space while the teacher displacement accumulates in parameter space. Understanding this makes it clear why reusing the teacher-student displacement is not information-theoretically inferior to a momentum buffer—the difference is that the former exists "for free" as an already-present state, while the latter requires additional memory allocation.

Persistent memory: This refers to GPU memory that remains resident throughout training and is not released with the computation graph, typically including optimizer states (AdamW's (m) and (v), Muon's momentum buffer) and model copies (EMA teacher). Unlike temporary memory such as activations, its size scales with the number of parameters rather than batch size, and therefore directly determines the model size that can fit on a single GPU. Eliminating one full-parameter-count momentum buffer means savings on the order of several GB for models above 7B.

Spectral normalization and Newton-Schulz orthogonalization: The core of Muon is to approximately orthogonalize the momentum matrix so that updates are balanced across singular value directions. The Newton-Schulz iteration approximates ((MM^\top)^{-1/2}) via matrix polynomials, replacing expensive SVD with several matrix multiplications, and is key to Muon's computational efficiency. Spectral normalization scales the matrix by its largest singular value so that subsequent iterations fall within a numerically stable range. Removing the radial component means stripping the component of the displacement matrix that points along the parameter's own direction (i.e., is aligned with the parameter), retaining only directions orthogonal to the parameter, to avoid updates that purely amplify the existing weight norm. Together, these two steps ensure that the borrowed teacher displacement, before being fed into the NS iteration, has its scale and directional distribution aligned as closely as possible with the statistical properties of the original gradient momentum; otherwise, the step size after orthogonalization would be distorted.

Theoretical Guarantees and Convergence Analysis

Under explicit regularity and alignment assumptions, the paper establishes that the expected objective-gap bound contracts linearly toward the residual error level. This means the optimization process has provable convergence behavior at the theoretical level, rather than being a purely empirical improvement.

For optimizer design, this theoretical characterization is particularly important—it shows that reusing the teacher displacement is not a heuristic trick but a formal method with convergence guarantees under specific assumptions.

Expected objective-gap bound: This refers to (\mathbb{E}[f(\theta_t) - f^\star]), the expected gap in objective function value between the current parameters and the optimal solution, a commonly used quantity in optimization theory for characterizing convergence rates. In this paper, "linear contraction to the residual error level" means that this gap shrinks by a fixed ratio each step until it bottoms out at a constant term (the residual error) arising from noise or relaxed assumptions. Such conclusions typically require two types of conditions: smoothness of the objective function (regularity) and the optimization signal not being completely misaligned with the true gradient direction (alignment assumption). The alignment assumption is precisely the risk point of this method: it requires that the inner product of the teacher displacement and the current gradient be non-negative. If the teacher lags too far behind and severely diverges from the student's current direction, the assumption may be violated.

Experimental Results: Accuracy Gains and Memory Savings

The paper validates Muon-RT on two models and two self-distillation methods.

SDFT and Qwen2.5-7B-Instruct

Under the Self-Distillation Fine-Tuning (SDFT) setting with Qwen2.5-7B-Instruct, Muon-RT improves average accuracy by 2.6 percentage points over AdamW while reducing persistent memory by 42%; compared with Muon with momentum, it improves average accuracy by 5.4 percentage points while reducing persistent memory by 27%. Interestingly, the comparison here is between the "momentum-free" version and the "with-momentum" version of Muon, and the improvement is even larger than the comparison with AdamW, suggesting that the signal quality from reusing teacher trajectories may be superior to an independent momentum buffer.

SDPO and Olmo-3-7B-Instruct

Under the Self-Distillation Policy Optimization (SDPO) setting with Olmo-3-7B-Instruct, Muon-RT improves average accuracy by 2.0 percentage points over AdamW. This result demonstrates that the method is applicable across different self-distillation paradigms, rather than being effective only for one particular fine-tuning pipeline.

Implications for Self-Distillation Training Practice

The value of Muon-RT lies in working with the existing structure of the self-distillation framework rather than against it. The EMA teacher already exists, the teacher-student displacement can already be computed, and the optimizer only needs to use this information differently to save substantial persistent memory. For large-scale model training, reducing persistent memory directly affects the trainable model size and batch size configuration.

Of course, the paper's convergence conclusions depend on regularity and alignment assumptions, and whether these assumptions always hold in real training still requires validation across more scenarios. In addition, the current experiments are concentrated on models around 7B, and performance at larger scales remains to be seen. Nevertheless, methodologically, the direction of "extracting optimization signals from teacher trajectories" offers a new entry point for memory-efficient self-distillation training.

Muon-RT paper source screenshot

Key Takeaways

Muon-RT eliminates Muon's independent momentum buffer by reusing the teacher-student displacement, provides a theoretical convergence guarantee of linear contraction, and validates dual benefits of accuracy improvement and persistent memory savings on Qwen2.5-7B-Instruct and Olmo-3-7B-Instruct. For researchers and engineers focused on self-distillation and optimizer memory efficiency, this is a direction worth following.

#Muon-RT#self-distillation#Muon optimizer#momentum buffer#EMA teacher#memory-efficient training#LLM fine-tuning
Share:

Related articles