11 분 소요

0. Introduction

Paper link

Draft-OPD는 speculative decoding에서 draft model을 어떻게 더 잘 학습할 것인가를 다루는 논문이다. 겉으로는 또 하나의 decoding acceleration 논문처럼 보이지만, 핵심은 속도보다 training distribution mismatch에 있다. 기존 EAGLE-3, DFlash 계열 draft model은 target model이 만든 trajectory 위에서 SFT를 한다. 그런데 실제 inference에서는 draft model이 자기 policy로 token block을 제안하고, target model은 그 block을 검증한다. 즉 학습 때 보는 prefix와 acceptance를 결정하는 prefix가 다르다.

이 논문이 던지는 질문은 단순하다.

“오프라인 SFT를 더 오래 돌리면 draft model이 계속 좋아질까?”

저자들의 답은 아니다에 가깝다. Figure 1에서 보듯이 초기 SFT 이후에는 accepted length가 빠르게 plateau에 걸리고, OPD prompt pool로 만든 target response에 다시 SFT를 해도 오히려 accepted length가 줄 수 있다. 문제는 데이터 양이 아니라, draft model이 실제로 틀리는 inference-time state를 학습하지 못한다는 점이다.

한 줄 요약: Draft-OPD는 target-assisted rollout으로 안정적인 continuation을 만들되, speculative verification에서 드러난 draft failure anchor를 replay하여 accepted token과 rejected token을 다르게 distill하는 on-policy draft-model post-training 방법이다.

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

  • speculative decoding 성능을 draft architecture만이 아니라 post-training objective 문제 로 본다.
  • OPD를 LLM student distillation이 아니라 speculative draft model 에 맞게 재설계한다.
  • accepted token과 rejected token에 서로 다른 KL 방향을 쓰는 설계가 꽤 실용적이다.
  • Qwen3 thinking mode처럼 긴 reasoning trace에서 draft model 학습이 왜 더 어려운지까지 연결한다.

이 논문의 좋은 점은 speculative decoding을 lossless inference trick으로만 보지 않는다는 데 있다. 속도를 좌우하는 것은 결국 target model과 draft model의 local distribution agreement인데, 그 agreement가 깨지는 지점을 학습 데이터로 다시 끌어오는 방식이 매우 직접적이다.

1. Problem Setting

1-1. Problem definition

Speculative decoding은 target model $p_theta$와 draft model $q_phi$를 함께 쓴다. draft model이 먼저 $K$개의 token block을 제안하고, target model이 이를 병렬로 검증한다. 검증을 통과한 prefix만 accept되므로, 실제 speedup은 target model이 한 번 verify할 때 평균 몇 개 token을 받아들이는지에 강하게 묶인다.

논문에서는 이 값을 accepted length, 즉 $tau$로 본다.

\[\hat{y}_{t+k} \sim q_phi( . \mid x_{<t}, \hat{y}_{t:t+k-1}), \quad k = 0, ..., K - 1\] \[\text{speedup} \propto \text{accepted length } tau\]

문제는 draft model을 보통 offline SFT로 학습한다는 점이다. 이 경우 모든 prefix는 target model이 만든 trajectory에서 나온다. 하지만 실제 speculative decoding에서는 draft model이 직접 token block을 만들고, 그 block이 target model에게 reject될 수 있다. 결국 acceptance를 결정하는 state는 target trajectory state가 아니라 draft-induced state다.

정리하면 이 논문의 문제 설정은 아래와 같다.

  • draft model은 target-generated trajectory에서 SFT된다.
  • inference에서는 draft model이 자기 token block을 제안한다.
  • accepted length는 draft-induced state에서 target model과 얼마나 잘 맞는지에 의해 결정된다.
  • 따라서 offline SFT만으로는 inference-time mismatch를 없애기 어렵다.

1-2. Why previous approaches are insufficient

기존 접근이 부족한 이유는 크게 3가지다.

