MoE의 공통 전문가가 다른 전문가의 학습을 가로막을 수 있을까?

배경과 문제

MoE는 트랜스포머의 피드포워드 네트워크(FFN)를 여러 개의 전문가로 대체하고, 각 토큰마다 일부 전문가만 활성화하여 계산량을 줄이면서 파라미터 수를 늘리는 기법이다. 최근 모델들은 두 가지 설계를 결합한다. 첫째, 미세 전문가 분할은 같은 활성화 비율에서 전체 전문가 수 E와 활성화 전문가 수 K를 함께 늘려 (예: (E,K)=(8,1) 대신 (64,8)) 가능한 전문가 조합 수를 조합적으로 크게 늘린다. 둘째, 공유 전문가는 모든 토큰에 대해 항상 활성화되어 공통 지식을 담당하고, 라우팅된 전문가는 특화된 지식을 학습하도록 돕는다. 논문은 DeepSeek와 Qwen 계열 등을 이런 설계를 채택한 사례로 든다.

문제는 이 두 설계가 탐색 관점에서 충돌한다는 점이다. 미세 분할은 라우팅 공간을 크게 만들므로 최적의 전문가 조합을 찾으려면 더 많은 탐색이 필요하다. 예를 들어 (E,K)=(8,1)에서는 조합 수가 8개이지만 (64,8)에서는 약 44억 개로 늘어난다. 반면 공유 전문가는 출력의 큰 부분을 차지하여 라우팅된 전문가의 영향력을 줄인다. 따라서 서로 다른 라우팅 결정을 내려도 표현 변화가 작아져 탐색이 비효율적이 된다. 이 긴장을 저자들은 '탐색 딜레마'라고 부른다.

기존 접근의 한계

이 논문에서 비교한 공유 전문가 없는 MoE에서는 라우팅이 학습됨에 따라 무작위 라우팅과의 성능 격차가 커지며, 이는 라우팅이 점점 전문화된다는 신호다. 그러나 공유 전문가가 있는 MoE에서는 이 격차가 학습 초기(20%)와 중기(50%) 사이에 거의 증가하지 않는다. 이는 공유 전문가가 표현을 지배하여 라우팅의 중요성을 떨어뜨리고, 라우팅 전문화를 늦춘다는 것을 시사한다.

저자들은 498M 전체 파라미터, 121M 활성 파라미터, E=66, K=8, 공유 전문가 2개인 모델을 20%, 50%, 100% 학습 시점에서 분석했다. 개별 전문가를 제거(출력을 0으로)했을 때 공유 전문가를 제거하면 PPL이 크게 증가하지만 라우팅 전문가를 제거하면 증가폭이 작았다. 또한 한 층씩 무작위 라우팅으로 바꾸는 실험에서 공유 전문가 모델은 초기·중간 층에서 PPL 증가가 훨씬 작았다. 이는 공유 전문가가 표현의 대부분을 차지하여 무작위 라우팅에도 강건하다는 뜻이다.

표현 다양성을 정량화하기 위해 두 가지 지표를 사용했다. 전문가 분산(ExpertVar)은 공유 전문가 출력은 유지한 채 라우팅 전문가 하나씩을 썼을 때 출력이 얼마나 달라지는지 측정한다. 라우터 분산(RouterVar)은 학습된 라우팅 분포에서 256개의 전문가 집합을 샘플링했을 때 출력의 변동을 측정한다. 두 지표 모두 공유 전문가가 있을 때 모든 층에서 더 작았다. 특히 RouterVar는 약 10배 감소했다. 이 설정에서는 공유 전문가가 클수록 다른 경로를 택해도 표현이 덜 달라진다는 해석을 뒷받침한다.

연구방법을 단계별로

MirrorMoE는 탐색을 장려하는 대신 탐색 공간을 줄여 모델의 탐색 능력에 맞추는 접근이다. 먼저 라우터 가중치 G_l에 토큰 표현을 곱해 얻은 점수, 즉 로짓의 부호를 뒤집은 분포를 만든다. 원래 소프트맥스 분포 ξ는 큰 로짓을 선호하고, 부호 반전 분포 ζ는 작은 로짓을 선호한다. 이 둘을 (1-ρ)와 ρ로 섞어 새로운 확률 p를 만든다. 여기서 ρ는 0에서 1 사이의 값으로, 초기에는 0.4로 시작해 코사인 스케줄로 0까지 감소시킨다.

