Atlas: El preentrenamiento conjunto de Recuperador-Lector supera a los LLM de 540B de parámetros con solo 11B de parámetros

Atlas es la continuación de Izacard y Grave a su propio artículo sobre Fusion-in-Decoder, extendiendo FiD a un sistema totalmente entrenado de forma conjunta donde el recuperador y el lector se coentrenan desde cero. Lo estoy leyendo ahora porque cierra el linaje arquitectónico desde el artículo original de RAG, pasando por FiD, hasta la recuperación entrenada conjuntamente —exactamente el espacio de decisión que cualquier sistema de QA para libros contables necesita navegar.

El artículo

"Atlas: Few-shot Learning with Retrieval Augmented Language Models" (Izacard et al., JMLR 2023) se pregunta si los modelos aumentados por recuperación pueden igualar a los LLM de parámetros masivos en tareas de pocos ejemplos (few-shot) intensivas en conocimiento. La contribución principal es un sistema aumentado por recuperación cuidadosamente preentrenado que entrena conjuntamente un recuperador denso basado en Contriever junto con un lector Fusion-in-Decoder basado en T5. La idea clave es que el preentrenamiento conjunto —no la arquitectura— es lo que impulsa el rendimiento del conocimiento en pocos ejemplos. El sistema recupera los 20 documentos principales, codifica cada uno de forma independiente en el codificador y luego los fusiona en la atención cruzada del decodificador, el mismo diseño FiD del artículo de los autores de 2021.

Ideas clave

• Atlas-11B logra una precisión del 42,4% en Natural Questions con solo 64 ejemplos de entrenamiento, superando a PaLM (540B de parámetros) por aproximadamente 3 puntos usando 50 veces menos parámetros.
• En TriviaQA (64-shot), Atlas-11B alcanza el 74,5% en el conjunto filtrado y el 84,7% en la prueba oculta no filtrada, lo que demuestra que el componente de recuperación compensa fuertemente la supervisión limitada de la tarea.
• Se evalúan cuatro objetivos de entrenamiento del recuperador: Destilación de Atención (ADist), EMDR2 (tratando los documentos recuperados como variables latentes), Destilación de Perplejidad (PDist) y LOOP (leave-one-out). Las diferencias de rendimiento entre ellos son pequeñas; se adopta PDist por eficiencia de cómputo.
• El preentrenamiento conjunto en texto no etiquetado es el factor más importante: todas las configuraciones de preentrenamiento aumentado por recuperación superan con creces la línea base de solo ajuste fino aumentado por recuperación.
• El índice de documentos puede actualizarse después del entrenamiento sin volver a entrenar el modelo, lo cual es arquitectónicamente importante para las bases de conocimiento dinámicas. Los índices con desajustes temporales degradan el rendimiento notablemente.
• En MMLU (5-shot), Atlas-11B alcanza el 47,9%, superando el 43,9% reportado de GPT-3, a pesar de tener aproximadamente 16 veces menos parámetros.

Lo que se mantiene — y lo que no

La afirmación principal —que la recuperación permite un rendimiento de conocimiento con pocos ejemplos a una fracción del recuento de parámetros— se mantiene de forma convincente. La cifra de 42,4% en NQ con 64 ejemplos es un resultado sorprendente, y la comparación con PaLM es justa porque PaLM era el punto de referencia de escala de vanguardia en ese momento.

Sin embargo, tengo tres reservas. Primero, la precisión de la recuperación no es excelente incluso después del entrenamiento conjunto: análisis independientes muestran que Contriever falla en al menos una declaración de referencia (gold statement) en aproximadamente el 85% de los casos, y logra alrededor del 47% de precisión de recuperación en QA. El entrenamiento conjunto mejora la recuperación sobre las líneas base no entrenadas conjuntamente, pero el lector está haciendo un trabajo enorme para compensar una recuperación imperfecta; las cifras principales de pocos ejemplos reflejan el techo del sistema, no la calidad del componente de recuperación. Segundo, el costo de la infraestructura es real: actualizar los índices de documentos durante el preentrenamiento añade aproximadamente un 30% de sobrecarga computacional, y el índice completo de Wikipedia+CommonCrawl requiere 587GB en fp16. Eso es manejable en un entorno de investigación, pero es una restricción operativa real para el despliegue en producción. Tercero, se reconoce la fuga de datos (data leakage) pero no se resuelve: el 2,8% de las preguntas de MMLU aparecen literalmente en el corpus CCNet utilizado para el preentrenamiento, inflando los resultados de MMLU por un margen desconocido.

También hay una limitación arquitectónica más sutil en la que el artículo no profundiza del todo: FiD codifica cada pasaje recuperado de forma independiente antes de la fusión, lo que ayuda al paralelismo pero significa que el codificador no tiene atención entre pasajes. Las cadenas largas de razonamiento de múltiples saltos que necesitan conectar información a través de los pasajes deben hacer todo ese trabajo en el decodificador —y con 20 pasajes recuperados, la atención cruzada del decodificador está soportando una carga pesada.

Por qué esto es importante para la IA en finanzas

Para el QA de libros contables de Beancount, la contribución más relevante de Atlas es la demostración empírica de que el entrenamiento conjunto recuperador-lector rinde frutos en configuraciones de pocos ejemplos, y su honesta rendición de cuentas de cuándo no es así. Un agente de Beancount que consulta un historial de transacciones de varios años se enfrenta exactamente al problema del índice dinámico: llegan nuevas entradas diariamente, y un índice con un mes de antigüedad produce respuestas incorrectas. Atlas muestra que el índice se puede intercambiar en caliente sin volver a entrenar, lo cual es arquitectónicamente alentador.

Sin embargo, las cifras de precisión de recuperación son aleccionadoras. Si Contriever falla al encontrar la entrada relevante del libro contable en el 53% de los intentos de recuperación incluso después del entrenamiento conjunto en texto general, un agente del dominio financiero que opere sobre libros de Beancount —con sus nombres de materias primas (commodities) específicos del dominio, jerarquías de cuentas y directivas de Beancount— necesitará un entrenamiento del recuperador adaptado al dominio o una recuperación aumentada por métodos de consulta estructurados (coincidencia exacta de cuentas, filtrado por fechas). La recuperación de estilo RAG por sí sola, incluso entrenada conjuntamente, no será suficiente para operaciones de libros contables de alta precisión.

La comparación con PaLM también aclara el compromiso arquitectónico: la recuperación permite comprimir el conocimiento en menos parámetros, reduciendo el costo de inferencia. Para un producto como Beancount.io, donde el costo de inferencia importa a escala, la filosofía de diseño de Atlas es atractiva. Pero el costo del índice de 587GB traslada la carga a la infraestructura de almacenamiento y recuperación, un tipo diferente de restricción operativa que no aparece en las cifras de los puntos de referencia.

Qué leer a continuación

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — el marco anterior de preentrenamiento conjunto recuperador-lector que Atlas extiende; esencial para entender qué mejora realmente Atlas y qué deja sin cambios.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — logra un rendimiento competitivo con Atlas utilizando el ajuste de instrucciones en lugar del preentrenamiento conjunto desde cero; sugiere que la brecha entre el entrenamiento conjunto e independiente puede cerrarse sin el costo de infraestructura.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — el enfoque de DeepMind para la recuperación durante el preentrenamiento a una escala diferente; completa el panorama de los enfoques de preentrenamiento aumentados por recuperación antes de tomar decisiones arquitectónicas para el QA de libros contables.