라우터는 전문가의 기하학적 구조를 학습한다: 희소 전문가 혼합 (Sparse Mixture-of-Experts) 모델에서의 기하학적 결합

희소 전문가 혼합 (Sparse Mixture-of-Experts, SMoE) 모델은 언어 모델을 효율적으로 확장할 수 있게 해주지만, 라우팅 (routing)이 소수의 전문가에게 집중되어 붕괴될 수 있고 보조 부하 분산 손실 (auxiliary load-balancing losses)이 전문화를 저해할 수 있어 학습이 여전히 어렵습니다. 이러한 난제들에 착안하여, 본 연구에서는 SMoE의 라우팅 결정이 기계론적으로 어떻게 형성되는지 연구합니다. 첫째, 우리는 라우터와 그에 대응하는 전문가 사이의 기하학적 결합 (geometric coupling)을 밝혀냅니다. 주어진 토큰에 대해, 선택된 전문가를 위한 라우터 가중치와 해당 토큰을 처리하는 전문가 가중치는 스칼라 계수만 다를 뿐 동일한 입력 방향을 따라 그래디언트 (gradients)를 받습니다. 따라서 일치하는 라우터-전문가 방향은 동일한 라우팅된 토큰 이력을 축적합니다. 이러한 이론적 결합은 라우팅 역학에서도 경험적으로 나타납니다. 처음부터 학습된 1B SMoE에서, 더 높은 라우터 점수는 더 강력한 전문가 뉴런 활성화를 예측하며, 이는 라우팅 결정이 선택된 전문가 내부에서 거울처럼 반영됨을 보여줍니다. 다음으로, 우리는 보조 부하 분산 (auxiliary load balancing)이 라우터-전문가 기하학적 결합에 미치는 영향을 분석하며, 이러한 손실이 입력 방향 그래디언트를 라우터 가중치 전체로 분산시킴으로써 이 구조를 깨뜨리고, 서로 다른 라우터 방향을 거의 3배 더 유사하게 만든다는 것을 보여줍니다. 마지막으로, 우리는 각 전문가가 자신에게 라우팅된 은닉 상태 (hidden states)의 이동 평균을 유지하고 코사인 유사도 (cosine similarity)를 기반으로 토큰을 할당하는 파라미터가 없는 온라인 K-Means 라우터를 통해, 효과적인 라우팅을 위한 기하학적 결합의 중심성을 입증합니다. 보조 손실 (auxiliary-loss) 및 손실 없는 균형 (loss-free balancing)과 비교했을 때, 이 라우터는 약간의 퍼플렉시티 (perplexity) 증가만으로 가장 낮은 부하 불균형을 달성하였으며, 이는 기하학적 결합이 라우터가 학습하는 내용의 상당 부분을 포착하고 있음을 나타냅니다. 종합적으로, 우리의 결과는 라우터가 어떻게 효과적인 분업을 지원하는 할당 기하학 (assignment geometry)을 형성하는지 설명합니다.