PaperlayerView exact revision on arXiv

Source-linked AI summary

ResiSpec: Enhancing Multi-Candidate Speculative Sampling via Residual Distribution Shaping

Zhi-Kai Chen, Jun-Jie Tao, Wei-Xiang Mao, De-Chuan Zhan, Han-Jia Ye

arXiv:2608.24411v1cs.AI

TL;DR

Multi-candidate speculative decoding is limited by Residual Drift, which makes later pre-sampled candidates ineffective after early rejections. ResiSpec reshapes verification residuals toward draft-model high-confidence regions while preserving exact sampling, achieving up to 1.92× speedup.

  • Problem

    Residual Drift shifts rejection residuals away from the draft distribution, making subsequent pre-sampled candidates ineffective in multi-candidate verification.

  • Method

    ResiSpec uses an auxiliary proxy distribution to reshape verification and align rejection residuals with draft-model high-density regions.

  • Results

    1.92× speedup over multi-candidate methods is achieved while maintaining exact, zero-bias sampling.

  • Takeaways & Limitations

    ResiSpec prevents candidate obsolescence during sequential verification and can complement existing candidate-generation and tree-construction methods.

  • Takeaways & Limitations

    Achieving the global L1 minimum under full draft support and disjoint support constraints forces the slack function and shaping mass to zero.

Abstract

from arXiv · show

The efficiency of Large Language Model (LLM) serving is fundamentally limited by the sequential nature of autoregressive decoding. Speculative Decoding (SD) mitigates this by using a lightweight draft model to speculate future tokens, which are then validated by the LLM in a single parallel forward pass. To further boost efficiency, multi-candidate schemes propose diverse candidate sets to increase the likelihood of token acceptance. However, we show that these schemes are bottlenecked by Residual Drift: a phenomenon where the rejection of initial candidates causes the residual target distribution to diverge from the draft model's predictions. This shift renders subsequent candidates ineffective and forces the system into expensive resampling. To resolve this, we propose ResiSpec, a framework that strategically reforms the proposal distribution during verification to anchor the residual target mass within the draft model's high-confidence regions. By mathematically re-aligning the verification process without compromising output exactness, ResiSpec prevents candidate obsolescence and achieves up to 1.92$\times$ speedup over state-of-the-art multi-candidate methods. Code is available at https://github.com/Czzzk/Resispec.

Introduction

Autoregressive decoding limits efficient LLM serving, while speculative decoding uses draft-model proposals and parallel target verification to reduce sequential overhead. Multi-candidate methods expose a residual-drift bottleneck, which ResiSpec addresses by reshaping verification while preserving exact sampling.

  • Autoregressive decoding generates tokens sequentially, creating I/O bottlenecks and underutilizing modern GPU compute.
  • Speculative decoding uses a lightweight draft model to propose tokens that the target model validates in one parallel forward pass.
  • Multi-candidate schemes expand the search space with diverse token choices to increase the probability of accepting a high-quality candidate.
  • Residual Drift occurs when sequential rejections shift the target distribution toward regions where the draft model has low confidence, invalidating pre-sampled candidates.
  • ResiSpec reshapes verification using an auxiliary proxy that keeps rejection residuals near draft-model high-density regions and preserves candidate viability.
  • 1.92× speedup over multi-candidate methods is demonstrated while maintaining zero-bias sampling across benchmarks.

Related Work

Prior work established speculative decoding and multi-candidate tree verification as acceleration strategies, but standard protocols scale poorly with candidate counts. ResiSpec targets this acceptance-length bottleneck by tailoring sampling to large branching trees.

  • Speculative decoding is a lossless acceleration paradigm that verifies draft-model candidate tokens in parallel with the target model.
  • SpecInfer introduced tree-based attention for verifying branching token sequences in a single forward pass.
  • Standard sampling protocols fail to scale efficiently with candidate counts, creating an acceptance-length bottleneck.
  • ResiSpec optimizes sampling logic for large branching trees to increase the marginal utility of additional candidates and expected token yield.

Preliminary

Speculative decoding drafts candidate tokens and verifies them against the target distribution, while tree-based verification applies sequential acceptance checks across branches and levels.

  • Speculative decoding operates in drafting and verification stages using target and draft conditional distributions.
  • The draft model generates γ candidate tokens, and the target model evaluates speculative tokens in parallel using one forward pass.
  • Each candidate is accepted with probability min(1, p_tgt(x)/p_dft(x)) to preserve exact sampling.
  • Tree proposals organize tokens by depth and branching factor, with b independent candidates proposed at each speculative position.
  • Candidates are verified sequentially, and rejection updates the target distribution for the next candidate to a residual distribution.
  • If all candidates at a position are rejected, the branch is truncated and a token is resampled from the cumulative residual distribution.
  • Tree attention enables verification of all paths within the proposal tree in a single forward pass.

Decoding

