Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

1. 문제 제기 (Problem)

• 기존 DLM 의 MoE 한계: 최근 대규모 언어 모델은 파라미터 수를 늘리면서도 추론 비용을 낮추기 위해 혼합 전문가 (Mixture-of-Experts, MoE) 아키텍처를 채택하고 있습니다. 기존 DLM 들은 자가회귀 (autoregressive) 모델에서 유래한 토큰 선택 (Token-Choice, TC) 라우팅 방식을 그대로 사용하는데, 이는 DLM 의 구조적 특성과 맞지 않습니다.
• 부하 불균형 (Load Imbalance): TC 방식에서는 각 토큰이 독립적으로 전문가를 선택하므로, 특정 전문가에 부하가 집중되고 다른 전문가는 유휴 상태가 되는 불균형이 발생합니다. 이를 해결하기 위해 보조 손실 함수 (auxiliary loss) 를 사용하지만, 이는 추가 연산 비용을 발생시키고 주 학습 목표와 경향성을 상충시킬 수 있으며, 부하 균형을 완벽하게 보장하지 못합니다.
• 고정된 계산 할당: DLM 은 비인과적 (non-causal) 인 어텐션을 사용하여 모든 토큰을 동시에 처리하므로, TC 방식의 '토큰이 전문가를 선택'하는 구조보다는 '전문가가 토큰을 선택'하는 구조가 더 적합함에도 불구하고, 이를 체계적으로 연구한 사례가 부족했습니다.

2. 방법론 (Methodology)

A. 전문가 선택 (Expert-Choice, EC) 라우팅 도입

• 원리: TC 방식과 반대로, 각 전문가가 고정된 수의 토큰을 선택하도록 합니다.
• 장점:
  • 결정론적 부하 균형: 설계 단계에서 각 전문가가 처리할 토큰 수를 고정하므로, 보조 손실 함수 없이도 완벽한 부하 균형을 달성합니다.
  • 처리량 향상: GPU 간 대기 시간 (straggler effect) 이 제거되어 훈련 및 추론 속도가 크게 향상됩니다.
  • DLM 에 최적화: DLM 은 모든 토큰을 한 번에 볼 수 있으므로, 전문가가 전역적으로 토큰을 선택하는 EC 방식과 구조적으로 완벽하게 호환됩니다.

B. 시간 단계 의존적 전문가 용량 (Timestep-Dependent Expert Capacity)

• 개념: 확산 과정의 각 시간 단계 (denoising step) 는 마스킹 비율 (masking ratio) 이 다르므로, 학습 난이도와 필요한 계산량이 다릅니다. EC 방식은 전문가의 용량 (capacity) 을 외부에서 제어할 수 있으므로, 시간 단계에 따라 할당되는 계산량을 동적으로 조절할 수 있습니다.
• 전략: 마스킹 비율이 낮은 단계 (대부분의 토큰이 이미 노출되어 있고, 소수만 예측해야 하는 단계) 에 더 많은 계산 자원을 할당하는 선형 역방향 (Linear-Reverse) 스케줄링을 제안했습니다.
  • $k(r) = k_{min} + (k_{max} - k_{min}) \cdot (1 - r)$
  • 여기서 $r$은 마스킹 비율이며, $r$이 작을수록 (마스킹이 적을수록) 더 많은 전문가 용량이 할당됩니다.

3. 주요 기여 (Key Contributions)

1. EC 라우팅의 우월성 입증: DLM MoE 훈련에서 EC 라우팅이 TC 라우팅보다 부하 균형, 처리량 (throughput), 수렴 속도 측면에서 압도적으로 우수함을 실험적으로 증명했습니다.
2. 적응적 계산 정책 제안: DLM 의 확산 단계별 특성을 분석하여, 마스킹 비율이 낮은 단계에 더 많은 계산 자원을 집중시키는 것이 최적의 성능을 낸다는 것을 발견하고 이를 구현했습니다.
3. 학습 효율성 메커니즘 규명: 저마스킹 비율 단계에서 토큰의 학습 효율성 (수렴 속도) 이 고마스킹 비율 단계보다 약 10 배 이상 높다는 것을 발견했습니다. 이는 문맥이 풍부할 때 예측을 정교화하는 데 추가 자원을 투입하는 것이 가장 큰 한계 수익 (marginal return) 을 준다는 것을 의미합니다.
4. 기존 모델의 호환성 (Retrofitting): 사전 훈련된 TC 기반 DLM 모델을 아키텍처 변경 없이 라우터 (router) 만 교체하여 EC 방식으로 전환할 수 있음을 보여주었습니다. 이를 통해 미세 조정 (fine-tuning) 시 더 빠른 수렴과 향상된 정확도를 달성했습니다.

4. 실험 결과 (Results)

• 훈련 효율성:
  • EC 방식은 TC 방식보다 약 2 배 빠른 속도 (10.6 시간 vs 20.7 시간) 로 동일한 손실 (loss 3.75) 에 도달했습니다.
  • GPU 당 처리량 (TFLOP/s) 이 TC 대비 1.5~2.1 배 향상되었습니다.
  • GPU 메모리 사용량의 표준 편차가 TC(3.6 GB) 에서 EC(0.0 GB) 로 균일해져, 병목 현상이 사라졌습니다.
• 동적 용량 스케줄링:
  • OpenWebText 및 Nemotron-CC 데이터셋 실험에서, 마스킹 비율이 낮은 단계에 자원을 집중하는 'Linear-Reverse' 스케줄러가 고정 용량 (Static) 이나 다른 스케줄러보다 최소 Perplexity를 기록했습니다.
  • 8B-A1B 규모 모델에서도 동적 EC 가 정적 EC 보다 모든 평가 지표 (Perplexity, MMLU, ARC-Challenge) 에서 우위를 보였습니다.
• 미세 조정 (Fine-tuning) 성능:
  • 사전 훈련된 TC 모델 (LLaDA-MoE) 을 EC 로 전환한 후 GSM8K, HumanEval, MedQA 등 다양한 태스크에서 미세 조정을 수행했습니다.
  • EC 는 TC 보다 1.3~1.5 배 빠른 추론 속도를 보였으며, 동적 EC 는 최종 정확도도 가장 높았습니다 (예: HumanEval Pass@1 58.6% vs TC 53.9%).

5. 의의 및 결론 (Significance)

이 논문은 확산 언어 모델 (DLM) 에서 **계산 자원 할당을 고정된 아키텍처 상수가 아닌 적응적 정책 (adaptive policy)**으로 다루어야 함을 강력하게 주장합니다.

• 아키텍처 패러다임 전환: DLM 의 비인과적 특성을 활용하여, TC 대신 EC 라우팅을 표준으로 삼음으로써 부하 불균형 문제를 근본적으로 해결하고 효율성을 극대화할 수 있음을 보였습니다.
• 계산의 지능적 분배: 확산 과정의 모든 단계가 동일한 계산량을 필요로 하지 않으며, 특히 문맥이 풍부한 저마스킹 비율 단계에 자원을 집중하는 것이 학습 효율성을 높인다는 통찰을 제공했습니다.
• 실용성: 기존에 배포된 TC 기반 DLM 모델을 라우터 교체만으로 쉽게 고도화할 수 있어, 실제 산업 적용에 매우 유용합니다.

결론적으로, 이 연구는 DLM 의 MoE 아키텍처 설계에 있어 Expert-Choice 라우팅과 시간 단계 기반 적응적 계산이 새로운 표준이 되어야 함을 제시하며, 더 빠르고 정확한 대규모 언어 모델 개발의 길을 열었습니다.