PaperlayerView exact revision on arXiv

Source-linked AI summary

ReTrace: Rejected-Trajectory Conditioning for Speculative Decoding

Luxi Lin, Zhanpeng Zeng, Shuang Peng, Songwei Liu, Rongrong Ji

arXiv:2608.29748v1cs.CL

TL;DR

Speculative decoding wastes rejected draft suffixes after an early mismatch, motivating reuse of their retained predictive information. ReTrace conditions each subsequent DFlash block on corrected rejected-suffix representations without extra forward passes, consistently improving acceptance length and speed over DFlash.

  • Problem

    After the first mismatch, standard prefix verification discards the remaining draft suffix even though its computation may retain information about the target continuation.

  • Method

    ReTrace carries rejected hidden representations into the next draft, aligns and target-corrects them, and injects them through gated residual conditioning.

  • Results

    ReTrace consistently improves macro-average acceptance length by 6.97% and macro-average speedup by 5.34% over DFlash across model scales and decoding regimes.

  • Takeaways & Limitations

    Rejected draft computation can condition later drafting while preserving target-side verification and the lossless property of speculative decoding.

Abstract

from arXiv · show

Speculative decoding accelerates autoregressive language model inference by having a lightweight draft model propose multiple candidate tokens, which are then verified in parallel by a larger target model. However, after the first rejection, standard prefix-based verification discards the remaining draft suffix, so the computation spent generating and verifying those positions does not contribute to decoding progress. Focusing on DFlash, we show that rejected positions in a rejected suffix may still align with the target continuation, indicating that the draft model can retain useful semantic and structural information despite local token-level errors. Motivated by this observation and inspired by conditional diffusion, we introduce~\textbf{ReTrace}, a rejected-trajectory conditioning method that conditions each draft block on the rejected suffix from the previous round rather than generating it from fresh mask placeholders alone. ReTrace retains the hidden representations of the rejected suffixes, aligns them with the next draft block, refines them using target-aware correction signals from the same verification pass, and admits them into the drafter's input embeddings through gated residual fusion. Because rejected tokens are never committed and target-side verification remains unchanged, ReTrace preserves the lossless property of speculative decoding without requiring an additional model forward pass. Experiments with Qwen3 models across mathematical reasoning, code generation, and open-ended dialogue demonstrate that ReTrace consistently improves average acceptance length and end-to-end decoding speed over its DFlash backbone. By introducing cross-round conditioning without modifying within-round proposal generation, ReTrace is largely orthogonal to existing drafting improvements and might be combined with them for further gains.

Introduction

Speculative decoding reduces autoregressive inference cost by drafting multiple tokens and verifying them in parallel, but early mismatches waste later draft computation. ReTrace reuses rejected hidden representations to condition subsequent DFlash proposals without extra forward passes.

  • Speculative decoding uses a lightweight drafter to propose tokens that a larger target model verifies in parallel.
  • Early mismatches reject the remaining draft suffix, preventing its computation from advancing decoding.
  • Later rejected positions can realign with the target continuation, indicating that rejected hidden representations retain predictive information.
  • ReTrace carries rejected hidden representations across rounds as auxiliary conditioning for the next DFlash proposal.
  • ReTrace preserves standard draft-then-verify decoding and requires no additional forward pass through either model.

Related Works

Related approaches improve proposal generation or verification within a decoding round, whereas ReTrace reuses rejected-suffix representations across rounds. Its evidence combines token-rank and representation-similarity analyses showing that rejected trajectories retain useful but imperfect continuation information.

  • Related Work: Existing speculative-decoding methods improve draft quality or verification through feature prediction, diffusion drafting, intra-block refinement, confidence scheduling, and tree construction.
  • ReTrace: ReTrace differs by carrying rejected-suffix representations across rounds, refining them with target states, and conditioning the next proposal.
  • Draft Reuse: Lyanna resamples candidate trees from reusable hidden trajectories, whereas ReTrace uses rejected states only as conditioning signals for the next regular draft.
  • Choice of Drafter: DFlash is suited to cross-round reuse because parallel block prediction means a local rejection need not invalidate all subsequent states.
  • Token-Level Predictive Signal: 30.4% of target tokens remain the drafter’s top prediction and 75.2% fall within its top five beyond the first mismatch.
  • Representation-Level Predictive Signal: Rejected draft states are more similar to their own target continuation than to random windows, with macro-average similarity rising from 0.128 to 0.351.
  • Bounded Reusability: The rejected suffix is not reliable enough to copy directly but can provide auxiliary conditioning for subsequent drafting.

Method