첫째, offline SFT는 data distribution이 맞지 않는다. SFT가 보는 token prefix는 대부분 target model이 만든 깨끗한 trajectory다. 반대로 speculative decoding에서 중요한 것은 draft model이 실제로 제안했다가 틀린 token 주변의 state다.

둘째, standard OPD를 그대로 적용하기 어렵다. 일반 OPD는 student가 자기 policy로 full sequence를 rollout할 수 있다고 가정한다. 하지만 DFlash나 EAGLE-style draft module은 독립 autoregressive generator가 아니라 target hidden state와 verification에 기대는 short-span proposer에 가깝다. draft-only rollout을 강제로 만들면 반복적이거나 낮은 품질의 sequence가 나오기 쉽다.

셋째, target-assisted rollout만 쓰면 on-policy signal이 사라진다. target model이 잘못된 draft token을 버리고 continuation을 고쳐 버리기 때문에, 최종 sequence는 다시 target distribution을 따르게 된다. 이러면 rejected token이 주는 가장 중요한 error signal이 training objective에서 빠진다.

“draft model이 실제로 어디서 틀리는가”

이 질문에 답하지 못하면, 더 많은 SFT data나 더 긴 SFT training은 accepted length plateau를 근본적으로 넘기 어렵다.

2. Core Idea

2-1. Main contribution

Draft-OPD의 핵심 기여는 OPD를 speculative draft model에 맞게 바꾼 것이다. 구성은 3단계다.

  1. target-assisted rollout으로 stable continuation을 만든다.
  2. speculative verification 중 draft block start position을 anchor로 기록한다.
  3. anchor에서 draft proposal을 replay하고, accepted token과 rejected token을 분리해 distillation한다.

이 설계 덕분에 Draft-OPD는 두 가지를 동시에 잡는다.

  • rollout 자체는 target-assisted라서 품질이 무너지지 않는다.
  • replay는 draft-induced token block을 다시 보므로 rejected proposal까지 학습 신호로 들어온다.

즉 논문은 OPD의 원래 목표였던 on-policy state supervision을 유지하면서도, draft model이 full sequence를 독립적으로 생성하지 않아도 되게 만든다.

2-2. Design intuition

직관은 상당히 깔끔하다. speculative decoding의 병목은 target model과 draft model이 얼마나 자주 어긋나는가다. 그런데 어긋난 token은 verification 단계에서 바로 버려진다. 버려졌다는 사실 때문에 inference는 lossless가 되지만, training 입장에서는 가장 중요한 error case가 데이터에서 사라진다.

Draft-OPD는 이 지점을 anchor replay로 복구한다.

  • rollout 시점에는 target model correction을 받아 stable sample을 만든다.
  • 동시에 draft block이 시작된 위치를 저장한다.
  • 학습 시점에는 그 anchor에서 draft model이 냈던 proposal을 다시 펼쳐 본다.
  • target model도 같은 draft-generated prefix 위에서 log probability를 계산한다.

이렇게 하면 final verified trajectory만 학습하는 것이 아니라, verification 과정에서 실제로 드러난 failure state를 직접 학습한다.

이 논문에서 가장 중요한 표현은 accepted token보다 rejected token이다. accepted token은 이미 target과 어느 정도 맞는 자리다. rejected token은 accepted length를 짧게 만드는 직접 원인이다. Draft-OPD는 바로 그 rejected token을 버리지 않고 학습 objective의 중심에 둔다.

3. Architecture / Method

3-1. Overview

Item Description
Goal speculative draft model을 offline SFT plateau 이후에도 더 개선하는 것
Base setting target model + lightweight draft model + speculative verification
Main mismatch SFT는 target trajectory에서 학습하지만 inference는 draft-induced block에서 평가됨
Key module target-assisted rollout, error-position anchor, replay, acceptance-aware KL
Baselines EAGLE-3, DFlash
Main metrics speedup ratio, accepted length $tau$
Key claim Draft-OPD는 thinking model에서 5x 이상 lossless acceleration을 달성하고 EAGLE-3와 DFlash보다 각각 23%와 13% 개선

