OneGen: LLMs를 위한 효율적인 One-Pass 통합 생성 및 검색

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

September 8, 2024

저자: Jintian Zhang, Cheng Peng, Mengshu Sun, Xiang Chen, Lei Liang, Zhiqiang Zhang, Jun Zhou, Huajun Chen, Ningyu Zhang

cs.AI

초록

최근 대형 언어 모델 (LLM)의 발전에도 불구하고, 다양한 자연어 처리 (NLP) 작업에 대한 생성 능력을 크게 향상시켰지만, LLM은 여전히 직접 검색 작업을 처리하는 데 제한을 겪고 있습니다. 그러나 많은 실용적인 응용 프로그램에서는 검색과 생성을 원활하게 통합해야 하는 요구가 있습니다. 본 논문은 새로운 효율적인 One-pass Generation and retrieval 프레임워크 (OneGen)를 소개하며, 이는 생성과 검색이 모두 필요한 작업에서 LLM의 성능을 향상시키기 위해 설계되었습니다. 제안된 프레임워크는 검색 토큰을 자기회귀적으로 생성하여 생성과 검색을 위한 기존에는 분리되어 있던 훈련 접근 방식을 연결합니다. 이를 통해 단일 LLM이 통합된 순방향 전달에서 두 작업을 동시에 처리할 수 있게 됩니다. 우리는 RAG와 Entity Linking 두 가지 다른 유형의 복합 작업에 대한 실험을 수행하여 OneGen의 훈련 및 추론에서의 연결성, 효과성 및 효율성을 검증합니다. 더 나아가, 결과는 생성과 검색을 동일한 맥락에서 통합하는 것이 LLM의 생성 능력을 보존하면서 검색 성능을 향상시킨다는 것을 보여줍니다. 우리의 최고 지식으로, OneGen은 LLM이 생성 중에 벡터 검색을 수행할 수 있도록 하는 첫 번째 방법입니다.

English

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.