1 Introduction↩︎

Recent advances in large language models (LLMs) have enabled strong performance across reasoning, coding, and general assistant tasks, but their growing model sizes and longer generations substantially increase inference cost [1], [2]. Speculative decoding (SD) mitigates this cost by using a lightweight draft model to propose tokens that are verified in parallel by a larger target model, preserving the target model’s output distribution [3], [4]. In SD, the speedup depends heavily on the quality of the draft model: when the drafter closely matches the target model, longer token spans can be accepted, reducing the number of expensive target-model decoding steps. Recent mainstream methods, such as EAGLE-3 and DFlash, have achieved strong acceleration by training lightweight draft models on target-generated trajectories. [5], [6].

However, we find that offline supervised fine-tuning quickly reaches its limit for draft-model training. Figure 1 shows a representative training curve. After an initial SFT warm-up, continuing SFT does not improve the draft model’s acceptance length on test data; instead, the accepted length fluctuates around a fixed plateau. Moreover, continuing SFT on the data used for OPD can even reduce the accepted length. This suggests that the key limitation does not lie in offline training compute. Instead, the plateau points to a mismatch between the states used for training and the states that determine speculative acceptance. In SFT, the drafter learns from fixed target-generated trajectories, so every prefix is produced by the target model. During speculative decoding, however, the target model verifies token blocks proposed by the drafter itself. The accepted length is therefore determined by draft-induced inference states, rather than only by the static states seen in offline target trajectories.

This offline-to-inference mismatch motivates on-policy distillation (OPD), where the target model supervises states induced by the current draft policy [7]. Directly applying OPD to draft models, however, is not straightforward. Standard OPD assumes that the student can roll out full sequences under its own policy. But mainstream draft models such as EAGLE-style or DFlash-style drafters are designed to propose short token spans under target-model guidance, rather than to act as standalone autoregressive generators. As a result, draft-only rollouts can easily become repetitive or low quality. Using target-assisted rollouts can produce usable sequences, but the verified continuation follows the target distribution: target model discards the incorrect tokens proposed by draft model, even though these errors are the most valuable training signals. Consequently, draft model still learns on trajectories generated by target model rather than on its own inference-time states.

To address this gap, we propose Draft-OPD, an on-policy distillation framework designed for speculative draft models. Draft-OPD uses target-assisted rollout to maintain stable continuations, and records error positions exposed by speculative verification. It then replays drafting from these positions and asks the target model to score the same draft-generated prefixes. This makes it possible to train on states where the draft model actually acted, including both accepted proposals and rejected proposals. Finally, Draft-OPD uses an acceptance-aware distillation objective that treats these two token groups differently: accepted tokens reinforce reliable agreement with the target model, while rejected tokens focus learning on draft-induced errors that limit speculative acceptance.

In summary, our contributions are:

• We identify a key limitation of offline SFT for draft models: SFT quickly plateaus because the drafter is trained on fixed target trajectories but evaluated on blocks induced by its own policy during speculative decoding.

• We explain why standard OPD does not directly apply to draft models: draft-only rollouts are unstable, while target-assisted rollouts remove the on-policy signal.

• We propose Draft-OPD, an on-policy distillation framework with error-position replay, making it feasible to efficiently post-train draft models on verification-time errors.

• Experiments show that Draft-OPD achieves over \(5\times\) lossless acceleration for thinking models and improves over EAGLE-3 and DFlash by \(23\%\) and \(13\%\) under matched FLOPs.

2 Related Work↩︎

2.1 Speculative Decoding↩︎

Speculative decoding accelerates autoregressive inference by using a lightweight draft model to propose multiple future tokens and a larger target model to verify them in parallel while preserving the target distribution [3], [4]. Recent work further improves speculative decoding by leveraging feature-level context from the frozen target model, as in EAGLE [5], [8], or by designing stronger block-level draft architectures such as DFlash [6]. These methods substantially improve decoding efficiency, but they still train draft models with offline SFT on target-generated trajectories, which limits further improvements in draft-model acceptance length. In contrast, our work studies how to move draft models beyond this offline training recipe.

2.2 Draft Model Distillation↩︎

