PaperlayerView exact revision on arXiv

Source-linked AI summary

Turning Trash into Treasure: Accelerating Inference of Large Language Models with Token Recycling

Xianzhen Luo, Yixuan Wang, Qingfu Zhu, Zhiming Zhang, Xuanyu Zhang, Qing Yang, Dongliang Xu

arXiv:2408.08696v3cs.CLcs.LG

TL;DR

Autoregressive decoding contributes to LLM inference latency, while existing speculative-decoding approaches face trade-offs in training, storage, retrieval cost, or adaptability. Token Recycling stores and reuses candidate tokens in an updated adjacency matrix. It reports nearly 2x speedup with less than 2MB of storage and over 31% improvement over previous train-free approaches.

  • Problem

    Autoregressive LLM decoding generates one token per step, and existing speculative-decoding methods face trade-offs involving extra training, storage, retrieval cost, or adaptability.

    Source: §1 · p. 1

  • Method

    Token Recycling stores candidate tokens in an adjacency matrix, retrieves them to construct a draft tree for tree-attention verification, and updates the matrix with new candidates.

    Source: §1 · p. 2, §7 · p. 9

  • Results

    Nearly 2x speedup with less than 2MB of storage, improving over 31% compared to previous train-free approaches.

    Source: §7 · p. 9

  • Takeaways & Limitations

    TR can be integrated seamlessly with existing LLMs and tasks as a train-free approach.

    Source: §7 · p. 9

    The draft tree uses a static structure; a dynamic tree could adapt better to each decoding step but would introduce additional complexity.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Massive parameters of LLMs have made inference latency a fundamental bottleneck. Speculative decoding represents a lossless approach to accelerate inference through a guess-and-verify paradigm. Some methods rely on additional architectures to guess draft tokens, which need extra training before use. Alternatively, retrieval-based training-free techniques build libraries from pre-existing corpora or by n-gram generation. However, they face challenges like large storage requirements, time-consuming retrieval, and limited adaptability. Observing that candidate tokens generated during the decoding process are likely to reoccur in future sequences, we propose Token Recycling. It stores candidate tokens in an adjacency matrix and employs a breadth-first-search (BFS)-like algorithm to construct a draft tree, which is then validated through tree attention. New candidate tokens from the decoding process are then used to update the matrix. Token Recycling requires \textless2MB of additional storage and achieves approximately 2x speedup across all sizes of LLMs. It significantly outperforms existing train-free methods by 30\% and even a widely recognized training method by 25\%.

1 Introduction

Autoregressive decoding makes LLM inference latency a bottleneck, while existing speculative-decoding approaches trade off training needs, storage, retrieval cost, or adaptability. Token Recycling reuses candidate tokens in an updated adjacency matrix and reports approximately 2x speedup across LLM sizes.

1 Introduction

  • Retrieval-based methods trade off static, storage-heavy or slow libraries, limited continuation options, and the extra inference generation used to build n-grams.The cited examples are REST, PLD, and Lookahead.Source: §1 · p. 2
  • TR stores candidate tokens in an adjacency matrix, retrieves them to build a draft tree with a BFS-like approach, verifies it using tree attention, and updates the matrix with newly generated candidates.The matrix reuses candidates that other speculative-decoding approaches discard.Source: Figure 1 · p. 1, §1 · p. 2
  • Experiments on SpecBench and code-domain MBPP with Vicuna and Code Llama report improvements over prior train-free methods across 7b, 13b, and 33b/34b models, and performance exceeding Medusa.Source: §1 · p. 2
  • TR requires less than 2MB of storage, has low retrieval cost, covers many new tokens, and updates its retrieval space dynamically.The update uses candidate tokens generated during decoding rather than additional generation.Source: §1 · p. 2
  • Approximately 2x speedup across LLM sizes, with over 31% improvement over prior train-free approaches and performance exceeding a training approach.The reported result covers all LLM sizes.Source: §1 · p. 2

2 Background

The background describes autoregressive decoding and the guess-and-verify basis of speculative decoding. Tree attention enables parallel verification of multiple candidate sequences while preserving their branching structure.

2.1 Auto-Regressive Decoding

  • At each autoregressive decoding step, the LLM computes a next-token probability distribution from the current content and selects the next token; this paper focuses on greedy decoding.Source: §2.1 · pp. 2–3

2.2 Speculative Decoding

  • Speculative decoding guesses subsequent draft tokens, checks them against the LLM's predictions, and confirms tokens until a draft token fails to match.Parallel processing lets one forward pass confirm multiple tokens.Source: §2.2 · p. 3
  • Tree attention verifies multiple draft sequences in parallel by merging them while preserving tree structure, so each child attends only to its parents; the longest correct sequence is selected.In rare equal-probability cases, it can select a different token than autoregressive decoding, with minimal reported quality impact.Source: §2.2 · p. 3

3 Methodology

Token Recycling stores candidate tokens in a hot-start adjacency matrix, expands them into draft trees for tree-attention verification, and updates the matrix with new candidates. It recycles candidates without additional model structures or training.