3-2. Module breakdown

1) Rollout with error-position collection

Draft-OPD는 먼저 speculative decoding 자체를 사용해 rollout을 만든다. prompt $x$가 주어지면 draft model $q_phi$가 block $d_m$을 제안하고, target model $p_theta$가 이 block을 verify한다.

\[d_m = (d_{m,1}, ..., d_{m,K}) \sim q_phi( . \mid x, y_{<=a_m})\]

여기서 $a_m$은 현재 verified prefix가 끝나는 위치다. Draft-OPD는 block을 제안하기 직전의 $a_m$을 anchor로 저장한다. target model이 $r_m$개의 token을 accept하면 다음 anchor는 $a_m + r_m$이 된다.

이 과정의 의미는 다음과 같다.

  • final rollout은 target-assisted라서 품질이 안정적이다.
  • anchor는 draft model이 실제로 action을 취한 지점을 보존한다.
  • rejected proposal도 나중에 replay할 수 있다.

2) Replay for log-probability computation

rollout과 anchor를 모은 뒤에는 각 anchor에서 draft block을 replay한다. anchor $a_m$의 context는 아래처럼 둘 수 있다.

\[c_m = (x, y_{<=a_m})\]

각 drafted token $d_{m,k}$에 대해 draft model과 target model의 log probability를 같은 draft-generated prefix 위에서 계산한다.

\[\log q_{m,k}(d_{m,k}) = \log q_phi(d_{m,k} \mid c_m, d_{m,<k})\] \[\log p_{m,k}(d_{m,k}) = \log p_theta(d_{m,k} \mid c_m, d_{m,<k})\]

핵심은 target model을 final verified token 위에서만 보지 않는다는 점이다. target model도 draft model이 실제로 낸 prefix 위에서 scoring한다. 그래서 rejected position까지 teacher feedback을 받을 수 있다.

3) Accepted and rejected token partition

verification 결과는 drafted token을 accepted set과 rejected set으로 나눈다.

\[I_{acc} = \{(m, k) : 1 <= k <= r_m\}\] \[I_{rej} = \{(m, k) : r_m < k <= K\}\]

accepted token은 target verifier를 통과한 자리다. 따라서 draft distribution이 target distribution을 잘 cover하도록 만들면 된다. rejected token은 draft model이 high probability로 냈지만 target과 맞지 않은 자리다. 따라서 draft distribution의 잘못된 mode를 직접 누르는 쪽이 중요하다.

4) Acceptance-aware distillation objective

Draft-OPD는 accepted token과 rejected token에 서로 다른 KL 방향을 쓴다.

accepted token에는 forward KL을 쓴다.

\[L_{acc} = \frac{1}{|I_{acc}|} \sum_{(m,k) \in I_{acc}} D_{KL}(p_{m,k} || q_{m,k})\]

rejected token에는 reverse KL을 쓴다.

\[L_{rej} = \frac{1}{Z} \sum_{(m,k) \in I_{rej}} w_k D_{KL}(q_{m,k} || p_{m,k})\]

여기서 rejected token weight는 block 안에서 뒤쪽으로 갈수록 작아진다.

\[w_k = \gamma^{k-1}\]

final objective는 두 항을 평균한다.

\[L_{Draft-OPD} = \frac{\lambda_{acc} L_{acc} + \lambda_{rej} L_{rej}}{\lambda_{acc} + \lambda_{rej}}\]

이 설계는 꽤 합리적이다. accepted token에서는 target-supported distribution을 넓게 맞추는 것이 중요하고, rejected token에서는 draft model이 스스로 크게 믿은 잘못된 mode를 낮추는 것이 중요하다. 그래서 forward KL과 reverse KL의 역할이 다르게 배치된다.

5) Why this is not just more SFT data

