Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale

Resumen Técnico: ¿Pueden los modelos de lenguaje recuperar realmente en contexto?

Planteamiento del problema
Aunque los modelos de lenguaje (LM) ofrecen una alternativa potencial a la recuperación basada en vectores al condicionarse a un corpus en contexto para generar directamente identificadores de documentos relevantes, esta capacidad sigue siendo ampliamente inexplorada a escalas realistas. El trabajo previo se ha centrado en sistemas propietarios o tareas de reordenamiento (reranking) a pequeña escala, dejando dos desafíos críticos sin resolver: (C1) la escalabilidad a corpus que contienen millones de tokens, y (C2) la generalización a tamaños de corpus muy superiores a los vistos durante el entrenamiento. Los LM actuales tienen dificultades tanto con los costes de atención de un millón de tokens como con la fragilidad de la extrapolación a contextos más largos. La cuestión central es si un LM puede actuar de forma fiable como un recuperador a escala de corpus bajo estas condiciones sin colapsar.

Metodología: BLOCKSEARCH y Análisis
Los autores presentan BLOCKSEARCH, un estudio sistemático y una arquitectura de modelo diseñada para abordar la recuperación en contexto (ICR) a escala de corpus.

• Arquitectura y Entrenamiento: Construido sobre un núcleo de 0.6B de parámetros (Qwen3-0.6B), BLOCKSEARCH emplea atención dispersa por bloques (block-sparse attention) para manejar documentos disjuntos de manera eficiente. Las modificaciones clave de arquitectura y entrenamiento incluyen:

  • Identificadores de Documentos Aleatorios: En lugar de IDs secuenciales, a los documentos se les asignan códigos aleatorios de cuatro dígitos por cada paso de entrenamiento para evitar el sobreajuste a posiciones absolutas.
  • Negativos en el Lote (In-Batch Negatives) y Pérdida On-Policy: El modelo se entrena con 1 positivo y 15 negativos difíciles por consulta. Para mitigar el sesgo de exposición (donde el modelo se condiciona a sus propios prefijos generados potencialmente incorrectos durante la inferencia), se añade una pérdida auxiliar on-policy. Esto implica realizar un despliegue (rollout) de un código de la propia distribución del modelo y entrenar al modelo para corregir su trayectoria basándose en las puntuaciones del lote.
  • Configuración de Evaluación: El modelo se evalúa en conjuntos de datos BEIR (MS MARCO, NQ, HotpotQA) con tamaños de corpus que oscilan entre 500 y 10,000 documentos (hasta ~1 millón de tokens), superando con creces la ventana de contexto nativa de 32k tokens del modelo.

• Análisis Mecanicista del Fallo: El estudio identifica que, si bien BLOCKSEARCH mejora la extrapolación, sigue colapsando a gran escala (recuperación cercana a cero a 1M de tokens). Los autores trazan esto no a un fallo en la identificación del documento relevante (ranking), sino a un efecto de dilución de la atención:

  • Brecha entre Recuperación y Generación: Incluso a 10,000 documentos, las cabezas de atención individuales a menudo siguen asignando la puntuación más alta antes del softmax al documento de oro ($R_{any} \approx 1.0$).
  • Colapso del Softmax: Sin embargo, a medida que el corpus ($N$) crece, la contribución agregada de los documentos irrelevantes al denominador del softmax aumenta más rápido que la puntuación del documento de oro. En consecuencia, la masa de atención normalizada (peso post-softmax) en el documento de oro colapsa, causando que la salida de la capa sea un promedio de los distractores en lugar del contenido relevante.

