技术摘要：CacheMuon

问题陈述

Muon 是一种针对深度学习中矩阵值隐藏层参数设计的几何感知优化器。与逐坐标优化器（如 Adam）不同，Muon 使用通过动量矩阵的极分解（polar factor）导出的正交化方向来取代标准的动量更新。尽管 Muon 在大规模语言模型中展示了强大的经验性能和可扩展性，但其计算效率受限于正交化步骤。

目前的实现使用 Newton–Schulz 迭代（具体为 Gram Newton–Schulz）来近似极分解，以避免高昂的精确奇异值分解（SVD）成本。然而，这种近似仍然需要在每个优化步骤中进行多次矩阵乘法。因此，Muon 的很大一部分计算成本都消耗在重复求解高度相关的矩阵正交化问题上。虽然最近的研究优化了单次正交化调用的效率（例如 PolarExpress、Gram Newton–Schulz），但这些方法在每一步都会从头开始重新计算正交化变换，忽略了训练过程中固有的时间相关性。

方法论：CacheMuon

作者提出了 CacheMuon，一种利用动量矩阵在训练过程中平滑演化的时域预处理方法。其核心假设是：如果动量矩阵变化缓慢，那么在先前步骤中计算的正交化变换对于当前步骤仍然具有信息量。

核心机制

CacheMuon 通过引入缓存机制修改了 Muon 的正交化程序：

1. 新鲜求解（刷新步骤）： 在特定的“刷新”步骤（由触发器决定）中，该方法运行完整的 Fresh Gram Newton–Schulz (FreshGNS) 求解器。这会产生近似极分解 $\tilde{D}_t$，并且至关重要的是，产生一个缓存的左变换矩阵 $Q_t$。
2. 缓存复用（中间步骤）： 在非刷新步骤中，该方法不再运行完整的求解器，而是复用最近一次刷新步骤的缓存变换 $Q_{a_t}$（其中 $a_t$ 是最近一次刷新的索引）来近似极分解：$\hat{D}_t = Q_{a_t} X_t$，其中 $X_t$ 是归一化后的动量矩阵。
3. 基于残差的触发器： 为了确定缓存变换何时不再准确，作者引入了一个轻量级的、基于残差的复用触发器。在复用缓存之前，算法会计算候选更新的正交性残差：
   $$\text{err}_t = \frac{\| \hat{D}^{\text{cand}}_t (\hat{D}^{\text{cand}}_t)^\top - I \|_F}{\|I\|_F}$$
   如果 $\text{err}_t \leq \gamma$（一个预设阈值），则接受缓存的变换。否则，将触发一次新鲜求解以更新缓存。

理论分析

论文将 CacheMuon 分析为一种不精确的 Muon 更新。缓存方向 $\hat{D}_t$ 与精确 Muon 方向 $D_t$ 之间的偏差由一个加性误差项 $\delta_t$ 约束：
$$\| \hat{D}_t - D_t \|_{\text{sp}} \leq \delta_t = L_Q \tau_t + \varepsilon_{\text{ref}}$$
其中：

• $\varepsilon_{\text{ref}}$ 是新鲜求解器的误差（由于有限的 Newton–Schulz 迭代引起）。
• $\tau_t$ 是缓存的 Gram 矩阵相对于当前矩阵的“陈旧度”（staleness）。
• $L_Q$ 是由求解器诱导的多项式映射导出的 Lipschitz 常数。

在标准假设（光滑性、有界方差、有界陈旧度）下，作者证明了其在非凸优化中的随机收敛保证。收敛速率取决于聚合不精确度参数 $\bar{\delta} = L_Q \bar{\tau} + \varepsilon_{\text{ref}}$。只要 $\bar{\delta} < 1$，该方法就会收敛，其速率按 $O((1-\bar{\delta})^{-1} N^{-1/4})$ 缩放。

主要贡献

1. CacheMuon 算法： 一种 Muon 的变体，通过在多个优化步骤之间复用缓存的正交化变换，利用时间平滑性来摊销极分解计算的成本。
2. 收敛理论： 一个将 CacheMuon 视为不精确优化器的理论框架，提供了受新鲜求解器误差和缓存陈旧度控制的随机收敛保证。
3. 基于残差的触发器： 一种实用的、低成本的启发式方法，根据候选更新的正交性动态决定何时刷新缓存，作为更新保真度的代理指标。
4. 经验验证： 在语言模型（GPT-2）和视觉（ResNet-18）工作负载上展示了可控的质量–效率前沿。

实验结果

作者在 GPT-2 Large/Small (OpenWebText) 和 ResNet-18 (CIFAR-10) 上评估了 CacheMuon。

• 质量–效率权衡： 通过改变残差阈值 $\gamma$，作者展示了一个平滑的前沿。
  • 保守复用 ($\gamma=2$)： 紧密追踪 Fresh Muon（GramMuon 和 PolarExpMuon）的性能，同时在 GPT-2 Large 上减少了 13.6%（对比 GramMuon）和 30.9%（对比 PolarExpMuon）的正交化 FLOPs。
  • 激进复用 ($\gamma=15$)： 在仅造成轻微验证损失下降（例如 GPT-2 Large 为 3.096 vs 3.037）的情况下，实现了高达 65.0% 的正交化 FLOPs 减少（对比 GramMuon）。
• 诊断洞察： 实验表明，残差触发器能有效与缓存更新与新鲜更新之间的对齐程度相关联。在训练早期，残差较高，迫使频繁刷新；随着训练趋于稳定，残差下降，从而允许更高的缓存命中率和更高的效率。
• 可扩展性： 该方法在不同模型规模（GPT-2 Small vs Large）以及不同模态（视觉 vs 语言）下均保持了竞争力的表现。

意义与主张

论文声称 CacheMuon 提供了一种在不改变底层几何更新规则的情况下，降低 Muon 计算开销的实用方法。通过将正交化视为一个时域问题而非纯粹的空间问题，该方法减少了冗余计算。

作者强调，虽然该方法显著降低了正交化 FLOPs，但由于系统级开销（例如用于残差检查的主机-设备同步），其在当前高层实现中的实际墙钟速度提升较为有限。他们认为，更优化的底层实现（例如将探测过程融合进设备内核）可以将更大比例的算术节省转化为实际的运行时间收益。

最终，这项工作表明，时间复用是优化几何感知优化器的一个有前景的方向，它提供了一个可调节的旋钮，用于平衡计算效率与训练质量。