PaperlayerView exact revision on arXiv

Source-linked AI summary

Cascade Speculative Drafting for Even Faster LLM Inference

Ziyi Chen, Xiaocong Yang, Jiacheng Lin, Chenkai Sun, Kevin Chen-Chuan Chang, Jie Huang

arXiv:2312.11462v5cs.LGcs.CL

TL;DR

Speculative decoding lowers LLM inference latency but is limited by autoregressive drafting and equal-cost generation of tokens with differing acceptance likelihoods. CS Drafting combines vertical and horizontal cascades to remove neural autoregressive drafting and allocate drafting effort across token positions. It achieves up to 81% additional speedup over speculative decoding while maintaining the target model’s output distribution.

  • Problem

    Speculative decoding remains bottlenecked by autoregressive draft generation and inefficiently allocates equal time to tokens with different acceptance likelihoods.

  • Method

    CS Drafting combines a recursive vertical cascade of smaller-model reviews with a horizontal cascade that uses smaller models for later draft tokens.

  • Results

    Up to 81% additional speedup over speculative decoding was achieved across tasks and settings while preserving the target model’s output distribution.

  • Takeaways & Limitations

    CS Drafting improves speculative-decoding latency through complementary cascades without sacrificing the target model’s output distribution.

  • Takeaways & Limitations

    Performance outcomes may differ on systems with different hardware configurations, motivating both standardized and raw walltime reporting.

Abstract

from arXiv · show

Introduced to enhance the efficiency of large language model (LLM) inference, speculative decoding operates by having a smaller model generate a draft. A larger target model then reviews this draft to align with its output, and any acceptance by the target model results in a reduction of the number of the target model runs, ultimately improving efficiency. However, the drafting process in speculative decoding includes slow autoregressive generation and allocates equal time to generating tokens, irrespective of their importance. These inefficiencies collectively contribute to the suboptimal performance of speculative decoding. To further improve LLM inference, we introduce Cascade Speculative Drafting (CS Drafting), a speculative execution algorithm that incorporates two types of cascades. The Vertical Cascade eliminates autoregressive generation from neural models, while the Horizontal Cascade optimizes time allocation in drafting for improved efficiency. Combining both cascades, CS Drafting achieves greater speedup compared to the baselines in our experiments, while preserving the same output distribution as the target model.

1 Introduction

Speculative decoding reduces LLM inference latency but remains limited by autoregressive drafting and equal time allocation across tokens. CS Drafting addresses these bottlenecks with vertical and horizontal cascades, achieving up to 81% additional speedup over speculative decoding.

  • Speculative decoding lets a smaller draft model propose multiple tokens for parallel target-model review, typically providing 2 to 3 times speedup without changing output distributions.The target model accepts aligned tokens and generates one additional token per review run.
  • Drafting latency limits speculative decoding because smaller draft models generate tokens autoregressively, despite larger drafts often producing higher-quality proposals.This cost motivates improving the drafting process rather than only enlarging the draft model.
  • Later draft tokens are progressively less likely to be accepted yet incur the same latency, making their equal-cost generation inefficient.Figure 2 examines acceptance probability by token position across FLAN-T5 model sizes on GSM8K and MMLU.
  • The vertical cascade uses smaller models to accelerate neural-model drafting, while the horizontal cascade assigns smaller models to later, high-rejection tokens.The two cascades are combined in CS Drafting as complementary efficiency mechanisms.
  • Up to 81% additional speedup over speculative decoding was achieved across tasks and settings while preserving generation quality.The reported comparison spans latency evaluations and maintains the target model’s output distribution.

2 Preliminary

Speculative decoding samples a draft autoregressively and has a larger target model review it, accepting aligned tokens while preserving the target model’s output distribution. This mechanism accelerates inference without task-specific retraining or architecture changes.

  • Speculative decoding samples k tokens autoregressively from a draft model, then has the target model review them in parallel.Accepted tokens are retained, rejected tokens are resampled from a recalibrated distribution, and following tokens are discarded.
  • The review rule accepts a draft token when its draft probability does not exceed the target probability; otherwise, rejection and resampling may occur.The target model generates one additional token after the accepted prefix.
  • The procedure guarantees the same output distribution as autoregressive sampling from the target model alone.This preserves generation behavior while reducing the number of target-model runs.
  • Speculative decoding has demonstrated 2x-3x faster inference without changing outputs, task-specific training, model architectures, or training procedures.The reported practical benefits cover varied tasks and model sizes.

