OneGen: Generación Unificada Eficiente en un Solo Paso y Recuperación para LLMs

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

September 8, 2024

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

cs.AI

Resumen

A pesar de los avances recientes en los Modelos de Lenguaje de Gran Tamaño (LLMs), que han mejorado significativamente las capacidades generativas para varias tareas de Procesamiento del Lenguaje Natural (NLP), los LLMs aún enfrentan limitaciones al manejar directamente tareas de recuperación. Sin embargo, muchas aplicaciones prácticas requieren la integración fluida de la recuperación y la generación. Este artículo presenta un marco novedoso y eficiente de Generación y Recuperación en un solo paso (OneGen), diseñado para mejorar el rendimiento de los LLMs en tareas que requieren tanto generación como recuperación. El marco propuesto une los enfoques de entrenamiento tradicionalmente separados para la generación y la recuperación al incorporar tokens de recuperación generados de forma autoregresiva. Esto permite que un solo LLM maneje ambas tareas simultáneamente en un solo pase unificado. Realizamos experimentos en dos tipos distintos de tareas compuestas, RAG y Vinculación de Entidades, para validar la capacidad de conexión, efectividad y eficiencia de OneGen en entrenamiento e inferencia. Además, nuestros resultados muestran que integrar generación y recuperación dentro del mismo contexto preserva las capacidades generativas de los LLMs al tiempo que mejora el rendimiento de la recuperación. Hasta donde sabemos, OneGen es el primero en permitir que los LLMs realicen recuperación vectorial durante la generación.

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.