DMax: Aggressive Parallel Decoding for dLLMs Review

0. Introduction

Paper link

Code link

DMax를 그냥 “dLLM decoding을 조금 더 빠르게 만든 논문”으로 읽으면 핵심을 놓치기 쉽다. 이 논문의 진짜 문제의식은 parallel decoding 자체가 아니라, aggressive parallel decoding에서 왜 성능이 무너지는가에 있다. 기존 masked diffusion language model은 mask가 token으로 바뀌는 순간 그 예측을 사실상 commit한다. 그러면 초반 오답이 다음 step의 context로 굳고, block 단위로 많이 decode할수록 error accumulation이 커진다. dLLM이 이론적으로는 parallel generation에 유리해 보여도, 실제 decoding에서는 threshold를 보수적으로 둘 수밖에 없는 이유가 여기에 있다.

DMax는 이 병목을 training과 inference 양쪽에서 동시에 건드린다. training에서는 On-Policy Uniform Training, inference에서는 Soft Parallel Decoding을 도입한다. 둘을 합치면 model이 자기 예측을 다시 볼 수 있고, 틀린 token을 다음 iteration에서 수정할 수 있다. 이 논문의 진짜 가치는 “한 번에 더 많이 뽑는다”가 아니라, revision-friendly decoding state를 처음부터 학습시킨다는 데 있다.

한 줄 요약: DMax는 masked dLLM의 일방향 mask-to-token decoding이 만드는 error accumulation 문제를, On-Policy Uniform Training과 Soft Parallel Decoding으로 완화해 공격적인 parallel decoding에서도 품질을 최대한 유지하려는 논문이다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같음.

• dLLM의 병목을 “parallel token을 얼마나 많이 열 수 있나”가 아니라 오류를 얼마나 되돌릴 수 있나로 다시 정의한다.
• training objective와 inference state를 같이 바꾸기 때문에, 단순 threshold tuning이나 decoding heuristic보다 설계 의도가 더 선명하다.
• math reasoning과 code generation 둘 다에서 TPF는 크게 올리고 accuracy는 비교적 보존하는 operating point를 보여준다.
• low-parallelism에서도 accuracy가 오르는 결과를 함께 보여줘서, 이 방법을 단순 speed trick으로만 보기 어렵다.

1. Problem Setting

1-1. Problem definition

이 논문이 겨냥하는 문제는 diffusion language model의 practical parallel decoding gap이다.

dLLM은 한 step에 여러 token을 병렬로 예측할 수 있으니 autoregressive LLM보다 빠를 것처럼 보인다. 하지만 실제로는 그렇지 않다. 기존 masked dLLM은 한번 mask를 token으로 바꾸면 그 값을 다음 step의 고정 context처럼 사용한다. 그래서 초반에 잘못 나온 token이 뒤쪽 token prediction을 계속 오염시킨다.

이 문제는 특히 아래 상황에서 심해진다.

• confidence threshold를 낮춰서 한 번에 더 많은 token을 decode하려 할 때
• math reasoning처럼 앞쪽 실수가 뒤의 chain-of-thought 전체를 무너뜨릴 때
• code generation처럼 local syntax error가 이후 completion 전부를 흔들 때

즉 이 논문의 문제 설정은 “dLLM을 더 빠르게 만들자”가 아니라, parallel decoding을 키울수록 커지는 error accumulation을 어떻게 제어할 것인가에 가깝다.

1-2. Why previous approaches are insufficient

기존 접근은 크게 세 방향으로 나뉜다.

Type	Main idea	Limitation
Default masked dLLM decoding	confidence threshold를 기준으로 mask를 token으로 바꾸고 commit	한번 나온 오답을 다시 revise하기 어렵다
Better decoding heuristic	divide-and-conquer, threshold scheduling 등	decode policy는 좋아질 수 있지만 self-revision 능력 자체를 만들지는 못한다
Uniform diffusion training	random token corruption에서도 복원하도록 학습	실제 inference rollout과 noisy input 분포가 달라 train-inference gap이 크다