Multi-candidate verification suffers from residual drift because later candidates sampled from the original draft distribution face shifted residual criteria. ResiSpec seeks to keep those residuals aligned with draft support, improving scalability.

  • After x1 is rejected, later candidates are verified against a residual distribution rather than the original target distribution.
  • The residual distribution assigns zero probability where q(x) meets or exceeds p(x), concentrating mass in the draft model’s blind spots.
  • All N candidates are sampled independently from q(x), producing a support mismatch as later verification criteria move away from q(x)’s high-density regions.
  • Subsequent candidate acceptance approaches zero after the first candidate, leaving later candidates with near-zero marginal utility across temperatures and target models.
  • Expanding tree width incurs linear verification cost but produces sub-linear acceptance gains because residual distributions diverge from the draft distribution.
  • ResiSpec proposes re-aligning verification criteria with q(x)’s high-density regions so subsequent candidates remain useful after earlier rejections.

Method: Residual Distribution Shaping

ResiSpec reshapes the residual distribution during multi-candidate verification by replacing the draft distribution with an auxiliary proxy. This keeps rejected-candidate residual mass near the draft model’s high-density support while preserving exact sampling.

  • Proxy-based verification: ResiSpec introduces an auxiliary proxy distribution k(x) that replaces the original draft distribution q(x) during verification.The proxy is constructed so rejected candidates produce a residual distribution p′(x) proximal to q(x).
  • Residual alignment: The proxy reshapes rejection residuals toward the draft model’s high-confidence regions, keeping pre-sampled candidates viable across the verification batch.This design aims to maximize acceptance probability for subsequent candidates.
  • Proxy mechanism: The proxy mechanism defines k(x) using an unnormalized residual mass r(x) together with a nonnegative slack function s(x).The residual distribution is p′(x) = r(x)/Z.
  • Support constraints: The slack function restores k(x)’s normalization while remaining separate from reserved residual mass.Its injected mass is permitted only where residual mass is absent, imposing a disjoint-support constraint.
  • Exactness: The constraints on k(x) are imposed to preserve exact sampling, meaning the marginal output distribution remains p(x).The method changes verification behavior without changing the target marginal distribution.

1. Probability Constraint: P

ResiSpec constructs a valid proxy distribution by reserving residual mass from the target and balancing it with slack mass, while aligning locally with the draft distribution. Because ideal residual matching can be infeasible in long-tailed or full-support settings, the framework uses constrained approximation and minimal feasible residual budgets while preserving unbiased sampling.

  • Proxy constraints: The proxy distribution k(x) must be normalized, nonnegative, and locally match the draft probability at the sampled point.The local alignment condition is k(x_i) = q(x_i), while the construction also requires k(x) ≥ 0 and total probability mass equal to one.
  • Mass conservation: The slack mass must equal the total residual mass reserved from the target distribution.Proposition 1 states that the slack budget satisfies Z = Σ_x r(x), provided r(x) ≤ p(x) for every x.
  • Residual shaping objective: Maximizing future acceptance is formulated as maximizing overlap between the shaped residual and the draft distribution, equivalently minimizing their L1 distance.The ideal residual satisfies r(x) = Zq(x), but this objective remains subject to the exactness and proxy-validity constraints.
  • Capacity bound: Ideal residual matching is bounded by the minimum density ratio between the target and draft distributions.The bound ensures that k(x) remains nonnegative under the disjoint-support constraint.
  • Practical boundary: Long-tailed distributions can make the feasible residual budget nearly zero, while full draft support makes perfect alignment incompatible with nonzero slack.These conflicts motivate a constrained approximation that retains a functional residual budget rather than enforcing the unattainable ideal exactly.
  • ResiSpec procedure: ResiSpec therefore selects the smallest feasible residual budget and uses proxy-based verification to preserve exact target-model sampling.The minimal-budget rule limits deviation from the draft trajectory, and the iterative proxy process is proved unbiased with zero distributional shift relative to autoregressive decoding.

Experiment

ResiSpec improves multi-candidate speculative decoding across datasets, temperatures, baselines, and tree configurations while preserving the target distribution. Its runtime overhead comes from distribution-level shaping after verification rather than an additional model forward pass.

  • Main Results: Up to 1.86× higher accepted length and 1.68× higher throughput than SpecInfer were observed across tested datasets, temperatures, and configurations.Experiments used a 3-branch complete tree with depth 5, comparing Llama-2-7B with JF68M on CNN/DM, OpenWebText, and C4 at T = 1.0 and T = 0.6.
  • Compatibility with Existing Methods: ResiSpec integrates with EAGLE and Sequoia as a verification-time add-on, improving accepted length and throughput under matched decoding settings.EAGLE accelerates candidate generation, while Sequoia optimizes the verification tree.
  • Computational Overhead: ResiSpec shaping adds distribution-level operations after verification and does not require an additional model forward pass.Most overhead comes from memory-bound GPU operations over token distributions; fused kernels could further reduce latency.
  • Ablation Study: ResiSpec shows larger gains than SpecInfer on deeper, narrower trees under comparable verification budgets.The pattern indicates stronger mitigation of accumulated Residual Drift over longer speculation horizons.
  • Distributional Exactness: The KL divergence between ResiSpec-generated and target distributions remains on the order of 10^-10, consistent with floating-point noise rather than systematic bias.The effective distribution is reconstructed by marginalizing all acceptance and rejection trajectories within a candidate batch.

