전문가 선택 경로 지정을 통한 확산 언어 모델의 적응형 연산 구현

초록

확산 언어 모델(DLM)은 병렬 비자기회귀 텍스트 생성을 가능하게 하지만, 기존 DLM 전문가 혼합(MoE) 모델은 자기회귀 시스템의 토큰 선택(TC) 라우팅 방식을 그대로 차용하여 부하 불균형과 경직된 계산 할당 문제를 야기합니다. 본 연구에서는 전문가 선택(EC) 라우팅이 DLM에 더 적합함을 보입니다: EC는 설계상 결정론적 부하 분산을 제공하여 TC보다 높은 처리량과 빠른 수렴을 달성합니다. EC의 전문가 용량이 외부적으로 제어 가능한 특성에 기반하여, 우리는 잡음 제거 단계에 따라 전문가 할당을 가변화하는 시간 단계 종속 전문가 용량을 도입합니다. 낮은 마스크 비율 단계에 더 많은 용량을 할당할 때 동일 FLOPs 기준으로 최고 성능이 일관되게 달성됨을 확인하였으며, 이에 대한 기계론적 설명을 제시합니다: 낮은 마스크 비율 컨텍스트의 토큰들은 학습 효율성이 수십 배 높아, 이러한 단계에 계산을 집중할 때 한계 수익이 가장 크게 나타납니다. 마지막으로, 기존 사전 학습된 TC DLM 모델이 라우터만 교체하여 EC 방식으로 개선 가능함을 보이며, 다양한 다운스트림 작업에서 더 빠른 수렴과 향상된 정확도를 달성합니다. 이러한 결과들은 EC 라우팅이 DLM MoE 모델의 우수한 패러다임임을 입증하고, DLM의 계산을 고정된 구조적 상수가 아닌 적응형 정책으로对待할 수 있음을 보여줍니다. 코드는 https://github.com/zhangshuibai/EC-DLM에서 이용 가능합니다.

English

Diffusion language models (DLMs) enable parallel, non-autoregressive text generation, yet existing DLM mixture-of-experts (MoE) models inherit token-choice (TC) routing from autoregressive systems, leading to load imbalance and rigid computation allocation. We show that expert-choice (EC) routing is a better fit for DLMs: it provides deterministic load balancing by design, yielding higher throughput and faster convergence than TC. Building on the property that EC capacity is externally controllable, we introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step. We find that allocating more capacity to low-mask-ratio steps consistently achieves the best performance under matched FLOPs, and provide a mechanistic explanation: tokens in low-mask-ratio contexts exhibit an order-of-magnitude higher learning efficiency, so concentrating compute on these steps yields the largest marginal return. Finally, we show that existing pretrained TC DLMs can be retrofitted to EC by replacing only the router, achieving faster convergence and improved accuracy across diverse downstream tasks. Together, these results establish EC routing as a superior paradigm for DLM MoE models and demonstrate that computation in DLMs can be treated as an adaptive policy rather than a fixed architectural constant. Code is available at https://github.com/zhangshuibai/EC-DLM.

전문가 선택 경로 지정을 통한 확산 언어 모델의 적응형 연산 구현

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

초록

Support