Because speculative speedup depends on how closely the draft distribution matches the target distribution, several methods train draft models through distillation. DistillSpec aligns a compact draft model with the target model and highlights the importance of on-policy data and task-specific divergence choices [9]; temperature-centric distillation studies further show that matching training and decoding configurations can improve speculative decoding under challenging sampling settings [10]. Online speculative decoding updates draft model from observed query distributions during deployment [11], while broader on-policy distillation trains students on self-generated sequences to reduce exposure bias [7]. However, almost all previous work only studies standalone autoregressive draft models, which are often selected as smaller models from the same family as the target model. Our Draft-OPD is designed for training-based draft models such as DFlash or EAGLE3. These draft models cannot independently generate full student trajectories, making online distillation difficult because training must obtain usable rollouts rather than meaningless repetitions.

3 Preliminary↩︎

3.0.0.1 Speculative Decoding.

Speculative decoding accelerates autoregressive generation by pairing a target model \(p_\theta\) with a lightweight draft model \(q_\phi\) [3], [4]. Given a prefix \(x_{<t}\), the draft model proposes a block of \(K\) candidate tokens, \[\begin{align} \hat{y}_{t+k} &\sim q_\phi(\cdot \mid x_{<t}, \hat{y}_{t:t+k-1}), \\ k &= 0,\ldots,K-1. \end{align}\] and the target model verifies these tokens in parallel. The verifier accepts the longest valid prefix of the drafted block and then continues generation from the verified prefix, preserving the target model’s output distribution while reducing the number of expensive target-model decoding steps.

Draft models can be categorized by how they are obtained. Model-based draft models use smaller autoregressive models from the same family as the target model [12]. A more mainstream approach trains a lightweight draft model specifically for a given target model [13]–[15]. These training-based draft models often share the target model’s embedding layer and LM head, and are designed to predict short draft spans rather than to generate complete sequences independently. In this work, we focus on training-based draft models because they can provide stronger acceleration.

A central metric for speculative decoding is the accepted length \(\tau\), i.e., the number of draft tokens accepted in each verification round. Higher \(\tau\) means that each target-model verification generates more tokens, which directly improves decoding efficiency. It also reflects the alignment between the draft and target models, with better alignment yielding longer accepted spans.

3.0.0.2 On-Policy Distillation.

On-policy distillation trains a student model on states induced by the student’s own policy, rather than only on fixed teacher-generated trajectories [7]. Given a prompt \(x\), the current student samples a trajectory and induces prefix states \[\tilde{y}\sim q_\phi(\cdot\mid x), \quad s_t=(x,\tilde{y}_{<t}).\] At each state \(s_t\), the teacher provides the next-token distribution. Writing \(p_\theta^t=p_\theta(\cdot\mid s_t)\) and \(q_\phi^t=q_\phi(\cdot\mid s_t)\), a standard OPD objective is \[\mathcal{L}_{\mathrm{OPD}} = \mathbb{E}_{s_t} \left[ D_{\mathrm{KL}}(p_\theta^t\,\|\,q_\phi^t) \right].\] Thus, OPD directly targets the exposure mismatch between offline supervised training and inference-time generation: the student learns from states that it actually visits under its own policy.

This perspective is particularly relevant to speculative decoding. As discussed above, the accepted length \(\tau\) is determined by how well the draft model aligns with the target model on the prefixes encountered during the draft-verify process. However, SFT trains \(q_\phi\) only on fixed target-generated trajectories, whereas speculative decoding evaluates blocks produced by \(q_\phi\) itself. OPD offers a natural way to reduce this mismatch by using the target model to supervise draft-induced states, particularly those induced by the draft model’s own deviations from the target.

4 Method↩︎

4.1 Challenges of Direct OPD↩︎

Applying OPD to draft models requires both stable rollouts and draft-induced training states. Standard OPD assumes that the student can roll out full sequences under its own policy, but EAGLE- and DFlash-style draft modules are designed to propose short token blocks under target-model verification, rather than to act as standalone autoregressive generators. As shown in Figure 2 (a), forcing such draft modules to self-rollout full trajectories can produce repetitive or degenerate samples, making the resulting supervision unreliable.