전문가 선택은 p에서 상위 K-S개를 고른다. K는 공유 전문가 S개까지 포함한 총 활성 수다. 따라서 원래 확률이 낮은 전문가도 선택될 수 있다. 그러나 집계(aggregation) 단계에서는 p가 아닌 원래 확률 ξ를 가중치로 사용한다. 이 때문에 저확률 전문가가 선택되어도 출력에 기여하는 가중치가 매우 작아 기능적으로 억제된다. 이렇게 하면 표현에 실질적으로 기여하는 정도를 나타내는 유효 전문가 수 K_eff가 줄어들어 탐색 공간이 축소된다.

K_eff는 선택된 전문가들의 ξ 가중치 합이 1이 되도록 다시 정규화한 뒤 계산한 지수 엔트로피로 정의된다. 모든 가중치가 균등하면 K_eff=K-S이고, 한 전문가만 1이면 K_eff=1이다. 100M 모델에서 학습 초기 MirrorMoE의 K_eff는 기준 모델보다 작았고, 선택된 전문가 중 최소 가중치는 0에 가까웠다. 이는 일부 전문가가 출력 기여가 거의 없어짐을 보여준다. 또한 층별로 K_eff가 다르게 나타나 토큰·층 적응적 억제가 일어남을 확인했다.

부하 균형 손실도 수정했다. 비교용 미러 라우팅 변형은 선택 확률 p를 균형화했지만, MirrorMoE는 원래 확률 ξ를 균형화한다. 이는 기능적으로 활성화된 전문가에 집중하도록 한다. 100M 모델의 제거 실험에서 이 수정이 PPL을 52.94에서 51.77로 낮췄다. 전체 훈련 절차는 각 반복에서 ρ를 코사인 어닐링하고, (3)과 (4)로 전문가를 선택하며, (5)로 집계한 후 작업 손실과 부하 균형 손실의 합으로 역전파한다.

결과와 숫자의 의미

주 실험은 Llama 스타일 MoE를 DCLM 데이터셋으로 처음부터 훈련했다. 모델 규모는 전체 파라미터 100M, 498M, 625M, 2.4B이고, 활성 파라미터는 각각 38M, 121M, 100M, 270M이다. K=8 (공유 2 + 라우팅 6), R=64 또는 128이다. 훈련 토큰은 4.4B~12.8B, 시퀀스 길이 2048, AdamW, 전역 배치 1024, H100 GPU 2대(100M) 또는 4대를 사용했다.

테스트 PPL은 모든 규모에서 MirrorMoE가 더 낮았다. 100M: 54.78→51.77, 498M: 28.31→27.36, 625M: 28.46→27.84, 2.4B: 21.17→20.50. PPL은 낮을수록 좋으며, 2.4B에서 0.67 감소했다. 하위 과제 6개(ARC-E, OBQA, HellaSwag, PIQA, BoolQ, SIQA) 평균 정확도는 2.4B에서 45.23→46.87로 1.64점 상승했다. 그러나 모든 개별 과제가 개선된 것은 아니다. SIQA는 38.33→38.28로 소폭 하락했다.

K를 6, 8, 10으로 바꾼 Figure 1(a)를 보면, 100M 모델에서 공유 전문가 없는 기준의 PPL은 56.49→54.71→54.26으로 낮아지고, 공유 전문가 있는 기준은 54.83→54.78→54.54로 개선 폭이 작다. MirrorMoE는 각 K에서 두 기준보다 낮지만 51.58→51.77→52.07로 오히려 올라간다. 이 세 설정 중 최저는 K=6이다. 따라서 이 방법이 기준보다 낫다는 결과와, K를 늘리면 항상 더 좋아진다는 주장은 구분해야 한다.

반복 시간의 평균은 MirrorMoE에서 약간 낮지만 오차 범위가 겹친다. 2.4B에서 25711.1±176.2ms → 25635.4±109.4ms, 100M에서 7871.3±65.9 → 7844.4±85.6ms이다. 통계 검정 없이 유의한 속도 향상이나 총 훈련 비용 절감을 결론 내릴 수 없다. 저자들은 추가 소프트맥스 연산에도 평균 시간이 늘지 않은 이유로 부하 균형 개선 가능성을 제시한다.

구체적인 예시

