Research map / Large language models / Retrieval and knowledge
One Token Can Help! Learning Scalable and Pluggable Virtual Tokens for Retrieval-Augmented Large Language Models
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 SPRING, a frozen-LLM tuning method that learns a small set of pluggable virtual token embeddings placed between retrieved passages and the question, scalable to any number of tokens and robust to retrievers and passage counts.
Problem
Fine-tuning LLMs for retrieval-augmented generation improves RAG performance but degrades general generation abilities in non-RAG settings and causes prompt overfitting, which is problematic for already-deployed LLMs.
Approach
Freeze the LLM and train only embeddings of n virtual tokens inserted between retrieved passages and the question; a random-prefix training scheme (using first k tokens) makes the number of tokens adjustable at inference, and the tokens can be removed for non-RAG queries.
Key results
- SPRING with 50 virtual tokens on Mistral-7b-instruct reaches average EM 34.35 / F1 65.00 over nine QA datasets versus 25.66 / 57.70 for the manual prompt, using only 0.2M trainable parameters.
- LoRA fine-tuning drops non-RAG performance (e.g., GSM8K 17.33 vs SPRING 31.89 with 8-shot; BoolQ 79.30 vs 82.97), while SPRING keeps the original model's behavior when virtual tokens are removed.
- Across retrievers (BM25, BGE-base, E5-base, E5-large), SPRING gives consistent gains over prompting and reduces variance across retrievers (1.75 vs 2.69 EM).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add 50 trainable virtual tokens instead of using a manually crafted prompt | Mistral-7b-instruct · 12 QA datasets, with retrieval | EM: 19.80 -> 28.40 | Helped |
| Add 50 trainable virtual tokens instead of using a manually crafted prompt | Mistral-7b-instruct · 12 QA datasets, with retrieval | F1: 51.30 -> 60.15 | Helped |
| Use SPRING rather than LoRA for RAG fine-tuning | Mistral-7b-instruct · 12 QA datasets, with retrieval | EM: 30.78 -> 28.40 | Hurt |
| Use SPRING rather than LoRA for RAG fine-tuning | Mistral-7b-instruct · GSM8K, 8-shot | accuracy: 17.33 -> 31.89 | Helped |
| Move virtual tokens from between retrieved passages and question to the beginning (ablation) | SPRING · Average on nine QA datasets | EM: 34.35 -> 6.53 | Hurt |
| Move virtual tokens from between retrieved passages and question to the end (ablation) | SPRING · Average on nine QA datasets | F1: 65.00 -> 63.97 | Hurt |
| Train on a mixture of QA datasets rather than only TQA (ablation) | SPRING · Average on nine QA datasets | EM: 26.66 -> 34.35 | Helped |
| Train on a mixture of QA datasets rather than only NQ (ablation) | SPRING · Average on nine QA datasets | F1: 61.40 -> 65.00 | Helped |
| Evaluate the E5-large-trained virtual tokens with BM25 rather than E5-large retrieval (ablation) | SPRING · Average on nine QA datasets | EM: 34.35 -> 30.94 | Hurt |
| Fine-tune LoRA for RAG rather than use SPRING tokens | Mistral-7b-instruct · 12 QA datasets, without retrieval | EM: 16.32 -> 0.02 | Hurt |
Surprising findings
- Prefix-tuning (tokens at sequence start) performs worst for RAG, contrary to the original prefix-tuning results, which the authors attribute to the different target of RAG tasks.
- LoRA-tuned models are highly sensitive to inference prompts and fall below the original model when no prompt is used.
Limitations
- Results on larger models (30B/70B) not evaluated; authors state this as future work.
- Performance of single-dataset training can fail to generalize across datasets, e.g., NQ-only training hurts MS MARCO below the prompt baseline.
- Gains depend on the number of virtual tokens and vary across LLMs (LLaMA-2-7b needs at least three tokens to beat the prompt baseline).
Experimental setup
Datasets: TriviaQA Natural Questions HotpotQA SQuAD 1.1 Web Questions 2WikiMultiHopQA CoQA MS MARCO PopQA Fermi Musique Bamboogle
Benchmarks: TriviaQA Natural Questions HotpotQA SQuAD 1.1 Web Questions 2WikiMultiHopQA CoQA MS MARCO PopQA Fermi Musique Bamboogle
Metrics: exact match F1
Baselines: Concat Prompt Prefix-tuning LoRA AAR LongLLMLingua RECOMP Self-RAG
Builds on: prefix-tuning LoRA prompt tuning
Closest papers by meaning
- A Theory for Token-Level Harmonization in Retrieval-Augmented Generation (ICLR 2025 · similarity 0.92)
- Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG (ICLR 2025 · similarity 0.92)
- RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs (NeurIPS 2024 · similarity 0.90)
- LaRA: Benchmarking Retrieval-Augmented Generation and Long-Context LLMs – No Silver Bullet for LC or RAG Routing (ICML 2025 · similarity 0.90)
- Speculative RAG: Enhancing Retrieval Augmented Generation through Drafting (ICLR 2025 · similarity 0.89)
- Follow My Instruction and Spill the Beans: Scalable Data Extraction from Retrieval-Augmented Generation Systems (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.