논문은 OPD prompt pool로 target response를 생성한 뒤 SFT를 더 하는 baseline과 비교한다. 결과적으로 Draft-OPD가 더 좋다. 즉 gain은 단순히 더 많은 prompt나 target response를 본 데서 나오지 않는다.

중요한 차이는 target distribution을 어디서 보느냐다.

“target-generated state” 에서 보는가.

“draft-induced error state” 에서 보는가.

Draft-OPD의 답은 두 번째다.

4. Training / Data / Recipe

4-1. Data

실험은 Qwen3 family에서 수행된다.

  • Qwen3-4B
  • Qwen3-8B
  • Qwen3-30B-A3B-Thinking-2507

평가는 3개 category로 나뉜다.

Category Benchmarks
Math GSM8K, MATH-500, AIME25
Code MBPP, HumanEval, SWE-Lite
Chat MT-Bench

SFT stage는 DFlash와 같은 training data mixture를 따른다. OPD stage에서는 16K prompt pool을 만든다.

Source Count
GSM8K train 2K
MATH corpus 5K
AoPS 4K
CodeAlpaca 5K

여기서 중요한 점은 response를 static reference answer에서 가져오지 않는다는 점이다. 논문은 questions or instructions만 사용하고, response는 OPD 중 target model이 online으로 생성한다고 설명한다.

4-2. Training strategy

Draft-OPD는 SFT-trained DFlash draft model 위에서 진행된다. draft model layer 구성은 다음과 같다.

Target model Draft model depth
Qwen3-4B 5 Transformer layers
Qwen3-8B 5 Transformer layers
Qwen3-30B-A3B-Thinking-2507 8 Transformer layers

공통 block size는 training과 inference 모두 $K = 16$이다. Appendix 기준 recipe는 다음과 같다.

  • Draft-OPD는 6 SFT epochs 이후 checkpoint에서 시작한다.
  • EAGLE-3와 DFlash baseline은 같은 data setup에서 10 SFT epochs를 계속 학습하고 best checkpoint를 보고한다.
  • OPD stage는 8 epochs를 수행한다.
  • thinking-enabled training maximum response length는 4096 tokens다.
  • thinking-disabled training maximum response length는 2048 tokens다.
  • optimizer는 AdamW다.
  • learning rate는 $3 * 10^{-4}$다.
  • cosine learning rate schedule과 warmup ratio 0.05를 쓴다.
  • rejected-token decay는 $\gamma = 0.8$이다.
  • final objective에서는 $\lambda_{acc} = \lambda_{rej} = 1$로 둔다.

4-3. Engineering notes

이 논문을 engineering 관점에서 읽을 때 중요한 부분은 학습과 inference가 같은 draft-verify loop를 공유한다는 점이다.

  • training rollout은 speculative decoding으로 만든다.
  • target hidden state는 rollout 과정에서 이미 계산되므로 replay에서 재사용할 수 있다.
  • anchor replay는 rejected token까지 포함하지만, full draft-only rollout을 요구하지 않는다.
  • speed metric은 vanilla autoregressive decoding 대비 actual test speedup이다.
  • generation quality는 별도 지표로 보고하지 않는다. speculative decoding 절차가 target distribution을 보존한다고 보기 때문이다.

이 점은 실무적으로 중요하다. Draft-OPD는 target model 자체를 바꾸는 방법이 아니라, target model의 lossless verification을 더 자주 통과하는 draft model을 만드는 방법이다. 따라서 output quality improvement paper가 아니라 inference efficiency paper로 읽어야 한다.

5. Evaluation

5-1. Main results

Table 1은 Qwen3-4B와 Qwen3-8B에서 EAGLE-3, DFlash, Draft-OPD를 비교한다. 조건은 thinking mode enabled/disabled, temperature 0/0.6으로 나뉜다.

Setting Result
Thinking enabled, temperature 0 Draft-OPD average speedup 4.88x, DFlash average tau 5.35 -> Draft-OPD average tau 5.85
Thinking enabled, temperature 0.6 Draft-OPD average speedup 4.17x
Thinking disabled Draft-OPD average tau 6.33, average speedup 5.17x
Overall claim thinking models에서 5x 이상 lossless acceleration, EAGLE-3 대비 23%, DFlash 대비 13% improvement

