Universal Assisted Generation: Faster Decoding with Any Assistant Model

TL;DR

Hugging Face and Intel Labs have developed Universal Assisted Generation (UAG), a technique that enables the use of any small language model as an assistant for speculative decoding, regardless of whether it shares the same tokenizer as the target model. This allows for inference speedups of 1.5x to 2.0x for models that previously lacked compatible small-scale variants.

The Challenge of Tokenizer Dependency in Assisted Generation

Assisted generation, or speculative decoding, accelerates LLM inference by using a small assistant model to predict a sequence of tokens, which a larger target model then verifies in a single forward pass. While this significantly reduces latency, standard assisted generation requires the target and assistant models to share the same tokenizer (i.e., they must be from the same model family).

This dependency creates a bottleneck because many high-performance models lack a sufficiently small version (typically 50-100 times smaller than the target) to provide meaningful speedups. For example, gemma-2-9b only has a 2B variant, and CodeLlama-13b has no smaller version available for this purpose.

Universal Assisted Generation (UAG) Mechanism

Universal Assisted Generation removes the tokenizer constraint by implementing a two-way tokenizer translation process. This allows a target model to be paired with any assistant model, such as using the tiny vicuna-68m to accelerate gemma-2-9b.

Two-Way Tokenizer Translation

  1. Assistant to Target: When the assistant model generates a sequence of tokens, those tokens are converted to text and then re-tokenized using the target model's tokenizer.
  2. Target to Assistant: After the target model verifies the tokens, the resulting target tokens are converted back into the assistant's token format and appended to the assistant model's context for the next iteration.

Handling Vocabulary Discrepancies

To ensure accuracy during re-encoding, UAG prepends a context window of previous tokens to the newly generated sequence. This entire sequence is re-encoded into the target format and aligned with the most recent target tokens to determine the exact insertion point for new tokens. Additionally, when translating from target to assistant, any mismatched tokens are discarded from the assistant model's key-value (KV) cache to maintain data integrity.

Performance Benchmarks

UAG provides significant latency improvements for models that lack suitable same-family assistants. The following speedups were observed across various tasks:

Target Model Assistant Model Task Speedup
CodeLlama-13b-Instruct-hf tiny_starcoder_py Code Generation 1.90x
Mixtral-8x22B-Instruct-v0.1 vicuna-68m Summarization 1.52x
gemma-2-9b vicuna-68m Summarization 1.76x
Mixtral-8x22B-Instruct-v0.1 Qwen2-0.5B-Instruct Long-context Summarization 1.78x
Llama-3.1-70B Qwen2-0.5B-Instruct Long-context Summarization 1.78x
Phi-3-medium-128k-instruct Qwen2-0.5B-Instruct Long-context Summarization 1.91x

Experiments were conducted on 100 randomly selected examples using hardware ranging from a single A6000 GPU to four A100 GPUs depending on the model size.

Implementation and Usage

Universal Assisted Generation is integrated into the 🤗 Transformers library starting with release 4.46.0. Users can implement UAG by passing both the tokenizer and assistant_tokenizer to the generate() method.

from transformers import AutoModelForCausalLM, AutoTokenizer

prompt = "Alice and Bob"
checkpoint = "google/gemma-2-9b"
assistant_checkpoint = "double7/vicuna-68m"

assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt")

model = AutoModelForCausalLM.from_pretrained(checkpoint)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)
outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)

Current Limitations and Future Work

UAG currently supports multinomial sampling when do_sample=True. Unlike speculative sampling, multinomial sampling automatically rejects a token if the target model does not sample the same token as the assistant, which can result in lower throughput compared to assistants that share the same tokenizer. Future updates plan to introduce support for speculative sampling and integrate UAG directly into 🤗 Transformers pipelines for a more streamlined user experience.

Sources