HTMuon: Improving Muon via Heavy-Tailed Spectral Correction

Each language version is independently generated for its own context, not a direct translation.

这篇论文介绍了一种名为 HTMuon 的新型“训练助手”，专门用来帮助大型人工智能模型（LLM）学得更好、更快。

为了让你轻松理解，我们可以把训练 AI 模型想象成教一个超级聪明的学生（模型）通过做大量的练习题（数据）来掌握知识。在这个过程中，我们需要一个“教练”（优化器）来告诉学生：“这道题你哪里做错了？下一步该怎么改？”

1. 现有的教练（Muon）有什么问题？

以前，大家最常用的教练叫 Muon。它很厉害，因为它不仅看学生每一道题的对错，还能看出题目之间的联系（比如，这道题做错了，可能意味着另一类题也做错了）。

但是，Muon 有一个致命的习惯：
它有一个“强迫症”，认为所有方向的错误都一样重要。

比喻：想象学生在跑步。有些方向是平坦的大道（信号强，学得快），有些方向是布满荆棘的泥潭（全是噪音，学得很慢）。Muon 这个教练会强行把学生在“大道”和“泥潭”上的步长调整得一模一样。
后果：学生在“泥潭”里（噪音方向）浪费了大量精力，甚至越跑越偏；而在“大道”上又没跑够。这就导致模型虽然学了不少，但最后的效果（泛化能力）不够好，就像学生虽然背了很多书，但遇到新题还是不会做。

2. 什么是“重尾”（Heavy-Tailed）？

论文里提到了一个听起来很学术的词：重尾（Heavy-Tailed）。

比喻：想象一下学生的成绩分布。
- 轻尾（Muon 的现状）：大部分学生成绩都差不多，都在平均分附近，很难出现特别拔尖或特别差的学生。这就像 Muon 把步长都拉平了，导致模型变得“平庸”。
- 重尾（HTMuon 的目标）：允许出现一些特别突出的“天才”方向（大步长），同时也允许一些特别弱的“噪音”方向（小步长，甚至忽略）。这种分布更符合真实世界的规律——真正的学习往往是由少数几个关键突破点驱动的。

3. HTMuon 是怎么做的？

HTMuon 就是 Muon 的“升级版”，它给教练加了一个智能调节器。

核心操作：它不再把步长强行拉平，而是根据方向的重要性进行**“幂次调整”**（把步长开一个小于 1 的方，比如 0.125 次方）。
通俗解释：
- 对于重要的方向（大道）：HTMuon 会放大步长，让学生跑得更快、更坚决。
- 对于不重要的方向（泥潭/噪音）：HTMuon 会缩小步长，甚至让学生“轻轻点一下”就过，避免在噪音里浪费时间。
结果：这种“重尾”的更新方式，让模型能更敏锐地捕捉到数据中的核心规律，而不是被噪音带偏。

4. 效果怎么样？

论文在多个任务上做了测试，效果非常显著：

语言模型（LLM）：
- 在训练像 LLaMA 这样的大模型时，HTMuon 让模型在同样的训练量下，“困惑度”（Perplexity，衡量模型懂不懂人话的指标，越低越好）降低了近 1 个点。
- 比喻：这就像同样的学习时间，用 HTMuon 教出来的学生，阅读理解能力直接从“大学生”提升到了“研究生”水平。
图像识别：
- 在识别图片（如 CIFAR、ImageNet）的任务中，HTMuon 也让模型的准确率有了明显的提升。
兼容性：
- 它不仅能独立使用，还能像“插件”一样，直接加在 Muon 的其他变种上，让它们变得更强。

5. 理论上的“为什么”

论文还从数学角度解释了为什么这招管用：

重尾自我正则化理论 (HT-SR)：研究发现，那些真正学得好的神经网络，它们的内部结构（权重分布）天然就是“重尾”的。也就是说，好的模型天生就喜欢“重尾”分布。
HTMuon 只是顺应了这种天性，而不是强行改变它。它让模型在训练过程中自然地形成了这种健康的“重尾”结构，从而变得更聪明、更稳定。

总结

HTMuon 就像是一个懂得因材施教的超级教练。
它不再搞“一刀切”，而是把力气花在刀刃上：

在关键方向上，它鼓励学生大胆冲刺；
在噪音方向上，它让学生轻描淡写。

这种策略不仅让 AI 模型学得更快、更准，还让它们在处理新任务时表现得更加出色。对于想要训练大模型的研究者和公司来说，这是一个既简单又强大的新工具。

