NuMuon: Nuclear-Norm-Constrained Muon for Compressible LLM Training¶

会议: NeurIPS 2026
arXiv: 2603.03597
领域: 模型压缩
关键词: 大模型优化器, 模型压缩, 低秩近似, 核范数约束, 谱正交化

一句话总结¶

针对 Muon 优化器全秩正交化更新在高压缩率下性能急剧恶化的问题，本文提出在谱范数 LMO 基础上引入核范数预算的 NuMuon 优化器，通过闭式 top-$k$ 奇异向量截断显式约束更新秩，在保持预训练收敛速度的同时将权重低秩可压缩性提升高达 55.8%。

研究背景与动机¶

随着大语言模型（LLM）参数规模向数十亿甚至数千亿迈进，极高的显存占用与推理延迟成为阻碍其实际部署的核心瓶颈。为降低部署成本，以奇异值分解（SVD）为代表的低秩分解压缩方案（如 ASVD、SVD-LLM 与 Dobi-SVD）被广泛应用。这类方案的压缩效果根本上取决于训练后权重矩阵的有效秩与谱能量集中程度。传统 AdamW 优化器具有隐式低秩偏置，但近年来兴起的新型矩阵优化器 Muon 通过对动量矩阵进行全秩正交化（极分解），极大地提升了大模型的预训练收敛效率。然而，Muon 在参数空间所诱导的几何结构与其下游可压缩性之间的关系此前仍未被探明。

对 Muon 训练过程的几何追踪揭示了一个反直觉的现象：尽管 Muon 施加的是严格全秩且各奇异方向均匀赋权的正交更新，其最终训练得到的模型权重在各层仍展现出高度显著的低秩特性（归一化稳定秩显著低于理论上限）。然而，这种自然涌现的低秩性在应对激进压缩时极为脆弱。在 20%–40% 的温和压缩率下模型尚能维持可用精度，但在 60%–80% 的高压缩率下，Muon 训练模型的困惑度急剧发散、下游任务准确率发生灾难性跌落。究其原因，Muon 的更新方向与权重矩阵的主奇异子空间缺乏动态对齐，无法在训练中主动抑制尾部奇异值能量。

这一困境的切入点在于将 Muon 的极分解操作置于无投影线性最小化神谕（LMO）的凸优化框架下审视。核心 idea：通过在 Muon 的谱范数球 LMO 中进一步相交引入核范数约束作为凸低秩代理，NuMuon 将更新方向的选择规约为主奇异方向的封闭解 top-$k$ 截断，并在训练中搭配分数秩调度，以极小训练开销诱导更低稳定秩，实现高压缩率下极限部署吞吐与精度的双赢。

方法详解¶

整体框架¶

NuMuon 建立在动量矩阵的线性最小化神谕（Linear Minimization Oracle, LMO）视角之上。在预训练阶段，对神经网络中任意二维矩阵参数（如 Transformer 的注意力投影与前馈层权重），优化器维护梯度的一阶动量 $M_t$。标准 Muon 在谱范数球上求解 LMO 得到全秩正交矩阵，而 NuMuon 则将容许更新方向限制在谱范数球与核范数球的交集区域中。

整个优化流水线遵循统一的时序演进：首先计算小批量梯度并更新指数移动平均动量；随后，层自适应秩调度器依据当前训练步数计算目标秩截断数 $k_t$；接着利用高效的随机分块 Krylov 迭代法在极低开销下提取动量矩阵的前 $k$ 对主奇异向量；最后构造 top-$k$ 极分解更新方向并更新权重参数。通过将更新方向限制在低秩极值点，参数轨迹在凸组合机制下持续被牵引至紧凑的低维几何子空间。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入小批量梯度与当前权重矩阵"] --> B["动量累积更新<br/>计算一阶指数移动平均矩阵 Mt"]
    B --> C["核范数约束 LMO 规划<br/>谱范数界与核范数预算凸交集"]
    C --> D["Top-k 奇异空间近似提取<br/>随机分块 Krylov 快速正交迭代"]
    D --> E["动态分数秩调度<br/>余弦退火调节截断比例 kt"]
    E --> F["闭式截断极因子参数更新<br/>低秩正交方向更新层权重"]
    F --> G["高压缩兼容的低稳定秩模型"]

关键设计¶

1. 核范数约束 LMO 规划：将更新方向规约为截断单纯形线性规划