• Mitigaciones Propuestas: Para contrarrestar la dilución de la atención, los autores proponen dos ajustes conscientes de la longitud:

  1. Ajustes de Softmax Conscientes de la Longitud:
     • Sumideros Aditivos (Additive Sinks): Anexar una constante aprendida al denominador del softmax para absorber la masa de atención difusa (inspirado en Streaming-LLM).
     • Reescalado Multiplicativo (SSMax): Escalar las puntuaciones pre-softmax por $s \cdot \log N$ para ampliar la brecha entre la puntuación del oro y los distractores, evitando que el peso del oro se desvanezca a medida que $N$ aumenta.
  2. Atención Dispersa a Nivel de Documento: Un mecanismo de enrutamiento en una capa intermedia (L16) que puntúa todos los documentos y restringe la atención densa a los top-$B$ (por ejemplo, 256) documentos para las capas subsiguientes, recortando efectivamente el contexto antes de la generación final.

Resultados Clave

• Rendimiento de la Línea Base: El modelo BLOCKSEARCH inicial iguala el rendimiento de la recuperación densa en benchmarks estándar (MS MARCO, NQ) en tamaños de corpus pequeños (~45k tokens), y se generaliza hasta ~500k tokens (10x el contexto de entrenamiento), pero colapsa cerca de 1M de tokens.
• Impacto de las Modificaciones:
  • SSMax: Mejora significativamente el rendimiento a lo largo de todo el barrido de $N$. En MS MARCO a 1M de tokens, logra un 16.5% de Recall@1 (frente al 0.2% del modelo sin modificar), cerrando la mayor parte de la brecha con la línea base densa (20.2%).
  • Enrutamiento (Routing): El enrutamiento a nivel de documento mejora aún más el rendimiento, logrando un 18.8% de Recall@1 en MS MARCO a 1M de tokens, casi igualando la línea base densa.
  • Enfoque Combinado: La combinación de SSMax y enrutamiento (BLOCKSEARCH-SSMax-routing) logra un 20.5% de Recall@1 en MS MARCO a 1M de tokens, superando ligeramente la línea base densa (20.2%) y superando al modelo concurrente MSA-4B, 7 veces más grande (16.0%).
• Fuera de Distribución (LIMIT): En el benchmark LIMIT, que requiere similitud léxica en lugar de similitud de embeddings semánticos (donde la recuperación densa tiene dificultades), el BLOCKSEARCH modificado supera significativamente a la línea base densa. A 850k tokens, BLOCKSEARCH-SSMax-routing logra un 14.9% de Recall@1 frente al 3.5% del recuperador denso (una mejora de ~3x).
• Techo de Atención: En todos los experimentos, el "techo de atención" (la fracción de consultas donde al menos una cabeza clasifica el documento de oro en primer lugar) permanece cerca del 100%, confirmando que la señal de recuperación está presente internamente pero se pierde durante la fase de lectura/generación debido a la dilución de la atención.

Significancia y Reivindicaciones
El artículo afirma proporcionar el primer estudio sistemático de la recuperación en contexto a la escala de un millón de tokens, estableciéndola como una alternativa viable a la recuperación densa clásica bajo condiciones específicas.

• Viabilidad: Con ajustes arquitectónicos (SSMax y enrutamiento), los LM pueden igualar o superar la recuperación densa en benchmarks estándar y superarlos significativamente en tareas que requieren similitud no semántica.
• Eficiencia: El modelo propuesto de 0.6B logra estos resultados con 7 veces menos parámetros que los modelos concurrentes de contexto largo (MSA-4B) y sin necesidad de reentrenamiento cuando el corpus cambia (a diferencia de la recuperación generativa).
• Cuello de Botella Fundamental: El trabajo identifica la dilución de la atención (dispersión del softmax) como el principal cuello de botella para la recuperación escalable en contexto. Aunque intervenciones simples pueden mitigarlo, el artículo enfatiza que controlar la atención bajo un crecimiento extremo del contexto sigue siendo un desafío fundamental para las arquitecturas de recuperación futuras.

Los autores mantienen la modestia, señalando que, aunque sus métodos extienden la recuperación funcional a corpus más grandes, no "resuelven" completamente el problema de la recuperación a escalas extremas, ya que el rendimiento sigue disminuyendo a medida que $N$ aumenta, y persiste una brecha residual respecto a la recuperación densa en algunos conjuntos de datos.