Self-RAG: Recuperación Adaptativa y Autocrítica para LLMs

El RAG estándar recupera información cada vez, independientemente de si la recuperación ayuda o no. Self-RAG, de Asai et al. (ICLR 2024 Oral), plantea una pregunta diferente: ¿qué pasaría si el propio modelo decidiera cuándo buscar algo y luego calificara el resultado? Resulta que eso importa bastante, y el mecanismo es lo suficientemente limpio como para que valga la pena estudiarlo con detenimiento.

El artículo

La queja principal con la Generación Aumentada por Recuperación (RAG) básica es que es indiscriminada: recupera un número fijo de pasajes para cada entrada, los antepone y genera la respuesta. Eso funciona bien cuando la recuperación ayuda, pero perjudica activamente cuando los pasajes son irrelevantes o cuando el modelo ya tiene la respuesta en sus pesos. El artículo presenta la Generación Aumentada por Recuperación Autorreflexiva (Self-RAG), cuyos autores son Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil y Hannaneh Hajishirzi (Universidad de Washington e IBM Research).

El mecanismo clave es un conjunto de cuatro tokens de reflexión especiales integrados en el vocabulario del modelo en el momento del entrenamiento. Retrieve decide si llamar al recuperador en absoluto. IsRel (relevancia) evalúa si un pasaje recuperado contiene realmente información útil para la consulta. IsSup (apoyo) comprueba si la afirmación generada está total, parcial o nada respaldada por el pasaje. IsUse (utilidad) califica la calidad general de la respuesta del 1 al 5. El modelo aprende a emitir estos tokens junto con su salida normal, de modo que critica su propia recuperación y generación en una sola pasada.

El entrenamiento consta de dos etapas: primero, un modelo crítico (LLaMA 2, 7B ajustado) se entrena con aproximadamente 4,000–20,000 ejemplos etiquetados por tipo de token, alcanzando más del 90% de acuerdo con las predicciones de GPT-4. Ese crítico luego anota un corpus de 150,000 ejemplos de instrucciones y resultados fuera de línea, y el generador se entrena con estos datos anotados, tratando los tokens de reflexión como vocabulario ordinario. No se requiere aprendizaje por refuerzo.

Ideas clave

• Los cuatro tokens de reflexión (Retrieve, IsRel, IsSup, IsUse) le otorgan al modelo un diálogo interno estructurado sobre si vale la pena confiar en la evidencia, no solo una decisión binaria de recuperar o no.
• Self-RAG 13B alcanza un 55.8% en PopQA, 69.3% en TriviaQA, 74.5% en PubHealth, 73.1% en ARC-Challenge y un FactScore de biografía de 80.2, superando a ChatGPT y a Llama2-chat con recuperación en cada una de ellas.
• Las ablaciones en PopQA muestran que eliminar la recuperación en el momento de la prueba cuesta 20.8 puntos porcentuales, mientras que eliminar solo al crítico cuesta solo 2.9 pp; el recuperador es el componente crítico, mientras que la crítica añade calibración adicional.
• En el momento de la inferencia, los pesos de los tokens de crítica se pueden ajustar para equilibrar la precisión de las citas con la fluidez sin necesidad de reentrenamiento. Esto hace que el comportamiento del modelo sea configurable para diferentes aplicaciones derivadas.
• El comité del programa ICLR 2024 otorgó a Self-RAG el estatus de oral (top 1%), lo que refleja un reconocimiento genuino de la contribución técnica por parte de sus pares.

Lo que se mantiene y lo que no

Los resultados de la ablación son convincentes. La brecha entre "siempre recuperar" y "nunca recuperar" es grande (20.8 pp); el modelo claramente aprendió a discriminar la recuperación útil del ruido. Los tokens IsRel e IsSup añaden un valor medible por encima de la recuperación adaptativa por sí sola. Ese es un resultado significativo, no solo un cambio de enfoque.

