SpecOPD: Branch-Parallel On-Policy Distillation for Speculative Decoding

SpecOPD is a branch-parallel on-policy distillation framework that supervises draft-sampled branches from reference anchors in parallel and trains speculative-decoding drafters to increase acceptance length.

Speculative decoding speedup depends on acceptance length, and on-policy distillation can push drafters beyond acceptance-length plateaus. Conventional SD-tailored OPD collects supervision through sequential draft-and-verify rollouts: the drafter proposes a block, the target verifies it, and the accepted prefix determines the next round. This serializes response collection before drafter updates and leaves low GPU utilization during rollout collection. The paper asks whether on-policy training must inherit these sequential dependencies.

SpecOPD samples B anchor positions from a pre-collected reference response and uses the reference sequence as a shared trunk; at each anchor the current drafter samples a W-token continuation independently, forming a proposal tree with branch parent map and semantic positions. It prefills the reference trunk once and batches branch scoring over cached representations, so B branches add B(W-1) new scoring rows after shared trunk processing. To preserve branch histories, it uses tree attention masks for full attention such as GQA, branch-consistent compressed indexer groups for sparse attention such as QSA, and indexed state replay for linear attention such as GDN. The drafter is trained with Bebop's E2E-TV acceptance-length surrogate: token acceptance probability is 1 minus total variation between draft and target, and the loss minimizes 1 minus the average prefix product over BW proposed tokens. Branches are resampled each update while sampled tokens and target distributions are held fixed; distributed normalization preserves the objective when ranks have unequal scored-position counts.

• SpecOPD improves five-benchmark mean acceptance length by up to 41.45%. §3
• Against sequential-rollout OPD, SpecOPD achieves higher update throughput while processing more supervised tokens per update, yielding higher supervised-token throughput. §3
• Measured generation speedup reaches 3.56× over autoregressive decoding. §3
• Agentic-task acceptance length improves by up to 81.51%. §3
• At matched optimizer steps and reference-sequence batches, SpecOPD improves mean AL by up to 15.71% over off-policy E2E-TV and 6.96% over off-policy CE; its best evaluated checkpoint gains 4.33% prefix survival at position seven. §3
• In the 125B evaluation, step 22,000 maximizes the six-cell mean acceptance length. §A.4

• For conditional proposals, the E2E-TV prefix products are a surrogate rather than an unbiased estimate of serving acceptance (§2.3).
• SWE and Terminal use fixed-history replay without live tools or task-success grading; their membership in official SWE-bench Verified and Terminal-Bench 2.0 subsets is unconfirmed, and 64 Terminal source trajectories are flagged as overlapping with training data (§A.1).
• Training-throughput comparisons exclude offline reference preparation and do not measure time to a matched acceptance threshold (§3, §B.2), leaving end-to-end time-to-quality open.
• Off-policy comparisons match sequence presentations but use different compute allocations and gradient paths (§A.6).

• A draft block can be supervised using only its reference prefix and draft-sampled history, so branches need not wait for verification outcomes. §1
• Branches use independent sampling streams, and the joint distribution over sampled continuations factorizes as a product over branches and positions. §2.1
• The proposal tree attaches draft branches to anchors on the reference trunk; the parent of the first branch token is the anchor token, later parents are preceding sampled tokens, and semantic position is anchor position plus offset. §2.1
• Prefilling the reference trunk once and batching branch scoring preserves each target query while amortizing overhead; B branches contribute B(W-1) new scoring rows after shared trunk processing. §2.2
• Full attention uses tree masks, sparse QSA uses branch-consistent compressed indexer groups, and linear GDN uses indexed state replay. §1
• The E2E-TV loss minimizes one minus the average prefix product of draft-target overlap, so an early mismatch affects all later prefix terms. §2.3
• Evaluation spans 8B full-attention GQA, 27B full-linear GQA-GDN, and 125B sparse-linear QSA-GDN targets, three draft architectures, and four task domains. §Abstract
• Off-policy E2E-TV and CE baselines decline from initialization despite decreasing training loss, while SpecOPD improves prefix survival at all seven draft positions. §3