Speculative decoding: draft ahead, verify together

Speculative decoding speeds up autoregressive generation by separating proposing tokens from checking them. A cheap draft mechanism suggests several tokens. The target model evaluates that proposed continuation in a wider pass, allowing multiple accepted tokens to advance the answer at once.

The target still controls which continuation is accepted. In a greedy version, tokens can be kept while they agree with the target's choices. At the first disagreement, the system takes the target's correction and discards the remaining speculative suffix. Sampling versions use an acceptance rule and a corrected sampling distribution to preserve the target distribution.

Try an assistant model

Transformers exposes this through assistant_model. This PyTorch example uses a smaller GPT-2-family draft with the target's tokenizer. The first run downloads both models.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()
draft = AutoModelForCausalLM.from_pretrained("distilgpt2").eval()
inputs = tokenizer("Far far away", return_tensors="pt")

output = model.generate(
    **inputs, assistant_model=draft, max_new_tokens=40,
    do_sample=False, pad_token_id=tokenizer.eos_token_id)
continuation = tokenizer.decode(output[0], skip_special_tokens=True)

The library handles draft generation, verification and cache bookkeeping. The important operation is validating several proposed positions together rather than making one expensive target call for every single token.

When speculation pays

Speed depends on how often the draft is accepted and how much it costs to produce and verify candidates. A poor draft spends work on tokens that are thrown away. A busy server with large batches can have different trade-offs from one interactive request.

Lossless speculative sampling preserves a distribution, not necessarily the same random sequence for a fixed seed. The acceleration comes from scheduling and parallel verification. It complements KV caching, which reuses earlier attention state.