标准 Muon 的更新等价于在谱范数球 $\mathcal{W}=\{\Delta W: \|\Delta W\|_2 \le \rho\}$ 上求解线性最小化 $\min_{\Delta W \in \mathcal{W}} \langle M, \Delta W \rangle$，其解析解为全秩极因子 $-\rho UV^\top$。NuMuon 为了引入秩控制，将可行域修正为谱范数与核范数的交集： $$ \mathcal{W}{*} := \bigl{\Delta W \in \mathbb{R}^{d \;\big|\; |\Delta W|}} \times d_{\text{in}}2 \le \rho,\; |\Delta W|* \le \tau \bigr} $$ 由于酉不变性与冯·诺依曼迹不等式，极值解必定与动量矩阵 $M$ 共享左右奇异向量，使得矩阵优化问题严格等价于奇异值上的受限线性规划（LP）：在 $0 \le s_i \le \rho$ 且 $\sum_i s_i \le \tau$ 的截断单纯形上最大化 $\sum_i \sigma_i s_i$。当设置核范数总预算 $\tau = k\rho$（$k$ 为正整数）时，贪心填水法给出了唯一的封闭解：前 $k$ 个奇异方向被分配最大上限 $\rho$，其余奇异值全部置为 0。最终更新方向严格退化为仅含前 $k$ 阶主奇异向量的外积之和： $$ \Delta W^\star = -\rho \sum_{i=1}^k u_i v_i^\top = -\rho U_k V_k^\top $$ 该设计在数学上以最简形式完成了从一阶 rank-1 更新到全秩 Muon 的平滑插值，无需昂贵的投影交替迭代。

2. Top-k 奇异空间近似提取：随机分块 Krylov 子空间投影

在大规模 LLM 训练中，对所有层的宽矩阵动量直接进行精确完全 SVD 分解是计算上无法承受的。为此，NuMuon 采用 Musco & Musco 提出的随机分块 Krylov 迭代法（Randomized Block Krylov Method）高效抽取前 $k$ 维子空间。算法从微小的高斯随机分块矩阵 $B_0 \in \mathbb{R}^{d_{\text{in}} \times b}$（其中分块大小 $b = k + 8$，提供少量过采样缓冲）出发，仅通过交替执行 $L=2$ 轮小规模矩阵连乘与 QR 分解： $$ T_i \leftarrow A B_{i-1}, \quad B_i \leftarrow \operatorname{qr}(A^\top T_i) $$ 即可跨越谱间隙快速收敛到主奇异子空间。随后将动量矩阵投影至构建的正交基 $Q$ 上计算极小规模的压缩 SVD，从而以近乎线性复杂度的矩阵向量乘法精确重构 $U_k$ 和 $V_k$。这一机制使得 NuMuon 相比原版 Muon 的单步时间开销仅增加约 18%，显存占用维持在相近水平。

3. 动态分数秩调度：匹配预训练动力学演进的谱退火策略

经验观测表明，Transformer 在训练初始阶段的各层权重自然具有较高的探索性稳定秩，过早强制压低更新秩会严重破坏网络表征的学习能力与初始收敛。为克服固定截断秩导致的欠拟合，NuMuon 将截断秩参数化为与矩阵形状无关的连续时间调度函数 $r(t) \in (0, 1]$，每层更新秩动态决定为 $k_t = \lceil r(t) \cdot \min(d_{\text{in}}, d_{\text{out}}) \rceil$。在预训练初期保留 10% 步数的全秩热身（$r_{\text{start}} = 1.0$），随后依据余弦函数平滑退火至设定下限（如 $r_{\text{end}} = 0.25$）： $$ r(t) = r_{\text{end}} + \frac{r_{\text{start}} - r_{\text{end}}}{2} \left(1 + \cos\left(\pi \frac{t - T_h}{T - T_h}\right)\right) $$ 动态退火策略确保了模型前期具有足够的谱探索广度，而在中后期迫使权重参数与更新方向的格拉斯曼距离（Grassmann Distance）显著收缩，将模型表征能量高度浓缩在最核心的低秩子空间中。

损失函数 / 训练策略¶

NuMuon 继承了现代化矩阵优化器的解耦训练管线：对于 Transformer 核心隐层中的二维线性变换矩阵（自注意力投影与 MLP 投影），统一采用 NuMuon 进行谱更新，学习率设为 $1\times 10^{-2}$，搭配 0.1 的权重衰减与 WSD（Warmup-Stable-Decay）学习率调度；对于嵌入层（Embedding）与输出预测头（LM Head），采用常规 AdamW 进行参数更新；对于一维标量和 LayerNorm/RMSNorm 参数，则采用 Lion 优化器进行更新。这种异构分治的优化策略与梯度裁剪阈值（1.0）协同工作，兼顾了谱优化的几何收敛性与非矩阵参数的稳定性。

实验关键数据¶

主实验¶

论文基于 FineWeb-EDU 数据集预训练了 Qwen3-0.6B、Olmo2-1.4B 与 Llama3-1.8B 三种架构模型，并采用 ASVD、SVD-LLM 与 Dobi-SVD 三种前沿低秩压缩方法在 20%–80% 压缩率下进行全面测试。在基准未压缩模型上，NuMuon 验证集困惑度与 Muon 基本相当，且明显优于传统 AdamW。下表展示了 Llama3-1.8B 模型在 40% 压缩率下的下游任务综合评测（包含 ARC-C、ARC-E、HellaSwag、LAMBADA、OpenbookQA、PIQA、Winogrande 七项基准平均准确率）与 WikiText2 验证集困惑度对比：

压缩方案	优化器	WikiText2 验证困惑度 ↓	下游任务平均准确率 (%) ↑	相比 Muon 相对变化
ASVD (40% 压缩)	AdamW	718.71	32.40	—
	Muon	1443.55	32.40	基线
	NuMuon	20.56	50.48	困惑度 -98.6% / 准确率 +55.8%