논문의 관점에서 가장 중요한 한계는 train-inference mismatch다. uniform diffusion training은 vocabulary에서 임의 token을 뽑아 noisy sequence를 만든다. 하지만 실제 decoding에서 model이 마주치는 noisy sequence는 랜덤 token 덩어리가 아니라, 자기 자신이 조금 틀리게 예측한 자연어 형태의 sequence다. 이 mismatch가 크면 model은 실제 rollout에서 필요한 self-correction을 잘 못 배운다.

또 하나 흥미로운 점은, 단순히 inference-time heuristic만으로는 문제가 안 풀린다는 것이다. token을 더 늦게 commit하거나 block을 더 잘 나누는 것은 도움은 되지만, model이 predicted token을 다시 입력으로 받았을 때 clean target으로 되돌리는 mapping 자체를 배운 것은 아니다. DMax는 바로 이 점을 푼다.

2. Core Idea

2-1. Main contribution

DMax의 핵심 기여는 두 가지다.

1. On-Policy Uniform Training
   • masked noisy sequence를 먼저 만들고,
   • 현재 model이 그 위에서 실제로 예측한 token들로 predicted noisy sequence를 만들고,
   • 두 입력 모두를 clean sequence로 복원하도록 학습한다.
2. Soft Parallel Decoding
   • intermediate state를 discrete token이 아니라,
   • predicted token embedding과 mask embedding의 confidence-weighted hybrid embedding으로 표현한다.
   • 이렇게 하면 이전 step의 uncertainty를 다음 step으로 넘길 수 있다.

이 두 설계를 합치면 model은 “mask만 잘 푸는 model”에서 “자기 예측이 섞인 상태에서도 clean token으로 되돌릴 수 있는 model”로 바뀐다.

2-2. Design intuition

이 논문의 직관은 매우 단순하다.

• 기존 masked dLLM은 token을 한번 commit하면 되돌리기 어렵다.
• 그래서 high parallelism에서 틀린 token이 누적된다.
• 그렇다면 model이 자기 자신의 noisy prediction을 입력으로 받아도 다시 복원하는 능력을 training에서 배워야 한다.
• 그리고 inference에서는 이 uncertainty를 discrete token이 아니라 soft state로 유지해야 한다.

이 관점에서 보면 DMax는 decoding trick보다 state design에 가깝다. model이 보는 intermediate state가 너무 hard하면 self-revision이 어렵고, 너무 random하면 generation이 불안정하다. DMax는 그 중간 지점으로 mask embedding과 token embedding 사이의 hybrid state를 선택한다.

3. Architecture / Method

3-1. Overview

Item	Description
Goal	aggressive parallel decoding에서도 dLLM accuracy drop을 최대한 줄이는 것
Base model	LLaDA-2.0-mini
Key training design	On-Policy Uniform Training
Key inference design	Soft Parallel Decoding
Core bottleneck diagnosis	binary mask-to-token commit가 error accumulation을 만든다
Key difference from prior work	decode heuristic만 바꾸지 않고 training distribution과 intermediate state를 함께 바꾼다

3-2. Module breakdown

1) Why binary commit is the bottleneck

기존 masked dLLM은 token이 한번 열리면 사실상 fixed context가 된다. parallel decoding을 공격적으로 하면 초반 오답도 여러 개가 동시에 열린다. 이후 step은 그 오답들을 condition으로 사용하므로 연쇄 오염이 생긴다.

이 논문이 좋은 이유는 이 failure mode를 흐릿하게 말하지 않는다는 점이다. 병목을 그냥 “parallel decoding이 어렵다”로 두지 않고, binary commit가 self-revision을 막는다고 명확하게 진단한다.

2) On-Policy Uniform Training

OPUT의 핵심은 noisy sequence를 random vocabulary token으로 만들지 않는 것이다. 먼저 clean sequence에서 일부 token을 [MASK]로 바꾼 masked noisy sequence를 만든다. 그다음 현재 model이 masked position에서 실제로 예측한 token을 sample해서 predicted noisy sequence를 만든다.

이 predicted noisy sequence는 개념적으로 아래처럼 쓸 수 있다.

\[x_t^{pred} = \begin{cases} x_t^{mask}, & \text{if token is not masked} \\ \hat{x}, & \text{if token is masked and sampled from the current model} \end{cases}\]

