💬 NLP

Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts

이 논문은 Mixture of Experts(MoE) 모델의 추론 과정에서 특정 전문가에게 부하가 집중되어 전체 속도가 저하되는 '스트래글러 효과(Straggler Effect)'를 해결하기 위해, 과부하된 토큰을 제한적으로 버리는 'Capacity-Aware Token Drop'과 유휴 전문가의 활용도를 높이는 'Capacity-Aware Expanded Drop' 기법을 제안하여 모델 성능 저하를 최소화하면서 추론 효율을 크게 향상시켰습니다.

원저자: Shwai He, Weilin Cai, Jiayi Huang, Ang Li

게시일 2026-02-10

📖 2 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Shwai He, Weilin Cai, Jiayi Huang, Ang Li

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

👨‍🍳 상황 설정: "맛집 거리의 불공평한 주문"

여러분이 아주 큰 쇼핑몰 안에 있는 '맛집 거리'를 관리하는 매니저라고 상상해 보세요. 이 거리에는 8개의 전문 식당(전문가, Experts)이 있습니다.

1번 식당은 '초밥 전문'
2번 식당은 '파스타 전문'
... 이런 식이죠.

손님(데이터, Tokens)이 들어오면, 안내 데스크(라우터, Router)가 손님의 취향을 보고 "초밥 드실 분은 1번으로 가세요!"라고 안내합니다.

문제는 여기서 발생합니다.
어느 날 갑자기 손님 100명 중 80명이 "오늘 초밥이 너무 먹고 싶어요!"라고 외칩니다. 그럼 1번 식당은 손님이 너무 많아 주문이 밀려 난리가 나고, 2번 파스타 식당은 손님이 한 명도 없어 직원들이 하품만 하며 놀고 있습니다.

이때, 전체 쇼핑몰의 운영 속도는 가장 느린 1번 초밥 식당의 속도에 맞춰집니다. 파스타 식당 직원들이 아무리 빨라도, 초밥 식당이 주문을 다 처리할 때까지 쇼핑몰 전체는 다음 단계로 넘어가지 못하고 멈춰 서 있게 되죠. 논문에서는 이 현상을 **'스트래글러 효과(Straggler Effect, 뒤처지는 자 효과)'**라고 부릅니다.

🛠️ 해결책 1: "적당히 거절하기" (Capacity-Aware Token Drop)

매니저인 여러분은 결단을 내립니다.
"1번 초밥 식당이 너무 힘들어하니, 너무 많이 밀리면 일부 손님은 그냥 돌려보내거나 다른 곳으로 안내하자!"

이것이 논문에서 말하는 **'Capacity-Aware Token Drop'**입니다.
식당마다 **'최대 수용 인원(Capacity)'**을 정해두는 거죠. 만약 초밥 식당이 감당할 수 있는 한계를 넘으면, 가장 중요도가 낮은 손님(가장 맛이 덜 궁금해 보이는 손님)을 정중히 거절합니다.

이렇게 하면 초밥 식당이 너무 오래 걸리는 것을 막을 수 있어, 쇼핑몰 전체의 회전율(인퍼런스 속도)이 엄청나게 빨라집니다. (논문에 따르면 속도가 약 30%나 빨라졌다고 해요!)

🛠️ 해결책 2: "빈자리 채워주기" (Capacity-Aware Expanded Drop)

그런데 단순히 손님을 돌려보내기만 하면, 손님들이 맛없는 경험을 할까 봐 걱정되죠? 그래서 두 번째 전략을 씁니다.

"초밥 식당이 꽉 찼다면, 바로 옆에 있는 '일식 퓨전 식당'이나 '해산물 식당'도 후보에 넣어보자!"

이것이 **'Expanded Drop'**입니다. 원래는 초밥만 먹으려던 손님에게 "초밥 식당이 지금 너무 바쁘니, 근처에 있는 다른 괜찮은 식당들도 한번 살펴보시겠어요?"라고 제안하는 것입니다.

이렇게 하면:

**놀고 있는 식당(저부하 전문가)**의 일감을 가져와서 효율을 높이고,
**손님(데이터)**은 아예 굶지 않고 다른 선택지를 가질 수 있어, AI의 성능(정확도)이 오히려 좋아지기도 합니다.

🚀 요약하자면?

이 논문은 AI가 일을 할 때 "어떤 놈은 너무 바쁘고, 어떤 놈은 너무 노는" 불균형 때문에 전체 속도가 느려지는 문제를 찾아냈습니다.

이를 해결하기 위해:

너무 바쁜 놈은 적당히 쳐내서(Token Drop) 전체 속도를 올리고,
노는 놈에게는 일을 나눠줘서(Expanded Drop) 효율과 실력을 동시에 잡았습니다.

결과적으로 AI는 더 빠르게(최대 1.85배), 그러면서도 더 똑똑하게 대답할 수 있게 된 것입니다! 🥳

[기술 요약] 용량 인식 추론: MoE 모델의 Straggler Effect 완화

1. 문제 정의 (Problem Statement)

Mixture of Experts (MoE) 구조는 입력 토큰마다 일부 전문가(Expert)만 활성화하여 효율성을 높이지만, 추론(Inference) 단계에서 심각한 토큰-전문가 할당 불균형(Imbalanced token-to-expert assignment) 문제를 겪습니다.