Target-assisted rollout avoids degenerate samples, but it removes the draft-policy signal that OPD is meant to capture. As illustrated in Figure 2 (b), strict speculative verification is lossless with respect to the target model, so the verified continuation follows the target distribution rather than the draft policy. Moreover, this brings back the offline-to-inference mismatch: target-assisted keeps accepted tokens and discards rejected proposals, even though these rejected tokens reveal the draft model’s most informative errors. We empirically verify this limitation in Section 5.3, where naive target-assisted rollout underperforms Draft-OPD. A suitable OPD method for draft modules must therefore keep target-quality rollouts while preserving the draft-induced errors exposed during verification.

4.2 Draft-OPD↩︎

Motivated by these failure modes, we design Draft-OPD to enable effective on-policy training for draft models. The key idea is to use target assistance to keep rollouts stable, while replaying draft-induced errors to preserve the on-policy training signal. Draft-OPD implements this idea through three coordinated designs: rollout with error-position collection, replay for log-probability computation, and an acceptance-aware distillation objective.

4.2.0.1 Rollout with error-position collection.

Let \(p_\theta\) denote the target model and \(q_\phi\) denote a lightweight draft model initialized from supervised draft-model training. Given a prompt \(x\), speculative decoding asks \(q_\phi\) to propose a block of \(K\) tokens and uses \(p_\theta\) to verify the block in parallel. Draft-OPD uses this interaction to collect a target-quality rollout while recording each draft block’s starting position as an anchor for later error replay. Let the verified rollout be \(y=(y_0,\ldots,y_{T-1})\). During rollout step \(m\), the current verified prefix ends at position \(a_m\), where \(a_m=-1\) denotes the beginning of the sequence. The draft model proposes a block \[d_m = (d_{m,1},\ldots,d_{m,K}) \sim q_\phi(\cdot \mid x, y_{\le a_m}),\] and the target model verifies the block. We record \(a_m\) as an anchor before moving to the next step. If the verifier accepts \(r_m\) tokens from the block, the next anchor is \(a_m + r_m\). This process is repeated until the rollout reaches the maximum generation length or the target model emits an end-of-sequence token.

Anchors preserve the draft model’s local actions without requiring the draft model to generate an entire sequence alone. The final rollout remains a high-quality target-model sample, but each anchor identifies a state where the draft model actually proposed a block during inference. Drafting from any anchor requires the target-model hidden states of the preceding tokens; since the rollout has already computed these hidden states on the same sample, Draft-OPD can reuse them during subsequent replay.

4.2.0.2 Replay for Log-Probability Computation.

After collecting the rollout and anchors, Draft-OPD replays drafting from all anchors to compute the token-level student and teacher log-probabilities. For anchor \(a_m\), define the replay context \[c_m = (x, y_{\le a_m}).\] Starting from \(c_m\), we replay the draft-generated block \(d_m\). For each drafted token \(d_{m,k}\), the draft model provides the student log-probability, while the target model provides the teacher log-probability on the same draft-generated prefix: \[\begin{align} \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}). \end{align}\] This replay step differs from training on the final rollout tokens: it evaluates the target model on the draft model’s proposed block, including positions that were rejected during verification.

The verification outcome naturally partitions drafted tokens into accepted and rejected sets. If \(r_m\) tokens are accepted from block \(m\), then \[\begin{align} \mathcal{I}_{\mathrm{acc}} &= \{(m,k): 1 \le k \le r_m\}, \\ \mathcal{I}_{\mathrm{rej}} &= \{(m,k): r_m < k \le K\}. \end{align}\] Accepted tokens reflect states where the draft model agrees with the target well enough to pass verification. Rejected tokens include the first failed token and later positions in the same draft block, which become less reliable because an earlier rejection invalidates the remaining drafted suffix.

4.2.0.3 Acceptance-Aware Distillation Objective.