SVD-LLM (40% 压缩)	AdamW	27.94	42.05	—
	Muon	26.39	39.50	基线
	NuMuon	19.18	48.91	困惑度 -27.3% / 准确率 +23.8%
Dobi-SVD (40% 压缩)	AdamW	28.83	43.79	—
	Muon	20.44	45.30	基线
	NuMuon	18.63	49.38	困惑度 -8.9% / 准确率 +9.0%

消融实验¶

在 Qwen3-0.6B 模型上，采用强度最高的可微低秩压缩方案 Dobi-SVD，对 NuMuon 在极限 80% 压缩率下的秩调度策略及固定秩预算开展消融对比：

调度类型与最终秩比例	单步耗时 (s)	未压缩困惑度 (0%)	80% 压缩困惑度	0% 下游准确率 (%)	80% 压缩下游准确率 (%)
分段调度 Piecewise (0.25)	9.94	20.01	33.39	45.87	37.92
余弦调度 Cosine (0.25)	10.49	20.20	36.10	45.69	35.89
固定秩 Fixed (0.05)	8.96	29.13	31.93	39.64	38.74
固定秩 Fixed (0.25)	9.81	21.85	44.69	44.30	34.36
固定秩 Fixed (0.50)	11.33	20.27	46.66	45.17	34.38
固定秩 Fixed (0.80)	11.49	19.31	43.70	46.33	34.42
全秩基线 Muon	8.86	19.20	50.31	46.81	33.32

关键发现¶

高压缩率下的断崖式免疫：随着压缩率从 40% 激增至 60% 甚至 80%，Muon 与 AdamW 模型的困惑度均迅速恶化至破千甚至数万，而 NuMuon 在各模型尺寸下均表现出平滑退化，80% 激进压缩下仍能保留接近 35%–38% 的下游零样本准确率。
子空间对齐的机理验证：格拉斯曼距离测量显示，Muon 在整个训练过程中与权重的主子空间夹角始终维持在高位且不收敛；而 NuMuon 的更新方向逐步向权重主奇异向量靠拢，格拉斯曼距离显著缩小，从机理上解释了其权重的奇异值分布为何更加陡峭集中。
退火调度的必要性：直接使用低固定秩（如 0.05）虽然压缩鲁棒性极强，但基座模型本身表征受损，未压缩困惑度高达 29.13；余弦或分段调度在前期给予全秩探索空间，不仅未压缩表现紧咬全秩 Muon（20.20 vs 19.20），更兼顾了极致的可压缩性。

亮点与洞察¶

将经验启发上升至严谨凸规划：本文没有停留在“直接在 SVD 分解后截断奇异向量”的工程直觉，而是从核范数正则化神谕的角度推导出了 top-$k$ 截断是谱范数与核范数双约束 LMO 的闭式最优解，为矩阵优化器的低秩修剪提供了扎实的理论支点。
非凸非光滑条件下的收敛性泛化：证明了在梯度外部能量（Tail Energy）满足次要奇异值衰减的弱假设下，NuMuon 能够继承并泛化 Muon 的非凸 Ky Fan $k$-范数平稳点收敛界，且实测验证了神经网络梯度在优化过程中的残差谱能量极度贴近零。
推动高效大模型全生命周期协同：以往“模型预训练”与“部署后剪枝压缩”彼此脱节，NuMuon 证明优化器在预训练阶段对几何流形的约束可以直接转化为下游部署推理的吞吐红利（在相同困惑度下将批推理吞吐提升数倍）。

局限与展望¶

计算开销的轻微妥协：尽管采用了高效的分块 Krylov 迭代，NuMuon 的单步迭代时间相较于高度硬件优化、无 SVD 调用的 Newton-Schulz 多项式迭代仍慢约 15%–20%，在大规模超算集群中的工程集成度仍有优化空间。
单矩阵独立截断的超参粒度：当前所有隐层矩阵共享统一的相对秩比例退火系数，未考虑不同层（如深层 MLP 与浅层 Attention）对奇异值稀疏度敏感性的内在异质性，未来可结合自动层敏感度分析实现自适应核范数预算分配。
向分布式低带宽通信的延展：NuMuon 产出的外积更新形式 $U_k V_k^\top$ 天生具备因子化低秩结构，这一特性极适于与去中心化训练或带宽受限的梯度通信算法（如 DiLoCo 或 DION）相结合，具有广阔的分布式系统探索潜力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次揭示 Muon 的隐式低秩几何偏置，并以规范的凸 LMO 框架提出闭式核范数截断优化器。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 0.6B 到 1.8B 三大模型家族，横跨 ASVD、SVD-LLM 与 Dobi-SVD 三大压缩管线及多尺度压缩率，论据坚实。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨清晰，图表完备，消融实验与理论假设验证层层递进。
价值: ⭐⭐⭐⭐⭐ 直击 LLM 预训练成本与端侧低成本部署的衔接鸿沟，对大模型轻量化落地具备极高实践指导价值。