Source-linked AI summary
D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting
Tianyu Wu, Yu Yao, Zhenting Qi, Han Zheng, Zhuohan Wang, Haoran Ma, Lawrence Liao, Himabindu Lakkaraju, Ju Li, Yilun Du
TL;DR
Speculative decoding can reduce inference latency, but existing position-weight schedules do not adapt to which draft positions currently limit acceptance. D-PACE derives per-position cross-entropy weights from a smooth accepted-length surrogate. Across the main Qwen3-4B configurations it improves average decoding speedup and emitted length, with no architecture or inference-procedure changes.
Problem
Existing multi-token drafter objectives use position-dependent weights that do not adapt as the positions limiting acceptance change during training.
Source: §1 · p. 2Method
D-PACE derives per-position cross-entropy weights from a differentiable surrogate of expected accepted draft length, smoothing weight inputs while leaving the cross-entropy term unsmoothed.
Source: §1 · p. 2, §3.3 · p. 4Results
On Qwen3-4B, D-PACE improves average decoding speedup by +8.0%–9.7% and average emitted length by +8.5%–10.7% across the main configurations.
Source: §1 · p. 2Takeaways & Limitations
D-PACE adds 2.3% training-time overhead and requires no changes to the drafter architecture or inference procedure.
Source: §8 · p. 9The derivation applies to parallel block drafters, and the evaluation covers only lossless speculative decoding.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Speculative decoding accelerates LLM inference by having a small drafter propose tokens that a larger target model verifies in parallel. Recent diffusion-based parallel drafters such as DFlash predict the full B-token block in one forward pass, enabling deeper drafters and longer accepted blocks. However, existing multi-token drafter objectives often use fixed position-dependent weighting schedules, such as head-dependent weights or block-position decays, which do not adapt as the positions limiting acceptance change during training. To address this, we derive per-position training weights from a differentiable surrogate of expected accepted draft length, matching the weight of each position to its log-probability gradient contribution. The resulting loss, D-PACE (Dynamic Position-Aware Cross-Entropy), shifts training signal toward positions that currently limit acceptance as the drafter improves. Across six benchmarks, two Qwen3-4B draft depths, two decoding temperatures, and two additional target models, D-PACE consistently improves both wall-clock speedup and average emitted length, with 2.3\% measured training-time overhead and no changes to the drafter architecture or inference procedure.
1 Introduction
Speculative decoding addresses the latency bottleneck of sequential token generation by letting a small drafter propose tokens for parallel target verification [2] [3] [39]. D-PACE replaces fixed position schedules with weights derived from a smooth accepted-length surrogate, improving decoding metrics across evaluated configurations.
1 Introduction
- Autoregressive decoding is a major LLM inference latency bottleneck because token-by-token generation exposes little parallelism and can leave GPUs underutilized [1] [20] [36] [37].The passage highlights long reasoning-model chains of thought as an especially relevant setting.Source: §1 · p. 1
- Fixed position-dependent weights can misallocate training signal because later positions matter only when the preceding draft prefix is accepted.When early positions are already accurate, later positions can become the acceptance bottleneck; an earlier error prevents later tokens from extending the accepted prefix.Source: §1 · pp. 1–2
- D-PACE differentiates a smooth surrogate of expected accepted draft length to derive adaptive position weights, and Figure 1 relates the surrogate to emitted length τ.The introduction reports that D-PACE improves τ across targets; the surrogate–τ relationship is also described for the DFlash baseline.Source: §1 · p. 2, Figure 1 · p. 2
- On Qwen3-4B, D-PACE raises average decoding speedup by +8.0%–9.7% and average emitted length by +8.5%–10.7% across the main configurations.Average emitted length also improves by about +12.5% on Llama-3.1-8B-Instruct and +12.8% on Qwen3-8B.Source: §1 · p. 2
2 Preliminaries
This section defines speculative decoding in terms of draft blocks and left-to-right verification, then identifies expected accepted draft length as a key speedup factor. It also introduces DFlash as the parallel block drafter used in the main experiments.
2 Preliminaries
- A draft model proposes a block of B candidate tokens, which the target verifies in one forward pass; verification stops at the first rejected position.Tokens before that rejection are accepted, defining the accepted draft length.Source: §2 · pp. 2–3
- For fixed draft and verification costs, expected accepted draft length is a key factor in speculative-decoding speedup over standard autoregressive decoding [2] [3].Source: §2 · p. 3
- DFlash, the parallel block drafter in the main experiments, predicts all B positions in one pass, and its baseline trains with decayed cross-entropy coefficients.The fixed coefficients are the baseline that D-PACE replaces.Source: §2 · p. 3
3 D-PACE: Dynamic Position-Aware Cross-Entropy
D-PACE uses draft confidences as a smooth proxy for per-position acceptance to form an accepted-length surrogate. It differentiates that surrogate to derive adaptive cross-entropy weights, then smooths the weights' inputs to prevent later-position weights from vanishing.
3.1 Accepted-Length Surrogate
- Because per-position acceptance rates are unavailable during training, the accepted-length surrogate uses draft confidence as a smooth proxy for them.The surrogate supplies a smooth basis for deriving per-position training weights.Source: §3.1 · p. 3
3.2 Per-Position Coefficient
- Differentiating the block-level surrogate yields per-position weights combining prefix confidence with the remaining accepted-length continuation value.The weights depend only on draft confidences, so computing them after target-generated training tokens are available requires no additional target-model forward pass.Source: §3.2 · pp. 3–4
3.3 Asymmetric Weight Smoothing
- Asymmetric smoothing interpolates each draft confidence toward 1 inside the weights while leaving the cross-entropy term unsmoothed.The smoothed cumulative product is bounded below by α^j, preventing later weights from vanishing while preserving the standard cross-entropy gradient.Source: §3.3 · p. 4
4 Experiments
Experiments compare D-PACE with related baselines across benchmarks, draft depths, temperatures, and target models. They report improvements in both decoding speedup and average emitted length for every reported Qwen3-4B setting.
4 Experiments
- The experiments address efficacy, ablation, cross-target generalizability, and sensitivity to smoothing and block size, comparing D-PACE with decayed-CE, Top-3 prefix mask, and Accept-rate.The Top-3 mask is adapted from GRIFFIN; the experiments adopt its mask construction, not the full method.Source: §4 · p. 5, §4.1 · p. 6
4.1 Experimental Setup
- Evaluation covers six math, code, and chat benchmarks, Qwen3-4B 3L and 5L drafters at T=0 and T=1, and 3L drafters for Qwen3-8B and Llama-3.1-8B.Both average emitted length (τ) and wall-clock speedup over autoregressive decoding (SR) are reported; higher is better.Source: §4.1 · p. 5
- Qwen3-4B drafter training uses the DFlash authors’ Qwen3-4B-Instruct-100K dataset, while the other targets use ShareGPT prompts with generated target responses.Training uses SpecForge and AdamW; results are reported at epoch 6 unless stated otherwise.Source: §4.1 · p. 5
4.2 Main Results
- D-PACE improves both decoding speedup (SR) and average emitted length (τ) over the DFlash decayed-CE baseline for every reported Qwen3-4B setting.On 3L MATH-500, it maintains roughly 12% higher τ than the baseline throughout a 9-epoch run.Source: Table 1 · p. 6, §4.2 · p. 6
- Top-3 prefix masking discards graded position-wise credit, while Accept-rate couples token likelihood with accepted-length contribution and lacks asymmetric smoothing.D-PACE separates these roles by deriving smoothed weights from the surrogate while retaining token-level cross-entropy gradients.Source: §4.2 · p. 6
5 Ablations and Additional Analyses
The experiments test transfer across targets, isolate weighting components, probe hyperparameters and surrogate choice, and examine position-weight dynamics and training overhead. D-PACE improves both reported metrics across targets and block sizes, with only 2.3% per-step training-time overhead.
5.1 Generalization to Other Targets
- D-PACE improves both τ and SR on every benchmark for the Llama-3.1-8B-Instruct and Qwen3-8B targets.Table 2 reports per-benchmark SR and τ for 3L drafters at T=0, epoch 6.Source: §5.1 · p. 6, Table 2 · p. 7
5.2 Component Ablation
- Continuation-only stays about 2% above DFlash but 7%–9% below cumulative-only and D-PACE; cumulative-only exceeds DFlash by 10%–12%, while full D-PACE adds a smaller consistent gain.The ablation compares MATH-500 τ over training epochs.Source: §5.2 · p. 7, Table 3 · p. 7
5.3 Hyperparameter Sensitivity
- D-PACE improves average SR and τ at every tested block size; tested α values outperform DFlash except α=0, with performance stable across α ∈[0.3, 0.7].At α=0, the cumulative product vanishes and τ stalls near 1.0 across the three configurations.Source: §5.3 · p. 7, Table 4 · p. 7, Table 5 · p. 8
5.4 Target-Probability Surrogate
- Replacing draft probability with target probability leaves both SR and τ much closer to DFlash than full D-PACE, because target probability cannot identify positions limiting acceptance for the drafter.Table 6 reports average SR and τ over six benchmarks for the target-probability variant.Source: §5.4 · p. 8, Table 6 · p. 8
5.5 Position-Weight Dynamics
- D-PACE shifts weight from low-index positions early in training toward higher-index positions later, recomputing weights for each block, example, and update.DFlash’s fixed exponential decay remains constant over training.Source: §5.5 · p. 8
5.6 Training-time overhead
- On a single H200, D-PACE increases per-step training time by only 2.3% over the DFlash decayed-CE baseline.The weight-computation step calculates block-level weights in parallel across positions.Source: §5.6 · p. 8
6 Related Work
Related work improves speculative decoding through drafter architectures, verification strategies, parallel block generation, and acceptance-aware training objectives. D-PACE addresses the allocation of learning signal across positions in parallel blocks.
Speculative decoding architectures
Parallel and diffusion-based drafters
- Parallel block drafters reduce sequential dependence on block length, with DFlash [11], DART [16], and DiffuSpec [43] among the related approaches.This line connects to diffusion language models and block-diffusion models.Source: §6 · pp. 8–9
Acceptance-aware drafter objectives
- Acceptance-aware objectives include GRIFFIN [10], SpecDiff-2 [41], and OPT-Tree [42]; D-PACE replaces fixed positional schedules such as DFlash [11] and DART [16] decays with accepted-length-proxy weights.D-PACE detaches the resulting weight from cross-entropy and applies asymmetric smoothing only to the weight.Source: §6 · p. 9
7 Limitations and Future Work
D-PACE’s derivation is scoped to parallel block drafters and its experiments to lossless speculative decoding. The paper also identifies limitations of its confidence proxy and leaves broader objective extensions unevaluated.
7 Limitations and Future Work
- The accepted-length surrogate is derived for parallel block drafters; sequential drafters can violate its fixed-prefix factorization because later drafting depends on earlier hidden states.The experiments evaluate D-PACE only for lossless speculative decoding.Source: §7 · p. 9
- Empirical evaluation covers token-level cross-entropy weighting; evaluation of the KL-distillation variant D-PAKL and other acceptance-aware dynamic weighting remains future work.D-PAKL is derived by replacing the soft acceptance analog with a Jensen lower-bound proxy.Source: §7 · p. 9
- The draft confidence proxy is not an exact acceptance probability, may align weakly with realized acceptance early in training, and that effect is not separately evaluated.The paper relies on the proxy empirically, citing correlation between the surrogate and reported accepted length.Source: §7 · p. 9
- The paper reports lower energy consumption and serving cost at scale, while leaving the target model’s distribution unchanged.It states that D-PACE introduces no novel risks beyond those inherent in the underlying LLMs.Source: §7 · p. 9
8 Conclusion
D-PACE replaces fixed positional decay with weights derived from an accepted-length surrogate. It improves emitted length and speedup across the reported settings, with modest training overhead and no architecture or inference changes.
8 Conclusion
- D-PACE improves average emitted length and speedup across DFlash drafts, reaching 4.47× speedup on MATH-500 for Qwen3-4B.The conclusion also reports transfer to Llama-3.1-8B and Qwen3-8B and 2.3% training-time overhead, without architecture or inference-pipeline changes.Source: §8 · p. 9
- D-PACE replaces DFlash’s fixed exponential decay with per-position weights derived from the gradient of an accepted-length surrogate.The weights tie training signal to accepted length, a key factor in speculative-decoding efficiency.Source: §8 · p. 9
- The same construction may apply beyond DFlash to other single-trajectory drafters when a differentiable proxy for acceptance is available.The conclusion further suggests replacing fixed position schedules with weights based on each position’s contribution to accepted length.Source: §8 · p. 9
A Confidence–acceptance correlation
The appendix reports that the cumulative confidence surrogate predicts accepted length better than a plain confidence sum, then outlines D-PACE loss ingredients and comparison baselines.
A Confidence–acceptance correlation
- The cumulative surrogate is more predictive of reported accepted length than the plain confidence sum; both correlations have p < 10−10.The analysis uses blocks from 128 MATH-500 prompts and compares per-block Spearman correlations.Source: A Confidence–acceptance correlation · p. 13, Table 8 · p. 13
B Loss formulations
- The loss formulations define smoothed draft confidence with default α = 0.5 and a smoothed continuation value for each position.The passages also introduce the block size and target-selected token, but their full definitions are not included here.Source: B Loss formulations · p. 13
B.1 D-PACE
- The appendix identifies the displayed objective as the main D-PACE loss from Section 3.3.Source: §B.1 · p. 13
B.2 Baselines
- Baselines include DFlash’s exponential position decay [11], a K=3 top-K prefix mask adapted from GRIFFIN [10], and direct ascent of the accepted-length surrogate.The GRIFFIN-derived baseline uses only the mask, not the full method; accept-rate uses neither smoothing nor detached weighted cross-entropy.Source: §B.2 · p. 13
B.3 Ablation variants of D-PACE
The ablations vary the separate factors in D-PACE’s per-position weights and test a target-probability surrogate while retaining draft cross-entropy.
B.3 Ablation variants of D-PACE
- Cumulative-only and continuation-only variants substitute each factor of the D-PACE weight separately into the loss.Source: §B.3 · p. 14
- The target-probability surrogate variant replaces draft confidence with target probability while keeping cross-entropy on the draft.Source: §B.3 · p. 14
- Its continuation value is computed analogously from smoothed target probabilities.Source: §B.3 · p. 14
C Training details
The appendix specifies the training setup and compute, then extends the position-weighted objective from hard labels to KL distillation. The extension builds on a soft-target accepted-length surrogate and yields D-PAKL, which reduces to D-PACE in the hard-CE limit.
C Training details
- Training uses sequence length 2048, 512 anchor positions per sequence, gradient clipping at max-norm 1.0, warmup ratio 0.04, and cosine decay after warmup.Within paired comparisons, only the training objective differs; reported configurations vary in draft depth, target model, and block size.Source: C Training details · p. 14
- The appendix consolidates the D-PACE, DFlash decayed-CE, Top-K prefix-mask, Accept-rate, and ablation loss formulations.The formulations are collected in Appendix B.Source: C Training details · p. 14
D Training compute
- A 3L draft run takes about 2:25 per epoch for 6 epochs; total reported training compute is about 465 H200 GPU-hours.All runs use one NVIDIA H200 GPU; Table 7 supplies the wall-clock comparison used to estimate overhead.Source: D Training compute · p. 14
E D-PAKL: KL-distillation variant
- D-PAKL substitutes a soft-target KL objective for hard cross-entropy, with the KL and soft-label cross-entropy sharing draft-parameter gradients.The hard-CE case is recovered when the target distribution is a Dirac distribution at the sampled token.Source: E D-PAKL: KL-distillation variant · p. 14
E.1 Accepted-length surrogate
- The soft-target acceptance proxy is the inner product of target and draft distributions, and Jensen’s lower bound yields a surrogate requiring no extra forward computation.The proxy collapses to the hard-label draft confidence for a Dirac target, and the bound is tight in that case.Source: §E.1 · pp. 14–15
E.2 Per-position coefficient
- Per-position coefficients are derived by differentiating the accepted-length surrogate and expressing the result through each position’s log-probability gradient.The derivation splits the outer sum at the position being differentiated and uses the continuation value.Source: §E.2 · p. 15
E.3 Weight smoothing
- Weight smoothing addresses coefficients that vanish when several estimated acceptance rates are small, while leaving the position loss unsmoothed.The passage specifies smoothing each estimated acceptance rate inside the weight only.Source: §E.3 · p. 15
E.4 D-PAKL loss
- D-PAKL uses the smoothed per-position weights detached from the gradient; replacing cross-entropy with KL preserves the draft-parameter gradient.In the hard-CE limit, the surrogate is tight and D-PAKL reduces to D-PACE.Source: §E.4 · p. 15