Muon-RT：复用教师轨迹的无动量 Muon 优化器用于大模型自蒸馏 MUON-RT Reusing the Teacher's Trajectory for Momentum-Free Muon for LLM Self-Distillation

每个优化步的完整流程（Algorithm 1）：输入学生 $\theta_S$、EMA 教师 $\theta_T$、损失 $\mathcal{L}$、混合权重 $\alpha$、EMA 系数 $\tau$、学习率 $\eta$、权重衰减 $\lambda$。第一步，对学生做前向/反向传播，得到每个隐层矩阵的梯度 $G_\ell=\nabla_{W_{S,\ell}}\mathcal{L}$。第二步，梯度支路：$O_\ell=NS_K(G_\ell)$，即 5 次五次 Newton–Schulz 迭代（系数 $\kappa=(3.4445,-4.7750,2.0315)$），输入先按 Frobenius 范数归一化。第三步，教师支路：计算位移 $D_\ell=W_{T,\ell}-W_{S,\ell}$；径向去除 $H_\ell=D_\ell-\frac{\langle D_\ell, W_{S,\ell}\rangle_F}{\|W_{S,\ell}\|_F^2+\epsilon}W_{S,\ell}$；用 8 次幂迭代（从上一步保存的右奇异向量热启动，初始随机向量会低估 $\sigma_{max}$ 约 8%，热启动 3 步内消除误差）估计 $\hat\sigma_{max}(H_\ell)$，得 $P_\ell=H_\ell/\max(\hat\sigma_{max},\epsilon)$。第四步，混合与更新：$U_\ell=(1-\alpha)O_\ell+\alpha P_\ell$（$\alpha=0.1$），就地更新 $W_{S,\ell}\leftarrow(1-\eta\lambda)W_{S,\ell}-\eta c_{A,B}U_\ell$，随后释放该矩阵的临时缓冲再处理下一个矩阵。第五步，非隐层参数（嵌入、输出头、norm 权重、bias）照常用 AdamW。第六步，EMA 更新教师 $\theta_T\leftarrow\tau\theta_T+(1-\tau)\theta_S$（SDFT 用 $\tau=0.99$，SDPO 用 $\tau=0.95$）。分布式实现基于 FSDP2：学生与教师用相同分片布局，位移本地计算，式 (9) 的内积跨 rank all-reduce，$G$ 和 $H$ 临时 gather 成全矩阵做 NS 和幂迭代，各 rank 只更新自己的分片；已验证分片与非分片结果一致、$\alpha=0$ 时与 Muon µ=0 逐位相同。

理论新颖性：这是首个系统刻画“EMA 教师位移信息内容”的工作——证明它是折扣更新历史（Lemma 1）、在温和条件下是降噪方向估计（Lemma 2），并用 Proposition 1 排除了一种过度解读（教师分支无法改进单步瞬时对齐），从而把增益机制严格归因于时间信息。这与 Defazio 等人“动量可写成迭代平均”的视角互补但不同：后者把平均当作专用缓冲，本文刻画的是两份活跃模型副本之间位移的语义。工程新颖性：三个具体设计都有实证支撑——(1) 混合必须在 NS 之后，文中测得 NS 前混合即使 $\alpha=0.02$ 造成的更新改变也超过其自身范数（NS 会把教师的弱奇异方向抬升到单位值），而 NS 后混合随 $\alpha$ 线性增长、影响可控；(2) 径向去除抑制权重缩放噪声，消融显示保留径向分量掉 1.3–4.5 点；(3) 谱归一化（而非 Frobenius）使 $\|\alpha P\|_2\le\alpha$ 有界，Frobenius 版本掉 0.9–3.3 点。实证闭环也很扎实：预测位移与实测相关 0.97、两分支余弦绝对值小于 0.015（教师几乎正交于当前梯度、补充梯度未覆盖的方向）、$\alpha=0$ 时实现与 Muon µ=0 逐位一致，理论声称都有对应的测量。

