기술 요약: DeMuon - 그래프 상의 행렬 최적화를 위한 분산형 Muon

문제 정의

본 논문은 완전히 분산된 통신 그래프 상에서의 행렬 최적화 문제를 다룹니다. 목표는 다음과 같은 다중 에이전트 최적화 문제를 해결하는 것입니다:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \frac{1}{N} \sum_{i=1}^N f_i(X)$$
여기서 $N$개의 에이전트는 유향 그래프 $\mathcal{G}=(\mathcal{V}, \mathcal{E})$를 통해 연결되어 있으며, 각 노드 $i$는 로컬 목적 함수 $f_i$를 보유합니다. 행렬 변수를 벡터화하는 경향이 있는 기존의 분산 최적화 방법들과 달리, 본 연구는 행렬 구조를 보존하는 데 중점을 두며, 특히 중앙 집중식 Muon 알고리즘(Jordan et al., 2024)의 기법을 활용합니다. 목표는 중앙 서버 없이도 에이전트 간의 전역 수렴과 합의(consensus)를 달야하는 것이며, 동시에 로컬 목적 함수의 이질성(heterogeneity)을 완화하는 것입니다.

방법론

본 논문은 두 가지 알고리즘인 DeMuon(DeMuon의 분산형 변형)과 DeMuon-A(가속화된 변형)를 제안합니다.

1. DeMuon: 분산형 Muon

DeMuon은 스펙트럼 노름(spectral-norm) 기반 업데이트를 그래디언트 트래킹(gradient-tracking) 프레임워크에 통합합니다. 이 알고리즘은 세 가지 시퀀스를 유지합니다: 로컬 그래디언트 추정치 $\{M^{(k)}\}$, 글로벌 그래디언트 추정치 $\{V^{(k)}\}$, 그리고 로컬 반복값(iterates) $\{X^{(k)}\}$.

• 로컬 모멘텀(Local Momentum): 각 노드는 확률적 그래디언트의 지수 가중 이동 평균을 사용하여 로컬 추정치 $M^{(k)}_i$를 업데이트합니다.
• 그래디언트 트래킹(Gradient Tracking): 노드들은 인접 노드로부터 정보를 집계하여 $V^{(k)}_i$를 업데이트함으로써 글로벌 그래디언트를 추적하며, 이를 통해 로컬 이질성에도 불구하고 그래디언트 방향에 대한 합의를 보장합니다.
• 행렬 직교화(Matrix Orthogonalization): 로컬 반복값 $X^{(k+1)}_i$는 이웃들의 업데이트를 집계하여 업데이트되며, 특히 추적된 그래디언트에 행렬 부호 함수($\text{msgn}(\cdot)$)를 적용합니다. 이 단계는 Muon 방법론의 핵심인 스펙트럼 노름($\|\cdot\|$)을 사용하여 노름 제약 선형 하부 문제를 해결합니다.

DeMuon에 대한 이론적 보장:
유계 노이즈(bounded noise)와 스펙트럼 및 핵 노름(nuclear norm)에 대한 그래디언트 립시츠 연속성(Lipschitz continuity) 가정 하에, 저자들은 다음을 입증합니다:

• 합의 오차(Consensus Error): 오차 $\|X^{(k)} - \mathbf{1} \otimes \bar{X}^{(k)}\|$는 $O(\frac{\sqrt{N}\lambda \eta}{1-\lambda})$로 제한되며, 여기서 $\lambda$는 그래프의 혼합 속도(mixing rate)를 나타냅니다.
• 정상성 오차(Stationarity Error): 평균 정상성 오차는 적절히 튜닝된 스텝 사이즈 $\eta$와 모멘텀 $\theta$에 대해 $O(K^{-1/4})$의 속도로 수렴합니다. 이는 모멘텀이 적용된 중앙 집중식 정규화 SGD의 최적 수렴 속도와 일치합니다.

2. DeMuon-A: 가속화된 분산형 Muon

DeMuon-A는 다중 외삽(multi-extrapolation) 기술(He et al., 2026b)을 분산형 행렬 설정으로 확장합니다.