그다음 masked noisy sequence와 predicted noisy sequence 둘 다를 clean sequence로 복원하게 학습한다. 핵심 objective는 아래처럼 정리된다.

\[L_{onpolicy} = L_{mask} + L_{pred}\]

여기서 중요한 점은 predicted noisy sequence가 current model의 실제 output distribution에서 나온다는 점이다. 이게 train-inference gap을 줄인다. paper에서도 OPUT만으로 GSM8K accuracy가 threshold 0.5 setting에서 78%에서 90%로 올라간다고 설명한다.

3) Soft Parallel Decoding

OPUT만으로도 self-correction은 생기지만, 매우 aggressive한 parallel decoding에서는 여전히 동시에 많은 오답이 생길 수 있다. 그래서 DMax는 inference state도 바꾼다.

SPD에서는 intermediate token state를 hard token으로 commit하지 않고, confidence에 따라 predicted token embedding과 mask embedding을 섞는다. 개념적으로 쓰면 아래와 같다.

\[h_j^{(t)} = p_j^{(t-1)} e(y_j^{(t-1)}) + (1 - p_j^{(t-1)}) e_{mask}\]

실제 paper의 Eq. 10은 여기에 renormalization을 추가한다. 단순 합산은 norm collapse를 만들 수 있기 때문이다. 즉 위 식은 이해를 위한 축약형으로 보는 편이 맞다.

이 설계의 의미는 분명하다.

• confidence가 높은 token은 token embedding 쪽에 더 가깝다.
• confidence가 낮은 token은 mask embedding 쪽에 더 가깝다.
• model은 다음 forward pass에서 어떤 위치가 아직 불확실한지 알 수 있다.

4) Contiguous prefix rule

SPD가 흥미로운 이유는 soft embedding 자체만이 아니다. block 안에서 token을 여는 방식도 꽤 중요하다.

각 decoding step에서 model은 masked region을 왼쪽에서 오른쪽으로 훑는다. 그리고 confidence가 tau_dec보다 높은 위치들 중 가장 긴 contiguous prefix만 token position으로 승격시킨다. 처음 low-confidence 위치를 만나면 그 오른쪽은 모두 mask로 남긴다. 만약 어떤 위치도 threshold를 못 넘으면, 가장 왼쪽 mask 하나는 강제로 연다.

이 rule의 의도는 명확하다.

• 오른쪽의 불안정한 미래 token이 왼쪽 prediction을 오염시키지 않게 한다.
• block 내부의 mask region을 contiguous하게 유지한다.
• decoding progress도 보장한다.

겉으로 보면 heuristic처럼 보이지만, ablation을 보면 꽤 중요한 설계다.

5) Block convergence criteria

block이 언제 끝났다고 볼지도 정의한다. paper의 기준은 두 가지다.

1. 모든 위치의 top-1 prediction이 두 step 연속 동일하다.
2. block의 모든 위치 confidence가 acceptance threshold보다 높다.

paper의 implementation에서는 tau_acc = 0.9를 사용한다. consistency가 primary signal이고, confidence criterion은 마지막 forward pass를 아끼는 쪽에 가깝다.

6) Why OPUT is a prerequisite for SPD

이 논문에서 가장 중요한 결론 중 하나는 SPD를 그냥 기존 model에 얹으면 안 된다는 점이다. Table 3을 보면 OPUT 없이 soft hybrid state를 쓰면 generation이 사실상 collapse한다.

이유는 간단하다. OPUT가 있어야 model이 mask embedding과 self-predicted token embedding 둘 다에서 clean target으로 가는 mapping을 배운다. 그래야 two-state interpolation이 의미가 있다. 그렇지 않으면 hybrid state는 그저 training 때 본 적 없는 이상한 input일 뿐이다.

4. Training / Data / Recipe

4-1. Data

DMax는 별도의 external high-quality response를 쓰지 않고 self-distillation만으로 training data를 만든다.

• math training prompt source
  • GSM8K trainset
  • PRM12K
  • Numina-Math subset
  • OpenThoughts subset
• code training prompt source
  • OpenCodeInstruct subset