3 Cascade Speculative Drafting

CS Drafting combines vertical and horizontal cascades to reduce neural drafting costs and allocate effort according to token acceptance likelihood. The method uses recursive smaller-model review and progressively smaller drafters while preserving the target distribution when lenience is not applied to target review.

  • 3.1 Vertical Cascade: The vertical cascade recursively uses a smaller model to draft for a larger draft model, continuing to a low-cost statistical model and avoiding neural autoregressive drafting.Each recursion is intended to reduce drafting latency without altering the output distribution.
  • 3.1 Vertical Cascade: Lenience can accelerate intermediate draft reviews, but withholding it when the target model reviews preserves the final output distribution.Lenience changes acceptance conditions and can otherwise allow outputs to differ from the target model.
  • 3.2 Horizontal Cascade: Horizontal Cascade assigns the largest draft model to early tokens and progressively smaller models to later tokens that are less likely to be accepted.Under a Bernoulli acceptance model, the probability of accepting the n-th token decreases as p^n.
  • 3.3 Max-Gram for Better Statistical Drafting: Max-Gram supports both cascades by greedily matching input or existing-generation suffixes and falling back to a Wikipedia-based bigram model.It serves as the smallest statistical draft model, avoiding autoregressive neural generation.
  • 3.4 Algorithm: The algorithm uses an upper-triangular K_nn matrix to set stopping criteria for recursive layers and assumes a universal lenience value except during target review.This design enables lenience-based speedups while preserving the target distribution at the final review.

4 Analysis

The analysis formalizes speculative decoding efficiency through expected walltime improvement and acceptance-cost parameters, then derives how vertical and horizontal cascades affect it. The results indicate that negligible-cost statistical drafting can improve efficiency, while horizontal cascading can outperform vanilla speculative decoding in simulation.

  • Analysis Setup: Expected acceptance rate α(Mt, Md) measures the probability that draft-model output is accepted by the target, while c(Mt, Md) is their single-run time ratio.These quantities parameterize the expected walltime improvement factor used throughout the analysis.
  • Vertical Cascade: Generating functions are used to derive the probability distribution and EWIF for speculative decoding and vertical cascades.The analysis introduces a probability generating function for the number of generated tokens, then applies it to systems with multiple draft models.
  • Vertical Cascade: A statistical model with negligible cost can almost always improve the efficiency of a speculative decoding system.The vertical-cascade analysis shows that adding a statistical model improves EWIF when its cost relative to the target is much less than one.
  • Horizontal Cascade: The horizontal-cascade analysis uses EWIF to study walltime improvement and the importance of tokens at different drafting positions.The associated derivative analysis examines how the expected improvement changes with token-level acceptance parameters.
  • Horizontal Cascade: Horizontal-cascade simulations achieved better EWIF than vanilla speculative sampling on CNN Dailymail and WMT EnDe.The simulation assumes Bernoulli-distributed target acceptance events and compares FLAN-T5-BASE and FLAN-T5-SMALL settings.

5 Experiments

Experiments evaluate CS Drafting across datasets, model architectures, cascade configurations, and walltime metrics, comparing it with speculative decoding and Medusa.

  • Metrics: SWI estimates speedup using constant per-forward-run costs, while walltime measures actual GPU throughput in tokens per second.SWI uses model size or prior reported timing data, whereas walltime is more user-representative but noisier and less reproducible.
  • Experimental Results: 44% additional speedup on GSM8K and 81% on MMLU over speculative decoding demonstrate CS Drafting’s main improvement under SWI.Table 2 reports these maximum gains across the evaluated FLAN-T5 settings.
  • Effectiveness of MaG: CS Drafting with one neural model achieved up to 70% speedup on MMLU and 32% on GSM8K against the fastest speculative-decoding setup.The accompanying MaG uses only a tokenizer-sized bigram model, with negligible memory cost and no additional deployment overhead.
  • Draft Model Size: With MaG assistance, FLAN-T5-BASE consistently outperformed FLAN-T5-SMALL as the draft model, reversing their usual speculative-decoding relationship.This suggests the preferred single-draft-model size may increase when MaG assists drafting.
  • Results of Decoder-only Models: On decoder-only Vicuna-7B experiments, CS Drafting improved walltime over speculative decoding and exceeded Medusa when combined with tree attention.The authors identify integration with other efficient decoding designs as a direction for future work.
  • Ablation Study: Removing the horizontal cascade reduced GSM8K performance from 56.16 to 53.55 with Vicuna-7B and K00 = 1.Performance declined only slightly under sub-optimal K00 settings, allowing simpler near-optimal configurations for users who do not require maximum performance.

