To Memorize or to Retrieve: Scaling Laws for RAG-Considerate Pretraining

1. Planteamiento del Problema

El artículo aborda una brecha fundamental en la comprensión de los modelos de lenguaje (LM): la relación entre el conocimiento paramétrico (adquirido durante el preentrenamiento y almacenado en los pesos del modelo) y el conocimiento no paramétrico (accedido mediante recuperación o retrieval en tiempo de inferencia, como en RAG).

Hasta ahora, las leyes de escalado (scaling laws) han tratado el corpus de entrenamiento como un bloque monolítico, asumiendo que todo el conocimiento debe comprimirse en los pesos del modelo. Sin embargo, con presupuestos de datos fijos, surge una pregunta crítica de asignación de recursos: ¿Cómo se debe distribuir un corpus de $N$ tokens entre el preentrenamiento (para aumentar el tamaño del modelo o la cantidad de tokens de entrenamiento) y la construcción de una base de datos externa para recuperación?

Actualmente, no existe una respuesta establecida sobre cómo optimizar esta compensación (trade-off), especialmente considerando que el aprendizaje paramétrico puede llevar a alucinaciones, mientras que la recuperación puede introducir errores por documentos irrelevantes.

2. Metodología

Los autores diseñaron un estudio empírico sistemático para modelar esta interacción bajo un presupuesto de datos fijo.

• Configuración Experimental:

  • Modelos: Se entrenaron modelos basados en la arquitectura OLMo-2 con tamaños que van desde 30M hasta 3 mil millones de parámetros.
  • Datos: Se utilizó el corpus DCLM (hasta 100 mil millones de tokens).
  • Estrategia de Datos: Para un presupuesto fijo, se variaron dos dimensiones:
    1. Escala de datos de preentrenamiento ($D$): Desde 1x hasta 150x el número de parámetros.
    2. Tamaño del almacén de recuperación ($R$): Desde 1x hasta 20x el número de parámetros, construido a partir de una porción retenida del mismo corpus DCLM.
  • Métricas: Se evaluó el rendimiento en una suite diversa de benchmarks (razonamiento, QA científica, QA de dominio abierto). Se priorizó la perplejidad (PPL) sobre la precisión (accuracy) para el análisis de escalado, ya que PPL ofrece una métrica continua y suave, evitando los comportamientos de "umbral" o escalón típicos de la precisión.

• Enfoque de Modelado:

  • Se extendieron las leyes de escalado clásicas (que relacionan pérdida con tamaño del modelo $N$ y tokens de datos $D$) para incluir un tercer eje: el tamaño del corpus de recuperación $R$.
  • Se propuso una ley de escalado tridimensional que modela la pérdida ($L$) como una función de $N$, $D$ y $R$.

3. Contribuciones Clave

1. Análisis de la Compensación Preentrenamiento-Recuperación: Demostraron que la relación no es trivial ni monótona; la recuperación ofrece beneficios que dependen de la escala del modelo, el tipo de tarea y el grado de saturación del preentrenamiento.
2. Marco de Leyes de Escalado 3D: Introdujeron un modelo matemático que describe el rendimiento en función de los tres factores ($N, D, R$), permitiendo cuantificar la sustituibilidad entre ambos mecanismos.
3. Punto de Cruce Dependiente de la Escala: Identificaron un umbral crítico donde la recuperación se convierte en un sustituto eficiente del preentrenamiento adicional.
4. Guía Práctica para Asignación de Recursos: Proporcionan una base cuantitativa para decidir cuándo es más eficiente invertir en más datos de entrenamiento vs. construir un almacén de recuperación más grande.

4. Resultados Principales

A. Leyes de Escalado y Sustituibilidad

El modelo de pérdida ajustado sigue la forma:
$$L(N, D, R) = A \left(\frac{N}{10^9}\right)^{-\alpha} + B \left(\frac{D}{10^9}\right)^{-\beta} - C \log\left(1 + \eta \frac{R}{10^9}\right) + L_0$$
Donde el término logarítmico negativo representa la ganancia de la recuperación.

• Sustituibilidad ($\sigma$): Se definió como el número de tokens de preentrenamiento que se pueden "ahorrar" por cada token de recuperación utilizado para alcanzar el mismo nivel de pérdida.
  • Hallazgo: En regímenes de pocos datos (baja relación $D/N$), la recuperación no es un sustituto efectivo.
  • Umbral: Existe un punto de cruce claro alrededor de 4.14 tokens de preentrenamiento por parámetro ($D/N \approx 4.14$). Por encima de este umbral, la recuperación se vuelve un sustituto eficiente: cada token de recuperación puede reemplazar a múltiples tokens de preentrenamiento.

B. Beneficio Marginal y Escala del Modelo

• Modelos Pequeños (ej. 30M): Se benefician enormemente de la recuperación. La ganancia marginal por token recuperado es alta, actuando como un "acelerador" de capacidades que les falta.
• Modelos Grandes (ej. 3B): Muestran rendimientos decrecientes. A medida que el modelo crece y el preentrenamiento se satura, el beneficio absoluto de añadir más datos de recuperación disminuye significativamente.
• Dependencia de la Tarea:
  • Tareas de Conocimiento (QA, hechos): La recuperación es muy efectiva y sustituye bien al preentrenamiento.
  • Tareas de Razonamiento (ej. PIQA, StrategyQA): La recuperación tiene un impacto limitado o nulo. El razonamiento parece depender más de la capacidad computacional interna (pesos) que del acceso a hechos externos.

C. Calidad de la Recuperación

Las pruebas con diferentes formulaciones de consultas (solo pregunta vs. pregunta + respuesta correcta) mostraron que, aunque una mejor recuperación mejora el rendimiento, no altera las tendencias de escalado fundamentales. El cuello de botella no es solo la calidad del recuperador, sino la naturaleza de la tarea y la escala del modelo.

5. Significado e Implicaciones

Este trabajo cambia la perspectiva de diseño de los sistemas de lenguaje:

1. Optimización Conjunta: El preentrenamiento y la recuperación no deben verse como elecciones de diseño separadas, sino como dos mecanismos competidores por un presupuesto de datos fijo que deben optimizarse conjuntamente.
2. Estrategia de Asignación: Para modelos pequeños o con recursos limitados, es más eficiente invertir en un almacén de recuperación robusto que en preentrenar masivamente. Para modelos grandes y altamente saturados, el enfoque debe volver al preentrenamiento o a arquitecturas más eficientes, ya que la recuperación ofrece retornos marginales bajos.
3. Cognición Artificial: El estudio valida la hipótesis de que, al igual que en la cognición humana (que internaliza habilidades abstractas pero externaliza la memoria factual), los LM deben asignar recursos de datos de manera diferenciada: preentrenamiento para razonamiento y patrones, y recuperación para hechos de "cola larga" y conocimiento factual.

En conclusión, el paper establece que la recuperación es un sustituto eficiente del preentrenamiento solo en regímenes específicos (modelos pequeños o tareas de conocimiento factual), proporcionando una hoja de ruta cuantitativa para la construcción de sistemas de lenguaje escalables y eficientes.