요즘 대형 LLM을 서빙하다 보면 MoE(Mixture of Experts) 얘기를 피해 가기 어렵다. Mixtral, DeepSeek-V2, Qwen-MoE까지 — 추론 비용을 줄이려는 모델들이 MoE 아키텍처를 택하는 데는 이유가 있다. 토큰 하나를 처리할 때 전체 파라미터를 쓰지 말고, 일부 전문가(expert)만 활성화해서 연산량을 줄인다는 논리다.
Mixtral 8x7B가 전형적인 예다. 전체 파라미터는 47B이지만 한 토큰을 처리할 때 활성화되는 파라미터는 약 13B다. 어텐션 레이어(~5B)는 공유되고, FFN 부분은 8개 전문가 중 2개(top-2)만 활성화된다. 연산량은 13B 덴스 모델 수준이지만 모델이 담고 있는 지식의 용량은 47B다. 메모리에는 여전히 47B 전체를 올려야 하지만, GPU 연산 시간만큼은 훨씬 싸다.
문제는 이 "2개만 활성화" 결정을 어떻게 내리느냐다. 이 결정이 만드는 구조적 긴장이 MoE 서빙을 다루기 어렵게 만드는 원인이고, Mixture of Tokens(MoT)라는 아이디어가 등장한 배경이다.
MoE top-k 라우팅이 만드는 구조적 긴장
MoE에서 라우팅의 주체는 토큰이다. 각 토큰이 라우터(게이트 네트워크)를 통과하면서 어느 전문가에게 갈지 점수를 계산하고, 상위 k개 전문가를 선택한다. 선택받지 못한 나머지 전문가는 해당 토큰과 무관하게 지나간다.
이 구조에서 반복되는 첫 번째 긴장은 load imbalance다. 비슷한 유형의 토큰들이 같은 전문가를 선호하는 경향이 생기면서, 배치 내에서 특정 전문가에 토큰이 몰리고 나머지는 논다. 극단적으로는 expert collapse — 일부 전문가가 아예 토큰을 받지 못하는 상태 — 까지 간다. Switch Transformer(Fedus et al., 2022)는 보조 손실 없이 top-1 라우팅을 적용하면 이런 collapse가 발생한다고 보고했다.
이를 막기 위해 훈련 중에 auxiliary load balancing loss를 추가한다:
$$\mathcal{L}{aux} = \alpha \cdot N \sum{i=1}^{N} f_i \cdot P_i$$
$f_i$는 전문가 $i$가 받은 토큰의 비율, $P_i$는 라우터가 전문가 $i$에 할당한 확률의 비율이다. 두 분포가 고르게 맞춰지도록 유도하는 손실이다. 그런데 이 보조 손실은 훈련 때만 작동한다. 추론 시점에 실제로 들어오는 입력 분포가 훈련 분포와 다르면, 라우팅 편향이 다시 발생한다.
두 번째 문제는 gradient flow 단절이다. top-k hard routing — k개를 골라 나머지를 버리는 결정 — 은 불연속(non-differentiable)이다. 선택받지 못한 전문가에는 gradient가 흐르지 않는다. 반복 학습에서 자주 선택받는 전문가만 강해지고 약한 전문가는 점점 쪼그라드는 양극화가 생긴다. STE(Straight-Through Estimator)나 noisy top-k 같은 우회책이 있지만 근본 구조를 바꾸진 않는다.
라우팅 방향을 뒤집기 — Mixture of Tokens
Mixture of Tokens(MoT)는 NeurIPS 2024에서 발표된 논문(Antoniak et al.)으로, 라우팅의 주체를 바꾼다. 토큰이 전문가를 고르는 대신, 각 전문가가 배치 전체에서 어떤 토큰 조합을 처리할지 결정한다.
메커니즘은 이렇다. 컨트롤러(controller)가 배치 내 모든 토큰에 대해 각 전문가의 가중치 행렬 $\Phi \in \mathbb{R}^{E \times T}$를 계산한다($E$ = 전문가 수, $T$ = 배치 내 총 토큰 수). 전문가 $e$가 받는 입력은 단일 토큰이 아니라 토큰들의 선형 조합이다:
$$\tilde{x}e = \sum{t} \Phi_{e,t} \cdot x_t$$
출력도 마찬가지로, 토큰 $t$의 업데이트는 모든 전문가 출력의 선형 조합이다:
$$y_t = \sum_{e} \Phi_{e,t} \cdot \text{Expert}_e(\tilde{x}_e)$$
"cross-example aggregation"이라는 이름이 붙은 이유가 여기 있다. 서로 다른 시퀀스(example)에서 온 토큰들이 한 전문가에게 혼합되어 들어간다.
이 소프트 집계는 앞서 말한 두 긴장을 구조적으로 다르게 다룬다. 모든 전문가가 항상 활성화되기 때문에 gradient가 전 전문가로 균등하게 흐른다. 가중치 $\Phi$가 소프트하게 정규화되어 있으므로 특정 전문가로 토큰이 편중되는 상황 자체가 만들어지지 않는다 — load imbalance를 보조 손실 없이 구조적으로 막는다. 연산 전체가 연속(differentiable)이라 학습 안정성도 높아진다. 논문은 덴스 Transformer 대비 훈련 속도 3배 향상을 보고하며, 기존 MoE와 동등한 성능을 달성한다고 밝혔다.
두 구조가 다루는 문제의 차이
| 항목 | MoE (top-k) | Mixture of Tokens |
|---|---|---|
| 라우팅 주체 | 토큰 | 전문가 |
| 라우팅 방식 | hard (discrete) | soft (continuous) |
| expert collapse | 보조 손실로 억제 | 구조적 불가 |
| gradient flow | 선택된 전문가만 | 전체 전문가 |
| load balance | 보조 손실 필요 | 자동 |
| 추론 시 imbalance | 잔존 가능 | 없음 |
MoE의 auxiliary loss는 훈련과 추론 사이의 분리 문제를 완전히 해결하지 못한다. 특히 prefill/decode 분리(disaggregation) 환경의 decode 단계에서 배치당 토큰 수가 극히 작아질 때 imbalance가 심해진다. 어떤 전문가 GPU는 과부하, 어떤 전문가 GPU는 거의 쉬는 상황이 벌어진다.
MoT는 이 긴장이 처음부터 없다. 다만 배치 크기가 작으면 cross-example 혼합의 다양성이 줄어든다. 배치에 시퀀스가 1개뿐이면 "서로 다른 예시의 토큰을 섞는" 효과가 사라지고 시퀀스 내 토큰들만 섞인다.
추론 서빙 관점 — all-to-all 통신의 실제 비용
서빙 엔지니어 입장에서 MoE를 다루기 어려운 부분은 파라미터 수가 아니라 expert parallelism에서 발생하는 통신 비용이다. 전문가들이 여러 GPU에 분산되어 있으면 토큰을 해당 GPU로 보내는 dispatch와, 처리 결과를 다시 모으는 combine이 필요하다. 두 단계 모두 all-to-all 집합 통신으로 구현된다. forward pass에서만 두 번, 학습 시에는 backward까지 포함해 MoE 레이어 하나당 총 네 번 all-to-all이 발생한다.
실측치를 보면 이 비용이 얼마나 큰지 체감된다. SGLang에서 DeepSeek-V2-Lite를 8-GPU expert parallelism으로 추론했을 때, MoE 레이어 forward pass 레이턴시의 최대 59.2%가 expert parallelism 통신이었다는 측정값이 있다. 전문가 연산 자체보다 토큰을 GPU 사이에서 셔플하는 데 더 많은 시간이 드는 상황이다.
MegaBlocks(Gale et al., 2022)는 sparse MoE의 GPU 커널 수준 최적화로 기존 대비 1.8x~2.4x 학습 속도를 달성했지만, all-to-all 통신 자체를 없애는 건 아니다. 통신을 연산과 겹치는(overlapping) 방식으로 비용을 줄이는 것이다. vLLM은 NCCL EP, DeepEP HT/LL 등 여러 all-to-all 백엔드를 플러그인 형태로 지원하면서 이 비용을 다루고 있다.
MoT 구조에서 이 패턴이 어떻게 달라지는지는 아직 대규모 서빙 환경의 실측이 충분하지 않다. cross-example 집계는 어텐션과 유사한 dense 행렬 연산으로 구현할 수 있어서, 토큰을 다른 GPU로 dispatch하는 방식 대신 FlashAttention 계열 커널을 재사용할 여지가 있다. 그게 가능하다면 MoE의 고질적인 all-to-all 비용을 구조적으로 우회하는 경로가 생긴다.
continuous batching 환경에서도 차이가 있다. MoE에서 특정 전문가 GPU가 straggler가 되면 전체 배치 레이턴시가 그 전문가에 묶인다. MoT는 전문가마다 처리하는 혼합 토큰 수가 균등하게 설계되어 있어 이 문제가 구조적으로 발생하지 않는다.
현재 한계와 실용 판단
MoT가 서빙 환경에서 바로 쓰이기 어려운 이유는 실전 배포 사례가 부족하다는 점만이 아니다. cross-example 집계가 배치 내 서로 다른 예시를 섞기 때문에, KV cache를 활용한 prefill-decode 분리 구현이 복잡해진다. 논문은 autoregressive 생성과 호환된다고 밝히지만, vLLM의 PagedAttention, continuous batching 최적화들과 어떻게 맞물리는지는 별도 검증이 필요하다.
메모리도 고려해야 한다. 소프트 집계는 배치 전체에 걸친 가중치 행렬 $\Phi$를 유지해야 하므로, 배치 크기에 비례해 활성화 메모리(activation memory)가 늘어난다. hard routing MoE는 활성화된 k개 전문가 경로만 유지하면 되므로 이 비용이 훨씬 작다.
현재 MoE를 운용하면서 expert utilization을 측정하고 싶다면, 레이어별 토큰 분배 히스토그램이 가장 직접적인 지표다. vLLM에서 --enable-expert-parallel 플래그로 expert parallelism을 활성화한 뒤 라우터 로짓을 커스텀 훅으로 수집하거나, SGLang 내부 통계에서 per-expert token count를 꺼내면 실제 운용 환경의 imbalance를 확인할 수 있다.
MoT가 서빙 최적화 생태계 — PagedAttention, FlashAttention, continuous batching — 를 얼마나 빠르게 흡수하느냐가 이 아이디어의 실용성을 결정할 것이다.