We use different KL directions for accepted and rejected draft tokens. For accepted tokens, we use forward KL to make the draft distribution cover the target distribution at states where the draft model is already close to the target: \[\mathcal{L}_{\mathrm{acc}} = \frac{1}{|\mathcal{I}_{\mathrm{acc}}|} \sum_{(m,k)\in\mathcal{I}_{\mathrm{acc}}} D_{\mathrm{KL}}\!\left(p_{m,k}\,\|\,q_{m,k}\right).\] For rejected tokens, we use reverse KL to penalize the draft model’s own high-probability modes when the target model disagrees: \[\mathcal{L}_{\mathrm{rej}} = \frac{1}{Z} \sum_{(m,k)\in\mathcal{I}_{\mathrm{rej}}} w_k D_{\mathrm{KL}}\!\left(q_{m,k}\,\|\,p_{m,k}\right),\] where \(Z=\sum_{(m,k)\in\mathcal{I}_{\mathrm{rej}}} w_k\) normalizes the rejected-token weights. We give the detailed rationale for this KL design in Appendix 8.

Rejected tokens at later positions in a draft block are less important than earlier rejected tokens. In speculative decoding, an early error prevents the verifier from using the remaining suffix, so mistakes near the beginning of a block have a larger effect on acceptance length. We reflect this point with an exponentially decaying weight over block positions: \[w_k = \gamma^{k-1}\] The final Draft-OPD objective averages the accepted-token and rejected-token losses: \[\mathcal{L}_{\mathrm{Draft\text{-}OPD}} = \frac{ \lambda_{\mathrm{acc}}\mathcal{L}_{\mathrm{acc}} + \lambda_{\mathrm{rej}}\mathcal{L}_{\mathrm{rej}} }{ \lambda_{\mathrm{acc}}+\lambda_{\mathrm{rej}} }, \label{eq:draft-opd-objective}\tag{1}\] We set \(\lambda_{\mathrm{acc}}=\lambda_{\mathrm{rej}}=1\) in all experiments. This objective trains the draft model on stable target-assisted rollouts while preserving the draft-policy errors that determine speculative acceptance.

5 Experiment↩︎

4pt

5.0.0.1 Models and tasks.

We conduct experiments on the Qwen3 family [1], including Qwen3-4B, Qwen3-8B, and Qwen3-30B-A3B-Thinking-2507. We evaluated three categories of benchmarks: mathematical reasoning, including GSM8K [16], MATH-500 [17], [18] and AIME [19]; code generation and software engineering, including MBPP [20], HumanEval [21] and SWE-bench Lite [22]; and out-of-domain benchmark MT-Bench [23].

5.0.0.2 Datasets.

For the SFT stage, we use the same training data mixture as DFlash [6]. For the OPD stage, we construct a 16K sample prompt pool randomly sampling 2K prompts from the GSM8K training set, 5K prompts from the MATH corpus after excluding MATH-500 held-out examples, 4K prompts from AoPS [24], and 5K prompts from CodeAlpaca [25]. We use only the questions or instructions from these datasets; responses are generated online by the target model during OPD rather than taken from static reference answers.

5.0.0.3 Implementation.

Unless otherwise specified, all experiments are conducted on NVIDIA H200 GPUs, with a batch size of 1, thinking mode enabled. We perform Draft-OPD on top of SFT-trained DFlash draft models: the draft models use 5 Transformer layers for Qwen3-4B and Qwen3-8B, and 8 layers for Qwen3-30B-A3B-Thinking [6]. We use a block size of 16 for both training and inference. Detailed hyperparameters for OPD training are provided in Appendix 7.

5.0.0.4 Baselines.

We compare against EAGLE-3 [5] and DFlash [6]. For a fair comparison, both EAGLE-3 and DFlash are trained with the data mixture introduced by DFlash, and their SFT training budget is matched to the total SFT plus OPD budget of our method so that all draft models are trained under approximately the same FLOPs budget.

5.0.0.5 Metrics.

We focus on efficiency-related metrics and do not report generation quality, since Draft-OPD only post-trains the draft model and does not change the speculative decoding procedure used at inference time, thereby preserving the exact output distribution of the target model.

• Speedup Ratio. The actual test speedup ratio relative to vanilla autoregressive decoding.

• Average acceptance length (\(\tau\)). The average number of draft tokens accepted by the target model in each verification cycle.

5.1 Main Results↩︎

Table 1 compares Draft-OPD with EAGLE-3 and DFlash on Qwen3 models under a matched training FLOPs budget. For EAGLE-3, we use a tree size of 16, with draft steps and top-\(k\) set to 8 and 4, respectively. We evaluate both greedy decoding and the recommended Qwen reasoning sampling setting, with temperature 0.6, top-\(p\) 0.95, and top-\(k\) 20.

