기술 요약: Orth-Dion

문제 정의

전체 분할 데이터 병렬성 (FSDP) 을 이용한 대규모 언어 모델 (LLM) 의 분산 학습에서 통신 비용은 주요 병목 현상입니다. Muon과 같은 스펙트럴 옵티마이저는 모멘텀을 직교화하고 스펙트럴 정규화된 방향으로 스텝을 밟음으로써 우수한 수렴성을 제공하지만, 정수행렬 연산 ($O(mn)$ 통신) 을 필요로 하므로 FSDP 분할과 양립할 수 없습니다.

Dion(Ahn et al., 2025b)은 단일 단계의 파워 반복을 통해 계산된 저랭크 인자 ($r \ll \min(m,n)$) 를 사용하여 Muon 을 근사함으로써 이를 해결했습니다. Dion 은 통신을 $O((m+n)r)$로 줄이지만, 정수행렬 스펙트럴 방법보다 훨씬 느리게 수렴합니다. 기존 주류 가설은 이 격차가 선택된 랭크-$r$ 부분공간 밖의 그라디언트 정보 손실로 인한 통계적 한계 때문이라고 여겨졌습니다.

본 논문은 수렴 격차가 저랭크 근사의 통계적 한계뿐만이 아니라 기하학적 불일치임을 규명합니다. 구체적으로 Dion 의 최종 정규화 단계 (열 정규화) 는 스펙트럴 기하학이 요구하는 올바른 랭크-$r$ 극인자 (polar factor) 를 생성하지 못합니다. 이로 인해 Ky Fan $r$-노름의 이중 노름 제약 조건을 위반하는 업데이트 방향이 발생하며, 저랭크 부분공간 자체가 정확하게 식별되더라도 수렴률에 최대 $\sqrt{r}$까지의 랭크 의존적 패널티 인자가 도입됩니다.

방법론

기하학적 통찰

저자들은 Ky Fan 이중 노름 인자 $\nu_t = \|\hat{D}_t\|_{(r)}$를 기준으로 업데이트 방향 $\hat{D}_t$를 형식화했습니다.

• 이상적인 경우: 랭크-$r$ 스펙트럴 업데이트의 경우, 최적 방향은 $\nu_t = 1$인 랭크-$r$ 극인자입니다.
• Dion 의 결함: Dion 은 파워 반복을 통해 오른쪽 인자 $W_t$를 계산한 후 열 정규화(각 열을 단위 길이로 스케일링) 를 적용합니다. 이는 열의 범위를 보존하지만 열 간의 직교성을 보장하지는 않습니다. 결과적으로 생성된 행렬의 연산자 노름은 $\sqrt{r}$까지 도달할 수 있어 $\nu_t \in [1, \sqrt{r}]$가 됩니다.
• 결과: 비유클리드 매끄러움 (non-Euclidean smoothness) 에 대한 수렴률 상한에는 $L_r \nu_t^2$에 비례하는 항이 포함됩니다. 열 정규화를 사용하면 이 값이 $O(\sqrt{L_r r / T})$로 부풀어 오르는 반면, 최적 수렴률은 $O(\sqrt{L_r / T})$입니다.

Orth-Dion

제안된 방법인 Orth-Dion은 열 정규화 단계를 오른쪽 인자의 QR 직교화로 대체합니다.

• 알고리즘: Dion 의 파워 반복, 잔여 버퍼, 저랭크 통신 패턴을 유지합니다. 유일한 수정은 업데이트 단계에서 ColNorm(Wt) 대신 orth(Wt)(QR 분해) 를 계산하는 것입니다.
• 기하학적 보장: 오른쪽 인자에 정규 직교 열을 강제함으로써 Orth-Dion 은 업데이트 방향이 부분 등거리 사상 (partial isometry) 이 되도록 보장하며, 구성상 $\nu_t = 1$을 보장합니다. 이는 업데이트를 스펙트럴 방법이 목표로 하는 랭크-$r$ 극 방향과 정확히 정렬시킵니다.

Ada-Orth-Dion

QR 단계의 계산 오버헤드 (각 단계당 $O(nr^2)$ 추가) 를 해결하기 위해 저자들은 Ada-Orth-Dion을 도입했습니다. 이 변형은 기하학적 수정을 레이어별 적응형 랭크 메커니즘과 결합합니다.