더 구체적으로, thinking mode enabled + temperature 0에서 Qwen3-4B는 Draft-OPD가 GSM8K 5.31x, MATH-500 5.55x, AIME25 5.28x, HumanEval 5.17x를 보인다. Qwen3-8B에서도 MATH-500 5.80x, AIME25 5.51x, HumanEval 5.19x를 보인다.

non-thinking mode에서는 수치가 더 높아진다. Qwen3-8B temperature 0에서는 Draft-OPD가 MATH-500 7.64x, AIME25 6.99x, HumanEval 6.02x를 보인다. 이는 shorter non-thinking generation에서도 draft-target alignment가 개선된다는 것을 보여준다.

5-2. SGLang serving result

논문은 SGLang FA3 backend에서도 deployment-time throughput을 본다. concurrency는 1, 4, 8, 16, 32로 바꾼다.

핵심 해석은 다음과 같다.

  • Draft-OPD는 모든 evaluated model, task, concurrency level에서 DFlash보다 throughput을 개선한다.
  • evaluated model-task pair 평균 accepted length는 11.2% 개선된다.
  • Qwen3-30B-A3B-Thinking에서는 최대 17% speedup을 보인다.
  • concurrency가 높아져도 gain이 사라지지 않는다.

이 결과는 단순히 offline metric인 $tau$가 올라간 것에 그치지 않고, serving engine 안에서도 실제 throughput gain으로 이어질 수 있음을 보여준다.

5-3. What really matters in the experiments

이 논문에서 가장 중요한 실험은 absolute speedup 수치보다 ablation이다. 이유는 Draft-OPD가 왜 먹히는지 비교적 선명하게 분해하기 때문이다.

Training-data ablation

OPD prompt pool로 target response를 만든 뒤 SFT를 더 하는 baseline보다 Draft-OPD가 좋다. 예를 들어 Qwen3-4B thinking mode에서 MATH-500은 EAGLE3 + OPD-data SFT 4.28x, DFlash + OPD-data SFT 5.14x, Draft-OPD 5.55x다. HumanEval은 3.79x, 4.67x, 5.17x이고, MT-Bench는 2.62x, 2.74x, 3.18x다.

즉 OPD prompt pool 자체가 중요한 것이 아니라, draft-induced state에서 distillation한다는 점이 중요하다.

KL type ablation

Table 3은 mixed KL objective가 all-forward KL과 all-reverse KL보다 낫다는 점을 보인다.

Method MATH-500 speedup HumanEval speedup MT-Bench speedup
Draft-OPD 5.55x 5.17x 3.18x
w/o Weight Decay 5.13x 4.96x 3.07x
All-reverse KL 5.11x 4.94x 3.08x
All-forward KL 5.34x 5.01x 3.09x
Random Anchors 5.04x 4.99x 2.96x

여기서 확인할 수 있는 것은 3가지다.

  • accepted/rejected token을 같은 방식으로 다루면 손해가 난다.
  • rejected position decay를 빼면 성능이 줄어든다.
  • random anchor보다 error-position anchor가 낫다.

이건 Draft-OPD의 설계가 단순 trick이 아니라, speculative verification의 failure structure에 맞춰져 있다는 근거다.

