MirrorMoE가 공유 전문가 MoE의 라우팅 탐색 딜레마를 줄인다

무엇인가

이 논문은 희소 전문가 혼합(Mixture-of-Experts, MoE) 모델에서 최근 표준이 된 두 설계, 즉 세분화된 전문가 분할(fine-grained expert segmentation)과 공유 전문가(shared experts)를 함께 쓸 때 생기는 라우팅 탐색(exploration)의 딜레마를 다룬다. 세분화된 분할은 같은 활성화 비율에서 더 많은 전문가 조합을 가능하게 해 표현력을 높이지만 라우팅 공간을 조합적으로 폭발시킨다. (E,K)=(8,1)에서는 전문가 부분집합이 8개뿐이지만 (E,K)=(64,8)에서는 C(64,8)≈4.4×10^9개로 몇 자릿수나 늘어난다. 반면 모든 토큰에 항상 활성화되는 공유 전문가는 학습된 토큰 표현을 지배해서, 라우팅된 전문가를 바꿔도 레이어 출력의 변화가 작아진다. 대안 경로들이 서로 구분되지 않으니 탐색이 약해지고, 세분화가 제공하는 추가 표현력이 실제로는 실현되지 못한다는 것이 저자들의 진단이다.

어떻게 동작하나

저자들은 이 진단을 세 가지 실험으로 뒷받침한다. 활성 121M·전체 498M 파라미터, E=66, K=8 모델에서 공유 전문가 2개(S=2)를 쓰는 MoE와 안 쓰는 MoE(S=0)를 같은 데이터·토크나이저·설정으로 비교하고, 사전학습 20%·50%·100% 체크포인트를 봤다. 첫째, 전문가 하나의 출력을 0으로 만들고 혼란도(perplexity) 증가를 재는 어블레이션에서 공유 전문가를 지웠을 때가 라우팅 전문가를 지웠을 때보다 훨씬 큰 증가를 낳았다. 둘째, 학습된 라우터를 균등 랜덤 라우팅으로 바꿔도 공유 전문가 MoE의 초·중간 레이어에서는 혼란도 증가가 표준 MoE보다 훨씬 작았다. 셋째, 전문가 분산(ExpertVar, 단일 전문가만 활성화했을 때의 출력 차이)과 라우터 분산(RouterVar, 라우팅 분포에서 256개 전문가 집합을 샘플링했을 때의 출력 차이)을 재보면 공유 전문가가 있으면 모든 레이어에서 분산이 더 작고, RouterVar는 약 한 자릿수 줄어든다. 논문은 어떤 라우팅 정책이 만들 수 있는 최대 변화가 R·ExpertVar로 상계된다는 것도 보인다.

무엇과 다른가

이 딜레마의 실질적 결과도 수치로 제시된다. 100M에서 2.4B 전체 파라미터까지 4개 규모에서 R∈{64,128}, S=2로 두고 활성 전문가 수 K를 {6,8,10}으로 바꿨다. K가 작을 때는 공유 전문가 MoE가 비공유 MoE를 앞서지만 K가 커지면 그 이점이 줄거나 사라진다. K를 6에서 10으로 늘렸을 때 혼란도 개선은 비공유 MoE에서 {2.23, 1.32, 1.08, 0.87}인데 공유 전문가 MoE에서는 {0.29, -0.03, 0.58, -0.03}로 떨어져 두 규모에서는 오히려 음수가 된다. 즉 전문가를 더 켜서 얻는 표현력이 공유 전문가 아래에서는 실현되지 않는다. 또 비공유 MoE는 학습이 진행될수록 랜덤 라우팅과의 혼란도 격차가 커지는데(라우팅 전문화), 공유 전문가 MoE는 20%와 50% 체크포인트 사이에서 이 격차가 거의 늘지 않아 전문화가 느리다.

어떻게 쓰나

제안 방법 MirrorMoE는 탐색을 늘리는 대신 탐색 공간을 줄이는 쪽을 택한다. 라우팅 확률을 표준 소프트맥스 ξ와 로짓 부호를 뒤집은 소프트맥스 ζ의 볼록결합 p=(1-ρ)ξ+ρζ로 정의하고(이 부호 반전이 이름의 유래), 이 p로 상위 K-S개 전문가를 고른다. 핵심은 집계 단계로, 선택된 전문가에 p가 아니라 ξ를 가중치로 준다. 그러면 로짓이 낮아 선택된 전문가들은 활성화되지만 표현에 거의 기여하지 못하고 기능적으로 억제된다. 이 억제는 토큰과 레이어마다 ξ에 따라 달라지므로 적응적이며, 논문은 활성 전문가 중 실질적으로 기여하는 수를 정규화된 ξ의 지수 엔트로피로 정의한 K_eff로 이를 정량화한다. K_eff가 K-S면 모든 전문가가 똑같이 기여하고, 1이면 한 전문가가 출력을 결정한다. 학습 중에는 ρ를 코사인 스케줄로 0까지 서서히 줄여, 전문가 전문화가 어느 정도 안정된 뒤에 탐색 공간을 다시 넓힌다. 로드 밸런싱 손실도 p 대신 ξ를 균형 맞추도록 바꾼다.

