용량 인식 추론으로 MoE 스트래글러 현상 완화

초록

Mixture‑of‑Experts(MoE) 모델에서 토큰이 특정 전문가에 과도하게 몰려 발생하는 “스트래글러 효과”를 해결하기 위해, 저자들은 두 가지 용량‑인식 기법을 제안한다. 첫 번째인 Capacity‑Aware Token Drop은 과부하 전문가의 토큰을 일정 용량 한도 내에서 버려 부하를 균등화하고, 두 번째인 Capacity‑Aware Expanded Drop은 저부하 전문가에게 추가 후보 전문가를 허용해 남은 용량을 활용한다. 실험 결과, 이 방법들은 언어·멀티모달 MoE 모델에서 평균 0.2% 성능 향상과 1.8배 이상의 추론 속도 개선을 달성한다.

상세 분석

본 논문은 전문가 병렬화 환경에서 MoE 레이어의 토큰‑전문가 할당 불균형이 전체 레이턴시를 좌우한다는 점을 “Straggler Effect”라는 용어로 정의한다. 저자들은 먼저 전문가당 최대 토큰 수를 γ·(\bar N) 로 제한하는 Capacity‑Aware Token Drop(CATD)을 제안한다. 여기서 (\bar N = tk/n) 은 평균 토큰 할당량이며, γ는 용량 팩터이다. 과부하된 전문가(J)에 대해, 각 토큰‑전문가 쌍에 대한 중요도 점수 (s_{ij}) 를 기반으로 상위 C(=γ·(\bar N)) 점수만 유지하고 나머지는 마스크한다. 점수 기반 선택은 “Score” 전략을 사용하며, 이는 토큰 순서에 민감하지 않아 안정적이다. 토큰을 버리는 비용은 전문가 연산에 비해 무시할 만큼 작으며, 실험에서는 OLMoE에서 30% 레이턴시 감소와 0.9% 성능 저하만을 보였다.

그러나 CATD만 적용하면 여전히 저부하 전문가가 용량 한도 이하로 남아 자원을 낭비한다. 이를 보완하기 위해 Capacity‑Aware Expanded Drop(CAED)을 도입한다. CAED는 동일 디바이스 내에서 각 토큰이 기존 top‑k 외에 추가 m개의 로컬 후보 전문가를 고려하도록 허용한다. 이때 각 후보 전문가에 대한 용량 제한은 동일하게 적용되어, 과부하 전문가에서 버려진 토큰이 저부하 전문가로 재분배될 수 있다. 결과적으로 전문가 활용률이 크게 상승하고, 모델의 표현력도 향상된다.

실험에서는 언어 모델(O​LMoE, Mixtral‑8×7B‑Instruct)과 멀티모달 모델 모두에 적용했으며, 특히 Mixtral에 CAED를 적용했을 때 평균 정확도 0.2% 상승과 1.85× 추론 속도 향상을 기록했다. 이미지 토큰에 대한 실험에서는 용량을 평균의 절반 수준으로 강하게 제한해도 성능 저하가 거의 없었으며, 이는 이미지 토큰의 중복성이 높다는 사실을 반영한다.

전체적으로 논문은 (1) 스트래글러 현상의 정량적 분석, (2) 용량 기반 토큰 드롭 메커니즘 설계, (3) 저부하 전문가 활용을 위한 확장 드롭 전략, (4) 다양한 베이스 모델과 태스크에 대한 광범위한 실증을 통해 제안 기법의 일반성을 입증한다. 또한, 기존 학습 단계에서의 균형 손실과 달리 추론 단계에서 추가 연산·통신 오버헤드 없이 부하를 조절한다는 실용적 장점이 강조된다.