OneGen: Geração Unificada Eficiente em Uma Única Passagem para Modelos de Linguagem de Longo Alcance (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

Resumo

Apesar dos avanços recentes em Modelos de Linguagem de Grande Escala (LLMs), que melhoraram significativamente as capacidades generativas para várias tarefas de Processamento de Linguagem Natural (NLP), os LLMs ainda enfrentam limitações ao lidar diretamente com tarefas de recuperação. No entanto, muitas aplicações práticas exigem a integração perfeita de tanto a recuperação quanto a geração. Este artigo apresenta um novo e eficiente framework de Geração e Recuperação em uma passagem (OneGen), projetado para melhorar o desempenho dos LLMs em tarefas que exigem tanto geração quanto recuperação. O framework proposto conecta as abordagens de treinamento tradicionalmente separadas para geração e recuperação, incorporando tokens de recuperação gerados de forma autoregressiva. Isso permite que um único LLM lide com ambas as tarefas simultaneamente em uma passagem unificada. Realizamos experimentos em dois tipos distintos de tarefas compostas, RAG e Vinculação de Entidades, para validar a adaptabilidade, eficácia e eficiência do OneGen no treinamento e inferência. Além disso, nossos resultados mostram que integrar geração e recuperação dentro do mesmo contexto preserva as capacidades generativas dos LLMs, ao mesmo tempo que melhora o desempenho de recuperação. Até onde sabemos, o OneGen é o primeiro a permitir que os LLMs realizem recuperação vetorial durante a geração.

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.