• 메커니즘: 이 알고리즘은 목적 함수의 고차 매끄러움(higher-order smoothness)을 활용하기 위해 매 반복마다 $q = p-1$번의 로컬 외삽 단계를 수행하여 점 $Z^{(k,s)}_i$를 생성합니다.
• 가정: 이는 스펙트럼 및 핵 노름에 의해 유도된 함수 $f_i$의 $p$차 도함수에 대한 고차 립시츠 조건을 요구합니다.
• 수렴: 이러한 조건 하에서, DeMuon-A는 개선된 $O(K^{-p/(3p+1)})$의 정상성 오차율과 $O(K^{-(2p+1)/(3p+1)})$의 합의 오차율을 달성합니다.

주요 기여

1. 알고리즘 제안: 저자들은 Muon의 스펙트럼 노름 기반 직교화와 그래디언트 트래킹을 결합한 최초의 완전 분산형 행렬 최적화 방법인 DeMuon을 도입했습니다.
2. 수렴 분석: 저자들은 정상성 및 합의 오차 측면에서 DeMuon과 DeMuon-A 모두에 대한 엄격한 전역 수렴 보장을 제공합니다. 특히, 이 분석은 행렬 차원(m, n)에 독립적인 차원 불변(dimension-free) 수렴을 확립하기 위해 행렬 마팅게일 모멘트 부등식(matrix martingale moment inequalities)을 사용합니다.
3. 가속화: 본 논문은 고차 매끄러움을 활용하기 위해 다중 외삽을 통합하여 기본 DeMuon 알고리즘보다 향상된 수렴 속도를 제공하는 DeMuon-A를 제안합니다.
4. 실험적 검증: 분산 딥러닝(Transformer 학습) 및 분산 행렬 회귀에 대한 수치 실험을 통해 제안된 방법론이 기존의 분산형 베이스라인들에 비해 갖는 실질적인 이점을 입증했습니다.

실험 결과

저자들은 DeMuon을 세 가지 베이스라인인 DSGD(순수 분산형 SGD), DSGD-C(클리핑 SGD), DSGD-N(그래디언트 트래킹이 적용된 정규화 SGD)과 비교 평가했습니다.

• 설정: 실험은 세 가지 그래프 토폴로지(Complete, Directed Exponential, Ring graph)에 걸쳐 Multi30k 데이터셋으로 학습된 7.4M 파라미터 규모의 GPT 모델을 대상으로 수행되었습니다.
• 성능: DeMuon은 모든 토폴로지에서 훈련 손실(training loss), 검증 손실(validation loss), 퍼플렉시티(perplexity) 측면에서 모든 베이스라인을 일관되게 능가했습니다.
• 합의: DeMuon의 합의 오차는 베이스라인들과 대등하거나 더 낮게 유지되었으며, 이는 에이전트 간의 효과적인 합의를 나타냅니다.
• 세부 사항: Complete Graph 토폴로지에서 DeMuon은 최종 검증 손실 0.240(퍼플렉시티 1.27)을 달성하였으며, 이는 DSGD-N(0.302) 및 DSGD(7.031)보다 현저히 낮은 수치입니다.

의의 및 주장

본 논문은 DeMuon이 중앙 집중식 딥러닝에서의 Muon의 경험적 성공과 분산 최적화의 요구 사항 사이의 간극을 성공적으로 메웠다고 주장합니다.

• 차원 불변 수렴(Dimension-Free Convergence): 문제의 차원에 명시적으로 의존할 수 있는 최근 연구들(예: Zhang et al., 2026a)과 달리, 저자들은 행렬 마팅게일 모멘트 부등식을 사용하여 행렬 차원 $(m, n)$에 명시적으로 의존하지 않는 수렴 보장을 달성했다고 주장합니다.
• 실질적 이점: 결과는 Muon의 특정 구조(스펙트럼 노름 제약 및 행렬 직교화)가 분산 환경, 특히 행렬 구조가 내재된 대규모 언어 모델 학습에서 명확한 실질적 이점을 제공함을 시사합니다.
• 이론적 엄밀성: 본 연구는 그래디언트 트래킹과 적절한 노름 기반 업데이트가 사용된다면, 중앙 코디네이터 없이도 행렬 최적화에서의 모멘텀 및 정규화의 이점이 완전히 분산된 환경에서도 유지될 수 있음을 입증합니다.

저자들은 본 연구가 분산 딥러닝 및 분산 PCA, 행렬 완결(matrix completion)과 같은 고전적 문제들을 동기로 하여, 고급 행렬 최적화 기술을 분산된 세계로 가져오는 단계라고 규정합니다.