[논문리뷰] Muon Makes the First Mixture-of-Experts Layer Harder to Balance

링크: 논문 PDF로 바로 열기

메타데이터

저자: Marcos Hernanz, et al.

1. Key Terms & Definitions (핵심 용어 및 정의)

• Muon: 각 가중치 행렬의 업데이트를 Orthogonalization하여 singular value를 1로 고정하는 최적화 기법으로, 최근 여러 open MoE 모델에서 사용됨.
• Span: 각 MoE 레이어 내 전문가(expert)들 사이의 가장 큰 bias와 가장 작은 bias의 차이로, 해당 레이어의 균형을 맞추기 위한 bias의 작업 부하를 측정하는 지표.
• Bias: MoE 레이어의 부하 균형을 맞추기 위해 라우팅 점수에 더해지는 값으로, 부하가 치우친 전문가의 bias를 조정하여 학습 안정성을 도모함.
• Shared Direction: 라우터 입력값들이 공통적인 방향을 공유하여 특정 전문가들로 라우팅이 쏠리게 만드는 현상.
• Overload: 가장 바쁜 전문가가 처리하는 토큰 비중을 fair share로 나눈 값으로, 부하 불균형의 심각도를 나타냄.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 최신 open MoE 언어 모델들이 학습 시 AdamW 대신 Muon을 채택함에 따라 발생하는 MoE 레이어의 부하 균형 문제에 주목한다. 기존 연구들은 Muon이 가중치 업데이트의 spectral norm을 일정하게 유지하여 부하 불균형을 완화할 것으로 기대했으나, 실제로는 특정 레이어에서 부하 불균형이 악화되는 현상이 관찰된다 [Figure 1]. 특히, 첫 번째 MoE 레이어가 모델의 다른 레이어들과 달리 Muon 사용 시 균형을 맞추기 훨씬 어렵다는 점이 문제로 제기된다. 본 연구는 이러한 현상의 원인을 규명하고, 라우터 및 백본(backbone) 학습 최적화 방식에 따른 영향을 분석하고자 한다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

저자들은 55M 및 108M 활성 파라미터를 가진 MoE 모델을 대상으로 AdamW와 Muon의 조합을 다양하게 변경하며 Span과 Overload 변화를 정량적으로 측정하였다. 실험 결과, 학습률이 튜닝된 환경에서 Muon은 다른 레이어와 달리 첫 번째 MoE 레이어의 Span을 평균적으로 약 2배 가까이 증가시키는 것으로 확인되었다 [Figure 1]. 이러한 현상은 라우터의 최적화 도구와 무관하게 백본이 Muon으로 학습될 때 발생하였으며, 라우터 입력값들이 Shared Direction을 형성하여 특정 전문가로의 쏠림 현상을 유도함이 밝혀졌다. 구체적인 수치로, Muon 환경에서 첫 번째 MoE 레이어의 가장 바쁜 전문가는 fair share 대비 4~7배의 부하를 기록하였으나, 입력값에서 batch mean을 제거함으로써 이러한 과부하를 획기적으로 줄일 수 있었다. 반면, 라우터를 AdamW로 옮길 경우 학습 초기 단계의 부하 변동성이 커지고 모델 성능이 저하되는 경향을 보여, 본 모델 규모에서는 라우터를 Muon으로 유지하는 것이 유리함을 확인하였다.

4. Conclusion & Impact (결론 및 시사점)

본 논문은 Muon 최적화 기법이 MoE 모델의 첫 번째 레이어에서 라우팅 쏠림과 부하 불균형을 심화시킨다는 사실을 실증적으로 증명하였다. 연구 결과는 MoE 구조에서 레이어별 최적화 방식이 학습 안정성에 미치는 중요성을 강조하며, 특히 첫 번째 레이어의 불균형 해소를 위해 별도의 입력값 처리나 최적화 전략이 필요함을 시사한다. 이 연구는 향후 고성능 MoE 모델 설계 및 학습 레시피를 구축함에 있어 최적화 도구 선택의 가이드라인을 제공할 것으로 기대된다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.