D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting

TL;DR

D-PACE introduces a dynamic, position-aware training method for parallel speculative decoding in large language models, improving speed and draft length without altering inference procedures.

Contribution

It derives a novel loss function that adapts training weights based on position-specific acceptance limits, enhancing draft quality and speed.

Findings

Consistently improves speedup and emitted draft length across multiple benchmarks.

Achieves 2.3% training-time overhead without changing inference architecture.

Effective across various models, depths, and decoding temperatures.

Abstract

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.