기술 요약: 세분화된 전문가 혼합을 위한 적응형 역색인 라우팅 (AIR-MoE)

1. 문제 제기

전문가 혼합 (MoE) 모델은 토큰당 전문가의 부분 집합만 활성화함으로써 확장 가능한 대규모 언어 모델 (LLM) 의 핵심 기반이 되었습니다. 최근의 경험적 확장 법칙은 세분화된 전문가(소수의 대형 전문가 대신 많은 수의 소형 전문가 사용) 를 사용할 때 성능이 향상됨을 시사합니다. 그러나 이러한 regimes 는 라우팅 비용이라는 상당한 계산 병목 현상을 도입합니다. 전문가 수 ($E$) 가 증가함에 따라 라우터는 모든 토큰 표현과 모든 전문가 중심점 사이의 내적을 계산하여 상위 $K$ 개의 전문가를 식별해야 합니다. 이 오버헤드는 전체 계산을 지배하여 매우 세분화된 모델의 훈련을 비현실적으로 만들 수 있습니다.

라우팅 비용을 줄이기 위한 기존 솔루션들은 일반적으로 모델의 표현력을 제한하는 구조적 제약을 부과합니다:

• 계층적 MoE는 전문가를 고정된 비겹치는 그룹으로 분할하여 해당 그룹 내에서의 라우팅만 허용합니다.
• 프로덕트-키/구조적 제약(예: He, 2024)은 정확한 검색을 가능하게 하기 위해 전문가 중심점을 구성 요소로 분할하지만, 허용 가능한 전문가 표현의 공간 (예: 데카르트 곱 구조 강제) 을 제한합니다.
• 고정 해싱 접근법은 학습 가능성이 부족하며 근본적으로 제한적임이 입증되었습니다.

과제는 전문가 매개변수에 구조적 제약을 부과하거나 전체 전문가 점수화를 요구하지 않고도 상위 $K$ 개 라우팅을 효율적으로 근사하는 것입니다.

2. 방법론: AIR-MoE

저자들은 정보 검색 및 근사 최근접 이웃 검색에 사용되는 역파일 색인에 영감을 받은 **MoE 를 위한 적응형 역색인 라우팅 (AIR-MoE)**을 제안합니다. 이 방법은 모든 전문가의 전체 점수화를 피하면서 두 단계의 coarse-to-fine 검색 과정을 통해 실제 상위 $K$ 개 전문가를 근사합니다.

2.1 두 단계 검색

각 입력 토큰 $h$에 대해 라우팅은 병렬로 진행됩니다:

1. Coarse Shortlisting (대략적 선별): 토큰은 코사인 유사도를 사용하여 학습된 코드북 $\mathcal{C} = \{c_1, \dots, c_G\}$에서 가장 가까운 코드워드 $c_g$로 매핑됩니다. 각 코드워드 $c_g$는 $M$개의 후보 전문가로 구성된 사전 계산된 shortlist $L_g$와 연관됩니다. 이러한 shortlist 들은 역색인의 "게시 목록 (posting lists)"과 본질적으로 동일하며, 코드워드 중심점에 대해 가장 높은 내적을 가진 전문가들의 인덱스를 저장합니다.
2. Fine Scoring (정밀 점수화): 라우터는 토큰 $h$와 검색된 shortlist $L_g$ 내의 전문가들 사이에서만 정확한 내적을 계산합니다. 그런 다음 상위 $K$ 개 전문가가 이 축소된 후보 집합에서 선택됩니다.

이 계층 구조는 활성 전문가 집합 $T$가 shortlist $L$의 부분집합이고, $L$이 전체 전문가 집합 $[E]$의 부분집합임을 보장합니다.

2.2 훈련 및 최적화

중요한 과제는 토큰을 코드워드에 할당하는 과정이 미분 불가능하다는 점입니다. AIR-MoE 는 이중 최적화 전략을 통해 이를 해결합니다:

• 기반 기반 훈련: 토큰 표현과 전문가 중심점은 하류 언어 모델링 손실로부터의 표준 역전파를 통해 업데이트됩니다.
• 기반 없는 코드북 훈련: 코드북 $\mathcal{C}$와 이에 연관된 shortlist 들은 적응형 구형 k-means 알고리즘을 사용하여 독립적으로 업데이트됩니다. 이 알고리즘은 다음을 수행합니다:
  • 토큰 표현의 분포를 추적하기 위해 지수 이동 평균 (EMA) 을 사용합니다.
  • 고차원 코사인 기하학을 효과적으로 처리하기 위해 단위 구면에서 작동합니다.
  • 붕괴를 방지하기 위해 "사망" 코드워드 (할당 횟수가 낮은 것들) 에 대한 재초기화 메커니즘을 포함합니다.
  • 최적화 단계 후에만 shortlist $L_g$를 업데이트하여 색인이 현재 전문가 중심점 상태와 일관되도록 유지합니다.

중요하게도, AIR-MoE 는 straight-through estimation 트릭을 사용하지 않습니다. 기울기는 양자화 단계를 거쳐 코드북으로 전달되지 않습니다.

