MoE는 부하 균형만으로 부족하다: DRMoET의 라우팅 강건성 개선

들어가며

Mixture-of-Experts(MoE) 트랜스포머는 토큰마다 소수의 전문가만 활성화해 계산량을 크게 늘리지 않고 모델 용량을 확장합니다. 그러나 희소 라우팅에는 숨은 신뢰성 문제가 있습니다. 라우터가 토큰을 전문가들에게 균등하게 분배하더라도, 각 전문가가 실제로 배정받은 입력을 충분히 학습했다는 뜻은 아니기 때문입니다.

「Distributionally Robust Mixture-of-Experts Training」은 이 문제를 겨냥해 DRMoET를 제안합니다. 핵심 아이디어는 전문가 사용량을 고르게 만드는 데서 멈추지 않고, 라우팅이 완벽하지 않을 때 손실이 커지는 경로를 직접 최적화하는 것입니다.

방법: 손실이 큰 라우팅 결과에 집중

기존의 부하 균형 기법은 전문가별 토큰 수를 조절해 특정 전문가에 요청이 몰리거나 일부 전문가가 거의 사용되지 않는 상황을 피합니다. 하지만 토큰 수의 균형은 전문가의 실제 처리 품질을 평가하지 않습니다.

DRMoET는 각 레이어에서 활성화된 전문가의 손실을 추적합니다. 입력 활성화에 따라 가중된 전문가 손실을 지수이동평균(EMA)으로 평활화하고, 엔트로피 정규화가 포함된 Softmax 규칙으로 전문가 분포를 업데이트합니다. 손실이 높은 라우팅 결과에 더 큰 학습 비중을 주는 방식입니다.

이 방식은 추론 때 더 많은 전문가를 활성화하는 방법이 아닙니다. 기존의 희소 Top-k 계산은 그대로 유지됩니다. 대신 학습 과정에서 현재 최상위 후보뿐 아니라 실제로 선택될 가능성이 있는 비Top 경로까지 더 잘 처리하도록 전문가를 훈련합니다. 따라서 라우터가 입력을 약간 잘못 분류해도 성능이 급격히 무너지지 않을 여지를 확보합니다.

실험에서 나타난 결과

연구진은 FLAME-MoE 학습 레시피를 사용해 총 746M과 10.3B 파라미터 규모에서 DRMoET를 평가했습니다. 두 규모 모두 표준 FLAME-MoE와 보조 손실을 사용하지 않는 균형 방식보다 다운스트림 평균이 개선됐다고 보고했습니다.

10.3B 파라미터 모델을 67B 토큰으로 학습한 설정에서는 7개 과제 평균이 기준값 0.6625에서 DRMoET 적용 후 0.6767로 상승했습니다. 보조 손실 없는 기준선은 0.6431이었습니다. 메커니즘 분석에서는 평균 손실을 거의 바꾸지 않으면서 전문가 손실의 분산을 낮춘 것으로 나타났습니다. 중간 순위 전문가를 강제로 선택해 오라우팅을 모사한 실험에서는 추가 손실이 4.3% 감소했고, 도메인별 전문가 전문화도 개선됐습니다.

의미와 남은 과제

이 결과는 MoE의 품질을 토큰 분배량만으로 판단하기 어렵다는 점을 보여줍니다. 트래픽이 균등해도 전문가 간 능력 차이가 크다면 라우터의 작은 판단 오류가 성능 저하로 이어질 수 있습니다. 따라서 이용률 균형과 함께, 후보 전문가들이 다양한 입력을 얼마나 안정적으로 처리하는지도 최적화할 필요가 있습니다.

DRMoET는 표준 희소 계산을 유지하는 플러그인형 학습 목표라는 점에서 실용적인 방향을 제시합니다. 다만 현재 보고된 결과는 특정 학습 레시피와 모델 규모, 과제에 기반합니다. 다른 라우팅 구조와 데이터 조합, 더 큰 학습 규모에서도 같은 효과가 나타나는지는 추가 검증이 필요합니다. 엔트로피 정규화와 손실 평활화 설정 역시 실제 안정성과 비용에 영향을 줄 수 있습니다.

출처: Hugging Face Daily Papers