Orth-Dion: Eliminating Geometric Mismatch in Distributed Low-Rank Spectral Optimization

TL;DR

Orth-Dion improves distributed low-rank spectral optimization by replacing column normalization with QR orthogonalization, closing the convergence gap to spectral methods while maintaining low communication costs.

Contribution

It introduces Orth-Dion, a novel method that corrects geometric mismatch in Dion, achieving spectral-level convergence rates with efficient communication.

Findings

Orth-Dion matches the convergence rate of exact spectral methods.

Experiments confirm Orth-Dion's effectiveness on large-scale language models.

Orth-Dion reduces the convergence gap while maintaining low communication overhead.

Abstract

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.