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

Résumé technique : Les modèles de langage peuvent-ils réellement effectuer de la récupération en contexte ?

Énoncé du problème
Bien que les modèles de langage (LM) offrent une alternative potentielle à la recherche vectorielle en se conditionnant sur un corpus en contexte pour générer directement des identifiants de documents pertinents, cette capacité reste largement inexplorée à des échelles réalistes. Les travaux antérieurs se sont concentrés sur des systèmes propriétaires ou des tâches de re-classement (reranking) à petite échelle, laissant deux défis critiques non résolus : (C1) le passage à l'échelle vers des corpus contenant des millions de jetons, et (C2) la généralisation à des tailles de corpus dépassant largement celles rencontrées lors de l'entraînement. Les LM actuels peinent tant avec les coûts d'attention d'un million de jetons qu'avec la fragilité de l'extrapolation vers des contextes plus longs. La question centrale est de savoir si un LM peut agir de manière fiable comme un moteur de recherche à l'échelle d'un corpus dans ces conditions sans s'effondrer.

Méthodologie : BLOCKSEARCH et analyse
Les auteurs introduisent BLOCKSEARCH, une étude systématique et une architecture de modèle conçues pour aborder la récupération en contexte (ICR) à l'échelle du corpus.

• Architecture et entraînement : Construit sur un backbone de 0,6 milliard de paramètres (Qwen3-0.6B), BLOCKSEARCH utilise une attention parcimonieuse par blocs (block-sparse attention) pour gérer efficacement les documents disjoints. Les modifications clés de l'architecture et de l'entraînement incluent :

  • Identifiants de documents aléatoires : Au lieu d'IDs séquentiels, les documents se voient attribuer des codes à quatre chiffres aléatoires par étape d'entraînement afin d'éviter le surapprentissage des positions absolues.
  • Négatifs au sein du lot (In-Batch Negatives) et perte on-policy : Le modèle est entraîné avec 1 exemple positif et 15 négatifs difficiles par requête. Pour atténuer le biais d'exposition (où le modèle se conditionne sur ses propres préfixes potentiellement incorrects générés lors de l'inférence), une perte auxiliaire on-policy est ajoutée. Cela implique de simuler un code issu de la propre distribution du modèle et d'entraîner le modèle à corriger sa trajectoire en se basant sur les scores du lot.
  • Configuration de l'évaluation : Le modèle est évalué sur les jeux de données BEIR (MS MARCO, NQ, HotpotQA) avec des tailles de corpus allant de 500 à 10 000 documents (jusqu'à environ 1 million de jetons), dépassant de loin sa fenêtre de contexte native de 32k jetons.

• Analyse mécaniste de l'échec : L'étude identifie que, bien que BLOCKSEARCH améliore l'extrapolation, il s'effondre toujours à grande échelle (rappel proche de zéro à 1M de jetons). Les auteurs attribuent cela non pas à un échec de l'identification du document pertinent (classement), mais à un effet de dilution de l'attention :

  • Écart entre récupération et génération : Même à 10 000 documents, les têtes d'attention individuelles assignent souvent le score pré-softmax le plus élevé au document de référence ($R_{any} \approx 1.0$).
  • Effondrement du softmax : Cependant, à mesure que le corpus ($N$) croît, la contribution agrégée des documents non pertinents au dénominateur du softmax augmente plus rapidement que le score du document de référence. Par conséquent, la masse d'attention normalisée (poids post-softmax) sur le document de référence s'effondre, provoquant le fait que la sortie de la couche devienne une moyenne des distracteurs plutôt que le contenu pertinent.

