Orth-Dion: 消除分布式低秩谱优化中的几何失配

# Orth-Dion: 消除分布式低秩谱优化中的几何失配

来源：https://arxiv.org/html/2605.16341
同等贡献：[email protected], [email protected] (https://arxiv.org/html/2605.16341v1/[email protected],[email protected])
通讯作者：[email protected] (https://arxiv.org/html/2605.16341v1/[email protected])

###### 摘要

低秩梯度压缩通过使用秩为 \(r\) 的因子表示更新，减少了分布式训练中的通信开销。Dion 是一种近期提出的方法，它近似了 Muon（一种对动量进行正交化的谱优化器），采用一步幂迭代后接列归一化（将右因子的每一列缩放到单位长度）。这使得它兼容全分片数据并行训练，但相比全秩谱方法收敛更慢。我们证明这一差距是几何性的：列归一化并不能产生 Muon 隐式目标中的秩 \(r\) 极分解因子，因此得到的更新方向违反了低秩谱几何的对偶范数约束，即使梯度本身在低秩近似下是准确的，收敛速率仍然会多出一个 \(\sqrt{r}\) 因子。同样的失配也会进入平滑项和误差反馈递归的分析中，对实际性能产生连锁效应。我们提出 Orth-Dion，将列归一化替换为右因子的 QR 正交化。在非欧几里得平滑条件下，设 \(L_r\) 为沿秩 \(r\) 方向的曲率常数，Orth-Dion 达到 \(O(\sqrt{L_r/T})\) 的速率，与精确谱方法相同，且每步通信成本与 Dion 一致。证明通过自洽不动点论证去除了以往误差反馈分析中常见的有界漂移假设，并使用时间平均收缩，只需误差序列在平均意义上收缩而非每一步都收缩。在大规模语言模型预训练上的实验验证了预测的 \(\sqrt{r}\) 缩放，并表明 Orth-Dion 在 Dion 通信成本下弥合了与 Muon 的收敛差距。

## 1 引言

训练大型语言模型越来越依赖于分布式分片方案，例如全分片数据并行（Zhao et al., 2023），在该方案中，优化器的设计同时受限于每步计算和通信。谱优化器如 Muon（Jordan et al., 2024）在这种情况下很有吸引力，因为它们利用了神经网络参数的矩阵结构。它们不是进行逐元素更新，而是对动量正交化，并沿着谱归一化的方向迈步。这种几何结构在实践中是有效的，但在分片下难以扩展，因为形成和正交化全动量矩阵需要额外的集体通信。Dion（Ahn et al., 2025b）通过用低秩分解更新替换全谱更新解决了这一瓶颈。每个矩阵更新通过秩 \(r\) 因子表示，这些因子通过一步幂迭代计算，通信成本与 \((m+n)r\) 成正比，而非 \(mn\)。乍看之下，Dion 与全秩谱方法之间剩余的差距似乎是使用秩 \(r\) 更新不可避免的代价（Ahn et al., 2025b; Carlson et al., 2015）。由于 Dion 只通信低秩因子，人们可能预期其较慢的收敛来自于所选子空间之外的梯度方向缺失。考虑到 PowerSGD 和 Dion 基于低秩、幂迭代压缩的演进脉络，这种解释是自然的（Vogels et al., 2019; Ahn et al., 2025b）。在这种观点下，问题是更新指向何方，Dion 之所以逊色是因为其秩 \(r\) 子空间太小。

我们证明这种解释是不完整的。Dion 的幂迭代能够识别出一个有用的低秩子空间，但后续的更新在该子空间内的形状仍然可能不正确。根源在于 Dion 的最终归一化步骤。在幂迭代之后，Dion 独立地归一化右因子的每一列（Ahn et al., 2025b）。这保住了因子的张成空间，但并没有产生与秩 \(r\) 谱更新相关联的正交右因子（Carlson et al., 2015; Bernstein and Newhouse, 2024b）。结果，Dion 本质上可以与理想的低秩极方向处于同一子空间，但使用的方向在谱下降中缩放不正确。我们通过对偶范数因子 \(\nu_t\) 形式化这一失配，表明即使低秩近似本身是准确的，列归一化也可能引入秩相关的惩罚。Orth-Dion 以尽可能小的改变修复了这一失配。它保持 Dion 的幂迭代、残差缓冲区和低秩通信模式不变，但将列归一化替换为右因子的 QR 正交化。这使得更新在几何上与谱方法所目标的秩 \(r\) 极方向对齐。在与 Dion 相同的低秩通信成本下，Orth-Dion 消除了秩相关的归一化惩罚，并恢复了精确秩 \(r\) 谱更新的主导收敛速率。

由于 Dion 系列方法使用误差反馈，证明必须控制残差缓冲区和被跟踪子空间的耦合演化。我们通过一个自洽残差控制论证和一个摊销收缩条件来实现这一点，避免了以往分析中使用的有界缓冲区漂移假设。我们在 Llama 3 320M 预训练期间直接测量了 \(\nu_t\)。Dion 的 \(\nu_t\) 在每一层和每个记录步骤中都随秩增加而上升；Orth-Dion 的 \(\nu_t\) 保持在 \(\nu_t \approx 1\)。这种失配并非最坏情况下的异常；它是 Dion 在其设计场景中的实际行为。固定 \(r\)，用 QR 替换列归一化可以在我们测试的每个秩下降低验证损失，这分离了 QR 修正与秩容量效应。由于 QR 增加了每步计算量，我们将其与自适应秩结合，得到 Ada-Orth-Dion：工作秩向每一层的内在维度缩小，在 17.1B 模型上恢复类似 Dion 的步长时间（图 1），同时保持收敛增益。组合方法提供了一种低秩分布式谱优化器，其收敛-通信权衡优于 Dion。

参照说明 (a) 在匹配秩下更低的验证损失。LLaMA 320M，后期收敛（均值 \(\pm\) 2 标准差，完整轨迹见图 3）。参照说明 (b) 大规模下更快的挂钟时间。LLaMA 17.1B 模型 50 步运行的每步时间（均值 \(\pm\) 2 标准差）。秩分数 \(r_f\) 是谱更新的秩除以参数矩阵的全秩。图 1：提出的方法同时改进了收敛性和挂钟时间。(a) 在匹配秩下，Orth-Dion 和 Ada-Orth-Dion 更早达到 Dion 的平台并继续下降。(b) 自适应秩吸收了 Orth-Dion 的 QR 开销，在 17.1B 模型上匹配了 Dion 的每步时间。（Ada-Orth-Dion 的秩固定在 \(0.93 r_f\)，这是在 320M 模型上达到的稳态秩；参见附录 L.2）。

#### 贡献。
我们 (i) 识别了 Dion 中一个几何来源的次优性：列归一化使 Ky Fan 对偶范数膨胀了高达 \(\sqrt{r}\)，产生了一个与低秩近似误差不同的秩相关收敛惩罚；(ii) 提出 Orth-Dion，一行代码将 ColNorm → QR，强制 \(\nu_t = 1\)，在 Dion 的通信成本下恢复精确的低秩谱速率；(iii) 通过自洽残差控制论证在非欧几里得平滑性下证明收敛，消除了有界缓冲区漂移并允许摊销而非每步收缩；(iv) 通过直接测量 \(\nu_t\)、在 Llama 3 320M 上匹配秩的改进以及 Ada-Orth-Dion 在更低验证损失下匹配 Dion 挂钟时间，验证了该机理和优化器层面的影响。

## 2 背景：秩约束优化的几何

范数下的最速下降。给定 \(f: \mathbb{R}^{m \times n} \to \mathbb{R}\)，在范数 \(\|\cdot\|\) 下的最速下降更新为 \(X_{t+1} = X_t - \eta D_t^*\)，其中 \(D_t^* = \arg\max_{\|D\| \leq 1} \langle \nabla f(X_t), D \rangle\)。SGD 使用 \(\ell_2\)；SignSGD（Bernstein et al., 2018）使用 \(\ell_\infty\)；谱方法使用算子范数。对于分布式训练，我们将更新约束为秩 \(\leq r\)，从而得到 Ky Fan \(r\)-范数几何。Ky Fan \(r\)-范数是前 \(r\) 个奇异值之和的对偶范数，可以写成 \(D_t^* = \arg\max_{\|D\|_{(r)} \leq 1} \langle M_t, D \rangle\)，其中 \(\|D\|_{(r)} = \max\{\sigma_1(D), \|D\|_F / \sqrt{r}\}\)。解是来自截断 SVD 的秩 \(r\) 极分解因子 \(P_r(M_t) = U_r V_r^\top\)，满足 \(\langle M_t, D_t^* \rangle = \|M_t\|_{\mathrm{KF}, r} = \sum_{i=1}^r \sigma_i(M_t)\)。

两种算法。两者都维护缓冲区 \(M_t = G_t + R_t\) 并使用一步幂迭代。唯一的区别在于右因子归一化：其中 \(W_t = M_t^\top U_t\) 且 \(U_t = \operatorname{orth}(M_t V_{t-1})\)。更新为 \(\hat{D}_t = U_t \bar{V}_t^\top\)，误差反馈为 \(R_{t+1} = \beta (I - U_t U_t^\top) M_t\)。QR 的计算量为 \(O(n r^2)\)，相比 \(O(m n r)\) 的幂步骤可以忽略不计，因为 \(r \ll m\)。完整过程见算法 1。

算法 1 Orth-Dion / 剥离版 Dion（一步）
0: 梯度 \(G_t \in \mathbb{R}^{m \times n}\)，残差 \(R_t\)，右因子 \(V_{t-1}\)，步长 \(\eta\)，EF 系数 \(\beta\)
1: \(M_t \leftarrow G_t + R_t\)  ▷ 缓冲区 = 梯度 + 误差反馈
2: \(U_t \leftarrow \operatorname{orth}(M_t V_{t-1})\)  ▷ 左因子通过 QR，\(O(m r^2)\)
3: \(W_t \leftarrow M_t^\top U_t\)  ▷ 右因子（未归一化），\(O(m n r)\)
4: \(\bar{V}_t \leftarrow \operatorname{orth}(W_t)\)  ▷ Orth-Dion: QR; Dion: \(\operatorname{ColNorm}(W_t)\)
5: \(\hat{D}_t \leftarrow U_t \bar{V}_t^\top\)  ▷ 低秩更新方向
6: \(X_{t+1} \leftarrow X_t - \eta \hat{D}_t\)  ▷ 参数更新
7: \(R_{t+1} \leftarrow \beta (M_t - U_t (U_t^\top M_t))\)  ▷ 误差反馈
8: \(V_t \leftarrow \bar{V}_t\)  ▷ 用于下一步的热启动

FSDP 通信。在 FSDP 下，参数在各设备间分片。每个训练步骤包括：(1) all-gather 以物化完整权重，(2) 前向/反向传播，(3) reduce-scatter 以重新分片梯度，(4) 本地优化器步骤。Muon 在正交化期间需要额外的 all-gather/reduce-scatter 对来获取完整动量矩阵。Dion 和 Orth-Dion 通过仅通信低秩因子（\(O((m+n)r)\) 对比 \(O(mn)\)）避免了这一点，从而以与秩成比例的成本兼容 FSDP 分片。

假设。我们使用以下假设（详情见附录 A）：
(A) 非欧几里得平滑性：\(f(X + \Delta) \leq f(X) + \langle \nabla f(X), \Delta \rangle + \frac{L_r}{2} \|\Delta\|_{(r)}^2\)，其中 \(L_r\) 捕捉沿低秩方向的曲率；
(A') Frobenius 梯度-Lipschitz：\(\|\nabla f(Y) - \nabla f(X)\|_F \leq L_F \|Y - X\|_F\)（与非欧几里得几何中的 A 无关）；
(B') 梯度谱间隙：\(\sigma_r(G_t) - \sigma_{r+1}(G_t) \geq \Delta_{\mathrm{gap}} > 0\)；
(B'') 小谱尾：\(\sigma_{r+1}(G_t) \leq \tau\)，在相关时间跨度上 \(\tau = O(\eta)\)；
(C') 梯度界：\(\|G_t\|_F \leq G_F\)，\(\kappa_r(G_t) \leq \kappa_G\)。

一步下降分解。在假设 A 下，任何满足 \(\|\hat{D}_t\|_{(r)} = \nu_t\) 的更新 \(\hat{D}_t\) 都产生（证明见附录 D）：

\[
f(X_{t+1}) \leq f(X_t) - \eta \|G_t\|_{\mathrm{KF}, r} + \eta \big( \delta_t + (1 + \nu_t) \|R_t\|_{\mathrm{KF}, r} \big) + \frac{L_r \nu_t^2}{2} \eta^2,
\tag{1}
\]

其中 \(\delta_t = \|M_t\|_{\mathrm{KF}, r} - \langle M_t, \hat{D}_t \rangle\) 是 oracle 缺陷。对偶范数因子 \(\nu_t\) 出现在两个关键位置：平滑性惩罚项 \(L_r \nu_t^2 \eta^2 / 2\) 和残差耦合项 \((1 + \nu_t) \|R_t\|_{\mathrm{KF}, r}\)。

## 3 \(\sqrt{r}\) 低效性

列归一化产生的 \(\bar{V}_t\) 具有单位长度的列，其 Gram 矩阵 \(\bar{V}_t^\top \bar{V}_t\) 的对角线元素为 1，但非对角线元素可能很大。由于这是一个相关矩阵：

**命题 3.1**（\(\sqrt{r}\) 界）。对于列归一化的 \(\bar{V}_t\)：\(\nu_t = \|\hat{D}_t\|_{(r)} = \|\bar{V}_t\|_{\mathrm{op}} \in [1, \sqrt{r}]\)。等式 \(\nu_t = 1\) 仅当 \(W_t\) 的列已经正交时才成立。

\(\nu_t\) 如何影响收敛速率。对 (1) 进行 telescoping 并优化 \(\eta\) 可得：

\[
\min_t \|G_t\|_{\mathrm{KF}, r} \leq \sqrt{\frac{2(f_0 - f_\infty) L_r \nu^2}{T}} + O(1/T), \quad \nu = \max_t \nu_t.
\tag{2}
\]

对于 \(\nu = \sqrt{r}\)（ColNorm）：速率 = \(O(\sqrt{L_r r / T})\)。对于 \(\nu = 1\)（Orth-Dion）：速率 = \(O(\sqrt{L_r / T})\)。\(\sqrt{r}\) 差距纯粹是几何性的。

几何解释。Ky Fan 对偶球 \(\{ D : \|D\|_{(r)} \leq 1 \}\) 同时约束了 \(\|D\|_{\mathrm{op}} \leq 1\) 和 \(\|D\|_F \leq \sqrt{r}\)。极分解因子 \(P_r(M_t)\) 位于其边界（\(\|\cdot\|_{\mathrm{op}} = 1\)，\(\|\cdot\|_F = \sqrt{r}\)）。ColNorm 保持了 \(\|\cdot\|_F = \sqrt{r}\)，但允许 \(\|\cdot\|_{\mathrm{op}}\) 高达 \(\sqrt{r}\)——将更新推出对偶球之外。