Orth-Dion:消除分布式低秩谱优化中的几何错位
机器学习
2026-05-19 v1
摘要
低秩梯度压缩通过以秩- 因子表示更新来减少分布式训练中的通信。Dion 是一种最近的方法,通过一次幂迭代后进行列归一化(即对右因子的每一列进行单位长度缩放)来逼近 Muon——一种对动量进行正交化的谱优化器。这使其能够与完全分片数据并行训练兼容,但收敛速度慢于全秩谱方法。我们指出,这一差距是几何性的:列归一化并未产生 Muon 隐式目标的秩- 极端因子,因此所得方向违反了低秩谱几何中的双范数约束,算法会额外引入一个 的因子,尽管低秩梯度本身的近似精度很高。相同的错位也进入平滑项和误差反馈递归中,这对实际性能产生连锁反应。我们提出 Orth-Dion,用右因子的 QR 正交化取代列归一化。在非欧几里得平滑性下,若 为沿秩- 方向的曲率常数,Orth-Dion 达到 的收敛率,匹配在 Dion 相同每步通信成本下的精确谱方法。该证明通过自洽的固定点论证消除了先前误差反馈分析中常见的有界漂移假设,并使用仅要求误差序列在平均意义上收缩的时序平均收缩。大规模语言模型预训练实验验证了预测的 比例,并表明 Orth-Dion 在 Dion 的通信成本下闭合了收敛差距。
引用
@article{arxiv.2605.16341,
title = {Orth-Dion: Eliminating Geometric Mismatch in Distributed Low-Rank Spectral Optimization},
author = {Tatsuhiro Nakamori and Laura Gomezjurado Gonzalez and Ganesh Talluri and Ansh Tiwari and Hideyuki Kawashima and Ioannis Mitliagkas and Guillaume Rabusseau and Hiroki Naganuma},
journal= {arXiv preprint arXiv:2605.16341},
year = {2026}
}
备注
24 pages, 3 figures, 11 tables