REFLEX: 확산 언어 모델에서의 MoE 추론 재고 - 정제(refinement) 인식 컴퓨팅 할당

REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models

Mixture-of-experts (MoE) 모델은 각 토큰에 대해 일부 전문가만 활성화하여 파라미터 용량을 증가시킵니다. 이러한 조건부 계산 방식 덕분에 자기 회귀 언어 모델은 토큰 당 계산량이 비례적으로 증가하지 않으면서도 모델 용량을 확장할 수 있었습니다. 그러나 확산 언어 모델(DLM)에서는 각 디노이징 과정에서 모든 토큰 위치를 동시에 처리하지만, 기본적으로 고정된 토큰 선택 라우팅 방식을 통해 모든 토큰에 동일한 전문가 할당량을 부여하므로, 전문가 계산량과 정제 요구 사항 간의 불일치가 발생합니다. 따라서 DLM에서의 MoE 추론은 다양한 토큰 정제 상태에 대한 정제 인지 컴퓨팅 할당으로 이해되어야 한다고 주장합니다. 본 연구에서는 학습이 필요 없는 방법인 REFLEX (REfinement-aware FLEXible expert allocation)를 제안합니다. 이 방법은 기본 라우터를 변경하지 않고, 진화하는 정제 과정을 중심으로 전문가 계산을 재구성합니다. 특히, REFLEX는 블록 상대적인 정제 역할에 따라 계산량을 할당하는 세분화된 계층 구조를 도입하고, Frontier-Progress Score를 사용하여 활성 블록의 우선순위를 결정합니다. 두 가지 대표적인 MoE 기반 DLM인 LLaDA-MoE와 LLaDA2.0-mini에 대한 여러 벤치마크에서 REFLEX는 평균적으로 할당된 전문가 계산량을 15% 줄이는 동시에, 대부분의 벤치마크에서 생성 품질을 유지하거나 향상시킵니다. 또한, 자기 회귀 방식의 가변 전문가 라우팅 방법과 비교했을 때, REFLEX는 더욱 일관된 품질-계산량 균형을 제공하며, 이는 각 디노이징 과정 내에 존재하는 다양한 정제 요구 사항에 따라 전문가 계산량을 할당하는 것이 중요하다는 점을 뒷받침합니다.

Mixture-of-experts (MoE) models increase parameter capacity by activating only a small subset of experts for each token. This conditional-computation paradigm has enabled autoregressive language models to scale model capacity without a proportional increase in per-token computation. In diffusion language models (DLMs), however, each denoising forward jointly revisits all token positions despite their sharply different refinement demands, while the default fixed token-choice routing assigns them a uniform expert budget, creating a mismatch between expert computation and refinement demand. We argue that MoE inference in DLMs should therefore be viewed as refinement-aware compute allocation across heterogeneous token refinement states. We propose REFLEX (\textbf{RE}finement-aware \textbf{FLEX}ible expert allocation), a training-free method that keeps the default router unchanged while reorganizing expert computation around the evolving refinement process. Specifically, REFLEX introduces a coarse-to-fine hierarchy for expert-budget allocation that aligns computation with block-relative refinement roles while using the Frontier-Progress Score to resolve active-block priorities. Across multiple widely used benchmarks on two representative MoE-based DLMs, LLaDA-MoE and LLaDA2.0-mini, REFLEX reduces allocated expert computation by 15\% on average while preserving or even improving generation quality on most benchmarks relative to default routing. Compared with autoregressive-style variable-expert routing methods, REFLEX also yields a more consistent quality--computation trade-off, further supporting the importance of allocating expert computation according to the heterogeneous refinement demands exposed within each denoising forward.