D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting
This work addresses the problem of suboptimal training signal distribution in multi-token speculative drafters, offering a simple yet effective loss modification that improves LLM inference efficiency without architectural changes.
D-PACE introduces a dynamic position-aware cross-entropy loss for parallel speculative drafting that adapts training weights to focus on positions limiting acceptance, achieving consistent improvements in wall-clock speedup and average emitted length across multiple benchmarks with only 2.3% training overhead.
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.