Multi-Mask Diffusion Language Models for Few-Step Generation

Background Masked diffusion language models (MDMs, the LLaDA/Dream family) generate by starting from an all-[MASK] sequence and progressively revealing tokens. Consistency distillation for few-step generation has succeeded in image and audio diffusion, and recent work has tried porting the same idea to language models.

배경 마스크 확산 언어모델(MDM, LLaDA/Dream 계열)은 시퀀스 전체를 [MASK] 상태에서 시작해 점진적으로 토큰을 드러내는 방식으로 생성한다. 최근 few-step(적은 스텝) 생성을 위한 컨시스턴시 증류가 이미지·오디오 확산에서 성공을 거두면서, 언어모델에도 같은 아이디어를 적용하려는 시도가 이어지고 있다.

Problem Every forward trajectory in an MDM collapses to the same single fully-masked state, regardless of which clean token it started from -- there is no way to distinguish terminal states. Consistency distillation depends on terminal-state diversity (entropy) to distill against, and MDMs structurally lack it.

문제 MDM의 모든 순방향(forward) 경로는 결국 동일한 완전-마스크 상태로 수렴한다 — 어떤 클린 토큰에서 출발했든 종결 상태가 구분되지 않는다는 뜻이다. 컨시스턴시 증류는 종결 상태의 다양성(엔트로피)에 의존하는데, MDM은 그 다양성을 원천적으로 갖지 못한다.

Prior limitations A recent alternative that avoids this degeneracy (uniform-state diffusion) keeps terminal states diverse, but makes it harder to tell clean tokens from noise, reportedly hurting modeling quality and training efficiency -- it pays for terminal diversity by giving up masking's clean/noise separability.

기존 한계 이 퇴화를 피하기 위한 최근 대안(uniform-state 확산)은 종결 상태를 다양하게 유지하지만, 대신 클린 토큰과 노이즈를 구분하기가 더 어려워져 모델링 품질과 학습 효율이 떨어진다고 보고된다 — 마스킹 구조가 주는 이점(클린/노이즈의 명확한 구분)을 포기하는 대가를 치른다.

Goal Keep the clean/noise separability that masking provides, while also recovering the terminal-state diversity that consistency-style few-step generation needs.

목표 마스킹 구조가 주는 클린/노이즈 구분의 장점은 유지하면서, 컨시스턴시 스타일 few-step 생성에 필요한 종결 상태 다양성을 함께 확보하는 것이 목표다.

Method MultiMDM expands the single mask token into a set of M masks (M much smaller than the vocabulary). Each clean token is pre-assigned its own designated mask; the forward process first pushes a token toward its own mask, then gradually mixes it across the full mask set. This decouples 'how fast a token leaves the visible vocabulary' (alpha_t) from 'how concentrated the masked probability mass stays on its own designated mask' (beta_t) -- a bit like not just dumping every item into one bin, but labeling each item so it first goes into its own labeled bin before gradually mixing with all the others. The paper also derives a closed-form ELBO objective supporting continual training from a pretrained single-mask MDM checkpoint.

방법 MultiMDM은 마스크 토큰을 하나가 아니라 집합(M개, M은 어휘 크기보다 훨씬 작음)으로 확장한다. 각 클린 토큰에는 지정된 마스크가 하나씩 배정되며, 순방향 과정은 먼저 토큰을 자신의 지정 마스크 쪽으로 밀었다가, 점차 전체 마스크 집합으로 섞이게 한다. 이는 '토큰이 얼마나 빨리 보이는 어휘에서 사라지는가'(α_t)와 '마스크된 확률질량이 자기 지정 마스크에 얼마나 집중되어 있는가'(β_t)를 분리하는 것과 같다 — 마치 물건을 그냥 하나의 큰 상자에 버리는 대신, 각 물건에 라벨을 붙여 처음엔 자기 라벨이 붙은 상자로, 나중엔 전체 상자로 서서히 섞이게 하는 것과 비슷하다. 사전학습된 단일-마스크 MDM 체크포인트로부터의 연속 학습(continual training)을 지원하는 닫힌 형태의 ELBO 학습 목적함수도 함께 유도한다.

Key idea The key move is that *which* designated mask a token resolves to is itself informative -- it becomes the terminal-state diversity a consistency distillation target needs, while the masking structure (clean/noise separability) is preserved. As a side effect, the backward process gains a drafting capability: predicting which mask a position will resolve to, before refining it into the actual clean token. For distillation, the paper also proposes a discrete-state consistency scheme using a shared-Gumbel coupling to reduce pathwise entropy.

핵심 아이디어 핵심은 '어느 지정 마스크로 향하는가'라는 정보 자체가 종결 상태의 다양성이 된다는 것이다 — 마스킹 구조(클린/노이즈 구분)는 그대로 유지하면서도, 컨시스턴시 증류가 필요로 하는 엔트로피가 생긴다. 부산물로 역방향(backward) 과정이 '어느 마스크로 귀결될지'를 먼저 예측한 뒤 클린 토큰으로 정제하는 드래프팅(drafting) 능력을 얻는다. 증류 단계에서는 경로별 엔트로피를 줄이기 위한 shared-Gumbel 결합을 쓰는 이산-상태 컨시스턴시 증류 기법도 함께 제안한다.

