# OneGen: Efficient One-Pass Unified Generation and Retrieval for LLMs

arXiv: https://arxiv.org/abs/2409.05152 · HTML: https://mlsanity.com/paper/2409.05152/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 59 of its 73 references that are in the map.
- First posted on arXiv 8 Sep 2024.

## Builds on (most cited first, 30 of 59)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Qwen Technical Report](https://mlsanity.com/paper/2309.16609/) (2023)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Unsupervised Dense Information Retrieval with Contrastive Learning](https://mlsanity.com/paper/2112.09118/) (2021)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Tools](https://mlsanity.com/paper/2406.12793/) (2024)
- [AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback](https://mlsanity.com/paper/2305.14387/) (2023)
- [Siren's Song in the AI Ocean: A Survey on Hallucination in Large Language Models](https://mlsanity.com/paper/2309.01219/) (2023)
- [Active Retrieval Augmented Generation](https://mlsanity.com/paper/2305.06983/) (2023)
- [Is ChatGPT a General-Purpose Natural Language Processing Task Solver?](https://mlsanity.com/paper/2302.06476/) (2023)
- [LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders](https://mlsanity.com/paper/2404.05961/) (2024)
- [Editing Large Language Models: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2305.13172/) (2023)
- [Autoregressive Entity Retrieval](https://mlsanity.com/paper/2010.00904/) (2020)

## Built on by (most cited first, 14 of 14)

- [LightThinker: Thinking Step-by-Step Compression](https://mlsanity.com/paper/2502.15589/) (2025)
- [Mitigating Lost-in-Retrieval Problems in Retrieval Augmented Multi-Hop Question Answering](https://mlsanity.com/paper/2502.14245/) (2025)
- [RetroLLM: Empowering Large Language Models to Retrieve Fine-grained Evidence within Generation](https://mlsanity.com/paper/2412.11919/) (2024)
- [Large Reasoning Embedding Models: Towards Next-Generation Dense Retrieval Paradigm](https://mlsanity.com/paper/2510.14321/) (2025)
- [MAGNET: Augmenting Generative Decoders with Representation Learning and Infilling Capabilities](https://mlsanity.com/paper/2501.08648/) (2025)
- [ImpRAG: Retrieval-Augmented Generation with Implicit Queries](https://mlsanity.com/paper/2506.02279/) (2025)
- [Large Language Model Can Be a Foundation for Hidden Rationale-Based Retrieval](https://mlsanity.com/paper/2412.16615/) (2024)
- [Hydra: Unifying Document Retrieval and Generation in a Single Vision-Language Model](https://mlsanity.com/paper/2603.28554/) (2026)
- [A Unified Model and Document Representation for On-Device Retrieval-Augmented Generation](https://mlsanity.com/paper/2604.14403/) (2026)
- [Efficient Rationale-based Retrieval: On-policy Distillation from Generative Rerankers based on JEPA](https://mlsanity.com/paper/2604.23336/) (2026)
- [The Periodic Table of LLM Reasoning: A Structured Survey of Reasoning Paradigms, Methods, and Failure Modes](https://mlsanity.com/paper/2606.11470/) (2026)
- [Co-LMLM: Continuous-Query Limited Memory Language Models](https://mlsanity.com/paper/2607.07707/) (2026)
- [Token-Native Storage: Read and Write in your Agent's Language](https://mlsanity.com/paper/2608.02376/) (2026)
- [Select, Don't Train: The Benefits of Modular Entity Disambiguation with LLM-Based Selection](https://mlsanity.com/paper/2608.27470/) (2026)

## Abstract (arXiv)

Despite the recent advancements in Large Language Models (LLMs), which have significantly enhanced the generative capabilities for various NLP tasks, LLMs still face limitations in directly handling retrieval tasks. However, many practical applications demand the seamless integration of both retrieval and generation. This paper introduces a novel and efficient One-pass Generation and retrieval framework (OneGen), designed to improve LLMs' performance on tasks that require both generation and retrieval. The proposed framework bridges the traditionally separate training approaches for generation and retrieval by incorporating retrieval tokens generated autoregressively. This enables a single LLM to handle both tasks simultaneously in a unified forward pass. We conduct experiments on two distinct types of composite tasks, RAG and Entity Linking, to validate the pluggability, effectiveness, and efficiency of OneGen in training and inference. Furthermore, our results show that integrating generation and retrieval within the same context preserves the generative capabilities of LLMs while improving retrieval performance. To the best of our knowledge, OneGen is the first to enable LLMs to conduct vector retrieval during the generation.