기술 요약: 2 단계 최적화를 위한 계층적 전문가 혼합 (Hi-MoE)

1. 문제 제기

희소 전문가 혼합 (MoE) 모델은 토큰을 소수의 전문가 네트워크로 라우팅하여 용량을 확장합니다. 그러나 이러한 모델을 학습시키는 과정에서는 **부하 균형 (load balancing)**과 전문가 특화 (expert specialization) 사이의 근본적인 트레이드오프에 직면합니다:

• 라우팅 붕괴 (Routing Collapse): 공격적인 다양성 메커니즘은 종종 "라우팅 붕괴"를 초래하여, 소수의 전문가가 토큰 선택을 지배하고 나머지는 유휴 상태로 남게 됩니다. 이는 비효율적인 용량 활용과 장치 간 부하 불균형 (stragglers) 을 유발합니다.
• 그룹화 딜레마 (The Grouping Dilemma): 최근 접근법은 전문가를 그룹화하고 그룹당 고정된 선택을 강제함으로써 불균형을 완화합니다 (예: $M$개의 그룹 각각에서 $K/M$개의 전문가 선택). 이는 장치 수준의 계산을 균일하게 보장하지만, 라우팅의 자유도를 제한합니다. 이러한 제약은 종종 **중복성 (redundancy)**으로 이어져, 서로 다른 그룹의 전문가들이 유사한 토큰 혼합물에 최적화됨에 따라 유사한 기능을 학습하게 되며, 이는 전문가 수준의 붕괴를 그룹 간 "다양성 붕괴"로 대체합니다.

기존 손실 기반 균형 조정 방법 (예: GShard) 은 종종 작업 그래디언트와 충돌하여 수렴을 늦추는 반면, 손실 없는 방법은 중복성을 방지하지 못할 수 있습니다. 하드웨어 정렬 균형을 동시에 강제하고 의미 있는 보완적 특화를 촉진하는 프레임워크가 필요합니다.

2. 방법론: Hi-MoE

저자들은 Hi-MoE를 제안합니다. 이는 라우팅 제어를 **그룹 간 균형 (inter-group balancing)**과 **그룹 내 특화 (intra-group specialization)**라는 두 가지 결합된 수준으로 분해하는 계층적 프레임워크입니다.

아키텍처

• 그룹화된 구조: 전문가들은 $M$개의 불연속 그룹 (종종 GPU 샤드에 정렬됨) 으로 분할됩니다.
• 라우팅 제약: 각 토큰에 대해 라우터는 각 그룹 $m$에서 고정된 수의 전문가 ($K_m$) 를 선택합니다.
• 2 단계 최적화: 프레임워크는 계층의 각 수준에 대한 특정 정규화 항을 도입합니다.

주요 구성 요소

1. 그룹 간 정규화 ($R_{inter}$):

   • 목표: 균형 잡힌 장치 활용을 보장하기 위해 그룹 간 공정한 트래픽 분배를 강제합니다.
   • 메커니즘: 그룹 간 선택 후 라우팅 가중치의 $\ell_2$ 집중도를 최소화합니다. 특정 그룹에 집중된 라우팅 질량을 패널티로 부과함으로써, 그룹 부하의 변동 계수 (CV) 최소화를 위한 대리 지표로 작용합니다.
   • 이론적 근거: 정리 4.1 은 이 항을 최소화하는 것이 그룹 수준 불균형에 대한 인증된 상한을 제공함을 증명합니다.

2. 그룹 내 정규화 ($R_{intra}$):

   • 목표: 그룹 내에서 보완적 특화를 촉진하고 전문가들이 중복된 기능을 학습하는 것을 방지합니다.
   • 메커니즘: 선택 전 라우팅 분포의 $\ell_2$ 집중도를 최대화합니다 (실제로는 겹침 방지 항). 이는 라우터가 많은 전문가를 낮은 확신으로 혼합하기보다 그룹 내 특정 전문가에 대해 결정적이고 희소한 약속을 하도록 장려합니다.
   • 이론적 근거: 정리 4.2 는 라우팅 겹침을 줄이면 전문가 간 그래디언트 결합이 감소하여 서로 다른 최적화 궤적을 따를 수 있음을 보여줍니다. 정리 4.3 은 균형 잡힌 주변부 하에서 라우팅 집중도를 높이면 입력과 선택된 전문가 간의 상호 정보를 최대화하여, 이는 특화의 대리 지표임을 입증합니다.

3. 안정화 메커니즘:

   • 그룹 내 다양성 신호가 퇴화 영역으로 이동하는 것을 방지하기 위해, 저자들은 라우터 로짓의 지수 이동 평균 (EMA) 을 사용한 편향 보정 라우팅 계산을 도입합니다. 이는 별도의 균형 조정 손실을 추가하지 않으면서 역사적으로 선호된 전문가의 지속적인 과선택을 억제합니다.

목적 함수

