Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

Guía de lectura

Problema y motivación

Los modelos de lenguaje preentrenados almacenan un conocimiento factual sustancial dentro de sus parámetros (memoria paramétrica). Sin embargo, presentan inconvenientes importantes: les cuesta expandir o actualizar su conocimiento sin reentrenamiento, no pueden proporcionar fácilmente procedencia o interpretabilidad para sus salidas y son propensos a sufrir alucinaciones.

Aunque los modelos híbridos paramétricos/no paramétricos como REALM y ORQA combinaron la recuperación diferenciable con el modelado de lenguaje enmascarado, estaban restringidos a la respuesta extractiva de preguntas. Las arquitecturas estándar de secuencia a secuencia (seq2seq) como BART y T5 permanecieron en gran medida como modelos de "libro cerrado" (closed-book). Este trabajo acorta esa distancia proponiendo la Generación Aumentada por Recuperación (RAG), un marco de ajuste fino de propósito general que acopla un recuperador neuronal denso preentrenado con un generador seq2seq autorregresivo, tanto para la generación en dominio abierto como para la clasificación.

Ideas clave

• Memoria híbrida paramétrica y no paramétrica: La memoria no paramétrica consiste en un índice de vectores densos sobre 21 millones de pasajes de 100 palabras de Wikipedia a los que se accede mediante búsqueda de producto interno máximo (MIPS). La memoria paramétrica es un generador BART preentrenado.
• Ajuste fino de extremo a extremo: El codificador de consultas y el generador se entrenan conjuntamente tratando los pasajes recuperados como variables latentes y minimizando la log-verosimilitud marginal negativa de los tokens objetivo sin supervisión explícita de recuperación.
• Dos formulaciones de condicionamiento:
  • RAG-Sequence: Se condiciona en el mismo documento recuperado a lo largo de toda la secuencia generada.
  • RAG-Token: Permite que cada token generado atienda y se marginalice sobre diferentes pasajes recuperados.
• Conocimiento del mundo intercambiable en caliente (hot-swappable): El índice externo se puede intercambiar en tiempo de inferencia (por ejemplo, cambiando un índice de Wikipedia de 2016 por uno de 2018) para actualizar el conocimiento del mundo sin modificar los pesos del modelo.

Método

Dada una secuencia de entrada $x$, el recuperador $p_{\eta }(z|x)$ produce una distribución sobre los documentos $z$, y el generador $p_{\theta }(y_i|x,z,y_{1:i-1})$ genera tokens objetivo condicionados tanto a $x$ como a $z$.

El recuperador utiliza Dense Passage Retriever (DPR): $$p_{\eta }(z|x)\propto \exp \left(\mathbf{d}(z{)}^{\top }\mathbf{q}(x)\right)$$ donde $\mathbf{d}(z)={\text{BERT}}_d(z)$ y $\mathbf{q}(x)={\text{BERT}}_q(x)$. Los candidatos top-$k$ se recuperan utilizando FAISS MIPS. El índice de documentos y el codificador de documentos ${\text{BERT}}_d$ se mantienen fijos, mientras que ${\text{BERT}}_q$ y los parámetros del generador BART $\theta$ se ajustan.

RAG-Sequence

El modelo trata el documento recuperado como una única variable latente a lo largo de toda la secuencia $y$: $$p_{\text{RAG-Sequence}}(y|x)\approx \sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)\prod _i^N{p}_{\theta }(y_i|x,z,y_{1:i-1})$$ La decodificación se puede evaluar utilizando Thorough Decoding (ejecutando una búsqueda en haz por documento y calculando pasadas hacia adelante para las hipótesis candidatas en todos los documentos) o Fast Decoding (aproximando las probabilidades faltantes del haz a 0).

RAG-Token

El modelo marginaliza sobre los documentos en cada paso de token, lo que permite la mezcla de contenidos a nivel de token entre varias fuentes: $$p_{\text{RAG-Token}}(y|x)\approx \prod _i^N\sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)p_{\theta }(y_i|x,z,y_{1:i-1})$$ Esto produce una probabilidad de transición autorregresiva estándar que se integra directamente en la búsqueda en haz (beam search).

Resultados

• QA de dominio abierto: RAG establece nuevos resultados de vanguardia en pruebas comparativas abiertas utilizando coincidencia exacta (EM):
  • Natural Questions: 44.5 (RAG-Seq) frente a 41.5 (DPR) y 36.6 (T5-11B+SSM).
  • TriviaQA (división abierta estándar): 56.8 (RAG-Seq); en la partición de prueba TQA-Wiki: 68.0 (RAG-Seq) frente a 60.5 (T5-11B+SSM).
  • WebQuestions: 45.5 (RAG-Token) frente a 41.1 (DPR).
  • CuratedTrec: 52.2 (RAG-Seq) frente a 50.6 (DPR).
  • En particular, en Natural Questions, RAG logró un 11.8% de EM cuando la respuesta correcta no aparecía en absoluto en el contexto recuperado (donde los modelos extractivos obtienen un 0%).
• Generación abstractiva: En Open MS-MARCO, RAG-Sequence alcanza 40.8 de Rouge-L y 44.2 de Bleu-1, superando a BART (38.2 Rouge-L, 41.6 Bleu-1).
• Generación de preguntas de Jeopardy: Los evaluadores humanos calificaron a RAG-Token como más factual que BART en el 42.7% de los casos, frente a un 7.1% en el que BART fue superior. Los evaluadores también encontraron que RAG era más específico (37.4% frente a 16.8%).
• Verificación de hechos: En FEVER de 3 clases, RAG alcanza una precisión de etiqueta del 72.5% (a menos del 4.3% de los sistemas de tubería que usan supervisión de evidencia oro); en 2 clases, RAG logra un 89.5% de precisión.
• Actualización del conocimiento mediante reemplazo en caliente: Utilizando una prueba temporal de 82 líderes mundiales entre 2016 y 2018, RAG obtuvo una precisión del 70% para los líderes de 2016 utilizando el índice de 2016 y del 68% para los líderes de 2018 con el índice de 2018, cayendo drásticamente (12% y 4%) cuando no coincidían.