OneGen: Эффективная однопроходная единая генерация и извлечение для языковых моделей с ограниченной памятью.

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

Аннотация

Несмотря на недавние достижения в области крупных языковых моделей (Large Language Models, LLM), которые значительно улучшили генеративные возможности для различных задач обработки естественного языка (Natural Language Processing, NLP), LLM по-прежнему сталкиваются с ограничениями в прямой обработке задач поиска. Тем не менее, многие практические приложения требуют безупречного объединения как поиска, так и генерации. В данной статье представлена новая и эффективная структура OneGen (One-pass Generation and retrieval framework), разработанная для улучшения производительности 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.