5.1.0.1 Thinking Mode Enabled.

With thinking mode enabled, Draft-OPD consistently improves acceptance length and decoding speed under matched training FLOPs. At temperature 0, it raises the two-model average \(\tau\) from 5.35 for DFlash to 5.85 and achieves 4.88\(\times\) average speedup across the seven benchmarks, improving over EAGLE-3 and DFlash by 23% and 13%.

At temperature 0.6, it remains the fastest method, averaging 4.17\(\times\) speedup. Although EAGLE-3 obtains the highest \(\tau\) on Qwen3-8B at temperature 0.6, its sequential drafting limits wall-clock speed, supporting our choice to apply OPD to DFlash-style parallel drafting.

5.1.0.2 Thinking Mode Disabled.

With thinking mode disabled, Draft-OPD preserves the same advantage across decoding temperatures and model sizes. It maintains an average acceptance length of 6.33 and achieves a 5.17\(\times\) average speedup. Together with the thinking-mode results, this shows that Draft-OPD improves draft-target alignment across both long reasoning traces and shorter non-thinking generations.

5.1.0.3 Performance on SGLang.

r0.5 3pt

We benchmark Draft-OPD and DFlash on SGLang [26] with the FA3 backend to evaluate deployment-time efficiency. Table 2 reports serving throughput under concurrency levels from 1 to 32, together with the average acceptance length \(\tau\).

Draft-OPD consistently improves SGLang serving efficiency across all evaluated models, tasks, and concurrency levels. It improves the acceptance length by 11.2% on average over the evaluated model-task pairs and achieves up to a 17% speedup on Qwen3-30B-A3B-Thinking. Notably, the throughput gains do not diminish under higher concurrency: at concurrency 32, the average relative gain is even higher than at concurrency 1. These serving results show that the higher acceptance lengths produced by Draft-OPD translate into practical throughput gains in an optimized inference engine.

5.2 Ablation Study↩︎

5.2.0.1 Training Data.

A key question is whether the gains of Draft-OPD come from on-policy distillation or simply from exposing the draft model to more prompts. Figure 4 compares Draft-OPD with EAGLE-3 and DFlash variants that continue supervised training on responses generated by the target model from the OPD prompt pool. Draft-OPD performs better, showing that its improvement is not explained by additional supervised data alone. Instead, the benefit comes from distilling target distributions on states induced by the draft model during speculative decoding.

5.2.0.2 KL Type.

The acceptance-aware KL objective is designed to match the asymmetric roles of accepted and rejected draft tokens. Table 3 studies this design by replacing the mixed objective with all-forward and all-reverse KL variants. The all-forward variant is consistently below Draft-OPD, while the all-reverse variant performs the worst among the KL variants. This supports using different KL directions for accepted and rejected positions rather than applying a single divergence to all draft tokens.

r0.5 4pt

5.2.0.3 Anchor Position.

Error-position replay is intended to focus training on the states where speculative decoding actually fails. Table 3 compares this design with randomly selected replay anchors. Random anchors lead to lower speedup and acceptance length, indicating that not all on-policy states are equally useful for draft-model post-training. Concentrating replay around rejected proposals better exposes the errors that determine speculative acceptance.

5.2.0.4 Weight Decay.

The rejected-token position decay down-weights later rejected tokens in the same drafted block, since they are more likely to be affected by earlier draft errors. As shown in Table 3, removing this decay reduces both speedup and acceptance length. This suggests that the rejected-token loss is most effective when it emphasizes the earliest and most informative failure positions.

5.3 Analysis Experiments↩︎

5.3.0.1 Target-Assisted Rollout.

The direct draft-only rollout in Figure 2 (a) quickly becomes repetitive and low quality, making it unsuitable for OPD training despite being on-policy. We therefore analyze the naive target-assisted rollout baseline illustrated in Figure 2 (b).

r0.5 4pt