Each language version is independently generated for its own context, not a direct translation.

1. 研究背景与问题 (Problem)

背景：

优化器的重要性： 优化器在大语言模型（LLM）训练中起着核心作用。传统的基于向量的优化器（如 Adam/AdamW）虽然流行，但往往忽略了参数之间的相互依赖关系（interdependencies）。
Muon 优化器： 最近提出的 Muon 是一种基于矩阵的优化器，它通过对动量矩阵进行正交化（Orthogonalization）更新，能够捕捉参数间的几何依赖关系，并在 LLM 训练中表现出优异性能。Muon 的更新规则相当于在 Schatten- $\infty$ 范数约束下的最速下降法。

核心问题：
尽管 Muon 表现良好，但作者发现其存在两个主要局限性，限制了模型的最终性能：

过度强调噪声方向： Muon 的正交化步骤将动量矩阵的所有奇异值强制设为 1。这意味着它对所有奇异向量方向赋予相同的权重。然而，小奇异值方向通常由噪声主导，均匀加权会导致训练在后期对噪声过于敏感，从而限制模型容量。
抑制重尾谱（Heavy-Tailed Spectra）： Muon 的更新规则导致动量更新矩阵和最终学习到的权重矩阵的谱分布呈现“轻尾”（Light-tailed）特性。根据重尾自正则化（HT-SR）理论，训练良好的神经网络通常具有重尾的权重谱分布，且重尾程度与模型质量强相关。Muon 的轻尾特性抑制了这种有益的重尾结构的形成，从而限制了模型的泛化能力。

2. 方法论 (Methodology)

为了解决上述问题，作者提出了 HTMuon（Heavy-Tailed Muon），一种基于矩阵的优化器。

核心思想：
HTMuon 旨在保留 Muon 捕捉参数相互依赖的能力，同时通过引入重尾谱校正，使更新方向和权重谱分布更加“重尾”。

算法设计：
HTMuon 的核心改进在于对动量矩阵 $M_t$ 的奇异值进行幂次变换：

Muon 更新： $O_t = U_t V_t^\top$ （即所有奇异值设为 1）。
HTMuon 更新： $O_t = U_t \Sigma_t^p V_t^\top$ $O_{t} = U_{t} Σ_{t}^{p} V_{t}^{⊤}$ ，其中 $p \in (0, 1)$ $p \in (0, 1)$ 。
- 当 $p=1$ 时，退化为 SGDM（向量优化器）。
- 当 $p=0$ 时，退化为 Muon。
- 作者选择 $p \in (0, 1)$ （默认 $p=0.125$ ），这使得较小的奇异值被相对放大（相对于 1 而言），从而产生重尾的更新分布。

理论依据：

HT-SR 理论： 通过拟合特征值分布的幂律指数 $\alpha$ 来量化重尾程度。较小的 $\alpha$ 意味着更重的尾部，通常对应更好的模型质量。
理论等价性： 作者证明 HTMuon 等价于在 Schatten- $q$ 范数约束下的最速下降法（Steepest Descent），其中 $q$ 与 $p$ 相关。这推广了 Muon 作为 Schatten- $\infty$ 范数约束下最速下降法的结论。
收敛性分析： 在平滑非凸设置下，HTMuon 的样本复杂度上界与 Muon 和 SGDM 一致，均为 $O(\epsilon^{-4})$ 。

加速实现：
为了降低 HTMuon 中 SVD 带来的计算开销，作者提出了两种加速方案：

间隔更新（Interval-based）： 仅在每 $k$ 步执行一次 HTMuon 更新，其余步骤使用 Muon。
HTMuon NS： 利用 Newton-Schulz 迭代近似计算矩阵的分数次幂（Fractional Power），避免显式 SVD，显著降低计算成本。

3. 主要贡献 (Key Contributions)

