Research map / Large language models / Language, safety and interpretability
Nearest Neighbor Speculative Decoding for LLM Generation and Attribution
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes NEST, a training-free semi-parametric LM decoding method that combines two-stage kNN token retrieval, relative retrieval confidence interpolation, and relaxed speculative span acceptance to improve factuality, attribution, and speed of Llama-2-Chat.
Problem
LLMs hallucinate and lack attribution; existing kNN-LM is slow and produces non-fluent text, while in-context retrieval augmentation does not reliably ensure faithful attribution.
Approach
Performs token-level kNN retrieval over an on-the-fly token index built from retrieved passages, uses a relative retrieval confidence score to set the interpolation weight, extends confident tokens into corpus n-gram spans, and accepts span prefixes via a relaxed speculative decoding rule so that generated spans are verbatim corpus text with source attribution.
Key results
- On WikiText-103 with Llama-2-Chat 70B, NEST lowers perplexity from 9.9 to 6.3 and raises ROUGE-1 from 22.9 to 32.6 (the abstract cites 42.3% relative ROUGE-1 gain).
- Llama-2-Chat 70B + NEST improves Biography FActScore from 34.2 to 41.6, and reaches 1.8x inference speedup on 8xA100 for the 70B model per the abstract.
- NEST attribution ratio ranges from 33.2% to 95.5% of generated tokens traceable to the corpus across NQ and Biography settings.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add Relative Retrieval Confidence to two-stage kNN-LM (ablation) | Llama-2-Chat 7B · WikiText-103 validation | ROUGE-1: 20.1 -> 24.7 | Helped |
| Add Relative Retrieval Confidence to two-stage kNN-LM (ablation) | Llama-2-Chat 7B · NQ validation | answer-level recall: 40.8 -> 44.4 | Helped |
| Add Relative Retrieval Confidence to two-stage kNN-LM (ablation) | Llama-2-Chat 7B · Biography validation | FActScore: 34.8 -> 41.6 | Helped |
| Add dynamic span selection after Relative Retrieval Confidence (ablation) | Llama-2-Chat 7B · WikiText-103 validation | ROUGE-1: 24.7 -> 24.5 | No clear effect |
| Add dynamic span selection after Relative Retrieval Confidence (ablation) | Llama-2-Chat 7B · NQ validation | answer-level recall: 44.4 -> 44.6 | No clear effect |
| Add relaxed speculative decoding after dynamic span selection (ablation) | Llama-2-Chat 7B · Biography validation | FActScore: 41.6 -> 46.8 | Helped |
| Add relaxed speculative decoding after dynamic span selection (ablation) | Llama-2-Chat 7B · WikiText-103 validation | ROUGE-1: 24.5 -> 26.8 | Helped |
| Add Nest to the base model | Llama-2-Chat 70B · TruthfulQA | ΔBLEU: 0.14 -> 0.17 | Helped |
| Add Nest to the base model | Llama-2-Chat 70B · Biography | FActScore: 34.2 -> 41.6 | Helped |
| Add in-context retrieval augmentation to Nest | Llama-2-Chat 70B · Biography | FActScore: 41.6 -> 59.2 | Helped |
| Add Nest to Llama-2-Chat 70B | Llama-2-Chat 70B · Inference | speedup: 1.8× speedup | Helped |
Surprising findings
- In-context retrieval augmentation hurts TruthfulQA scores while NEST, which interpolates at output level, improves them.
- Combining NEST with RA does not always help; RA-NEST underperforms on some QA and TruthfulQA settings.
Limitations
- Output may still contain factual errors depending on first-stage passage retrieval and token retrieval accuracy.
- Without fine-tuning, the plug-and-play integrated system may be sub-optimal.
- Semi-parametric LMs may not improve in-context learning since demonstrations rarely appear in the corpus.
Experimental setup
Datasets: Wikipedia 2021 dump (~33M passages) Pile of Law (~15M passages) WikiText-103 Natural Questions TriviaQA HotpotQA MedMCQA Biography (FActScore) TruthfulQA MMLU
Benchmarks: WikiText-103 Pile of Law Natural Questions TriviaQA HotpotQA MedMCQA Biography TruthfulQA MMLU
Metrics: perplexity MAUVE ROUGE-1 ROUGE-2 ROUGE-L answer-level recall FActScore delta BLEU delta ROUGE-1 macro accuracy
Baselines: Llama-2-Chat kNN-LM (two-stage) in-context retrieval augmentation (RA) REST
Builds on: kNN-LM speculative decoding Copy Generator RA-DIT
Closest papers by meaning
- Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies (ICML 2025 · similarity 0.91)
- RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding (ICML 2025 · similarity 0.91)
- Scaling Up, Speeding Up: A Benchmark of Speculative Decoding for Efficient LLM Test-Time Scaling (ICLR 2026 · similarity 0.89)
- SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.89)
- Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment (ICLR 2025 · similarity 0.89)
- Mixture of Attentions For Speculative Decoding (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.