Since speculative verification keeps the target-verified continuation and discards rejected draft proposals, this baseline reduces to KL-loss SFT on target-distributed trajectories. As shown in Table 4, naive target-assisted rollout underperforms Draft-OPD, reducing the average speedup from 4.63\(\times\) to 4.29\(\times\), a relative drop of 7.3%. This confirms that stable target-assisted rollouts alone are insufficient; preserving draft-induced errors is important for effective online draft-model training.

5.3.0.2 Thinking-Mode Draftability Gap.

We also observe that, on math and code tasks, draft models trained under thinking mode are substantially less effective than their non-thinking counterparts. We attribute this gap to the higher uncertainty of target-model responses in thinking mode, where long reasoning traces often allow multiple plausible next steps, making it more difficult for the draft model to fit the target model’s distribution. Appendix 9 provides a slightly more detailed analysis.

6 Conclusion↩︎

We propose Draft-OPD, an on-policy distillation framework for training-based draft models. By using target-assisted rollouts with an error-position replay mechanism, Draft-OPD keeps training samples stable while preserving the draft-policy errors that determine speculative acceptance. We further introduce an acceptance-aware distillation objective that treats accepted and rejected draft tokens differently, enabling the draft model to learn from both reliable proposals and informative failure modes. Experiments on Qwen3 models show that Draft-OPD improves average acceptance length and end-to-end decoding speedup over strong draft-model baselines under a matched training budget, with additional SGLang results confirming practical serving gains. Overall, Draft-OPD shows that draft models benefit from post-training on verification-time errors, offering an effective direction for improving training-based draft model.

Limitations↩︎

6.0.0.1 Training Length.

Due to compute constraints, for thinking-mode Draft-OPD training we cap the maximum response length at 4096 tokens, while evaluation uses 8192 tokens. Although Draft-OPD still achieves strong gains under this longer evaluation setting, the training rollouts may not fully cover late-stage states in very long generations. Scaling OPD training to longer rollouts could expose the draft model to a broader range of verification-time errors and may further improve draft-target alignment for long outputs.

6.0.0.2 Evaluation Scope.

This work focuses on post-training draft models for lossless speculative decoding. Our main experiments are conducted on Qwen3 models, and the Draft-OPD implementation follows the DFlash-style parallel draft architecture. Although we include comparisons with EAGLE-3 and evaluate deployment in SGLang, further evaluation is needed to understand how well Draft-OPD transfers to other model families, draft architectures, and inference backends.

6.0.0.3 Lossless Decoding.

Because speculative decoding preserves the target model distribution, Draft-OPD is designed to improve decoding efficiency rather than generation quality; extending on-policy draft-model training to settings with approximate or lossy verification is an interesting direction for future work.

Acknowledgments↩︎

This work was supported by the Shanghai Artificial Intelligence Laboratory. We are grateful to the authors and open-source communities whose work made this project possible. In particular, DFlash [6] and EAGLE-3 [5] provided strong foundations and reference points for training-based speculative decoding, while SGLang [26], SpecForge [27], and verl [28] offered practical infrastructure for serving and training experiments. We especially appreciate SpecForge for its open-source implementation of DFlash training, which provided a useful reference for our training setup. We also thank Runzhe Zhan for sharing valuable experience on OPD training and for discussions that helped us stabilize the training workflow.

7 Training Details↩︎

7.0.0.1 SFT stage.

We follow the SFT configuration of DFlash [6] for most training settings and use SpecForge [27] as the training framework. For Draft-OPD, we initialize the OPD stage from the draft-model checkpoint after 6 SFT epochs. For the EAGLE-3 and DFlash baselines used in our comparisons, we continue supervised draft-model training for 10 epochs under the same data setup and report the checkpoint with the best evaluation performance.

7.0.0.2 OPD stage.

For Draft-OPD training, we use the rejected-token position weights \[w_k = \gamma^{k-1},\] with \(\gamma=0.8\). We train on the OPD data mixture described in Section 5 for 8 epochs, using a maximum response length of 4096 tokens for thinking-enabled and 2048 tokens for thinking-disabled. Optimization uses AdamW with a learning rate of \(3\times 10^{-4}\) and a cosine learning-rate schedule with a warmup ratio of 0.05. We implement the OPD stage with verl [28]. For the final Draft-OPD objective in Equation 1 , we set \(\lambda_{\mathrm{acc}}=\lambda_{\mathrm{rej}}=1\) in all experiments.