(1) 32 位教师依赖：作者证明 bfloat16 教师会使 79% 的 EMA 条目永不更新（每步变化约为权重的 $10^{-5}$，而 bf16 只能表示 $4\times10^{-3}$ 以上的相对变化），因此位移信息会被舍入彻底破坏。这隐含每参数 2 字节的成本（7B 上每卡 7.6 GB），削弱了内存叙事；改进方向是研究随机舍入/误差补偿式低精度 EMA，让位移在 16 位下可读。(2) 与稳定化技术的冲突：QK-Clip 只作用于学生，会在 $W_T-W_S$ 中引入非学习偏移，实测首次裁剪使平均位移一步上升 12%；即便对教师同步裁剪（恢复到 67.3/70.1）仍低于无裁剪版本，因为裁剪还扰动了教师的注意力模式即蒸馏目标本身。改进方向是设计“位移感知”的稳定化：例如在径向去除后额外投影掉被裁剪头对应的子空间（文中显示 97% 的 $\|H\|_F^2$ 集中在被裁剪头上，说明这样的投影是可行的）。(3) 径向去除的副作用：它同时滤掉了权重衰减带来的合法缩放信号，实验中 SDFT 设 $\lambda=0$、SDPO 仅 0.01，强权重衰减的预训练场景未验证；改进方向是先把已知 $\eta\lambda$ 的缩放分量解析地扣除，再对残差做径向去除。(4) 幂迭代精度：8 次迭代低估 $\sigma_{max}$ 约 8%，依赖跨步热启动在 3 步内修正；训练早期或梯度突发突变时归一化可能暂时失准，$\|\alpha P\|_2\le\alpha$ 的界随之松动；改进方向是给 $P$ 加一个基于前几步 $\sigma_{max}$ 轨迹的保守上界。(5) 适用面：预训练从零开始时教师与学生都从同一点出发、位移从 0 缓慢生长，早期几乎没有历史信息可复用，方法的收益边界（微调/后训练 vs 预训练）需要在论文之外小心外推。

作者提出的方向：扩展到更大模型（7B 之外）和更多领域/任务；在 LLM 上实际估计收敛理论的常数（$\mu_{PL}$、$L$、对齐系数 $a_R,b_R$），把经验观察（梯度范数更低、KL 更低）与理论衔接得更紧。基于成果可延伸的方向：(1) 推广“复用模型副本作优化器状态”的思路到其他天然维护 EMA/权重平均的场合——mean-teacher 半监督、DINO/BYOL 自监督、AdEMAMix 的慢动量、Lookahead 外层权重、WARM 平均奖励模型，这些场景里副本间的位移同样可能承载动量信息；(2) 让教师分支与 logit 稳定化兼容（子空间投影或对教师-学生做对称裁剪的新方案），从而能进入 Muon 在大规模预训练中的主战场（Kimi K2 式训练）；(3) $\alpha$ 的自适应调度：理论里教师漂移 $\delta_s$ 和位移范数都是可在线测量的，可以按 $\|H\|_F/\|W\|_F$ 或目标漂移自动调节混合权重，替代固定 0.1；(4) 部分替代 AdamW 的二阶矩：作者明确说位移不能恢复 AdamW 的二阶统计，但能否用它构造廉价的逐坐标尺度近似（例如利用位移各坐标的分布）值得探索，这可能把内存优势扩展到嵌入/输出头等仍用 AdamW 的参数上；(5) 在更多 on-policy 目标（GRPO、DPO、on-policy distillation）中验证，并研究位移信息在持续学习序列（多个任务依次蒸馏）中的积累与衰减规律。

复现条件在论文中给出得相当完整。代码基础：官方 SDFT 发布代码（TRL + colocated vLLM 生成），只需替换优化器；分片与非分片、$\alpha=0$ 与 Muon µ=0 均已做逐位一致性验证。超参数：表 4（SDFT）与表 5（SDPO）完整列出——Qwen2.5-7B-Instruct / Olmo-3-7B-Instruct、2×H100 80GB + FSDP2、每步 32 prompt、SDFT 252/166 步、学习率 $\{3,5\}\times10^{-5}$ cosine、forward KL / Jensen–Shannon、教师 32 位存储 + bf16 前向、EMA $1-\tau=0.01/0.05$、NS 5 步及系数 $(3.4445,-4.7750,2.0315)$、$\alpha=0.1$、8 次热启动幂迭代、每配置 3 个随机种子。算力门槛：单次 SDFT 运行 1.1–1.8 小时（2×H100，生成主导），整套主表加消融估计需数百 GPU 小时，对学术实验室可负担；理论部分（图 4b/c、常数估计）只需在单个权重矩阵上做合成实验，普通单卡即可复现。数据全部公开（ToolAlpaca、SciKnowEval），AdamW 基线数字引自 SDFT/SDPO 原文并在作者的精度设置下复验过（复现差异在随机波动内）。未明确说明 Muon-RT 自身代码是否开源，这是复现的主要不确定性；不过方法只有约 30 行核心逻辑（Algorithm 1），配合论文的实现细节（A.4 节，含 FSDP2 的 gather/all-reduce 安排和 $5\times10^{-7}$ 的单矩阵数值验证）自行实现难度中等。