Lo que me convence menos es la afirmación de generalización. Las cinco tareas de evaluación (PopQA, TriviaQA, PubHealth, ARC-Challenge, ASQA) son de respuesta corta o de opción múltiple, el entorno exacto donde un solo pasaje recuperado puede proporcionar una señal decisiva. La generación de formato largo sobre contextos de múltiples documentos, que es donde residen las tareas financieras, recibe menos escrutinio. El FactScore de biografía (80.2) es la aproximación más cercana, pero las biografías están relativamente bien estructuradas en comparación con un libro mayor de gastos multianual en expansión.

También hay un inconveniente de reproducibilidad: las etiquetas de entrenamiento del modelo crítico provienen de GPT-4. Esto hace que la calidad de las etiquetas dependa de un sistema propietario e introduce costes de API que no se reportan. CRAG (arXiv:2401.15884) demostró más tarde que un evaluador de recuperación de 0.77B —mucho más ligero que el crítico de 7B de Self-RAG— podía corregir la calidad de la recuperación y ganar 19.0 pp sobre el RAG estándar en PopQA, lo que sugiere que el pesado crítico ajustado podría no ser necesario. Ese es un desafío significativo al diseño, incluso si la idea central sobre la recuperación selectiva se mantiene.

Finalmente, el punto de referencia de comparación importa. Superar a ChatGPT (probablemente GPT-3.5-turbo, finales de 2023) y a Llama2-chat es un listón razonable para un modelo abierto de 13B, pero los modelos de frontera han avanzado sustancialmente desde entonces. No se aborda si la recuperación adaptativa de Self-RAG superaría a un GPT-4o con un buen "prompt" en una configuración simple de "recuperar siempre" en estos mismos puntos de referencia.

Por qué esto es importante para la IA financiera

Los agentes financieros que operan sobre libros mayores de Beancount se enfrentan exactamente al problema de discriminación de recuperación que aborda Self-RAG. Cuando un usuario pregunta "¿cuál es mi ingreso neto este mes?", el agente puede calcularlo a partir de su contexto cargado; la recuperación podría simplemente añadir ruido. Cuando el mismo usuario pregunta "¿registré la factura del contratista del tercer trimestre?", el agente necesita escanear potencialmente años de entradas. Recuperar siempre desperdicia contexto y corre el riesgo de inyectar transacciones antiguas irrelevantes; no recuperar nunca pierde la búsqueda.

Los tokens IsRel e IsSup se mapean limpiamente a la lógica de validación del libro mayor. IsRel: ¿la entrada de transacción recuperada se relaciona realmente con la consulta? IsSup: ¿el contexto recuperado respalda realmente la cifra de saldo generada, o es un número alucinado? La puntuación de utilidad (1–5) podría informar la confianza en el asiento contable: solo confirmar una propuesta de asiento cuando el modelo califique su propio razonamiento con un 4 o 5, y marcar el resto para revisión humana.

La preocupación por la reproducibilidad también importa aquí. Para un agente contable de producción, depender de GPT-4 para generar etiquetas de entrenamiento es una limitación operativa. Si un evaluador más ligero (al estilo de CRAG) puede lograr una recuperación selectiva comparable, ese es el camino más viable para el despliegue. Los principios de diseño de Self-RAG —decidir antes de recuperar, criticar después de recuperar— siguen siendo valiosos incluso si se reemplaza la receta específica de entrenamiento de tokens.

Qué leer a continuación

• CRAG: Corrective Retrieval Augmented Generation (arXiv:2401.15884) — se basa en la idea de recuperación adaptativa de Self-RAG con un evaluador más ligero y una opción de búsqueda web cuando falla la recuperación local; vale la pena compararlo directamente con Self-RAG en puntos de referencia coincidentes.
• RQ-RAG: Learning to Refine Queries for Retrieval Augmented Generation (arXiv:2404.00610) — se centra específicamente en la descomposición de consultas para preguntas y respuestas complejas de múltiples saltos, que es el escenario que Self-RAG maneja con menos elegancia.
• FRAMES: Retrieval and Augmentation for Multi-Hop Evaluation (arXiv:2409.12941) — punto de referencia de Google DeepMind para RAG de múltiples documentos que requiere encadenar varios hechos recuperados; una prueba natural más difícil para modelos de estilo Self-RAG.