PaperlayerView exact revision on arXiv

Source-linked AI summary

SAM Decoding: Speculative Decoding via Suffix Automaton

Yuxuan Hu, Ke Wang, Xiaokang Zhang, Fanjin Zhang, Cuiping Li, Hong Chen, Jing Zhang

arXiv:2411.10666v3cs.CLcs.AI

TL;DR

Retrieval-based speculative decoding is limited by narrow retrieval sources, retrieval-efficiency limitations, and specialized-domain fit. SAM-Decoding uses suffix automata over corpus and generated text for exact longest-suffix drafting, and is 18%+ faster than previous retrieval-based methods on Spec-Bench.

  • Problem

    Retrieval-based speculative decoding often relies on incomplete retrieval resources, inefficient retrieval techniques, and domains where it yields limited acceleration.

    Source: §1 · pp. 1–2

  • Method

    SAM-Decoding uses suffix automata over a text corpus and current generated sequence to retrieve drafts from exact longest-suffix matches, with adaptive integration of other speculative-decoding methods.

    Source: §1 · p. 2, §3.1 · p. 4, §3.2 · p. 4

  • Results

    On Spec-Bench, SAM-Decoding is 18%+ faster than previous retrieval-based speculative decoding methods; with EAGLE-2, it adds 3.28%–11.13% speedup on MT-Bench across various LLM backbones.

    Source: §1 · p. 2

  • Takeaways & Limitations

    When combined with state-of-the-art techniques, SAM-Decoding enhances performance in multi-turn conversation, summarization, retrieval-augmented generation, and context Q&A tasks.

    Source: §6 · p. 9

    Performance depends on task type and retrieval-source quality, and the current corpus may be insufficiently diverse or may differ from text generated by other LLMs.

    Source: §7 · p. 9

Abstract

from arXiv · show