전체 학습 목적 함수는 작업 손실, 표준 부하 균형 손실, 그리고 제안된 계층적 항을 결합합니다:
$$\mathcal{L} = \mathcal{L}_{task} + \mathcal{L}_{load} + R_{intra} + R_{inter}$$
저자들은 이를 제약 최적화 문제로 규정하며, 여기서 계수 $\lambda_{intra}$와 $\lambda_{inter}$는 각각 다양성과 균형 제약에 대한 라그랑주 승수 역할을 합니다.

3. 주요 기여

• 긴장 관계의 식별: 이 논문은 엄격한 그룹 제약을 통한 장치 수준 균형 강제가 활용도를 안정화시킬 수 있지만 그룹 간 중복성을 증가시킬 수 있음을 식별합니다.
• 계층적 프레임워크: Hi-MoE 는 이러한 우려를 분리하는 모듈식, 플러그인 목적 함수를 도입합니다:
  • 하드웨어 정렬 균형을 위한 그룹 간 목적 함수.
  • 기능적 다양성과 특화를 위한 그룹 내 목적 함수.
• 이론적 보장: 이 논문은 정규화 항을 그룹 불균형 상한, 그래디언트 분리, 그리고 증가된 라우팅 정보 (특화) 와 연결하는 공식적 증명을 제공합니다.
• 실증적 검증: 비전 및 언어 도메인 전반에 걸쳐 일관된 개선을 보여주며, Hi-MoE 가 비분자/정확도와 부하 균형 (CV) 사이의 파레토 프론티어를 확장함을 입증합니다.

4. 실험 결과

저자들은 Hi-MoE 를 세 가지 다른 사전 학습 설정에서 평가했습니다:

A. 컴퓨터 비전 (Tiny ImageNet 의 Swin-MoE)

• 설정: 64 개 전문가 (16 개 그룹 4 개) 를 가진 10 억 파라미터 모델.
• 결과: Hi-MoE 는 Vanilla GShard, 손실 없는 균형 조정, MoGE 와 비교하여 가장 높은 Top-1 정확도 (59.13%) 와 최고의 부하 균형 (CV 0.29) 을 달성했습니다.
• 관찰: 어텐션 맵의 시각화는 Hi-MoE 그룹이 서로 다른 시각적 특화 (예: 한 그룹은 신체 부위에, 다른 그룹은 배경 표면에 집중) 를 발전시킨 반면, 베이스라인은 더 많은 중복성을 보임을 드러냈습니다.

B. 언어 모델링 (OpenWebText 의 nanoGPT)

• 설정: 8 개 전문가 (2 개 그룹 4 개) 를 가진 10.6 억 파라미터 모델.
• 결과: Hi-MoE 는 가장 낮은 비분자 (2.947) 와 최고의 균형 (CV 0.23) 을 달성했습니다.
• 애블레이션: $R_{inter}$를 제거하면 균형이 저하되었고, $R_{intra}$를 제거하면 품질이 저하되었습니다. 편향 보정은 안정제 역할을 했습니다.
• 효율성: 계층적 라우터는 레이어당 1.3ms 에서 1.6ms 의 최소 오버헤드를 추가하면서 평탄한 라우팅 대비 동기화 오버헤드를 약 3 배 감소시켰습니다.

C. 대규모 사전 학습 (OLMoE-7B 대 Hi-MoE-7B)

• 설정: 585.4 억 토큰 (Dolma 혼합물) 으로 학습된 70 억 파라미터 모델.
• 비분자: Hi-MoE-7B 는 표준 OLMoE-7B 베이스라인 대비 11 개 검증 도메인 전반에서 평균 5.6% 비분자를 감소시켰습니다.
• 하류 작업: Hi-MoE 는 MMLU 슬라이스 (예: STEM 에서 +22.2%, 인문학에서 +11.9%) 및 상식 추론 작업에서 상당한 향상을 보였습니다.
• 부하 균형: Hi-MoE 는 전문가 균형 (CV) 을 40% (0.153 대 0.257) 개선하여 대규모 학습의 자원 효율성 병목 현상을 직접 해결했습니다.

5. 중요성과 주장

이 논문은 Hi-MoE 가 희소 MoE 학습에 내재된 "균형–다양성" 긴장 관계를 해결한다고 주장합니다. 라우팅의 계층 (그룹 대 전문가) 을 명시적으로 모델링함으로써, 이 방법은 모델이 다음을 가능하게 합니다:

1. 효율적 확장: 모델의 조건부 용량을 희생하지 않으면서 균일한 GPU 활용도를 유지합니다.
2. 특화된 표현 학습: 그룹화된 전문가들이 기능적으로 중복되는 "다양성 붕괴"를 방지합니다.
3. 품질 향상: 비전 및 언어 작업 전반에서 비분자와 하류 작업 성능의 일관된 개선을 달성합니다.

저자들은 Hi-MoE 를 기존 안정화 메커니즘의 대체제가 아니라, 전문가 그룹화를 단순한 시스템 제약이 아닌 원칙적인 최적화 메커니즘으로 만드는 보완적이고 모듈적인 향상으로 위치시킵니다. 이 연구는 계층적 조정이 수조 파라미터 모델에서 신뢰할 수 있는 품질 향상을 unlocking 하는 실현 가능한 경로임을 시사합니다.