6 Related Work

Related work improves speculative decoding by lowering rejection rates, reusing target models, generating candidate trees, or accelerating drafting.

  • Efficient Inference: Efficient inference also includes pruning, knowledge distillation, and quantization, which reduce parameters, storage, or computation in other ways.These methods provide broader alternatives for lowering inference cost.
  • Speculative Decoding: Prior approaches reduce rejection rates through generalized knowledge distillation, self-drafting, or tree attention that generates multiple candidate continuations.These methods target acceptance or drafting structure from different directions.
  • Speculative Decoding: Spector et al. use speculative decoding for drafting, but their method has only two layers and does not model recursive vertical cascades or draft-model lenience.The paper identifies these omitted properties as important to vertical-cascade performance.

7 Conclusion

CS Drafting combines vertical and horizontal cascades to accelerate language-model inference while preserving the target model’s output distribution.

  • Conclusion: CS Drafting’s vertical cascade removes neural autoregressive drafting, while its horizontal cascade allocates drafting cost across token positions.Together, the cascades provide additional speedup over baselines across various settings without changing the target output distribution.

8 Limitations

The evaluation may vary across hardware configurations, so the reported improvements are accompanied by standardized and raw walltime measurements.

  • Hardware differences may change experimental outcomes, motivating reports of both standardized walltime improvement and raw walltime.The theoretical analyses are intended to remain independent of hardware configuration.

A Max-Gram Implementation

The implementation uses Max-Gram and fixed cascade hyperparameters to control drafting across model types and datasets.

  • A Max-Gram Implementation: Max-Gram generates tokens by locating matching subsequences and returning the continuation after the matched index.The implementation searches progressively longer suffixes and returns a fixed-length continuation when a match is found.
  • B Hyperparameters: MaG generates 10 tokens at once, while decoder-only experiments use fixed k-matrices across datasets to avoid hyperparameter tuning.With tree attention, the decoder-only configuration uses a reduced k-matrix and limits branching to the highest-probability lead node.
  • B Hyperparameters: The implementation avoids lenience for the target-model reviewer and between MaG and its reviewer to preserve output distribution and manage performance gaps.These constraints leave at most four hyperparameters in the described setup.

C.1 Proof for Theorem 4.1

The proof characterizes token-acceptance probabilities and relates them to the expected improvement formulation used for the experimental comparison.

  • C.1 Proof for Theorem 4.1: The probability of accepting exactly i tokens is αi − αi+1, while accepting k + 1 tokens has probability αk.The latter event requires all first i tokens to be accepted and the next token to be rejected.
  • C.1 Proof for Theorem 4.1: Table 5 reports FLAN-T5 speedups under model-size and prior-work latency assumptions alongside the cascade hyperparameters.The table defines Speedup (MS) as standardized walltime improvement and Speedup (PW) using latency data from previous work [14].
  • C.1 Proof for Theorem 4.1: The proof rearranges the acceptance expression into a form that is easier to work with.

C.2 Proof for Theorem 4.3

The proof models vertical-cascade acceptance with probability-generating functions and derives expected token acceptance and walltime improvement from cascade calls.

  • C.2 Proof for Theorem 4.3: The coefficient of xj in ϕn(x) gives the probability that n speculative steps accept a total of j tokens before target-model generation.The generating function therefore represents the acceptance distribution immediately before Mt runs.
  • C.2 Proof for Theorem 4.3: An operator maps the acceptance probability-generating function to the expected number of accepted tokens.The construction begins with a single polynomial term and extends by linearity to any polynomial.
  • C.2 Proof for Theorem 4.3: The vertical cascade makes n calls to Md1 and nk calls to Md2, yielding time cost 1 + ncd1 + nkcd2 in the EWIF expression.EWIF is the expected time improvement under the i.i.d. token-acceptance assumption.
  • C.2 Proof for Theorem 4.3: With nkcd2 = 0, the derived EWIF exceeds the corresponding speculative-decoding EWIF for step size n.
Loading 2312.11462v5…