이해를 돕기 위한 비유로, 공통 업무를 거의 다 처리하는 담당자와 세부 업무를 나누는 여러 담당자가 있는 팀을 생각해 보자. 공통 담당자의 영향이 너무 크면 세부 담당자 조합을 바꿔도 결과 차이가 잘 드러나지 않을 수 있다. MirrorMoE는 초기에 일부 세부 담당자의 의견 반영 비중을 아주 작게 두고, 안정적으로 역할을 익힌 뒤 그 제한을 풀어 가는 것에 비유할 수 있다. 실제 알고리즘에서는 낮은 원래 라우팅 확률을 가진 전문가를 선택하되 합산 가중치를 작게 유지한다. 이 비유가 직원 수나 실제 계산량을 줄인다는 뜻은 아니다.

실제 실험에서 100M 모델, K=8일 때 학습 초기 MirrorMoE의 유효 전문가 수는 약 3개였고, 기준 모델은 5~6개였다. 층별로 보면 3번째 층은 유효 전문가가 1개에 가까워져 매우 공격적으로 억제했지만, 7~8번째 층은 4개 정도를 유지했다. 이는 모든 토큰과 층에 동일한 K를 적용하는 비적응적 스케줄(예: K를 4에서 8로 선형 증가)보다 MirrorMoE가 이 설정에서 더 나은 성능을 낸 기제를 이해하는 데 도움이 된다. 비적응적 스케줄의 PPL은 54.91(선형) 또는 54.19(코사인)로 MirrorMoE의 51.77보다 나빴다.

한계와 오해하기 쉬운 점

이 연구는 프리프린트이며 동료 심사 완료 여부는 확인하지 못했다. 실험은 최대 2.4B 파라미터, 짧은 사전 학습, 6개 상식 추론 과제에 국한된다. 초대형 상용 모델이나 후학습, 실제 질의응답 전반에 대한 개선은 검증되지 않았다. 또한 진단은 498M 모델 하나에서만 수행되었으므로 모든 현대 MoE에서 반드시 발생한다고 단정할 수 없다.

Figure 1(a)의 100M 모델에서 MirrorMoE의 PPL은 K=6, 8, 10에서 각각 51.58, 51.77, 52.07로 K=6에서 최저다. 본문 4.2절은 'K가 증가함에 따라 PPL이 일관되게 감소한다'고 주장하지만, 이 그림과 모순된다. 따라서 K 증가 효과가 모든 규모에서 회복되었다고 일반화하면 안 된다.

MirrorMoE는 전문가를 삭제하거나 건너뛰어 FLOP를 절약하는 방법이 아니다. 실제로는 K-S개의 전문가를 항상 선택하고 계산한다. 다만 집계 가중치가 작아 기능적으로 억제될 뿐이다. 또한 K_eff는 가중치 엔트로피 기반의 연속 지표로, 실제 실행되는 전문가 수나 계산량을 나타내지 않는다.

공유 전문가 없는 모델과의 비교 조건도 구분해야 한다. Figure 1의 비교에서는 전체 파라미터와 훈련 토큰을 대략 맞췄지만 Megatron-LM의 분할 제약으로 GPU 수와 전역 배치 크기가 달라졌다. Table 7은 별도의 비공유 전문가 실험으로 층 깊이도 약간 조정했으므로 Table 1과 수치를 직접 비교하기 어렵다. top-p 라우팅과의 비교(Table 8)도 활성 전문가 수가 학습 중 변하여 완전히 동일한 계산 비용은 아니다.

실제로 어떻게 활용할까

공유 전문가를 쓰는 MoE를 새로 사전 학습할 때 시험해 볼 방법이다. 논문은 구현 코드 링크를 제공한다. 적용에는 선택 확률과 집계 가중치를 구분하는 변경, ρ를 줄이는 스케줄, 부하 균형 손실 조정이 함께 필요하다. ρ0=0.4와 코사인 스케줄은 저자들의 100M 모델 비교에서 가장 좋았던 설정이므로, 새로운 모델에도 그대로 최적이라고 보기보다는 자체 검증의 출발점으로 삼을 수 있다.

공유 전문가가 없는 MoE에서는 이점이 균일하지 않았다. Table 7에서 97M 모델은 PPL 52.01→51.23으로 개선됐지만 614M은 28.42→28.49로 나빠졌고, 나머지 두 규모의 차이는 작았다. 그러므로 공유 전문가 유무와 모델 규모별로 기준 성능을 비교해야 한다. 현재 근거로는 훈련 속도 향상을 기대하기보다, 같은 활성 전문가 수에서 예측 성능을 높일 가능성을 검증하는 활용이 타당하다.