전제와 한계

이론적으로도 적응성을 특성화한다. 어떤 전문가의 ξ가 ε→0으로 충분히 작고 ρ가 c/(1+c)보다 크면(여기서 c/R은 표준 라우팅이 뽑았을 최하위 전문가의 확률), MirrorMoE는 그 전문가를 활성화하되 집계에서 사실상 제거한다. 반대로 ξ_max/ξ_min ≤ κ=√((1-ρ)/ρ)이면 활성 집합이 표준 라우팅과 같아지고 어떤 선택된 전문가도 억제되지 않는다. 즉 라우터가 확신이 없을 때는 기존 라우팅처럼 동작한다. 저자들은 이를 탐색 외의 관점, 곧 토큰이 스스로 활성 전문가 수를 정하는 적응적 용량 배분으로도 해석한다.

실험은 DCLM 데이터셋과 Megatron-LM, NVIDIA H100에서 수행했고 코드는 공개돼 있다. 100M~2.4B 전체 파라미터의 공유 전문가 MoE 4종(R∈{64,128}, 라우팅 6개+공유 2개, K=8, 활성 파라미터 비중 16~38%)에서 MirrorMoE가 모든 규모에서 더 낮은 테스트 혼란도를 기록했다. 100M 모델 비교표에서는 MoE 54.78, 선형 스케줄 K̃ 54.91, 코사인 스케줄 K̃ 54.19에 대해 MirrorMoE가 51.77을 냈다. 2.4B 모델에서는 사전학습 성능 0.67포인트, 6개 상식 추론 다운스트림 과제 평균 1.64포인트를 개선했다. 레이어마다 소프트맥스가 하나 늘었는데도 반복당 학습 시간은 오히려 조금 줄었는데, 저자들은 로드 밸런스 개선 때문이라고 본다. 레이어 수와 전문가 수만 다른 MoE-498M-A121M과 MoE-625M-A100M 중 라우팅 공간이 더 큰 쪽에서 이득이 더 컸고, K∈{6,8,10} 실험에서는 MirrorMoE가 모든 경우에 비공유 MoE를 앞서면서 K가 커질수록 혼란도가 계속 낮아졌다. ρ 스케줄 어블레이션에서는 ρ0=0.4의 코사인 스케줄이 최고였고, 토큰·레이어 적응성 없이 K를 4에서 8로 늘리는 선형·코사인 베이스라인은 MirrorMoE에 못 미쳤다. 공유 전문가가 없는 MoE에 MirrorMoE를 적용하면 기존 학습과 비슷한 성능을 냈는데, 이는 공유 전문가가 없으면 딜레마 자체가 약하다는 분석과 일치한다.

개발자 관점에서 이 논문의 실용적 메시지는 두 가지다. 첫째, 공유 전문가를 쓰는 MoE에서 K를 키우는 것이 항상 이득이 아니라는 현상을 라우터 튜닝이나 로드 밸런싱 문제로만 보지 말고, 공유 전문가가 표현을 지배해 라우팅 탐색이 죽은 상태일 수 있다고 의심해볼 근거를 준다. 둘째, 처방이 값싸다. 라우팅 확률을 한 번 더 계산하고 집계 가중치만 ξ로 바꾸는 수준이라 추가 파라미터가 없고, 논문 기준으로 학습 시간도 늘지 않았다. 다만 ρ 스케줄과 로드 밸런싱 손실 수정이 함께 들어가야 하며, 공유 전문가가 없는 MoE에는 적용할 이유가 없다는 점도 실험으로 확인돼 있다.

한계로 저자들이 명시한 것은 현재 접근이 탐색 공간을 줄이는 방향이라는 점이다. 더 나은 라우팅 전략 자체를 설계해 탐색을 효율화하는 것을 향후 과제로 남긴다. 또한 이 방법이 공유 전문가 MoE에 맞춰 설계됐다는 전제가 있으며, 비공유 MoE에서는 이득이 없다는 실험 결과가 그 근거다. 논문은 별도의 한계 절을 두지 않고 결론에서 이 향후 방향만 밝히고 있다.