ReTrace carries rejected-suffix representations across drafting rounds, aligns them with the next proposal block, corrects them using target states, and injects them through gated residual fusion. It reuses representations from standard passes while preserving the original target-side verification and commitment procedures.

  • Overview: ReTrace treats the rejected hidden-state suffix as auxiliary conditioning for the next masked draft block, extending DFlash across drafting rounds.The rejected trajectory supplements, rather than replaces, the verified context and current masked input.
  • Rejected-Trajectory Construction and Alignment: After a rejection, ReTrace excludes accepted-prefix and first-rejected representations, carrying forward only states strictly after the first rejected position.The first rejected position is resolved by the target-side correction token, while later representations may retain semantic direction, formatting, or local structure.
  • Rejected-Trajectory Construction and Alignment: ReTrace shifts retained states to the beginning of the next proposal block while preserving their original order; with no usable remainder, it reduces to original drafting.The alignment is defined relative to the next round’s first unresolved position and applies only to the available suffix positions.
  • Target-Aware Correction: Target-aware correction fuses each aligned draft state with the corresponding target representation already produced during verification, requiring no additional target-model forward pass.A lightweight projection maps the pair into the draft representation space, and the residual formulation preserves the original suffix representation while adjusting it.
  • Gated Residual Conditioning: A learned vector-valued gate injects the corrected trajectory into the drafter’s input embeddings as a residual, while unconditioned positions retain their original inputs.The residual path is initialized to zero for stable optimization, and the original drafter performs the remaining computation unchanged.
  • Training and Inference: At inference, ReTrace reuses representations from standard drafting and verification, retains each signal for one round, and leaves target-side acceptance and output commitment unchanged.Training detaches cached representations across rounds, jointly optimizes the draft-side modules with the target model frozen, and introduces no extra forward pass.

Experiments

Experiments show that ReTrace improves DFlash-based speculative decoding across models, benchmarks, and decoding regimes, while adding little runtime overhead. Ablations attribute the gains to rejected-suffix conditioning, and compatibility results show that ReTrace can complement intra-block drafting improvements.

  • Main Results: ReTrace improves macro-average acceptance length by 6.97% and macro-average speedup by 5.34% over DFlash across four model–temperature configurations.The acceptance improvement appears across every benchmark and under both greedy and stochastic decoding.
  • Component Ablation: Target-aware correction achieves the best ablation result, reaching 6.90 average acceptance and 3.86× speedup, compared with 6.79 for gated fusion.Direct overwriting reaches only 6.51 average acceptance, while the unconditioned comparison starts at 6.45.
  • Effect of Additional Training: Continued DFlash training under the same data and optimization budget improves only marginally and saturates quickly, whereas ReTrace steadily increases average acceptance length.The comparison isolates rejected-suffix conditioning from additional training.
  • Runtime Efficiency Analysis: ReTrace raises GSM8K speedup from 3.37× to 3.73× and average accepted length from 6.36 to 7.36, with 0.13 ms conditioning overhead per cycle.The overhead is 0.40% of DFlash cycle latency and requires neither an additional target forward pass nor autoregressive drafting.
  • Compatibility with Domino: Adding ReTrace to Domino raises macro-average throughput from 287.17 to 301.39 tok/s and mean acceptance length from 7.43 to 7.79 tokens.The result supports compatibility between cross-round conditioning and Domino’s intra-block causal refinement.

Conclusion

The conclusion presents ReTrace as a cross-round conditioning method that reuses rejected draft representations without changing target verification. Across settings, it improves acceptance and speed while remaining compatible with within-block drafting refinements.

  • Conclusion: Rejected-suffix hidden representations retain predictive information that benefits subsequent drafting, even though the corresponding tokens cannot be committed.ReTrace uses this information to condition the next draft block.
  • Conclusion: ReTrace aligns rejected states, applies target-side correction, and injects the resulting trajectory through gated residual fusion.The correction uses representations already available from verification.
  • Conclusion: ReTrace improves acceptance length and end-to-end speedup over DFlash across model scales, temperatures, and diverse benchmarks at negligible latency cost.It also remains compatible with intra-block improvements such as Domino.
  • Conclusion: Each ReTrace round leaves target verification and sampling unchanged, while resetting the conditioning signal after full acceptance or a final-position rejection.The algorithm adds correction and input-fusion steps to the underlying speculative-decoding loop.

B Training Data

ReTrace is trained on a 40K-example prompt-only pool spanning dialogue, mathematical reasoning, code generation, and instruction following.

  • Training Data: The training pool contains 40K prompt-only examples drawn from Open-PerfectBlend and uses random seed 42.Open-PerfectBlend is described as an open reproduction of an instruction mixture spanning four task categories.

C Model and ReTrace Configuration