6. Limitations

  1. Training length limitation
    • thinking-mode training rollout은 4096 tokens로 제한되지만, evaluation은 8192 tokens까지 사용한다. 매우 긴 reasoning generation의 late-stage state를 충분히 학습하지 못할 수 있다.
  2. Model family and draft architecture scope
    • 주요 실험은 Qwen3 family와 DFlash-style parallel draft architecture에 집중되어 있다. 다른 target model family, EAGLE-style 세부 변형, 다른 inference backend로 얼마나 잘 이전되는지는 추가 검증이 필요하다.
  3. Lossless decoding scope
    • Draft-OPD는 target distribution을 보존하는 speculative decoding 효율 개선 방법이다. generation quality를 개선하는 학습이 아니고, lossy verification이나 approximate decoding으로 확장했을 때는 별도 설계가 필요하다.
  4. Thinking mode uncertainty
    • Appendix C는 thinking mode response가 non-thinking response보다 next-token NLL이 높다는 점을 보인다. 즉 reasoning draft model은 더 flat하거나 불확실한 target distribution을 맞춰야 한다. Draft-OPD가 이를 완화하지만, reasoning-specific drafter 설계는 여전히 남은 과제다.
  5. Artifact verification
    • PDF 첫 페이지에는 project page, code, models 버튼이 보이지만, 현재 리뷰 초안에서는 URL을 직접 추출하지 못했다. 최종 게시 전에는 PDF embedded link 또는 저자 공개 페이지를 직접 눌러 공개 artifact 범위를 확인하는 것이 좋다.

7. My Take

7-1. Why this matters for my work

이 논문은 speculative decoding을 inference algorithm만이 아니라 draft-model post-training problem 으로 재정의한다. 이 관점이 중요하다. 지금까지 speculative decoding을 보면 tree 구조, parallel drafting, verification kernel, cache reuse 같은 inference-time 설계가 주로 보였다. Draft-OPD는 거기서 한 걸음 더 들어가서, draft model이 어떤 state에서 target과 어긋나는지를 학습 문제로 다시 잡는다.

특히 reasoning model serving에서는 이 관점이 꽤 중요하다. thinking mode는 output length가 길고 target distribution도 덜 sharp하다. 그러면 작은 draft model이 target을 맞추기 더 어렵고, accepted length가 속도 병목이 된다. Draft-OPD는 그 병목을 training data mismatch로 설명하고, verification-time error replay로 직접 건드린다.

7-2. Reuse potential

재사용 가능성이 큰 부분은 아래 3가지다.

  1. Error-position replay
    • 어떤 online inference loop든 실패가 관측되는 위치가 있다면, 그 위치를 replay anchor로 삼는 방식은 다른 distillation 문제에도 옮길 수 있다.
  2. Acceptance-aware loss design
    • accepted case와 rejected case를 같은 KL로 다루지 않는 설계는 일반적인 preference, routing, retrieval, tool-use distillation에도 참고할 만하다.
  3. Evaluation metric separation
    • generation quality가 아니라 lossless speedup과 accepted length를 보는 점이 명확하다. output quality와 inference efficiency를 섞어서 해석하지 않는 태도는 실험 설계에 도움이 된다.

7-3. Follow-up papers

  • EAGLE-3: speculative decoding에서 feature-level draft model을 학습하는 strong baseline.
  • DFlash: DFlash-style parallel draft architecture와 training recipe를 이해하기 위한 직접 배경.
  • On-Policy Distillation of Language Models: OPD의 기본 framing과 exposure mismatch 관점.
  • Rethinking On-Policy Distillation of Large Language Models: OPD가 언제 잘 되고 왜 실패하는지 보는 동역학 분석.
  • Domino: causal modeling과 autoregressive drafting을 분리한 speculative decoding 접근.

8. Summary

  • Draft-OPD는 speculative draft model의 offline SFT plateau를 training distribution mismatch로 본다.
  • 핵심은 target-assisted rollout으로 stable sequence를 만들되, verification에서 드러난 draft failure anchor를 replay하는 것이다.
  • accepted token에는 forward KL, rejected token에는 reverse KL을 쓰며, rejected block position에는 decay weight를 둔다.
  • Qwen3 family에서 EAGLE-3, DFlash와 비교해 accepted length와 lossless speedup을 개선한다.
  • 이 논문은 speculative decoding 연구에서 inference-time algorithm과 draft-model post-training을 연결하는 좋은 예시다.

댓글남기기

참고