Response는 모두 LLaDA-2.0-mini가 생성한 것을 target으로 쓴다. generation setting은 confidence threshold 0.95, block size 32, max length 2048이다. length budget 안에 끝나지 않는 incomplete generation은 버린다. 이렇게 얻은 training data 규모는 아래와 같다.

Split	Size
Math trajectories	0.7M
Code trajectories	1.0M

이 부분도 꽤 중요하다. paper의 주장 중 일부는 “외부 teacher를 붙이지 않아도, model 자신의 rollout만으로 parallel decoding trade-off를 개선할 수 있다”는 데 있기 때문이다.

4-2. Training strategy

paper에 나온 주요 recipe는 아래와 같다.

Item	Setting
Base model	LLaDA-2.0-mini
Training objective	OPUT
Mask ratio	0.75
Epochs	2
Batch size	8
Initial learning rate	2e-6
LR schedule	cosine
Block size	32
Training compute	8 H200 GPUs
Inference decoding	SPD + block diffusion
Acceptance threshold	0.9

추가로 구현상 중요한 디테일이 하나 더 있다. masked noisy sequence와 predicted noisy sequence를 한 iteration에서 jointly optimize하지 않고, 같은 epoch 안의 separate iteration으로 처리한다. paper 설명에 따르면 이건 extra memory overhead를 줄이기 위한 선택이다.

또한 model도 하나가 아니라 task별로 두 개를 둔다.

• DMax-Math
• DMax-Coder

즉 이 paper의 main result는 single general-purpose model보다는 task-specialized post-training recipe에 가깝다.

4-3. Engineering notes

1. train-inference gap을 data distribution으로 직접 줄인다
   • random corruption을 더 clever하게 만드는 게 아니라, actual rollout-like corruption을 training input으로 삼는다.
2. speed와 accuracy를 동시에 보려면 state까지 바꿔야 한다
   • OPUT만으로도 improvement가 있지만, aggressive regime에서는 SPD가 없으면 accuracy가 급격히 꺾인다.
3. block scheduler가 생각보다 중요하다
   • contiguous prefix rule은 단순 heuristic처럼 보이지만, right-side unstable token interference를 줄여준다.
4. evaluation metric도 하나가 아니다
   • paper는 TPF, TPS, accuracy 외에 AUP score도 함께 본다. 병렬 decoding은 속도 하나만 보면 오해하기 쉽기 때문이다.

5. Evaluation

5-1. Main results

대표적인 결과만 압축해서 보면 아래와 같다.

Benchmark	Base TPF / Acc.	dParallel-SFT TPF / Acc.	DMax TPF / Acc.
GSM8K	2.04 / 92.6	2.79 / 92.3	5.48 / 92.1
MATH500	2.58 / 75.8	3.42 / 75.8	5.94 / 75.4
HumanEval-Instruct	4.38 / 84.2	5.12 / 76.8	7.36 / 83.5
MBPP-Instruct	2.71 / 80.6	3.66 / 74.7	5.86 / 79.2

이 표만 봐도 메시지는 꽤 분명하다. DMax는 대부분의 benchmark에서 TPF를 거의 2배 수준으로 밀어 올리면서 accuracy는 base model에 가깝게 유지한다. 특히 dParallel-SFT가 TPF를 올리는 대신 정확도를 더 크게 잃는 구간과 비교하면, DMax의 설계가 왜 self-revision capability에 초점을 두는지 잘 보인다.

throughput 쪽 claim도 강하다. abstract 기준으로 paper는 2 H200 GPUs에서 batch size 1일 때 평균 1,338 TPS를 보고한다. main text에서는 평균 TPF가 2.8에서 6.2로 올라간다고 정리한다.

또 하나 흥미로운 결과는 low-parallelism에서도 accuracy가 오른다는 점이다.

Benchmark	Base TPF / Acc.	DMax TPF / Acc.
GSM8K	2.04 / 92.6	3.54 / 93.4
MATH500	2.58 / 75.8	3.45 / 78.0
Minerva-Algebra	3.01 / 91.4	4.96 / 93.6
HumanEval-Instruct	4.38 / 84.2	4.58 / 87.2
MBPP-Instruct	2.71 / 80.6	3.58 / 83.4

