[Nature Machine Intelligence 潜力级] NuMuon：通过核范数约束，让 LLM 赢在压缩的起跑线上

TL;DR

随着 LLM 规模的无节制扩张，如何将其“塞进”边缘设备或降低云端成本成了工业界的头等大事。本文提出的 NuMuon 优化器通过在训练时显式控制权重的奇异值分布（核范数约束），让模型在出生时就具备了“易压缩”的基因。实验表明，它在 80% 的极高压缩率下，性能依然大幅领先于 AdamW 和 Muon 训练的模型。

背景定位：从隐式偏置到显式控制

在 Deep learning 的学术语境中，大家都知道模型权重通常具有隐式的低秩性（Low-rank bias）。即使是最近大火的 Muon 优化器——它通过矩阵正交化（Full-rank updates）来加速训练——在实际训练中也会产生低秩的权重分布。

然而，这种“天然”的低秩结构是脆弱的。当你尝试用 SVD-LLM 等工具强行砍掉 60% 以上的参数时，模型的逻辑能力往往会发生灾难性的退化。NuMuon 的核心动机很简单：既然我们最终都要进行低秩压缩，为什么不直接在训练更新方向上加上秩的约束？

痛点深挖：Muon 的野性更新

Muon 优化器的精髓在于它不看梯度的大小，只看梯度的方向（通过 Newton-Schulz 正交化）。这种“光谱平权”虽然加速了收敛，但其更新步是满秩的（Full-rank）。这导致模型权重包含了很多细碎、冗余的奇异值。对于后压缩而言，这些细碎的 singular values 就像是杂讯，砍掉它们会损伤模型的表达力。

核心方法论：NuMuon 的数学直觉

作者巧妙地将 Muon 的更新过程转化为一个 线性最小化算子 (LMO) 问题。

1. 约束的交集

NuMuon 不再仅仅约束更新步的谱范数（Spectral Norm），而是同时引入了**核范数（Nuclear-norm）**约束：

• 谱范数约束：限制最大的更新强度。
• 核范数约束：作为秩的凸代理，限制奇异值的总和，诱导稀疏性。

2. Top-k 奇异向量更新

数学推导证明（详见论文 Proposition 3.1 & 3.2），这个受约束的线性规划问题存在一个优雅的闭式解：NuMuon 的更新步本质上是 Momentum 矩阵的前 k 个奇异向量的组合。 这意味着 NuMuon 每一行、每一列的更新都在最关键的特征方向上精准发力，而舍弃了那些长尾的、对模型贡献微弱的方向。

(图a展示了 FFN Gate 投影层中更新方向的演变)

3. 动态秩调度 (Rank Scheduler)

为了防止训练早期过度约束导致收敛减慢，作者设计了 Rank Scheduler。模型在初期可以使用较高的秩进行探索，随着训练进行，秩约束逐渐收紧（例如从 1.0 降到 0.25），从而平稳地过渡到压缩友好型结构。

实验与结果：全线碾压

作者在 Qwen-0.6B 和 Llama-1.8B 等多个模型上进行了验证。

压缩质量的质变

在使用 SVD-LLM 或 Dobi-SVD 进行后压缩时，NuMuon 的优势被放大到了极致。

• 极高压缩率 (80%)：相比传统的 Muon，NuMuon 训练的模型在 WikiText2 上的困惑度下降了 90% 以上（数值越低越好），下游任务平均分数提升了显著的一截。
• 推理效率：在相同的精度门槛下，NuMuon 允许模型携带更少的参数运行，直接转化为更高的推理吞吐量。

(图 2 显示：在相同的 Perplexity 下，NuMuon 训练的模型实现了最快的生成速率)

深度洞察：为什么 NuMuon 有效？

作者通过 Grassmann Distance（格拉斯曼距离） 进行了解析。由于 NuMuon 在训练时显式选择了 Top-k 方向，其更新步与模型权重的 Top 谱子空间的对齐（Subspace Alignment）远好于 Muon。这种“持续对齐”保证了权重的关键信息始终集中在极少数维度上，从而实现了真正的“结构性冗余压缩”。

总结与局限

NuMuon 为 LLM 训练提供了一种新的思路：压缩不应该是后处理的补丁，而应该是训练时的目标。

• 贡献：提供了兼顾收敛效率（Muon 级）和极致压缩性（NuMuon 独有）的优化方案。
• 局限：计算 Top-k 奇异向量（使用 Randomized Block Krylov 方法）虽然比全 SVD 快，但相比 AdamW 仍有约 15%-20% 的每步计算开销。这在超大规模预训练中是否划算，仍需根据最终部署次数进行成本核算。

未来的研究可能会探索这种受限更新在多模态模型或更长上下文模型中的广义性。毫无疑问，NuMuon 在通往“轻量化 LLM”的道路上迈出了扎实的一步。