2.3 이론적 보장

이 논문은 Mass Recall(shortlist 에 유지된 실제 상위 $K$ 개 전문가의 총 확률 질량) 에 대한 하한을 유도합니다. 이 하한은 다음에 의존합니다:

1. 양자화 오차 $\epsilon(h) = \|h - c(h)\|$: 오차가 작을수록 유지율이 향상됩니다.
2. shortlist 질량 $\rho_M(c)$: 코드워드에 대한 상위 $M$개 전문가가 포착하는 확률 질량.
   이론적으로 AIR-MoE 는 토큰이 잘 양자화되고 shortlist 크기 $M$이 충분히 클 때 대부분의 라우팅 질량을 유지합니다.

3. 주요 기여

1. AIR-MoE 아키텍처: 전문가 표현에 구조적 제약 (예: 그룹화 또는 매개변수적 제한) 을 부과하지 않고 세분화된 MoE 를 위한 상위 $K$ 개 라우팅을 근사하는 역색인 영감 라우팅 메커니즘.
2. 이중 최적화: 코드북의 기반 없는 적응 (적응형 구형 k-means 를 통해) 과 모델 매개변수의 기반 기반 훈련을 분리하는 훈련 전략으로, straight-through 추정이 필요하지 않습니다.
3. 이론적 분석: 근사가 라우팅 품질을 보존하는 조건에 대한 직관을 제공하는 Mass Recall 에 대한 유도된 하한.
4. 경험적 성능: 기존 효율적 세분화 라우팅 방법보다 우수한 perplexity–FLOPs 트레이드오프를 달성함을 입증.

4. 실험 결과

저자들은 WikiText-103, C5, OpenWebText2 에서 65,536 개의 전문가를 사용하여 소형, 중형, 대형 모델 규모 (최대 약 0.45B 매개변수) 로 AIR-MoE 를 평가했습니다.

• 성능 vs 베이스라인: AIR-MoE 는 PEER(He, 2024) 및 계층적 MoE를 포함한 기존 효율적 세분화 베이스라인보다 일관되게 우수한 성능을 발휘했습니다.
  • 유사한 FLOP 비용에서 PEER 베이스라인 대비 최대 **10%**의 perplexity 개선을 달성했습니다.
  • 표준 세분화 라우팅 (모든 전문가 점수화) 은 가장 낮은 perplexity 를 달성했지만, prohibitive 한 계산 비용 (예: WikiText-103 Small 에서 AIR-MoE 의 324.1P 대비 약 467.5P FLOPs) 을 발생시켰습니다.
• 효율성: AIR-MoE 는 경쟁력 있는 perplexity 를 유지하면서 정확한 세분화 라우팅에 비해 훈련 FLOP 를 크게 줄였습니다.
• Ablation 연구:
  • 중심점 정규화를 제거하거나 구형 k-means 대신 유클리드 거리를 사용하면 성능이 약간 저하되었습니다.
  • 정적 코드북(비적응형) 을 사용하는 경우 가장 낮은 성능을 보였으며, 이는 적응형 코드북의 필요성을 확인시켜 주었습니다.
  • 전문가 선택 게이팅은 높은 중복을 초래했지만 심각한 전문가 기아 (78.6% 사망 전문가) 를 유발하여, 중복만으로는 라우팅 품질을 위한 충분한 지표가 아님을 시사합니다.
• 전문가 사용: 고정된 구조로 인해 라우팅 아티팩트를 초래하는 계층적 또는 PEER 라우터와 달리, AIR-MoE 는 구조적 제약 없이 유연한 전문가 사용을 가능하게 했습니다.

5. 중요성 및 주장

이 논문은 AIR-MoE 가 희소 전문가 혼합 아키텍처에서 라우팅을 위한 효과적인 메커니즘을 제공하며, 특히 전체 라우팅의 계산 패널티 없이 세분화된 전문가의 사용을 가능하게 한다고 주장합니다.

• Drop-in Replacement: 이 방법은 기본 모델 아키텍처나 손실 함수에 수정이 필요 없이 표준 라우터를 대체하는 drop-in 솔루션으로 작동합니다.
• 유연성: 전문가 매개변수에 구조적 제약을 피함으로써 AIR-MoE 는 고정된 그룹화나 프로덕트-키 구조가 필요한 방법보다 토큰 - 전문가 종속성을 더 유연하게 학습할 수 있게 합니다.
• 확장성: 이 접근법은 coarse shortlisting 비용이 분산되고 라우팅된 토큰 수와 무관하기 때문에 큰 유효 배치 regimes 에서 유리하게 확장됩니다.

저자들은 한계를 인정하며, FLOP 감소는 상당하지만 실제 월클럭 속도 향상은 메모리 트래픽과 세분화 MoE 고유의 색인 오버헤드를 완화하기 위한 하드웨어 지원 및 최적화된 구현에 달려 있다고 지적합니다. 또한 현재 구현은 차원 축소를 사용하지 않으며, 이는 향후 연구의 직교적인 방향이라고 덧붙였습니다.