Conclusion

ResiSpec reshapes rejection residuals through a proxy distribution while preserving exact sampling. Its proof shows that locally aligning the proxy with the draft distribution makes accepted and rejected paths sum to the target distribution.

  • Acceptance Path: The acceptance path contributes the intersection of the target distribution p and proxy k when local alignment gives k(x′) = q(x′).This identifies the accepted mass used in the marginal-distribution proof.
  • Rejection Path: The rejection path samples from a residual distribution whose mass complements the acceptance mass after normalization cancellation.The final marginal combines mutually exclusive acceptance and rejection events using the law of total probability.
  • Conclusion: Because the acceptance and rejection contributions sum for every x′, the output distribution is identically the target distribution p(x).This establishes the claimed universal exactness of ResiSpec.
  • Correctness and Exactness: ResiSpec couples residual shaping to a proxy distribution k(x) while preserving exact sampling when k and q align at the verification point.The proxy may differ from the draft distribution elsewhere, provided local alignment holds at the sampled candidate.

A.2 Monotonicity of Residual Shaping Loss with Redistribution

The analysis shows that residual shaping loss is zero up to a critical mass budget Z∗, then increases strictly as redistribution becomes necessary. ResiSpec’s path-based verification nevertheless preserves the target distribution, with measured KL divergence near zero.

  • Monotonicity of Residual Shaping Loss: The minimum L1 distance is zero for Z ∈ (0, Z∗], where the target can mirror the draft distribution’s shape.Here Z∗ is defined by the bottleneck capacity threshold.
  • Redistribution Mechanism: The residual optimization sets r(x)=p(x) in bottleneck regions and requires r(x)≥Zq(x) elsewhere, with local deviations balancing to zero globally.Bottlenecks occur where p(x)<Zq(x).
  • Monotonicity of Residual Shaping Loss: The minimum L1 distance increases strictly monotonically for Z ∈ (Z∗, 1].Once residual mass exceeds Z∗, the draft distribution cannot be matched without distortion.
  • Redistribution Mechanism: Beyond Z∗, capacity clipping at bottleneck tokens and redistribution to non-bottleneck tokens contribute equally to the increase in L1 distance.The result motivates using a minimal residual-mass budget to preserve distribution fidelity.
  • Exactness Verification: ResiSpec reconstructs the marginal output distribution by summing mutually exclusive accepted-candidate paths and the final residual fallback path.After each rejection, the residual distribution is reshaped before evaluating subsequent candidates.
  • Exactness Verification: Measured DKL remained typically below 10^-8, confirming no distributional bias within floating-point numerical error.The reconstructed distribution is mathematically equivalent to the target distribution under the verification procedure.

C Reproducibility Details

The reproducibility protocol standardizes execution conditions and randomization across matched baseline and ResiSpec comparisons. It also excludes auxiliary diagnostics from throughput measurements.

  • Execution Environment: Throughput experiments use the same Docker environment, model checkpoints, and matched configurations for baselines and ResiSpec variants.This controls differences from host-side Python, CUDA, and library versions.
  • Randomness Control: Experiments initialize Python, NumPy, PyTorch CPU, and CUDA random states with the same base seed and enable deterministic cuDNN mode.Prompt i uses seed s+i, avoiding identical random streams while preserving matched schedules.
  • Randomness Control: The same per-prompt seed schedule is applied to ResiSpec and its baseline under each matched model, dataset, temperature, and tree configuration.This makes the comparisons reproducible across evaluated prompts.
  • Timing and Diagnostics: Throughput runs disable auxiliary KL diagnostics and other distribution-checking instrumentation.The reported throughput therefore excludes the overhead of these diagnostics.

D Ablation Experiment Results

The ablation evidence indicates that adding candidates yields only modest accepted-length gains while efficiency declines sharply. Across model-pair configurations, ResiSpec outperforms SpecInfer in accepted length and throughput.

  • Scaling Inefficiency: With speculation depth fixed at 4, increasing candidate count N produces only modest EAL gains while efficiency score η drops sharply.Results are averaged over CNN/DM with Llama-2-7B.
  • Reported Metrics: The reported table metrics include temperature, prediction-tree shape, accepted length, and tokens per second.These abbreviations are defined as T, Tree, Acc., and Thru.
  • Model-Pair Ablation: ResiSpec significantly outperforms SpecInfer in accepted length and throughput across different target–draft model configurations.The ablation uses Llama-2-7B and Sheared-LLaMA-2.7B targets with four draft-model choices.
Loading 2608.24411v1…