Speculative decoding (SD) has been demonstrated as an effective technique for lossless LLM inference acceleration. Retrieval-based SD methods, one kind of model-free method, have yielded promising speedup, but they often rely on incomplete retrieval resources, inefficient retrieval methods, and are constrained to certain domains. This paper presents a novel retrieval-based speculative decoding method that adapts suffix automaton (SAM) for efficient and accurate draft generation by utilizing common text corpus and dynamic text sequence. Unlike existing $n$-gram matching methods, SAM-Decoding finds the exact longest suffix match, achieving an average time complexity of O(1) per generation step of SAM update and suffix retrieval. It can also integrate with existing methods, adaptively selecting a draft generation strategy based on match length to generalize to broader domains. Extensive experiments on Spec-Bench show that our method is $18\%+$ faster than other retrieval-based SD methods. Additionally, when combined with advanced EAGLE-2, it provides an additional speedup of $3.28\%$ -- $11.13\%$ across various-sized LLM backbones. Our code is available at our \href{https://github.com/hyx1999/SAM-Decoding}{repository}.

1 Introduction

Retrieval-based speculative decoding faces limits in retrieval coverage, efficiency, and domain range. SAM-Decoding uses suffix automata over corpus and generated text to draft tokens, reporting faster performance than retrieval-based baselines and added speedups when combined with EAGLE-2.

1 Introduction

  • Existing retrieval-based methods rely on a single source, have retrieval-efficiency limitations, and are often suited to specialized domains.PLD retrieves from current text, while REST uses a text corpus; the paper notes limited acceleration in other domains.Source: §1 · pp. 1–2
  • SAM-Decoding retrieves from both a common text corpus and the current sequence, using suffix automata to find exact longest-suffix matches with average O(1) update and retrieval time.The method records substring positions in static and dynamic automata and uses the matching positions and lengths to retrieve drafts.Source: §1 · p. 2
  • SAM-Decoding can integrate with other speculative-decoding methods and select a draft strategy according to match length.This allows auxiliary decoding techniques to be selected when retrieval is not applicable.Source: §1 · p. 2
  • 18%+ faster than previous retrieval-based speculative decoding methods on Spec-Bench; with EAGLE-2, SAM-Decoding adds 3.28%–11.13% speedup on MT-Bench.The MT-Bench comparison covers various LLM backbones.Source: §1 · p. 2

2 Background

A suffix automaton indexes substrings and supports fast matching of the longest suffix of a progressively generated sequence. Speculative decoding drafts multiple tokens for simultaneous target-model verification, with tree-based verification as an extension.

2.1 Suffix Automaton

  • A suffix automaton indexes substrings, with nodes representing states and edges encoding character extensions or links to states for shorter suffixes.The structure can be built incrementally in O(L) time for a string of length L.Source: §2.1 · p. 2
  • For a progressively generated sequence, the suffix automaton finds the longest matching suffix at each step with average O(1) time complexity.The extension and suffix-link edges support these transitions.Source: §2.1 · p. 2

2.2 Speculative Decoding

  • Speculative decoding drafts multiple candidate tokens quickly, then the target LLM evaluates them simultaneously and accepts tokens aligned with its output distribution.The draft has n tokens and the accepted sequence has m tokens.Source: §2.2 · pp. 2–3
  • Tree-mask verification lets the target LLM evaluate multiple branches of a candidate-token tree in parallel to increase draft acceptance length.Source: §2.2 · p. 3

3 SAM-Decoding

SAM-Decoding addresses retrieval-based speculative decoding's source, efficiency, and domain limits with suffix automata over corpus and live text. It selects drafts using match information and can pair with other speculative decoding methods.

3 SAM-Decoding

  • SAM-Decoding uses diverse text sources and suffix automata to address retrieval-source, efficiency, and n-gram-length limits, while integrating with model-free and model-based methods.The method also targets performance beyond specialized domains such as summarization and RAG.Source: §3 · p. 3

3.1 Suffix Automaton Construction

  • A static automaton is built offline for the common corpus, while a dynamic automaton is expanded incrementally for the current text; nodes record earliest substring positions as min_endpos.The current text includes user prompts and already generated tokens.Source: §3.1 · pp. 3–4

3.2 Drafting with Suffix Automaton

  • Dynamic drafts use n tokens after min_endpos; static drafts use a Prim-built tree based on top-k successor states.The dynamic automaton's matched state supplies the min_endpos used to locate the draft in reference text.Source: §3.2 · p. 4
  • Dynamic-automaton drafts are prioritized; a static draft is used only when its match length l1 exceeds the dynamic match length l2 by lbias.The paper reports that dynamic-automaton drafts often outperform drafts from the static text corpus.Source: §3.2 · p. 4
  • Average state-transfer complexity is O(1) (worst case O(L)); SAM finds exact longest suffix matches without a length limit, versus PLD's O(n^2L) and REST's O(n^2 log L).Here, n is the predefined maximum matching length and L is the current text or concatenated corpus length.Source: §3.2 · p. 4

3.3 Update of Suffix Automaton

  • After LLM verification, accepted tokens update the static automaton by state transfer; the dynamic automaton transfers its match state and then expands.The dynamic automaton represents the generated text sequence.Source: §3.3 · p. 5

3.4 Adaptive Draft Selection

  • When automaton drafts are unsatisfactory, an auxiliary speculative method can be used, with draft selection based on automaton match length; options include Token Recycling and EAGLE-2.The auxiliary method receives a fixed virtual match length as its selection threshold.Source: §3.4 · p. 5

4 Experiments

Experiments compare SAM-Decoding with retrieval-based and auxiliary speculative-decoding baselines across Spec-Bench, HumanEval, HAGRID, and multiple model sizes. Ablations examine parameter settings, draft size, and the static and dynamic suffix-automaton modules.

Experiment Results

  • 1.84× on Spec-Bench, 2.29× on HumanEval, and 2.24× on HAGRID: SAM-Decoding beats retrieval-based baselines on inference speedup.Table 1 reports these speedup ratios for Vicuna-7B-v1.3.Source: Table 1 · p. 6, §4 · p. 6
  • Combining SAM-Decoding with Token Recycling or EAGLE-2 further improves speedups on Spec-Bench and HAGRID.Spec-Bench: Token Recycling rises from 1.84× to 2.27× and EAGLE-2 from 2.38× to 2.58×; HAGRID: Token Recycling rises from 1.93× to 2.57× and EAGLE-2 from 2.41× to 2.81×.Source: §4 · p. 7
  • On HumanEval, EAGLE-2 changes slightly after integration, while Token Recycling rises from 1.94× to 2.45×.The authors attribute EAGLE-2's slight change to code generation being less likely to copy generated text.Source: §4 · p. 7
  • SAM-Decoding beats retrieval-based baselines across Spec-Bench tasks, and integration improves Token Recycling and EAGLE-2 on three retrieval-amenable tasks.For Multi-turn Conversation, Summarization, and Retrieval-augmented Generation, Token Recycling rises from 1.92×, 1.96×, and 1.68× to 2.48×, 2.86×, and 2.14×; EAGLE-2 rises from 2.87×, 2.33×, and 2.03× to 3.02×, 2.76×, and 2.23×, respectively.Source: Figure 4 · p. 6, Figure 5 · p. 6, §4 · p. 7
  • On MT-Bench, SAM-Decoding outperforms retrieval-based baselines across Vicuna-7B, Vicuna-13B, and Vicuna-33B, and can further improve model-free and model-based methods.The passage refers readers to Appendix B for additional experimental results.Source: §4 · p. 7

Ablation Experiments

  • Both mean accepted tokens and speedup ratio rise as lbias and lthreshold increase up to 5, then decline when both exceed 5.Figure 6 presents the parameter sweep for SAM-Decoding[T].Source: §4 · p. 7, Figure 6 · p. 7
  • Throughput rises with draft size up to 40, then declines beyond 40, with the decline more pronounced at 70.The authors attribute the slowdown beyond the average acceptance-length threshold to GPU-capacity strain; Figure 7 shows throughput across draft sizes.Source: §4 · pp. 7–8, Figure 7 · p. 7
  • Removing either suffix-automaton module reduces acceleration, with the dynamic module having the greater impact than the static module.The authors suggest that drafts from dynamic context are often more effective than retrieval from a pre-existing text corpus; Table 2 reports the module comparison.Source: §4 · p. 8, Table 2 · p. 8

5 Related Work

Related work covers model-based and retrieval-based speculative decoding, other model-dependent approaches, and architecture-focused inference acceleration.

Speculative Decoding

  • Model-based speculative decoding uses smaller neural networks to draft tokens; examples include pre-existing small LLMs, Medusa's decoding heads, and EAGLE-2's shallow Transformer layers.The passage notes that model-based methods can struggle to generate long drafts because of drafting overhead or declining prediction accuracy.Source: §5 · p. 8
  • Retrieval-based methods draft through retrieval and n-gram matching, while Token Recycling drafts from previously generated token distributions.Named retrieval-based examples include Lookahead Decoding, PIA, PLD, and REST.Source: §5 · p. 8
  • Other speculative-decoding research relies on the model itself or on sub-models within a larger architecture.Source: §5 · p. 8

Efficient LLM Architecture

  • Architecture-focused inference work includes distillation, quantization, and pruning, which reduce model size, parameter storage, or computation and memory overhead.Structured pruning can be combined with distillation; semi-structured pruning can reduce memory-access and computing overhead when combined with special hardware.Source: §5 · p. 8

6 Conclusion

SAM-Decoding uses suffix automatons built from generated text and a text corpus to retrieve drafts and can switch to other draft-generation methods when retrieval is infeasible. The authors report that it outperforms retrieval-based baselines and improves existing methods on several tasks.

6 Conclusion

  • SAM-Decoding constructs suffix automatons from both generated text and a text corpus to retrieve drafts.Source: §6 · p. 8
  • The authors report that SAM-Decoding outperforms retrieval-based speculative-decoding baselines.Source: §6 · pp. 8–9
  • When combined with state-of-the-art methods, SAM-Decoding improves performance in Multi-turn Conversation, Summarization, Retrieval-augmented Generation, and Context Q&A.Source: §6 · p. 9

7 Limitation

SAM-Decoding’s performance depends on task type and retrieval-source quality, while its current corpus is limited in diversity and may differ from text generated by other LLMs. Its method-selection heuristic does not fully use exact match lengths, and existing datasets may not reflect real usage scenarios for retrieval-based methods.

7 Limitation

  • Performance depends on task type and retrieval-source quality; the current corpus is not diverse enough and may diverge from other LLMs’ generated text.The corpus was collected from Vicuna-7B results on Stanford-alpaca, GSM8K, and Python-instruct-18k; the authors call for more specialized, diverse corpora.Source: §7 · p. 9
  • The match-length-based heuristic for selecting combined decoding methods does not fully use SAM’s exact match lengths.The authors plan to train a classifier to select decoding methods at each generation round.Source: §7 · p. 9
  • Existing datasets may not reflect real usage, and retrieval-based performance is highly correlated with usage scenarios.The authors call for datasets more compatible with real scenarios to evaluate retrieval-based methods.Source: §7 · p. 9

A Suffix Automaton

The appendix describes suffix-automaton construction and a frequency-based drafting procedure that uses a maximum spanning tree to obtain top-k successors and transition probabilities.

A.1 Construction Process of Suffix Automaton

  • Build-SAM constructs the suffix automaton and Expand extends it; nodes track substring frequency, initialized to 1 for expansion nodes and 0 for clone nodes.The root starts with link −1, empty next transitions, length 0, and min_endpos 0.Source: §A.1 · p. 11

A.2 Drafting via Prim’s Algorithm

  • Prim’s-algorithm drafting uses substring frequencies to compute top-k successors and transition probabilities for static-automaton nodes.Each transition probability divides the target state’s substring frequency by the current state’s frequency.Source: §A.2 · p. 11

A.3 Time Complexity of State Transfer

The analysis bounds the average cost of suffix-automaton state transfer by O(1), while a single transition can take O(L) in the worst case.

A.3 Time Complexity of State Transfer

  • The amortized analysis defines potential as matching length, then bounds each step’s amortized cost by 2 because link edges reduce potential and next edges increase it.Summing the amortized costs yields the average-time bound.Source: §A.3 · p. 11
  • O(1) average and O(L) worst-case time characterize each state transition over a sequence of L tokens.In the worst case, one operation may traverse up to l_i link edges before taking a next edge.Source: §A.3 · pp. 11–12

B Additional Experiment Results

Across the reported benchmarks and models, SAM-Decoding outperforms retrieval-based baselines and improves speedups when paired with Token Recycling or EAGLE-2.

B Additional Experiment Results

  • 18%+ faster than prior retrieval-based speculative decoding methods on Spec-Bench; Token Recycling’s overall Spec-Bench speedup rises from 1.91× to 2.05× with SAM-Decoding.For Llama3-8B-instruct, Token Recycling speedups also rise on HumanEval from 1.99× to 2.16× and HAGRID from 2.17× to 2.30×.Source: B Additional Experiment Results · p. 12
  • With SAM-Decoding, EAGLE-2 speedups across five tasks rise from 2.08×, 1.85×, 1.87×, 2.37×, and 2.18× to 2.36×, 1.98×, 2.11×, 2.54×, and 2.35×.The task order is Multi-turn Conversation, Summarization, Retrieval-augmented Generation, Code Generation, and Context Q&A.Source: B Additional Experiment Results · p. 12
  • On Vicuna-13B-v1.3 and Vicuna-33B-v1.3, SAM-Decoding remains faster than retrieval-based baselines, while combined methods further improve speed on several tasks.The reported tasks include Multi-turn Conversation, Summarization, Retrieval-augmented Generation, and Context Q&A.Source: B Additional Experiment Results · p. 12
  • Figures 8 and 9 report inference-time shares by module and draft-module usage alongside mean accepted tokens, respectively.Source: Figure 8 · p. 12, Figure 9 · p. 12

C Additional Ablation Experiments

The additional ablations profile SAM-Decoding’s runtime and draft-module use on Spec-Bench, and report its combination with Token Recycling and EAGLE-2 on A800.

C Additional Ablation Experiments

Related papers

Loading 2411.10666v3…