Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory¶
会议: NeurIPS 2026
arXiv: 2603.26554
OpenReview: https://openreview.net/forum?id=oNlQTzdbf5
论文: https://arxiv.org/abs/2603.26554
领域: 其他
关键词: 谱优化器, Muon, 联想记忆, 缩放定律, 临界批大小
一句话总结¶
本文针对非正交高斯嵌入与幂律频次分布下的线性联想记忆任务，严格给出了谱优化器 Muon 的锐化容量缩放定律，证明单步 Muon 即可达到与二阶牛顿法一致的 \(\widetilde{\Theta}(d^{1+\frac{1}{2\alpha}})\) 记忆容量并饱和于更大的临界批大小。
研究背景与动机¶
在大语言模型（LLM）的千亿参数预训练中，以 Adam 为代表的自适应一阶优化器长期占据统治地位。然而，近年来直接作用于参数矩阵几何结构的谱优化器（如 Muon）在大规模预训练基准测试中展现出超越传统一阶方法的性能，尤其在大批大小训练中优势显著。Muon 通过近似极分解（polar decomposition）将参数矩阵梯度沿谱正交化方向更新，在动力学层面对应于谱范数意义下的最速下降。尽管经验效果拔群，理论界对于 Muon 到底在何种数据分布与结构特性下能够产生超越一阶梯度的本质优势，依然缺乏严谨且细致的动力学解析。
现有针对谱优化器的联想记忆（associative memory）理论分析通常依赖于输入与输出嵌入相互严格正交的简化假设。这种正交性假设固然简化了极映射的矩阵分析，却也强制约束了模型容量上限——存储的记忆事实总数 \(N\) 绝不能超过嵌入维度 \(d\)。但在真实的 Transformer 知识记忆中，参数维度通常允许通过“特征叠加（superposition）”存储远超维度数量的信息，即在 \(d \times d\) 矩阵中信息论最优容量可达 \(\widetilde{\Theta}(d^2)\)。同时，自然语言中的词频与事实激活普遍遵循齐普夫定律（Zipf's law）所主导的重尾幂律分布，使得真实训练梯度呈现出极强的非各向同性。
如何在破除正交假设的高维高斯叠加态与幂律长尾分布下，定量刻画 Muon 的记忆容量扩展机制，构成了连接现代谱优化器实践与理论的关键鸿沟。本文的核心 idea 是：将基于重尾幂律分布的线性联想记忆建模为非正交特征叠加态下的分类动力学，揭示 Muon 的谱前条件化本质上是通过奇异谱分辨率对梯度的长尾微弱信号实施差异化放大，从而证明仅利用一阶信息的单步 Muon 即可在容量与损失衰减上严格逼近二阶牛顿法。
方法详解¶
整体框架¶
本文考虑输入与输出词表大小均为 \(N\) 的线性联想记忆模型。每个原子事实由各向同性高斯嵌入 \(v_i \sim \mathcal{N}(0, \frac{1}{d}\mathbf{I}_d)\) 与反嵌入（unembedding）向量 \(u_i \sim \mathcal{N}(0, \frac{1}{d}\mathbf{I}_d)\) 构成，目标是通过权重矩阵 \(\mathbf{W} \in \mathbb{R}^{d \times d}\) 存储关联事实 \(u_i \approx \mathbf{W}v_i\)。在预测时，模型通过多分类交叉熵计算第 \(i\) 个事实命中正确反嵌入的 logit 概率，样本采样频次严格服从幂律分布 \(p_i \propto i^{-\alpha}\)（其中指数 \(\alpha > 1\)）。
在优化过程中，优化器从批大小为 \(B\) 的小批量样本中计算经验负梯度 \(\mathbf{G} = -\nabla_{\mathbf{W}} L(\mathbf{W}; \mathcal{B})\)。Muon 优化器对梯度矩阵进行奇异值分解 \(\mathbf{G} = \mathbf{U}\mathbf{S}\mathbf{V}^\top\)，并引入基于截断分辨率参数 \(\lambda\) 的平滑极映射滤波器 \(h_\lambda(z) = \frac{z}{\sqrt{z^2 + \lambda^2}}\)，生成参数更新量 \(\mathbf{W}_{t+1} = \mathbf{W}_t + \eta_t h_\lambda(\mathbf{G}_t)\)。该更新在动力学上实现对梯度矩阵中主体小奇异值方向的自适应信号提升，使得模型能够在极短步数内检索出处于叠加态中的海量低频事实。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入与输出高斯嵌入<br/>u_i, v_i ~ N(0, I/d), 幂律频次 p_i ~ i^(-α)"] --> B["小批量交叉熵梯度构建<br/>负梯度 G = ∑ q_i(u_i - u_bar)v_i^T"]
B --> C["谱截断平滑极映射<br/>h_λ(G) = U h_λ(S) V^T, λ ~ d^(-α)"]
C --> D["单步与多步权重矩阵更新<br/>W_{t+1} = W_t + η h_λ(G_t)"]
D --> E["联想记忆召回容量与临界批大小判定<br/>单步容量 d^(1+1/(2α))，多步指数逼近 d^2"]关键设计¶
1. 谱截断平滑极映射：平衡极小奇异值放大与算子范数利普希茨连续性 传统符号极映射 \(\text{polar}(\mathbf{G}) = \mathbf{G}(\mathbf{G}^\top \mathbf{G})^{-1/2}\) 在数值上极易受到接近零的退化奇异值扰动。本文采用平滑正则化函数 \(h_\lambda(z) = \frac{z}{\sqrt{z^2 + \lambda^2}}\)，其核心机制在于将 \(\lambda\) 作为谱分辨率阈值：当奇异值 \(s \gg \lambda\) 时算子近似正交化方向 \(h_\lambda(s) \approx 1\)；当 \(s \ll \lambda\) 时则保留线性比例衰减 \(h_\lambda(s) \approx s/\lambda\)。通过将正则化参数设定为 \(\lambda \asymp \max\left\{\frac{(\log d)^{2\alpha+2}}{d^{\alpha}}, \frac{(\log d)^2}{B}\right\}\)，Muon 能够针对处于梯度主体（bulk）区域的微小奇异值赋予高达 \(\lambda^{-1} \sim d^\alpha\) 的乘性放大系数，从理论上保证算子具备 \(\lambda^{-1}\)-Lipschitz 连续性，有效抑制了伪特征维度的噪声发散。
2. 留一法扰动与非对角干扰界的精细控制：突破非正交叠加态的分析瓶颈 当嵌入向量失去相互正交性时，对角信号 logit \(u_i^\top \mathbf{W} v_i\) 与非对角交互噪声 \(u_j^\top \mathbf{W} v_i\) 之间产生复杂的耦合关联。本文采用留一法梯度分解 \(\mathbf{G}_{-i} = \mathbf{G} - q_i u_i v_i^\top\) 并构建标量辅助增长函数 \(\phi(q) = u_i^\top h_\lambda(\mathbf{G}_{-i} + q u_i v_i^\top) v_i\)。借助算子微积分中的 Daleckii-Krein 公式计算 Fréchet 导数，证明了当信号尖峰逐渐加回时，对角信号导数 \(\phi'(0)\) 受控于主体奇异值截断 \(\frac{h_\lambda(s_{d/2})}{s_{d/2}} \gtrsim \lambda^{-1} \sim d^\alpha\)，得到对角 logit 增长界： $\(\gamma_{ii} \gtrsim \min\left\{\frac{q_i}{\lambda}, 1\right\} - \sqrt{\frac{\log d}{d}}\)$ 针对非对角交叉项，结合块分解预解式（block resolvent）复积分展开与随机矩阵矩量扰动法，严格证明任意 \(j \neq i\) 的交互噪声满足 \(|\gamma_{ij}| \lesssim \frac{(\log d)^3}{\sqrt{d}}\)。当对角信号显著压倒交互界时，成功确立了单步记忆恢复的精确充要判据。
3. 隐式二阶白化机制与双向预条件等价性：一阶谱优化器匹敌牛顿法的内在根源 对于同一交叉熵损失，初始化处的全 Hessian 矩阵可精确分解为 Kronecker 积 \(\mathcal{H} = \mathbf{M}_v \otimes \mathbf{\Sigma}_u\)。牛顿步在梯度两侧分别进行二阶白化预条件，反求得 \(\mathbf{W}_1^{\text{Newton}} \propto \mathbf{\Sigma}_u^{-1} \mathbf{G}_0 \mathbf{M}_v^{-1}\)。相比之下，Muon 仅通过右乘近似白化矩阵 \(\mathbf{G}_0(\mathbf{G}_0^\top \mathbf{G}_0)^{-1/2}\)。在输入与输出嵌入呈现各向同性高斯的前提下，反嵌入协方差 \(\mathbf{\Sigma}_u\) 高度趋近于各向同性白噪声，使得左侧白化操作的必要性大大降低；而右侧谱归一化自适应对齐了样本输入协方差矩阵的主轴方向。这种结构巧合使得仅依赖一阶梯度的 Muon 在大样本极限下完全复现了牛顿法在一阶步长下的渐近恢复容量 \(\widetilde{\Theta}(d^{1+\frac{1}{2\alpha}})\)，揭示了 Muon 具备类牛顿曲率感知的本质机理。
损失函数 / 训练策略¶
模型在小批量 \(\mathcal{B} = \{i_1, \dots, i_B\}\) 上采用标准的多分类交叉熵损失进行训练： $\(\mathcal{L}(\mathbf{W}; \mathcal{B}) = -\sum_{i \in [N]} q_i \left( u_i^\top \mathbf{W} v_i - \log \sum_{j \in [N]} \exp(u_j^\top \mathbf{W} v_i) \right)\)$ 其中 \(q_i\) 为样本 \(i\) 在当前批次中的经验出现频次。Muon 的单步学习率设定为 \(\eta \asymp (\log d)^{-4}\sqrt{d}\)，在该学习率下模型不仅保证了高容量恢复，还实现总体损失上界 \(\mathcal{L}(\mathbf{W}_1^{\text{Muon}}) \le \widetilde{O}(\max\{d^{\frac{1}{2\alpha}+\frac{1}{2}-\alpha}, B^{\frac{1}{\alpha}-1}\})\)。在多步训练动力学中，根据阈值化梯度近似假设，每一轮更新自动剔除已恢复高频事实的残差梯度，保持恒定步长 \(\eta \asymp \sqrt{d}\)，即可实现恢复容量指数级向理论上限 \(\widetilde{\Theta}(d^2)\) 逼近。
实验关键数据¶
主实验¶
论文在合成高斯联想记忆任务以及基于双层 Transformer 的上下文联想召回（In-Context Recall）任务上系统验证了容量扩展理论。以下第一表汇总了总体分布下各优化器在单步与多步恢复容量上的理论标度指数与关键临界量对比；第二表记录了真实二维 Transformer 在诱导头联想召回中（词表大小 \(N=256\)，序列长度 \(T=24\)，批大小 \(B=256\)）的上下文分布内（ID）与分布外（OOD）预测准确率。
| 优化算法 | 单步容量理论阶 \(C(d)\) | 单步 SGD 对比倍率 | 单步临界批大小 \(B^\star\) | 多步早期加速阶段递归率 | 多步渐近理论极限 |
|---|---|---|---|---|---|
| SGD | \(\widetilde{\Theta}(d^{\frac{1}{2\alpha}})\) | 基准 (\(1\times\)) | \(\widetilde{\Theta}(\sqrt{d})\) | \(d_{t+1} \approx d^{\frac{1}{2\alpha}} d_t\) (线性缓慢) | \(\widetilde{\Theta}(d^2)\) |
| Newton | \(\widetilde{\Theta}(d^{1+\frac{1}{2\alpha}})\) | 提升 \(\sim d\) 倍 | \(\widetilde{\Omega}(d^\alpha)\) | 需要计算/求逆二阶全 Hessian | \(\widetilde{\Theta}(d^2)\) |
| Muon (本文) | \(\widetilde{\Theta}(d^{1+\frac{1}{2\alpha}})\) | 提升 \(\sim d\) 倍 (匹配 Newton) | \(\widetilde{\Theta}(d^{\alpha+\frac{1}{2}})\) | \(d_{t+1} \approx d^{2-(1-\frac{1}{2\alpha})^t}\) (单步越阶) | \(\widetilde{\Theta}(d^2)\) |
消融实验¶
下表展示了 Transformer 上下文召回任务中，当输出词分布的幂律重尾衰减系数 \(\alpha_o\) 逐渐增大（事实频次衰减极度陡峭、学习难度加剧）时，模型在固定维度 \(\text{dim}=64\)、迭代步数 512 时的 OOD 泛化召回准确率响应情况：
| 频次分布指数配置 | Muon (OOD 准确率) | AdamW (OOD 准确率) | SGD (OOD 准确率) | 现象说明 |
|---|---|---|---|---|
| \(\alpha_o = 0.0\) (均匀分布) | ~98.5% | ~97.8% | ~84.2% | 均匀无偏场景下 Muon 与 AdamW 性能持平 |
| \(\alpha_o = 0.5\) (弱重尾) | ~97.2% | ~91.5% | ~68.0% | AdamW 性能出现可察觉滑落 |
| \(\alpha_o = 1.0\) (标准齐普夫分布) | ~94.8% | ~76.4% | ~42.3% | AdamW 显著劣化，Muon 维持 90% 以上高召回 |
| \(\alpha_o = 1.5\) (强重尾分布) | ~90.3% | ~52.1% | ~18.5% | Muon 展现极强鲁棒性，领先 AdamW 达 38.2% |
| \(\alpha_o = 2.0\) (极重尾分布) | ~82.6% | ~31.0% | ~5.2% | SGD 趋近崩塌，Muon 仍具备稳定知识保留能力 |
关键发现¶
- 大临界批大小的内在成因：实验证实 SGD 容量饱和于 \(B \approx 100\) 的极小批大小，而 Muon 的性能随着批大小扩展至 \(10^4\) 以上仍保持近乎线性的信息论最优收益，与理论推导的 \(B^\star \asymp d^{\alpha + 1/2}\) 完美相符。
- 训练初期的超常规加速机制：在多步轨迹中，SGD 需要耗费 \(\lceil 2\alpha \rceil\) 步才能使记忆量突破维度量级 \(d_t \gtrsim d\)，而 Muon 仅需第 1 步即直接跨入超维度叠加态存储阶段，精准阐明了实践中 Muon 早期 Loss 陡降的机制。
- 各向异性破坏下的理论边界：消融协方差矩阵特征谱衰减指数 \(\kappa\) 发现，当反嵌入向量呈现强各向异性（\(\kappa > 0\)）时，牛顿法由于双侧白化仍能维持容量，而 Muon 的单侧谱前条件无法消除左侧奇异失真，容量发生衰减。
亮点与洞察¶
- 一阶谱正交化能够无损平替二阶牛顿法：在各向同性特征叠加的联想记忆场景下，Muon 纯粹利用梯度矩阵的谱分解实现了通常需要显式求逆二阶 Hessian 才能达到的曲率自适应收益。
- 严密阐明了 Muon 专属于大 Batch 训练的数理成因：给出了现代大模型预训练中“Muon 只有在大批大小下才显著战胜 Adam”的严谨动力学解释，量化了谱前条件器对批量信息容量的高吞吐利用率。
- 数学工具链具备强大的理论迁移价值：将随机矩阵理论、Daleckii-Krein Fréchet 微分以及长尾谱特征分析整合到注意力权重记忆解析中，为深度学习优化器在非独立同分布、非正交表示下的泛化分析提供了标杆范式。
局限与展望¶
- 各向同性高斯假设与真实词嵌入存在偏差：理论严格要求输入与输出嵌入服从各向同性高斯分布；当特征空间存在较强各向异性协方差结构时，Muon 与牛顿法之间出现性能分歧，未来亟需引入双侧前条件器（如 Shampoo 或 SOAP）作理论补全。
- 多步分析依赖阈值化梯度的解耦启发式假设：多步动力学通过逐步剥离已记忆项目的梯度进行简化，严格考虑中间预测概率 \(0 < \hat{p}(i \mid i) < 1\) 带来的高阶相关性仍作为公开猜想留待证明。
- 未建模动量项与更复杂的 Transformer 架构拓扑：本项工作聚焦于单矩阵无动量 Muon，未来可扩展至带动量与 Nesterov 加速的混合架构，并推广至多层注意力交叉推理与多跳事实检索任务。
相关工作与启发¶
- vs. 正交联想记忆分析 (Wang et al., 2025; Li et al., 2025)：以往理论假设特征向量完全正交，人为将容量上限死锁在 \(N \le d\)；本文破除正交假设，在高维高斯叠加态下导出了突破维度的 \(\widetilde{\Theta}(d^{1+\frac{1}{2\alpha}})\) 与多步 \(\widetilde{\Theta}(d^2)\) 容量定律。
- vs. 经验谱优化与结构自适应算法 (Bernstein et al., 2024; Vyas et al., 2024)：以往多停留在实证观察或凸函数谱范数平滑假设；本文从具体任务数据分布（Zipf 幂律）出发，为 Muon 在长尾语言建模中的优越性补齐了微观机理。
- vs. 二阶曲率优化器 K-FAC / Shampoo (Martens & Grosse, 2015; Gupta et al., 2018)：证明了在特征各向同性情境下单侧谱处理与双侧曲率白化等价，提示工业界在某些 Transformer 线性投影层上完全可用轻量级 Newton-Schulz 迭代代替昂贵的双侧逆矩阵更新。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次打破正交假设，在非正交高斯特征叠加态下严格证明了 Muon 的锐化容量标度法则。
- 实验充分度: ⭐⭐⭐⭐⭐ 理论推导完备，并在高维线性记忆仿真与两层 Transformer 诱导头任务上双重验证了标度指数与临界批大小。
- 写作质量: ⭐⭐⭐⭐⭐ 论证结构极其严密，从单步到多步层层递进，数学工具使用精湛且几何直觉清晰。
- 价值: ⭐⭐⭐⭐⭐ 为当前大模型训练中迅速普及的谱优化器奠定了至关重要的坚实理论基石。