8 Loss Design for Accepted and Rejected Draft Tokens↩︎

We provide a local training-objective justification for using different KL directions on accepted and rejected replay positions. The goal is not to show that the mixed KL objective universally dominates all-forward or all-reverse KL. For a fixed replay state and an unconstrained draft distribution, these objectives share the same optimum \(q=p\). Their difference lies in how the local loss weights target-supported tokens and draft-proposed tokens during finite-capacity optimization.

Fix a replay state \(s\) and let \(p(\cdot\mid s)\) and \(q(\cdot\mid s)\) denote the target and draft next-token distributions over vocabulary \(\mathcal{V}\). In the following derivation, all probabilities are conditioned on \(s\), so we write \(p(y)\) and \(q(y)\) for brevity. For an accepted replay position, the draft proposal has passed target verification, and the local objective should match the target distribution at this reliable state. This gives the target-weighted cross-entropy \[\mathcal{J}_{\mathrm{acc}}(q) = \mathbb{E}_{y\sim p}[-\log q(y)] .\] Expanding the expectation yields \[\begin{align} \mathcal{J}_{\mathrm{acc}}(q) &= -\sum_{y\in\mathcal{V}} p(y)\log q(y) \\ &= H(p) + D_{\mathrm{KL}}(p\,\|\,q), \end{align}\] where \(H(p)\) is independent of \(q\). Thus, the accepted-position objective is equivalent, up to an additive constant, to minimizing forward KL.

For a rejected replay position, the relevant signal is draft-induced: the token belongs to a draft-proposed suffix that failed or was invalidated by target verification. The loss should therefore place weight on tokens that the draft model itself is likely to propose, especially when those modes are not supported by the target. This leads to the draft-weighted disagreement objective \[\mathcal{J}_{\mathrm{rej}}(q) = \mathbb{E}_{y\sim q}\!\left[\log \frac{q(y)}{p(y)}\right].\] Since \[\mathcal{J}_{\mathrm{rej}}(q) = \sum_{y\in\mathcal{V}} q(y)\log \frac{q(y)}{p(y)} = D_{\mathrm{KL}}(q\,\|\,p),\] the rejected-position objective corresponds to reverse KL. Unlike the accepted-position objective, this loss is weighted by the draft distribution and therefore directly penalizes high-probability draft modes that disagree with the target distribution.

Aggregating these local objectives over replay positions gives the acceptance-aware form used in Draft-OPD: \[\sum_{i\in\mathcal{I}_{\mathrm{acc}}} D_{\mathrm{KL}}(p_i\,\|\,q_i) + \sum_{i\in\mathcal{I}_{\mathrm{rej}}} w_i D_{\mathrm{KL}}(q_i\,\|\,p_i),\] up to normalization and scalar weights. This decomposition clarifies the role of the two KL directions. Forward KL is appropriate for accepted positions because the supervision is target-weighted on verified states, while reverse KL is appropriate for rejected positions because the supervision is draft-weighted on states that expose draft errors. Applying a single KL direction to all positions ignores this distinction: all-forward KL treats rejected draft errors like reliable accepted states, whereas all-reverse KL treats verified accepted positions like draft-error states. The component ablations in Table 3 are consistent with this local objective view.

9 Thinking-Mode Drafting Gap↩︎

To examine this gap, we run Qwen3-4B in both thinking and non-thinking modes on prompts from ShareGPT [29], AoPS [24], the math and code splits of Nemotron-Post-Training-Dataset-v2 [30], and compute the next-token NLL at each token position for the generated responses. Figure 5 shows that thinking-mode responses have higher next-token NLL than non-thinking responses across the evaluated datasets. This suggests that, in thinking mode, the target model itself has higher uncertainty over the next token. A draft model trained for thinking-mode decoding is therefore asked to predict a less concentrated target distribution, making it harder to match the target model’s subsequent tokens accurately.

Draft-OPD partially mitigates this challenge by post-training the draft model on verification-time errors, improving acceptance for reasoning-oriented decoding; however, since our method is designed as a general post-training framework across decoding settings, approaches tailored specifically to reasoning draft models remain an important direction for future work.

References↩︎