3 Methodology

  • TR uses a hot-start adjacency matrix and BFS-like retrieval to build draft trees for tree-attention verification, then recycles candidates without extra model structures or training.Source: §3 · p. 3, §3.3 · p. 5

3.1 Adjacency Matrix Initialization

  • M stores each vocabulary token’s top-k candidates in a parallel-friendly matrix; hot-starting it leverages prior knowledge and broadens initial coverage.Source: §3.1 · pp. 3–4

3.2 Draft Tree Retrieval

  • A static, imbalanced tree keeps child counts fixed for parallel traversal while giving earlier, higher-probability nodes more children and depth.Source: §3.2 · p. 4
  • Starting from the current content’s last token, BFS-like retrieval expands matrix candidates layer by layer, exploring multiple continuations to form a draft tree.The resulting tree can be verified through tree attention.Source: §3.2 · pp. 4–5

3.3 Verification and Update

  • Tree attention verifies draft sequences and keeps the longest correct one; parallel output-distribution updates overwrite M, adapting candidates to current content without extra update operations.Source: §3.3 · p. 5
  • Parallel updates can mix top-k lists for duplicate tokens with differing distributions; the authors leave these variations unresolved because controls reduce performance.Source: §3.3 · p. 5

4 Experiment

Experiments compare Token Recycling (TR) with retrieval-based train-free methods and Medusa on SpecBench and MBPP. TR delivers near-2x speedups across model sizes, with results varying by task and memory trade-off.

4.1 Experimental Setup

  • Experiments use SpecBench with Vicuna and MBPP with Code Llama at 7B, 13B, and 33B, comparing TR with three train-free retrieval methods and Medusa.The evaluation focuses on greedy decoding with batch size one; the largest Code Llama model is 34B, referred to as 33B for consistency.Source: §4.1 · pp. 5–6

4.2 Main Results

  • On SpecBench, TR provides nearly 2x speedup across model sizes, with nearly 30% higher speedup than prior train-free methods and almost 25% improvement over Medusa on 7B.For 13B and 33B models, the reported speedup remains nearly 2x with a 30% acceleration advantage.Source: Table 1 · p. 6, §4.2 · p. 6
  • TR achieves its best speedup on most SpecBench subtasks, slightly trailing PLD on Sum; its RAG improvement rises from under 3% on 7B to over 10% on 33B.On Trans, TR exceeds 1.75x speedup across model sizes, while PLD is near 1x and REST has its lowest speedup across tasks.Source: §4.2 · p. 6
  • On MBPP, TR provides approximately 2.3x speedup across all model scales and achieves state-of-the-art performance.The comparison uses Code Llama; the table excludes Medusa because it lacks a Code Llama variant.Source: §4.2 · p. 6, Table 1 · p. 6
  • TR uses minimal additional memory while outperforming REST and Lade in speedup; PLD uses no additional memory but has limited speedup.Medusa involves training, whereas TR is training-free.Source: §4.2 · pp. 6–7, Table 2 · p. 7

5 Analysis

The analysis examines tree size, initialization, update choices, sampling temperature, comparisons with Eagle, and where Token Recycling spends time. The results show trade-offs in tree construction and update strategy, while retrieval and updating add negligible latency.

5.1 Tree Structure

  • Above breadth 80, Tokens/s decreases as added computational overhead outweighs the MAT gains from covering more possibilities.The experiment expands breadth while holding tree depth at six layers.Source: §5.1 · p. 7
  • With 80 nodes, increasing depth raises MAT rapidly before gains diminish, while Tokens/s fluctuates as adjacent-step storage weakens distant-token connections.The authors attribute reduced effectiveness at greater depth to the matrix storing candidates only for adjacent steps.Source: §5.1 · p. 7

5.2 Ablation Study

  • Updating from all draft tokens, including rejected ones, significantly improves MAT over updating only from accepted tokens.The authors interpret rejected tokens as carrying information useful for subsequent decoding.Source: Table 4 · p. 8, §5.2 · p. 8
  • Fixed, shuffled, and inherited matrices improve over Zero, while Random sharply decreases performance; the small differences among the stronger strategies suggest generalizable patterns.The initialization strategies include a matrix from 100 unrelated AlpacaEval queries, a shuffled test set, random vocabulary tokens, and an all-zero matrix.Source: §5.2 · pp. 7–8, Table 3 · p. 7
  • The last-occurrence update yields the highest MAT, but the current parallel-CUDA merge achieves higher Tokens/s by avoiding token-position processing.MAT differences among the three update strategies are minimal.Source: §5.2 · p. 8
  • At temperature 0.3, performance remains unaffected versus greedy decoding; temperatures 0.5 and 1 cause slight degradation, while speedup stays around 2.0.The authors suggest higher temperatures increase selection of tokens outside the stored top-k candidates.Source: Figure 4 · p. 8, §5.2 · p. 8

