改进分布式训练中的梯度压缩方法,提升收敛速度。
Orth-Dion: Eliminating Geometric Mismatch in Distributed Low-Rank Spectral Optimization
- 用QR正交化替代列归一化,解决几何不匹配问题。
- 理论证明收敛速率提升至与全秩方法相当,且通信开销相同。
- 适合大规模模型训练,尤其关注通信效率的场景。
低秩梯度压缩通过用秩-$r$因子表示更新来减少分布式训练中的通信量。Dion是一种近期方法,它通过一次幂迭代和列归一化近似穆昂(Muon)——一种对动量进行正交化的谱优化器。这使其兼容完全分片的数据并行训练,但收敛速度慢于全秩谱方法。我们发现该差距源于几何不匹配:列归一化未产生穆昂隐含的目标秩-$r$极分解因子,导致方向违反低秩谱几何的对偶范数约束,使收敛率额外增加$\\(sqrt{r}$因子,尽管梯度的低秩近似本身准确。同样的不匹配也出现在光滑性项和误差反馈递归中,进而影响实际性能。我们提出Orth-Dion,将右因子的列归一化替换为QR正交化。在非欧光滑性下,当$L_r$为沿秩-$r$方向的曲率常数时,Orth-Dion达到$O(\\(sqrt{L_r/T})$的速率,与精确谱方法一致,且每步通信成本与Dion相同。证明摒弃了以往误差反馈分析中的有界漂移假设,采用自洽固定点论证,并引入时间平均收缩机制,仅要求误差序列平均收缩而非每步都收缩。大规模语言模型预训练实验验证了预期的$\\(sqrt{r}$缩放关系,并显示Orth-Dion在保持Dion通信成本的前提下,弥补了与Muon的收敛差距。
原文摘要 · Abstract (English)
Low-rank gradient compression reduces communication in distributed training by representing updates with rank-$r$ factors. Dion is a recent method that approximates Muon, a spectral optimizer that orthogonalizes momentum, using one step of power iteration followed by column normalization (rescaling each column of the right factor to unit length). This makes it compatible with fully sharded data parallel training, but it converges more slowly than full-rank spectral methods. We show that this gap is geometric: column normalization does not yield the rank-$r$ polar factor that Muon implicitly targets, so the resulting direction violates the dual-norm constraint of the low-rank spectral geometry, and the rate picks up an extra factor of $\sqrt{r}$ even though the low-rank approximation of the gradient itself is accurate. The same mismatch enters the smoothness term and the error-feedback recursion in the analysis, which has a knock-on effect on empirical performance. We propose Orth-Dion, which replaces column normalization with QR orthogonalization of the right factor. Under non-Euclidean smoothness, with $L_r$ the curvature constant along rank-$r$ directions, Orth-Dion attains rate $O(\sqrt{L_r/T})$, matching exact spectral methods at the same per-step communication cost as Dion. The proof removes the bounded-drift assumption common in prior error-feedback analyses via a self-consistent fixed-point argument, and uses a time-averaged contraction that only requires the error sequence to contract on average rather than at every step. Experiments on large-scale language model pre-training validate the predicted $\sqrt{r}$ scaling and show that Orth-Dion closes the convergence gap to Muon at Dion's communication cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。