이건 단순하다. DMax가 속도만 올리는 게 아니라, 자기 prediction을 다시 검토하는 능력 때문에 일부 benchmark에서는 accuracy도 오른다.

5-2. What really matters in the experiments

1) 이 paper는 single-point result보다 trade-off curve가 핵심이다

DMax의 장점은 최고 accuracy를 새로 세웠다는 데 있지 않다. 핵심은 TPF를 밀어 올렸을 때 accuracy가 얼마나 천천히 떨어지는가다.

paper의 Figure 4에서 이 차이가 더 잘 보인다. 예를 들어 MATH500에서는 비슷한 TPF 6.5 근처에서 DMax가 71.6% 이상 accuracy를 유지하는 반면, original model은 15.2%까지 떨어진다. MBPP도 비슷하다. 유사한 TPF에서 DMax는 79.2%인데 original model은 2.3%까지 무너진다.

즉 이 논문은 best accuracy paper가 아니라 parallelism-accuracy operating point paper로 읽는 것이 맞다.

2) OPUT가 진짜 핵심이다

paper는 Uniform Diffusion Training baseline도 같이 넣는다. 결과는 꽤 냉정하다. decoding speed도 좋아지지 않고, accuracy는 심하게 떨어진다. Table 1에서 GSM8K 68.7, MATH500 33.6, HumanEval-Instruct 15.2처럼 quality drop이 매우 크다.

이 결과가 말하는 것은 명확하다. UDLM처럼 보이는 objective를 쓴다와 실제 rollout을 닮은 noisy input으로 학습한다는 전혀 다른 이야기다. DMax의 핵심 novelty는 uniform corruption 그 자체가 아니라 on-policy rollout이다.

3) SPD는 cosmetic add-on이 아니다

Table 3이 정말 중요하다. GSM8K에서 tau_dec = 0.0 같은 매우 공격적인 setting을 보면,

• original model은 7.86 TPF지만 accuracy는 0.9
• OPUT만 쓰면 5.89 TPF, 68.2 accuracy
• OPUT + contiguous prefix는 5.98 TPF, 69.6 accuracy
• OPUT + hybrid embedding은 6.01 TPF, 90.4 accuracy
• full DMax도 6.01 TPF, 90.4 accuracy

즉 hybrid embedding이 들어가야 aggressive regime이 실제로 살아난다. 그리고 OPUT 없이 SPD만 올리면 0.0 accuracy로 collapse한다. 이 실험은 paper의 method story를 거의 그대로 증명한다.

4) contiguous prefix rule도 실제로 도움이 된다

겉보기엔 가장 단순한 heuristic인데, 결과는 무시하기 어렵다. OPUT only 대비 contiguous prefix를 넣으면 tau_dec = 0.5에서 5.14 / 90.1이 5.28 / 91.3으로 좋아진다. 큰 차이는 아니지만, SPD 같은 핵심 설계 옆에서 block scheduler도 성능에 유의미하게 기여한다는 뜻이다.

5) low-parallelism gain이 생각보다 중요하다

Table 2가 이 논문의 숨은 강점이다. 많은 decoding 가속 논문은 speedup을 얻는 대신 원래 quality를 최대한 덜 잃는 데 집중한다. 반면 DMax는 low-parallelism regime에서 accuracy까지 오른다. HumanEval-Instruct는 84.2에서 87.2, MBPP-Instruct는 80.6에서 83.4로 올라간다.

이 결과는 DMax를 단순 decoding policy가 아니라 iterative self-revision post-training으로 볼 수 있게 만든다.

6) convergence criterion 해석도 실용적이다

Table 4를 보면 consistency criterion만으로 GSM8K 5.13 / 92.1, confidence criterion만으로는 2.28 / 92.2, 둘 다 쓰면 5.48 / 92.1이 나온다. 즉 consistency가 primary termination signal이고, confidence criterion은 마지막 step을 더 일찍 끊어 TPF를 추가로 확보하는 보조 장치에 가깝다.

6. Limitations

1. main evidence가 LLaDA-2.0-mini 한 family에 묶여 있다
   • 결과가 꽤 설득력 있지만, 다른 dLLM backbone에서도 같은 개선 폭이 그대로 나올지는 아직 열려 있다.