Straggler Effect (지연자 효과): 특정 전문가에게 토큰이 과도하게 몰리는 반면, 다른 전문가는 유휴 상태(Idle)가 됩니다. 전문가 병렬화(Expert Parallelism) 환경에서 전체 레이어의 지연 시간은 가장 부하가 큰 전문가(Straggler)의 계산 시간에 의해 결정됩니다. 즉, 부하가 적은 GPU가 부하가 많은 GPU를 기다리며 자원 낭비와 전체 지연 시간 증가를 초래합니다.
기존 방식의 한계: 학습 시 보조 손실 함수(Auxiliary loss)를 사용해 균형을 맞추려 하지만, 이는 추론 시 발생하는 동적인 불균형을 해결하지 못합니다. 또한, 부하가 높은 전문가를 복제하는 방식(예: DeepSeek-V3)은 추가적인 컴퓨팅 자원 오버헤드를 발생시킵니다.

2. 제안 방법론 (Methodology)

본 논문은 추가적인 연산 오버헤드 없이 추론 시점에 토큰 할당을 조절하는 Capacity-Aware Inference를 제안합니다.

① Capacity-Aware Token Drop (고부하 전문가 대응)

과부하된 전문가의 지연 시간을 제어하기 위해 전문가별 **최대 용량(Capacity, $C$ )**을 설정합니다.

메커니즘: 설정된 용량( $C = \gamma \bar{N}$ , 여기서 $\gamma$ 는 용량 계수)을 초과하는 토큰은 버립니다(Drop).
선택 기준 (Scoring Function): 단순히 순서대로 버리는 것이 아니라, **게이팅 점수(Gating Score)**를 기준으로 중요도가 낮은 토큰을 선별하여 버립니다. 이를 통해 모델의 성능 저하를 최소화하면서 지연 시간을 효과적으로 줄입니다.

② Capacity-Aware Expanded Drop (저부하 전문가 활용)

Token Drop만 적용할 경우, 여전히 용량이 남은 저부하 전문가들이 유휴 상태로 남는 문제가 있습니다. 이를 해결하기 위해 Expanded Drop을 도입합니다.

메커니즘: 토큰이 원래 선택한 Top- $k$ 전문가 외에도, **동일한 디바이스 내의 다른 로컬 전문가들( $m$ 개)**을 후보군에 추가합니다.
효과: 과부하된 전문가에서 버려질 뻔한 토큰들이 용량이 남은 로컬 전문가에게 재할당되도록 유도하여, 전문가 활용도(Utilization)를 높이고 모델의 표현력을 유지합니다.

③ Device-Level Capacity-Aware Inference (고급 변형)

전문가 개별 단위가 아닌 디바이스 단위로 용량을 관리합니다. 특정 전문가가 용량을 초과하더라도 디바이스 내 다른 전문가에게 여유 공간이 있다면 토큰을 버리지 않고 허용함으로써, 더 유연하고 성능 저하가 적은 추론을 가능하게 합니다.

3. 주요 기여 (Key Contributions)

현상 규명: MoE 추론 시 발생하는 'Straggler Effect'를 명확히 정의하고, 이것이 전체 지연 시간에 미치는 영향을 분석했습니다.
효율적 알고리즘 제안: 추가적인 모델 학습이나 복잡한 자원 복제 없이, 추론 시점에 토큰을 선별적으로 드롭하거나 재배치하는 두 가지 핵심 기법(Token Drop, Expanded Drop)을 제안했습니다.
멀티모달 확장성: 언어 모델뿐만 아니라 멀티모달 MoE 모델에서도 이미지 토큰의 중복성을 활용한 'Image-First' 전략 등을 통해 효과적으로 적용 가능함을 입증했습니다.

4. 실험 결과 (Results)

다양한 MoE 모델(OLMoE, Mixtral, DeepSeek-V2-Lite, Qwen 등)을 대상으로 광범위한 실험을 수행했습니다.

추론 속도 향상:
- Mixtral-8×7B-Instruct: Expanded Drop 적용 시 평균 성능은 0.2% 향상되면서도, 추론 속도는 1.85배 빨라졌습니다.
- OLMoE: Token Drop 적용 시 성능 저하를 0.9%로 억제하면서 30%의 속도 향상을 달기했습니다.
성능 유지 및 개선: Expanded Drop은 단순히 토큰을 버리는 것보다 저부하 전문가를 활용함으로써 모델의 정확도를 오히려 높이는 결과(예: Qwen1.5-MoE에서 0.9% 향상)를 보여주었습니다.
멀티모달 성능: 이미지 토큰의 중복성을 이용해 공격적인 용량 제한( $\gamma=0.5$ )을 가해도 성능을 유지하며 효율성을 극대화할 수 있음을 확인했습니다.

5. 의의 (Significance)

본 연구는 MoE 모델의 실전 배포(Serving) 시 가장 큰 병목 중 하나인 불균형한 연산 부하 문제를 매우 가볍고(lightweight) 효과적인 방식으로 해결했습니다. 특히, 하드웨어 자원을 최대한 활용하면서도 모델의 지능(Performance)을 희생하지 않는 실용적인 추론 최적화 경로를 제시했다는 점에서 높은 가치를 지닙니다.