BEAM: Dynamic MoE 라우팅을 위한 Binary Expert Activation Masking

MoE에서 고정 Top-K 라우팅은 토큰 복잡도에 따라 불필요한 전문가를 활성화시키며 계산 비효율을 초래한다. BEAM은 Lightweight한 mask router를 도입해 Top-K 후보에서 불필요한 전문가를 비활성화하고, STE로 학습하는 엔드-투-엔드 구조를 통해 토큰별로 활성화된 전문가 수를 동적으로 조절한다. vLLM에 CUDA 커널로 플러그인 가능하며, 중간 수준의 Sparsity에서 정확도 손실을 2% 내외로 유지하고 MoE 레이어 FLOPs를 최대 85%까지 감소시킨다. 이는 대규모 MoE 모델의 실무 배포에서 추론 속도와 처리량을 크게 향상시킨다.

출발점: MoE에서의 비효율은 고정 Top-K 라우팅으로 인해 토큰의 난이도에 따라 불필요하게 많은 전문가가 활성화되는 점이다. BEAM은 Mask Router를 추가해 Top-K 후보군에서 활성화 여부를 이진화하고, 활성화 수를 토큰별로 다르게 결정하도록 학습한다. 이로써 토큰 정보량이 높은 경우에만 더 많은 전문가를 사용하고, 간단한 토큰은 적은 계산으로 처리한다. 학습은 STE를 통해 이진 마스크를 미분 가능하게 만들고, Lreg를 통해 TK 내에서의 희소성을 직접 조절한다. 결과적으로, 중간 수준의 sparsity에서도 성능 손실을 크게 줄이고, 극단적 sparsity에서도 실용적인 가속을 달성한다.

Step 1. Standard Top-K Routing: R(x)로 계산된 로짓 r에서 Top-K를 유지하고 Softmax를 적용하여 gi를 구한다. Step 2. Raw Mask Generation: m_hat = σ(xWm)로 입력 x에서 마스크 후보를 생성한다. Step 3. Binary Masking: mi = 1 if m_hat_i ≥ 0.5, 그렇지 않으면 0으로 이진화한다. Step 4. Masked Aggregation: g_hat = g ⊙ m, y = ∑ g_hat_i Ei(x)로 최종 출력을 얻는다. 학습 전략: L = Llm + αLbal + βLreg, Lreg = (1/K) ∑_{i∈TK} |m_hat_i|. STE를 이용해 역전파에서 m_hat의 그래디언트를 흘려보낸다. 초기화는 mask router를 0으로 하여 시작 시 Top-K 동작을 보존하고, 훈련 중 점진적으로 sparsity를 유도한다.

주요 벤치마크에서 BEAM은 mid sparsity에서 모든 모델에서 원래 성능의 98% 이상을 유지하고 Avg-K를 47–61% 감소시켰다. High sparsity에서 Avg-K가 0.14~0.56 범위까지 감소하였으며, Extreme sparsity에서도 대체로 우수한 성능 유지가 관찰되었다. 속도 측면에서 TPOT은 최소 1.3×, TTFT는 최대 1.1× 이상 개선되었고, Throughput은 1.4× 이상 향상되었으며, DeepSeekV2-Lite에서 decoding은 2.5× 가속되었다.

BEAM은 MoE Block 내부에 Mask Router를 추가하여 Top-K 후보에 대한 Binary Mask를 학습한다. 4단계 파이프라인으로 구성되며, 1) Top-K Routing, 2) Raw Mask 생성, 3) Binary Masking, 4) Masked Aggregation으로 출력 y를 얻는다. Mask Router의 파라미터는 Wm이며, m_hat = σ(xWm)으로 계산하고, m_i = 1{m_hat_i ≥ 0.5}로 이진화한다. 최종 g_hat = g ⊙ m이며, y = ∑ g_hat_i Ei(x)이다. 학습은 L = Llm + αLbal + βLreg 형태의 손실을 사용하며, Lreg은 TK 내의 |m_hat|의 L1 노름이다. STE를 통해 역전파에서 mask의 그래디언트를 흘려보내고, 초기엔 마스크를 0으로 초기화해 BEAM의 Top-K 동작을 보존한다. 이로써 Primary Router의 로드밸런싱과 분리된 Sparsification이 가능해져 Conflicts를 피하고 토큰-특이적 활성화 패턴을 학습한다.

A. BEAM은 세 MoE 아키텍처에서 평가되었으며, 다른 gating 구조나 전문가 세분화(more granular MoE)에서의 일반화는 추가 검증이 필요하다. B. 마스크 학습을 위한 post-training SFT가 필요해 추가 비용이 발생한다. C. Shared-expert 비율이 높은 모델에서는 BEAM의 속도 향상 여지가 상대적으로 작을 수 있다. D. 단일-GPU 설정에서의 벤치마크이므로 다GPU 환경에서의 상호작용은 추가 연구가 필요하다.