• Mitigations proposées : Pour contrer la dilution de l'attention, les auteurs proposent deux ajustements sensibles à la longueur :

  1. Ajustements de Softmax sensibles à la longueur :
     • Puits additifs (Additive Sinks) : Ajouter une constante apprise au dénominateur du softmax pour absorber la masse d'attention diffuse (inspiré de Streaming-LLM).
     • Rescaling multiplicatif (SSMax) : Mise à l'échelle des scores pré-softmax par $s \cdot \log N$ pour élargir l'écart entre le score de référence et les distracteurs, empêissant le poids de la référence de s'évanouir à mesure que $N$ augmente.
  2. Attention parcimonieuse au niveau du document : Un mécanisme de routage à une couche intermédiaire (L16) qui évalue tous les documents et restreint l'attention dense aux $B$ meilleurs documents (par exemple, 265) pour les couches suivantes, effectuant ainsi un élagage du contexte avant la génération finale.

Résultats clés

• Performance de base : Le modèle BLOCKSEARCH initial égale la performance de la recherche dense sur les benchmarks standards (MS MARCO, NQ) à de petites tailles de corpus (~45k jetons), et se généralise jusqu'à ~500k jetons (10x le contexte d'entraînement), mais s'effondre près de 1M de jetons.
• Impact des modifications :
  • SSMax : Améliore significativement les performances sur toute la courbe de $N$. Sur MS MARCO à 1M de jetons, il atteint un Rappel@1 de 16,5 % (contre 0,2 % pour le modèle non modifié), comblant la majeure partie de l'écart avec la ligne de base dense (20,2 %).
  • Routage : Le routage au niveau du document améliore davantage les performances, atteignant un Rappel@1 de 18,8 % sur MS MARCO à 1M de jetons, rejoignant presque la recherche dense.
  • Approche combinée : La combinaison de SSMax et du routage (BLOCKSEARCH-SSMax-routing) atteint un Rappel@1 de 20,5 % sur MS MARCO à 1M de jetons, dépassant légèrement la ligne de base dense (20,2 %) et surpassant le modèle concurrent MSA-4B, 7 fois plus grand.
• Hors-distribution (LIMIT) : Sur le benchmark LIMIT, qui exige une similitude lexicale plutôt qu'une similitude d'embedding sémantique (là où la recherche dense peine), le modèle BLOCKSEARCH modifié surpasse significativement la recherche dense. À 850k jetons, BLOCKSEARCH-SSMax-routing atteint un Rappel@1 de 14,9 % contre 3,5 % pour le récupérateur dense (une amélioration d'environ 3x).
• Plafond d'attention : À travers toutes les expériences, le "plafond d'attention" (la fraction de requêtes où au moins une tête classe le document de référence en premier) reste proche de 100 %, confirmant que le signal de récupération est présent en interne mais perdu lors de la phase de lecture/génération à cause de la dilution de l'attention.

Signification et revendications
L'article affirme fournir la première étude systématique de la récupération en contexte à l'échelle du million de jetons, l'établissant comme une alternative viable à la recherche dense classique sous certaines conditions.

• Viabilité : Avec des ajustements architecturaux (SSMax et routage), les LM peuvent égaler ou dépasser la recherche dense sur des benchmarks standards et surpasser nettement la recherche dense sur des tâches nécessitant une similitude non sémantique.
• Efficacité : Le modèle de 0,6B proposé atteint ces résultats avec 7 fois moins de paramètres que les modèles à contexte long concurrents (MSA-4B) et sans nécessiter de réentraînement lorsque le corpus change (contrairement à la recherche générative).
• Goulot d'étranglement fondamental : Ce travail identifie la dilution de l'attention (dispersion du softmax) comme le principal goulot d'étranglement pour la récupération en contexte à grande échelle. Bien que des interventions simples puissent l'atténuer, l'article souligne que le contrôle de l'attention sous une croissance extrême du contexte reste un défi fondamental pour les futures architectures de recherche.

Les auteurs restent modestes, notant que bien que leurs méthodes étendent la récupération fonctionnelle à des corpus plus larges, elles ne "résolvent" pas totalement le problème de la recherche aux échelles extrêmes, car les performances déclinent à mesure que $N$ augmente et qu'un écart résiduel avec la recherche dense subsiste sur certains jeux de données.