Experiments use Qwen3-4B and Qwen3-8B targets with corresponding DFlash-b16 drafters. The targets remain frozen, and correction representations come from ordinary verification without an additional target forward pass.

  • Model and ReTrace Configuration: The evaluated target models are Qwen3-4B and Qwen3-8B, initialized with corresponding public DFlash-b16 drafter checkpoints.Each drafter has five draft layers and block size 16.
  • Model and ReTrace Configuration: The target model remains frozen, and all correction representations are taken from the ordinary verification pass.Consequently, ReTrace requires no additional target-model forward pass.
  • Model and ReTrace Configuration: Table 6 lists the model and ReTrace configuration.The supplied configuration passage identifies the table as the configuration summary.

D Optimization

ReTrace training retains the DFlash block-diffusion objective and introduces no auxiliary losses. Its correction is warmed up gradually, while inference uses the fully warmed coefficient and initially reproduces the original draft input.

  • Optimization: ReTrace uses the same block-diffusion objective as DFlash without adding target-model, counterfactual gate, or auxiliary conditioning losses.Correction and gated-fusion parameters are optimized through the current-round prediction loss, jointly with the draft backbone.
  • Optimization: Training runs on 32 NVIDIA A800 40GB GPUs across four eight-GPU nodes.
  • Optimization: The correction coefficient β is linearly warmed up, remains inactive at training start, and reaches β = βmax = 1 at inference.Because Wv is zero-initialized, ReTrace exactly matches the original draft input before learning begins.

E Evaluation Protocol

The evaluation standardizes prompts, checkpoints, tokenization, stopping, sampling, and generation limits across methods. It reports benchmark-level acceptance and speed under controlled runtime measurements, including a specified single-GPU decomposition.

  • Evaluation Protocol: All methods use the same prompts, target checkpoints, tokenizer, stopping criteria, and sampling parameters within each comparison.Unless otherwise stated, evaluation is single-request with block size 16 for DFlash and ReTrace, and at most 8,192 new tokens.
  • Evaluation Protocol: The evaluation covers GSM8K, MATH-500, LiveCodeBench, Alpaca, and the full HumanEval test set, using controlled temperature-specific sampling settings.The random seed is 42, with a warmup batch and runtime-cache clearing before benchmark measurement.
  • Evaluation Protocol: End-to-end speedup is measured against vanilla autoregressive decoding of the same target model, while average acceptance length τ is computed over target verification cycles.Latency includes proposal generation, ReTrace conditioning, and the remaining decoding-cycle components.
  • Evaluation Protocol: Runtime decomposition uses Qwen3-4B on one NVIDIA A800 GPU with batch size 1, context length 1,024, greedy decoding, and draft block size 16.Component medians are measured with CUDA events, and the total sums non-overlapping verification, draft, LM-head, and conditioning times.

F Ablation Settings

The ablation compares alternative rejected-state conditioning choices under matched training and evaluation conditions. A separate baseline and compatibility study isolate the role of rejected-suffix conditioning and its cross-block scope.

  • F Ablation Settings: Every ablation row uses Qwen3-4B, block size 16, greedy decoding, and the same initialization, training data, optimization budget, and evaluation prompts.
  • F Ablation Settings: Raw overwrite replaces active mask inputs with aligned rejected states, whereas raw gated retains the mask pathway and applies gated residual fusion.
  • F Ablation Settings: Corrected gated conditioning adds target-aware correction before applying the same gated fusion used by raw gated conditioning.
  • F Ablation Settings: The SFT baseline matches ReTrace's training data and optimization budget but omits rejected-suffix conditioning, while Domino operates within blocks and ReTrace carries information between blocks.The Domino compatibility study uses greedy decoding, block size 16, and a maximum generation length of 1,024 tokens.

G Concurrent-Serving Experiment

The concurrent-serving experiment tests whether ReTrace remains advantageous when speculative decoding competes with target-model batching. Across the tested concurrency levels, ReTrace is fastest and retains an acceptance-length advantage over DFlash.

  • G Concurrent-Serving Experiment: Concurrent serving evaluates Qwen3-4B on 128 GSM8K prompts with no-think greedy decoding, BF16, FlashInfer, and at most 1,024 new tokens.DFlash and ReTrace use 16-token blocks; EAGLE-3 uses seven draft steps with tree sizes 16 and 60.
  • G Concurrent-Serving Experiment: Client concurrency is swept over c ∈{1, 2, 4, 8, 16, 32} with each method running in an isolated SGLang server on one NVIDIA A800-SXM4-40GB GPU.
  • G Concurrent-Serving Experiment: Table 8 reports end-to-end speedup over target-only decoding and average acceptance length τ for each concurrency level, with bold marking the best result per level.
  • G Concurrent-Serving Experiment: ReTrace is fastest at every concurrency level and gains about one token of acceptance length over DFlash.The advantage persists under batching because conditioning is reused inside the drafter without adding verification work.
  • G Concurrent-Serving Experiment: At concurrency 32, ReTrace reaches 1.62× while both EAGLE-3 configurations fall below target-only decoding.
Loading 2608.29748v1…