5.3 Compare with Eagle

  • TR reaches 91.23% of Eagle2’s acceleration performance using 0.39% of its memory, and outperforms Eagle1 despite being training-free.Eagle2 uses a dynamic tree, whereas TR uses a static tree.Source: Table 5 · p. 8, §5.3 · pp. 8–9

5.4 Time Allocation

  • Model forward consumes most response time and verification the second most, while preprocessing, retrieval, and update add negligible latency.Verification includes extracting and validating feasible paths.Source: Figure 5 · p. 9, §5.4 · p. 9

6 Related Work

Related work includes speculative decoding methods that train additional draft structures and training-free methods that retrieve draft tokens. Token Recycling is presented as training-free, with low storage and retrieval costs and an adaptable retrieval space.

Training-dependent methods

  • Draft-model and parameter-efficient speculative methods require additional training or model-specific adaptation, which the paper describes as resource-intensive and limiting scalability.The approaches include smaller draft models and parameter-efficient structures.Source: §6 · p. 9

Train-free retrieval methods

  • Lookahead generates n-grams during decoding, PLD retrieves from prior content, and REST uses existing corpora; their stated trade-offs include computation, limited new-content coverage, storage, retrieval time, or query adaptation.PLD is described as providing little acceleration for tasks requiring new content, while REST's library cannot adapt to each query.Source: §6 · p. 9

Token Recycling

  • Token Recycling is characterized as requiring no additional generation, covering broader continuations, using little storage with low retrieval costs, and adapting its retrieval space through updates.The related-work passage presents these as contrasts with existing retrieval-library approaches.Source: §6 · p. 9

7 Conclusion

Token Recycling accelerates LLM inference by storing candidate tokens in an adjacency matrix, retrieving and verifying a draft tree, and updating the matrix with newly generated candidates. The paper reports nearly 2x speedup with less than 2MB of storage and improvement over previous training-free approaches.

Conclusion

  • Token Recycling stores candidate tokens in an adjacency matrix, retrieves a draft tree for tree-attention verification, and updates the matrix with candidates generated during decoding.The method is described as seamlessly integrable with existing LLMs and tasks.Source: §7 · p. 9
  • Nearly 2x speedup with <2MB of storage is reported for this training-free approach.Source: §7 · p. 9
  • Token Recycling improves over previous training-free approaches by more than 31%.Source: §7 · p. 9

Limitations

The stated limitation is that Token Recycling uses a static draft tree; the authors identify a dynamic tree as a possible alternative that could better adapt to each decoding step, at added complexity.

Limitations

  • The current draft tree is static, rather than dynamic.Source: Limitations · p. 10
  • Dynamic trees add complexity, according to the authors.Source: Limitations · p. 10
  • Dynamic trees could adapt better to each decoding step and potentially improve performance by tailoring structure to query requirements.Source: Limitations · p. 10

Ethical Considerations

The experiments use publicly accessible project resources, datasets, and baseline models, and acknowledge the original authors.

Ethical Considerations

  • The proposed-method data, experimental datasets, and baseline models are publicly accessible, and the original authors are acknowledged.The data comes solely from publicly accessible project resources; the authors report that no sensitive information is included.Source: Ethical Considerations · p. 10

A.1 Identical Probability Tokens

This appendix explains how Token Recycling handles near-tied probabilities, constructs its verification tree, reuses candidates, and updates its matrix. It also reports a case-study example and a trade-off between matrix-update strategies.

A.1 Identical Probability Tokens

  • Tree Attention can change the argmax for tied probabilities because mask-token softmax weights are tiny but nonzero; the rare issue negligibly affects accuracy and performance.The appendix attributes identical probabilities to floating-point precision and reports minimal impact on answer accuracy and quality.Source: §A.1 · p. 13, Table 6 · p. 13

A.2 Draft Tree Algorithm and Structure

  • The static, unbalanced verification tree is iteratively pruned using validation-set correctness, ending with 81 nodes across 6 layers and 79 draft inputs per forward.A BFS-like algorithm constructs the current-content draft tree from this structure.Source: §A.2 · p. 13

A.3 Reuse Mechanism Analysis

  • On SpecBench, candidates from rejected tokens yield significantly greater performance gain than accepted-only candidates; Token Recycling stores candidates for all draft tokens.The paper connects this reuse to common patterns in earlier generations that often meet local semantic needs.Source: §A.3 · p. 14

A.4 Case Study

  • In the MT-Bench case study, a rejected-token candidate predicts “a” after “be,” while an accepted-token candidate predicts “great” after “a.”Both candidate matches occur in the illustrated decoding rounds.Source: §A.4 · pp. 14–15

A.5 Explanation of Three Update Strategies

  • Position-controlled updates slightly affect Mean Accepted Tokens, but extracting positions adds latency and reduces overall acceleration versus uncontrolled updates.The appendix contrasts uncontrolled updates with first-occurrence and last-occurrence position selection.Source: §A.5 · p. 15

Related papers

Loading 2408.08696v3…