SignMuon: Communication-Efficient Distributed Muon Optimization

Distributed training of large neural networks is bottlenecked by full-precision gradient communication and by coordinatewise optimizers that ignore the matrix structure of weight tensors. We propose Sign-Muon, a 1-bit, matrix-aware optimizer that combines majority-vote sign aggregation from signSGD with the polar-step framework of Muon. Each worker forms a Muon-style direction by taking the polar factor of its momentum via a Newton--Schulz iteration, transmits only the entrywise signs, and aggregates by majority vote; an optional local polar step further enforces orthogonality at no extra communication cost. Under spectral-norm smoothness and bounded-variance stochastic gradients, the spectral-norm normalized sign step yields an $\mathcal{O}(1/\sqrt{T})$ nonconvex rate for an $\ell_1$-based stationarity measure. With unimodal symmetric noise, majority vote across $M$ workers cuts the stochastic term by $1/\sqrt{M}$, matching signSGD. In the $\alpha$-$\beta$ model, distributed Sign-Muon needs only one integer sum-allreduce per iteration; all orthogonalization is local, giving a $32\times$ bandwidth reduction over float32 ($4\times$ for int8). Across 330 CIFAR-10/ResNet-50 configurations Sign-Muon attains the best validation accuracy (92.15\%); its 4-GPU majority-vote variant reaches 92.02\% with 37\% less training time at matched effective batch. On nanoGPT, Sign-Muon achieves lower perplexity and better anytime performance than other sign-based baselines, with favorable weak-scaling up to 16 GPUs.

📋 한 줄 요약

**[분산 최적화]** signSGD의 다수결 집계와 Muon의 polar-step을 결합한 1-bit, 행렬 인식 분산 최적화기 Sign-Muon을 제안하여 float32 대비 32배 대역폭 절감과 함께 안정적인 수렴을 보였다.

🎯 핵심 기여도

1-bit 통신 + 행렬 구조 인식의 새로운 분산 최적화기 Sign-Muon 제안
Newton–Schulz 반복으로 polar factor를 로컬에서 계산 후 부호만 전송하는 설계
Spectral-norm smoothness 하에서 $O(1/\sqrt{T})$ 비볼록 수렴률 증명
$M$ 워커 다수결로 stochastic 항을 $1/\sqrt{M}$ 만큼 감소시키는 신호 분석 (signSGD 매칭)
분산 환경에서 iteration당 정수 sum-allreduce 1회만 필요 (모든 직교화는 로컬)

💡 핵심 아이디어

대규모 분산 학습은 (1) full-precision gradient 통신과 (2) 가중치 텐서의 행렬 구조를 무시하는 좌표별 최적화기 두 가지가 병목이다. Muon은 polar step으로 (2)를 해결하고 signSGD는 1-bit 통신으로 (1)을 해결한다. 두 아이디어를 결합해 워커가 모멘텀의 polar factor를 로컬로 만든 뒤 부호만 보내고 다수결로 집계하면, 행렬 구조를 보존하면서도 통신을 극단적으로 줄일 수 있다.

🔬 기술적 접근법

Newton–Schulz 반복으로 워커별 polar factor 계산
엔트리별 sign만 전송, majority vote로 집계
선택적 local polar step으로 추가 통신 없이 직교성 강제
분산에서는 integer sum-allreduce 1회/이터레이션
$\alpha$-$\beta$ 모델: float32 대비 32×, int8 대비 4× 대역폭 절감

📊 주요 결과

CIFAR-10/ResNet-50 330개 구성: Sign-Muon이 최고 검증 정확도 92.15%
4-GPU 다수결 변형: 동일 effective batch에서 92.02%, 학습 시간 37% 단축
nanoGPT: 다른 sign 기반 베이스라인 대비 더 낮은 perplexity와 우수한 anytime 성능
16 GPU까지 weak scaling 양호

💭 의의 및 한계

"1-bit + 행렬 구조"라는 두 효율화 축을 이론·실험 모두에서 결합한 점이 인상적이며, 분산 학습 통신 병목을 근본적으로 줄일 수 있는 실용적 옵션을 제시한다. 다만 대규모 LLM 사전학습 등 보다 큰 스케일에서의 안정성과 실효 성능 검증은 추가 연구가 필요하다.

🚀 실용적 활용

대역폭이 제한된 분산 학습(엣지 클러스터, 인터커넥트 약한 멀티노드, 통신 비용이 비싼 클라우드 환경)에서 Muon 수준 품질을 유지하면서 통신 비용을 극단적으로 줄이는 옵티마이저로 직접 채택 가능하다.