Validation Pretraining (from-scratch vs. continual training from a pretrained MDLM checkpoint) is run on a 170M-parameter DiT backbone on OpenWebText and LM1B, comparing unconditional GenPPL against DUO, MDLM, and CANDI baselines across a range of step budgets K (e.g. 4/8/16/32/64) under both entropy-aligned and temperature-aligned conditions.

검증 170M 파라미터 DiT 백본으로 OpenWebText·LM1B에서 사전학습(무작위 초기화 vs. 사전학습된 MDLM으로부터의 연속 학습, 두 레시피)을 수행하고, 스텝 예산 K(예: 4/8/16/32/64)에 걸쳐 entropy-aligned·temperature-aligned 두 조건에서 비조건부 GenPPL을 DUO·MDLM·CANDI 베이스라인과 비교한다.

Results The continual-training variant (MultiMDM-cont, 50K additional steps on a 150K-step pretrained MDLM) wins most clearly in the low-step-budget (K) regime -- e.g. on OpenWebText entropy-aligned at K=4, MDLM scores 721.1 versus MultiMDM-cont's 558.8. The gap narrows or reverses as K grows -- meaning the advantage concentrates exactly in the few-step regime this paper targets.

결과 연속 학습 변형(MultiMDM-cont, 15만 스텝 학습된 MDLM에서 5만 스텝만 추가 학습)이 특히 적은 스텝(K) 구간에서 가장 뚜렷하게 앞선다 — 예: OpenWebText entropy-aligned K=4에서 MDLM 721.1 대비 MultiMDM-cont 558.8. K가 커질수록(스텝이 많아질수록) 방법들 간 격차는 좁혀지거나 뒤집히는 경향을 보인다 — 이 논문이 목표로 하는 정확히 그 영역(few-step)에서 이득이 가장 크다는 뜻이다.

Comparison Compared against DUO (the uniform-state alternative), MDLM (standard single-mask), and CANDI under matched 170M DiT backbone, tokenizer, and training budget (CANDI is omitted from the LM1B table since its official implementation has no LM1B setup). Of the two training recipes tested, continual training consistently outperforms training from scratch.

비교 동일한 170M DiT 백본·토크나이저·학습 예산 아래 DUO(uniform-state 대안)·MDLM(표준 단일-마스크)·CANDI와 비교한다(CANDI는 LM1B 설정을 공식 구현이 지원하지 않아 LM1B 표에서는 제외). 두 학습 레시피(처음부터 vs. 연속 학습) 중 연속 학습 쪽이 일관되게 더 강한 결과를 보인다.

Significance From an efficient-language-model-inference standpoint, this paper pinpoints the exact root cause (no terminal entropy) behind masked diffusion's structural inability to support consistency-style few-step generation, and fixes it with a minimal structural change (a set of masks instead of one) -- a contribution on the decoding-step-count axis rather than compression, with the potential to cut serving cost roughly in proportion to steps.

의의 효율적 언어모델 추론 관점에서, 이 논문은 마스크 확산의 컨시스턴시 few-step 생성 불가능성이라는 구조적 결함의 근본 원인(종결 엔트로피 부재)을 정확히 지목하고 최소한의 구조 변경(마스크를 집합으로)으로 해결한다 — 압축(경량화)이 아니라 디코딩 스텝 수 자체를 줄이는 축의 기여이며, 서빙 비용을 스텝 수에 정비례해 낮추는 잠재력이 있다.

Limitations How the number of designated masks M is chosen, and how sensitive results are to it, is not pinned down with specific numbers within the scope this review draws from (abstract, introduction, method, and experiments sections) -- an appendix ablation is referenced in the table of contents but not read here -- not stated in the paper at this review's depth. All results are at 170M scale; whether the advantage transfers to LLaDA/Dream-scale (7-8B) or instruction-tuned/reasoning models is untested here (reviewer judgment).

한계 지정 마스크 개수 M을 어떻게 정하는지, 결과가 M에 얼마나 민감한지는 이 리뷰가 참조한 초록·서론·방법·실험 섹션 범위에서는 구체적 수치로 확인되지 않는다(부록에 관련 소거 실험이 있다고 목차에 언급되어 있으나 본문은 미확인) — 논문 내 명시 없음(이 리뷰 범위 기준). 모든 결과가 170M 규모이며, LLaDA/Dream급(7-8B) 또는 instruction-tuned/추론 모델로 전이되는지는 다뤄지지 않는다(리뷰어 판단).

Future work No explicit future-work items are identifiable within the scope this review draws from -- not stated in the paper.

향후 과제 이 리뷰가 참조한 범위에서 저자들이 명시적으로 밝힌 future-work 항목은 확인되지 않는다 — 논문 내 명시 없음.

Released resources No public code repository link is identifiable within the scope this review draws from (the abstract page) -- no public release verified.

자원 공개 공개 코드 저장소 링크는 이 리뷰가 참조한 범위(초록 페이지)에서는 확인되지 않았다 — 공개 링크 확인 안 됨.