기술 요약: 라우터는 자신의 전문가들의 기하학을 학습한다

문제 정의

희소 혼합 전문가 (SMoE) 모델은 토큰당 일부 파라미터만 활성화함으로써 언어 모델의 효율적인 확장을 가능하게 합니다. 그러나 이러한 모델을 학습시키는 데에는 상당한 어려움이 따릅니다. 개입이 없을 경우, 라우팅 메커니즘은 종종 붕괴되어 트래픽이 점점 더 축소된 전문가들의 하위 집합에 집중되는 (표현 붕괴) 현상이 발생합니다. 이를 완화하기 위해 실무자들은 균일한 전문가 활용을 강제하기 위해 보조 부하 균형 손실 (auxiliary load-balancing losses) 을 사용합니다. 부하 균형에는 효과적이지만, 이러한 손실은 전문가의 전문화를 감소시키는 것으로 알려져 있습니다. 본 논문이 다루는 핵심 문제는 SMoE 에서 라우팅 결정이 어떻게 형성되는지에 대한 기계적 이해의 부재와 보조 손실과 같은 일반적인 학습 개입이 라우터와 전문가 사이의 근본적인 역학을 어떻게 변화시키는지에 대한 이해 부족입니다.

방법론

저자들은 SMoE 아키텍처 내의 기울기 역학에 대한 기하학적 분석을 통해 문제에 접근합니다. 그들의 방법론은 다음 세 가지 주요 단계로 구성됩니다:

1. 이론적 기울기 유도: 저자들은 라우터 가중치 ($r_i$) 와 전문가 입력 측 가중치 ($W^{gate}_i$) 에 대한 역전파 업데이트를 유도합니다. 그들은 SMoE 레이어에 적용된 연쇄 법칙을 분석하여 매칭된 라우터 - 전문가 쌍과 선택되지 않은 전문가들에 대한 기울기 업데이트의 방향을 결정합니다.
2. 실증적 검증: OLMoE-mix-0924 데이터셋 (~50B 토큰) 에서 처음부터 학습된 1B 파라미터 SMoE 모델을 사용하여, 저자들은 라우터 점수와 선택된 전문가들의 내부 활성화 간의 상관관계를 측정합니다. 특히 라우터의 선호도가 반영되는지 확인하기 위해 "게이트 뉴런" 활성화 ($\sigma(W^{gate}_i x)$) 를 구체적으로 검토합니다.
3. 애블레이션 및 개입:
   • 보조 손실의 영향: 동일한 아키텍처를 가지지만 다른 균형 규칙을 사용하여 학습된 두 모델을 비교합니다. 하나는 표준 보조 부하 균형 손실 (예: Switch Transformer 스타일) 을 사용하고 다른 하나는 "손실 없는" 편향 기반 균형 방법을 사용합니다. 기하학적 붕괴를 평가하기 위해 라우터 가중치 벡터 간의 코사인 유사도를 측정합니다.
   • 센트로이드 라우터: 라우터가 본질적으로 입력 분포를 요약하도록 학습된다는 가설을 검증하기 위해, 파라미터가 없는 온라인 K-Means 라우터를 구현합니다. 이 설계에서 라우터 가중치는 라우팅된 은닉 상태의 지수 이동 평균인 비학습 가능한 센트로이드로 대체되며, 토큰은 코사인 유사도에 따라 할당됩니다. 이 라우터는 학습 가능한 라우팅 파라미터나 균형 기울기 없이 작동합니다.

주요 기여