理论洞察： 指出 Muon 的正交化更新规则倾向于噪声主导的方向，并抑制了权重矩阵中重尾谱分布的出现，这违背了 HT-SR 理论关于高质量模型的特征。
提出 HTMuon： 设计了一种简单但有效的优化器，通过奇异值幂次变换（ $p \in (0,1)$ ）引入重尾更新，同时保留矩阵优化器对参数耦合的建模能力。
广泛的实验验证：
- 在 LLM 预训练（LLaMA 系列、GPT-2）和图像分类（ResNet, ViT）任务上，HTMuon consistently 优于 Muon、AdamW 及其他 SOTA 优化器（如 Cautious, Mars, SOAP, COSMOS 等）。
- 具体数据： 在 C4 数据集上训练 LLaMA-60M，HTMuon 相比 Muon 降低了 0.92 的困惑度（Perplexity）；在 LLaMA-135M 上降低了 0.98。
- 兼容性： HTMuon 可作为插件与现有的 Muon 变体（如 NorMuon, AdaMuon）结合，进一步提升性能。
理论分析： 建立了 HTMuon 与 Schatten- $q$ 范数约束下最速下降法的理论联系，并提供了非凸环境下的收敛性证明。
高效实现： 提出了两种加速版本（间隔更新和 NS 近似），在保持性能优势的同时显著降低了运行时间开销。

4. 实验结果 (Results)

LLM 预训练 (C4 & OpenWebText)：
- HTMuon 在所有规模的 LLaMA 模型（60M, 135M, 350M, 1B）上均取得了最低的困惑度。
- 相比 Muon，LLaMA-60M 的 PPL 降低了 0.92，LLaMA-135M 降低了 0.98。
- 相比 AdamW，性能提升更为显著（例如 LLaMA-60M 降低 4.33 PPL）。
- 在下游任务（7 个常识推理基准）的零样本评估中，HTMuon 预训练的 LLaMA-1B 模型平均得分最高，比次优的 Muon 高出 1.05 分。
图像分类 (CIFAR & ImageNet)：
- 在 CIFAR-10/100 上的 ResNet 和 ImageNet-1K 上的 ViT-tiny 任务中，HTMuon 的准确率均优于 Muon、NorMuon 和 SGDM。
- 例如，在 CIFAR-100 上，HTMuon 相比 Muon 提升了 0.31% 的准确率。
谱分析验证：
- 实验测量了训练后权重矩阵的经验谱密度（ESD）的幂律指数 $\alpha$ 。
- 结果显示，HTMuon 产生的 $\alpha$ 值显著低于 Muon（即谱分布更重尾），这与 HT-SR 理论预测的“更重尾对应更好性能”一致。
- HTMuon 还产生了更小的谱范数和 Frobenius 范数，这通常与更好的泛化性相关。
效率分析：
- 虽然原始 HTMuon 的每一步计算开销略高于 Muon，但通过“间隔更新”或"HTMuon NS"策略，可以在仅增加少量时间开销（如 6%-11%）的情况下，获得显著的性能提升，甚至在总训练时间上优于 Muon（因为收敛更快）。

5. 意义与影响 (Significance)

重新审视优化器设计： 该工作将优化器设计与随机矩阵理论（Random Matrix Theory）及重尾自正则化（HT-SR）理论紧密结合，为设计下一代优化器提供了新的理论视角：即优化器应促进权重谱的重尾化，而非强制均匀化。
提升 LLM 训练上限： 证明了在 Muon 这一强基线之上，通过简单的谱校正（Spectral Correction）仍能带来显著的性能提升，为大模型训练提供了新的优化策略。
理论与实践的统一： 不仅提供了实证结果，还从 Schatten 范数约束的角度给出了理论解释，并证明了其收敛性，增强了方法的可靠性。
实际可用性： 提出的加速方案解决了矩阵优化器计算成本高的痛点，使得 HTMuon 能够应用于大规模模型（如 1B 参数以上）的训练，具有极高的实用价值。

总结： HTMuon 通过引入重尾谱校正机制，成功克服了 Muon 优化器在噪声处理和谱分布上的局限性，在保持矩阵优化器优势的同时，显著提升了 LLM 和图像分类模型的性能，是优化器领域的一项重要进展。

HTMuon: Improving Muon via Heavy-Tailed Spectral Correction

1. 现有的教练（Muon）有什么问题？

2. 什么是“重尾”（Heavy-Tailed）？

3. HTMuon 是怎么做的？

4. 效果怎么样？

5. 理论上的“为什么”

总结

1. 研究背景与问题 (Problem)

2. 方法论 (Methodology)

3. 主要贡献 (Key Contributions)

4. 实验结果 (Results)

5. 意义与影响 (Significance)

类似论文

Comparison of Outlier Detection Algorithms on String Data

Structure-Aware Epistemic Uncertainty Quantification for Neural Operator PDE Surrogates

Interventional Time Series Priors for Causal Foundation Models

Fingerprinting Concepts in Data Streams with Supervised and Unsupervised Meta-Information

Graph Tokenization for Bridging Graphs and Transformers