• 메커니즘: 파워 반복 인자에서 파생된 경량 엔트로피 기반 추정기를 사용하여 각 레이어 그라디언트 스펙트럼의 "내재적 차원"을 추정합니다.
• 목표: 작동 랭크 $r$을 레이어의 신호 - 잡음 전이 지점에 맞게 동적으로 조정하여 QR 단계의 계산 비용을 줄이면서도 기하학적 이점을 유지합니다.

이론적 기여

본 논문은 비유클리드 매끄러움 가정 하에서 Orth-Dion 에 대한 수렴 증명을 제공합니다. 주요 이론적 혁신은 다음과 같습니다:

1. 자기 일관된 드리프트 제거: 증명에서는 기존 오차 피드백 분석에서 흔히 사용되던 "유계 버퍼 드리프트" 가정을 제거합니다. 대신 추적 오차와 잔여 비율의 결합 시스템에 대한 자기 일관된 고정점 논증을 사용합니다.
2. 상각된 수축 (Amortized Contraction): 모든 단계에서 수축이 필요하지 않도록 (워밍업 기간 중에는 실패할 수 있음) 증명에서는 시간 평균 수축 조건을 수립하여, 창 (window) 에 걸쳐 기하학적 수렴을 보장하면서도 일시적인 확장을 허용합니다.
3. 오차 피드백 분석: 본 논문은 부분공간 밖 그라디언트의 단계 간 상관관계에 기반하여 오차 피드백 역학의 세 가지 영역을 특징짓습니다. 이는 확률적 환경에서 오차 피드백 계수 $\beta$를 1 미만으로 줄이는 것이 암묵적인 헤비볼 (heavy-ball) 모멘텀으로 작용할 수 있음을 보여줍니다.

주요 결과

실험적 검증

실험은 Llama 3모델 (320M 및 17.1B 파라미터) 과 GPT-2 45M에서 수행되었습니다.

1. 기하학적 수정: Llama 3 320M 사전 학습 중 $\nu_t$의 직접 측정은 이론을 확인했습니다. Dion 의 $\nu_t$는 랭크가 증가함에 따라 단조 증가하여 ($r=384$에서 $\approx 2.9$ 도달) 반면, Orth-Dion 은 모든 랭크와 레이어에서 $\nu_t \approx 1.0$을 유지했습니다.
2. 수렴성 개선: 일치하는 랭크에서 Orth-Dion 은 Dion 보다 일관되게 우수한 성능을 보였습니다.
   • Llama 3 320M 에서 Orth-Dion 은 Dion 의 최선 검증 손실을 12–13% 적은 단계 (약 10% 월클락 감소) 에 도달했습니다.
   • 개선은 모든 테스트된 랭크 ($r \in \{96, 192, 384\}$) 에서 관찰되어, 이 이득이 증가된 랭크 용량이 아닌 기하학적임을 확인했습니다.
3. 월클락 효율성: 17.1B 모델에서 고정 랭크 Orth-Dion 은 QR 오버헤드로 인해 Dion 보다 약간 느렸습니다. 그러나 Ada-Orth-Dion은 평균 작동 랭크를 고정 랭크의 약 93% 로 줄임으로써 Dion 의 단계당 시간을 성공적으로 회복하면서도 더 낮은 검증 손실을 달성했습니다.
4. 오차 피드백 튜닝: GPT-2 45M 에 대한 스윕 (sweep) 실험은 오차 피드백 계수 $\beta$를 1.0 에서 $\approx 0.1$로 줄이는 것이 추가 컴퓨팅 없이 검증 손실을 0.315 nats만큼 개선했음을 보여주어, 확률적 영역에서 암묵적 모멘텀 축적에 대한 이론적 예측을 검증했습니다.

의의

본 논문은 Orth-Dion 이 저랭크 스펙트럴 최적화의 근본적인 기하학적 비효율성을 해결한다고 주장합니다. 열 정규화를 QR 직교화로 대체함으로써 Dion 과 동일한 통신 비용으로 정수행렬 스펙트럴 방법 (Muon) 의 정확한 수렴률을 회복합니다. 이 기하학적 수정과 적응형 랭크 (Ada-Orth-Dion) 의 결합은 기존 저랭크 방법보다 엄격히 더 나은 수렴 - 통신 트레이드오프를 제공하는 분산 스펙트럴 옵티마이저를 제공하며, 통신 오버헤드를 발생시키지 않고 정수행렬 스펙트럴 옵티마이저와의 격차를 해소합니다.