2. task-specialized model 중심이다
   • paper의 main model은 DMax-Math와 DMax-Coder다. 하나의 general DMax가 여러 domain을 동시에 커버하는 결과는 본문에서 충분히 다루지 않는다.
3. throughput claim은 hardware와 framework 의존적이다
   • TPS 수치는 dInFer, 2 H200 GPUs, batch size 1 기준이다. 실제 serving stack이나 batch regime이 바뀌면 절대 수치는 달라질 수 있다.
4. self-distillation ceiling이 있다
   • training target이 전부 base model의 own generations이므로, quality ceiling 역시 base model behavior에 묶일 가능성이 있다.
5. comparison scope가 dLLM 내부에 집중되어 있다
   • 이 논문은 dLLM decoding trade-off를 푸는 paper이므로, strong AR speculative decoding과의 종합적인 serving cost 비교는 직접 다루지 않는다.
6. paper 자체도 working in progress 성격이 남아 있다
   • arXiv abstract page comments에 “Working in progress”가 명시되어 있으므로, 후속 version이나 code update에서 일부 detail이 바뀔 수 있다.

7. My Take

7-1. Why this matters for my work

DMax의 가장 중요한 메시지는 간단하다. parallel decoding의 병목은 threshold가 아니라 revision capability라는 점이다.

실무에서 iterative generator를 빠르게 만들려 할 때 흔히 하는 실수는 decode policy만 더 공격적으로 바꾸는 것이다. 하지만 model이 틀린 intermediate state를 다시 복원하는 능력을 학습하지 않았다면, aggressive decoding은 거의 항상 collapse risk를 동반한다. DMax는 이 문제를 아주 정직하게 보여준다.

또 하나 중요한 포인트는 training objective design이다. random corruption보다 on-policy corruption이 더 낫다는 메시지는 dLLM 밖에서도 재사용 가치가 크다. model이 deployment에서 보게 될 intermediate state와 training input distribution이 다르면, self-correction은 잘 안 생긴다.

7-2. Reuse potential

1) on-policy noisy state로 학습하기

이 아이디어는 dLLM에만 갇히지 않는다. iterative editing, refinement, planner-reviser 구조에서도 model의 actual rollout state를 training input으로 쓰는 편이 더 자연스럽다.

2) hard state 대신 soft state를 유지하기

SPD의 핵심은 token을 늦게 commit하는 것이 아니라, uncertainty를 state 안에 보존한다는 점이다. 이건 sequence generation뿐 아니라 iterative structured prediction에도 쓸 수 있는 아이디어다.

3) block scheduler 설계

contiguous prefix rule은 simple하지만 유용하다. partial state를 왼쪽에서 오른쪽으로 점진적으로 안정화하고, 불안정한 미래 위치가 현재 prediction을 망치지 않게 한다는 아이디어는 꽤 일반적이다.

4) low-parallelism quality gain을 같이 보기

가속 논문을 평가할 때 speedup만 보지 않고, low-parallelism에서 accuracy가 오르는지 보는 습관도 중요해 보인다. DMax는 이 지점을 잘 보여준다.

7-3. Follow-up papers

• LLaDA-2.0
• dParallel
• Hierarchical Decoding
• Dependency-Aware Parallel Decoding via Attention for Diffusion LLMs

8. Summary

• DMax는 masked dLLM의 binary mask-to-token commit가 만드는 error accumulation을 핵심 병목으로 본다.
• OPUT는 random corruption 대신 model의 own rollout에서 나온 noisy sequence로 학습해 train-inference gap을 줄인다.
• SPD는 predicted token embedding과 mask embedding의 hybrid state를 통해 uncertainty를 다음 step으로 전달한다.
• 실험에서 DMax는 GSM8K, MATH500, HumanEval-Instruct, MBPP-Instruct 전반에서 TPF를 크게 올리면서 accuracy를 비교적 잘 보존한다.
• 이 paper의 진짜 가치는 “더 공격적인 parallel decoding”보다, revision을 전제로 intermediate state를 다시 설계했다는 데 있다.