1. 기하학적 결합의 발견: 이 논문은 라우터와 해당 전문가 사이의 근본적인 기하학적 결합을 밝혀냈습니다. 전문가 $i$ 로 라우팅된 토큰의 경우, 라우터 가중치 벡터 $r_i$ 와 전문가의 입력 측 가중치 모두 동일한 입력 방향 $x$ 에 비례하는 기울기 업데이트를 받습니다. 결과적으로 매칭된 라우터 - 전문가 쌍은 공유된 라우팅된 토큰 역사의 결합된 적분체로 진화합니다.
2. 결합의 실증적 확인: 학습된 1B SMoE 에서 저자들은 특정 전문가에 대한 라우터 점수가 높을수록 해당 전문가의 게이트 뉴런 활성화가 강해지는 것과 강한 상관관계가 있음을 보여줍니다. 이는 라우팅 결정이 단순한 외부 할당이 아니라 선택된 전문가의 내부 계산에 반영됨을 확인시켜 줍니다.
3. 기하학적 붕괴의 메커니즘: 이 연구는 보조 부하 균형 손실이 이러한 자연스러운 결합을 방해한다는 것을 보여줍니다. 이러한 손실은 (선택 여부와 관계없이) 모든 토큰에 대해 모든 라우터 가중치 벡터로 입력 지향 기울기를 전송함으로써, 서로 다른 라우터 방향이 서로 거의 3 배 더 유사해지게 만들어 전문가별 방향을 통합하고 전문화를 약화시킵니다.
4. 센트로이드 기반 라우팅: 저자들은 라우팅된 은닉 상태의 센트로이드를 명시적으로 추적하는 단순한 파라미터 없는 K-Means 라우터를 구현합니다. 이 라우터는 (손실 없는 및 보조 손실 기준 모델을 포함한) 모든 테스트된 변형 중 가장 낮은 부하 불균형을 달성하며, 퍼플렉시티는 약간만 증가합니다.

결과

• 기울기 정렬: 이론적 유도는 선택된 토큰에 대해 라우터와 전문가 가중치가 스칼라 계수만 다를 뿐 동일한 입력 방향 $x$ 를 따라 업데이트됨을 확인합니다.
• 활성화 상관관계: 1B 모델에서 라우터 점수와 전문가 게이트 뉴런 활성화는 단조로운 관계 ($\rho = 0.43$) 를 보여주어, 라우터의 기하학적 선호도가 전문가의 계산에 반영됨을 검증합니다.
• 보조 손실의 효과: 보조 부하 균형 손실로 학습된 모델은 라우터 가중치 벡터 간의 쌍별 코사인 유사도가 ($\mu \approx 0.63$) 손실 없는 모델 ($\mu \approx 0.13$ ~ $0.32$) 에 비해 현저히 높았습니다. 이는 보조 손실이 라우터의 기하학적 다양성을 붕괴시킴을 나타냅니다.
• K-Means 라우터 성능:
  • 부하 불균형: K-Means 라우터는 MaxVio (최대 상대 과부하) 가 0.037로, Loss-Free 기준 (0.084) 과 Aux-Loss 기준 (0.526) 보다 현저히 낮았습니다.
  • 퍼플렉시티: K-Means 라우터는 학습된 Loss-Free 라우터에 비해 퍼플렉시티가 약간 증가했습니다 (학습 PPL: 15.40 대 15.01), 약 2.6% 의 비용을 의미합니다.
  • 파라미터: K-Means 라우터는 0개의 학습 가능한 라우팅 파라미터를 사용하며, 실행 통계량 (센트로이드 및 편향) 에만 의존합니다.

중요성과 주장

이 논문은 기하학적 결합이 기울기 기반 학습 라우터가 학습하는 것의 상당한 구성 요소이며 단순한 부산물이 아니라고 주장합니다. 연구 결과는 다음과 같은 시사점을 제공합니다:

• 기계적 통찰: 라우터는 임의의 할당 규칙을 학습하는 것이 아니라, 효과적인 노동 분업을 지원하는 공유 기하학을 형성하기 위해 전문가들과 함께 진화합니다.
• 현재 관행에 대한 비판: 일반적인 보조 부하 균형 손실은 부하를 균형 있게 유지하지만, 라우터 방향을 통합하는 "간섭 기울기"를 주입하여 전문가 전문화를 가능하게 하는 기하학적 구조를 능동적으로 저하시킵니다.
• 설계 시사점: 향후 라우팅 방법은 이러한 자연스러운 라우터 - 전문가 기하학을 보존하는 데서 혜택을 볼 수 있습니다. 파라미터 없는 K-Means 라우터의 성공은 복잡한 기울기 기반 라우팅 가중치나 균형 손실 없이도 명시적인 센트로이드 추적이 라우터의 기능 대부분을 회복할 수 있음을 보여줍니다.

저자들은 그들의 작업이 라우터가 할당 기하학을 어떻게 형성하는지 설명하며, 라우터와 전문가 사이의 자연스러운 결합을 보존하는 것이 효과적인 SMoE 학습에 중요하다고 결론지었습니다. 그들은 한계를 인정하며, 실증적 결과는 단일 1B 구성에 기반하며, K-Means 라우터는 구성적 개념 증명이지만 남은 퍼플렉시티 격차로 인해 아직 학습된 라우터의 실용적인 대체제가 아니라고 지적합니다.