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

Technische Samenvatting: Kunnen taalmodellen daadwerkelijk in-context ophalen?

Probleemstelling
Hoewel taalmodellen (LM's) een potentieel alternatief bieden voor vectorgebaseerde retrieval door te conditioneren op een in-context corpus om direct relevante document-identificatoren te genereren, blijft deze capaciteit grotendeels onverkend op realistische schaal. Voorgaand werk heeft zich gericht op propriëtaire systemen of kleinschalige reranking-taken, waardoor twee kritieke uitdagingen onbeantwoord blijven: (C1) schalen naar corpora die miljoenen tokens bevatten, en (C2) generaliseren naar corpusgroottes die veel groter zijn dan die tijdens de training gezien zijn. Huidige LM's worstelen met zowel de miljoenen-tokens dure aandachtskosten als de broosheid van extrapolatie naar langere contexten. De centrale vraag is of een LM betrouwbaar kan fungeren als een retriever op corpus-schaal onder deze omstandigheden zonder in te storten.

Methodologie: BLOCKSEARCH en Analyse
De auteurs introduceren BLOCKSEARCH, een systematische studie en modelarchitectuur ontworpen om in-context retrieval (ICR) op corpus-schaal aan te pakken.

• Architectuur & Training: Gebouwd op een 0.6B-parameter ruggengraat (Qwen3-0.6B), maakt BLOCKSEARCH gebruik van block-sparse attention om disjointe documenten efficiënt te verwerken. Belangrijke architecturale en training-aanpassingen zijn onder meer:

  • Gerandomiseerde Document-identificatoren: In plaats van sequentiële ID's worden documenten per trainingsstap toegewezen aan willekeurige viercijferige codes om overfitting op absolute posities te voorkomen.
  • In-Batch Negatives & On-Policy Loss: Het model wordt getraind met 1 positieve en 15 harde negatieven per query. Om exposure bias te mitigeren (waarbij het model conditioneert op zijn eigen potentieel incorrecte gegenereerde prefixes tijdens inferentie), wordt een on-policy hulpverlies (auxiliary loss) toegevoegd. Dit houdt in dat een code wordt gerold uit de eigen distributie van het model, waarbij het model wordt getraind om zijn pad te corrigeren op basis van in-batch scores.
  • Evaluatie-opzet: Het model wordt geëvalueerd op BEIR-datasets (MS MARCO, NQ, HotpotQA) met corpusgroottes variërend van 500 tot 10.000 documenten (tot ~1 miljoen tokens), wat de natuurlijke 32k token contextwindow van het model ver overstijgt.

• Mechanistische Analyse van Falen: De studie identificeert dat hoewel BLOCKSEARCH extrapolatie verbetert, het nog steeds instort op grote schaal (bijna-nul recall bij 1M tokens). De auteurs herleiden dit niet tot een falen in het identificeren van het relevante document (ranking), maar tot een attention dilution effect (aandachtsverwatering):

  • Retrieval vs. Generation Gap: Zelfs bij 10.000 documenten wijzen individuele attention heads vaak nog steeds de hoogste pre-softmax score toe aan het gouden document ($R_{any} \approx 1.0$).
  • Softmax Collapse: Echter, naarmale het corpus ($N$) groeit, neemt de geaggregeerde bijdrage van irrelevante documenten aan de softmax-noemer sneller toe dan de score van het gouden document. Bijgevolg stort de genormaliseerde aandachtsmassa (post-softmax gewicht) op het gouden document in, waardoor de output van de laag een gemiddelde van distractoren wordt in plaats van de relevante inhoud.

• Voorgestelde Mitigaties: Om attention dilution tegen te gaan, stellen de auteurs twee lengte-bewuste aanpassingen voor:

  1. Length-Aware Softmax Aanpassingen:
     • Additive Sinks: Het toevoegen van een geleerde constante aan de softmax-noemer om diffuse aandachtsmassa op te vangen (geïnspireerd door Streaming-LLM).
     • Multiplicatieve Rescaling (SSMax): Het schalen van pre-softmax scores met $s \cdot \log N$ om de kloof tussen de gouden score en de distractoren te vergroten, waardoor wordt voorkomen dat het gouden gewicht verdwijnt naarmale $N$ toeneemt.
  2. Document-Level Sparse Attention: Een routingmechanisme op een intermediaire laag (L16) dat alle documenten scoort en de dichte aandacht beperkt tot de top-$B$ (bijv. 256) documenten voor de daaropvolgende lagen, wat effectief de context filtert vóór de uiteindelijke generatie.

Belangrijkste Resultaten

• Baseline Performance: Het initiële BLOCKSEARCH-model komt overeen met de prestaties van dense retrieval op standaard benchmarks (MS MARCO, NQ) bij kleinere corpusgroottes (~45k tokens), en generaliseert tot ~500k tokens (10x de training context), maar stort in nabij 1M tokens.
• Impact van Aanpassingen:
  • SSMax: Verbetert de prestaties aanzienlijk over de volledige N-sweep. Op MS MARCO bij 1M tokens bereikt het 16,5% Recall@1 (tegenover 0,2% voor het ongewijzigde model), waarmee het grootste deel van de kloof naar de dense baseline (20,2%) wordt gedicht.
  • Routing: Document-level routing verbetert de prestaties verder en bereikt 18,8% Recall@1 op MS MARCO bij 1M tokens, wat bijna de dense baseline evenaart.
  • Gecombineerde Aanpak: De combinatie van SSMax en routing (BLOCKSEARCH-SSMax-routing) bereikt 20,5% Recall@1 op MS MARCO bij 1M tokens, wat de dense baseline (20,2%) licht overtreft en het gelijktijdige 7x grotere MSA-4B model (16,0%) verslaat.
• Out-of-Distribution (LIMIT): Op de LIMIT-benchmark, die lexicale gelijkenis vereist in plaats van semantische embedding-gelijkenis (waar dense retrieval moeite mee heeft), presteert de aangepaste BLOCKSEARCH aanzienlijk beter dan de dense baseline. Bij 850k tokens bereikt BLOCKSEARCH-SSMax-routing 14,9% Recall@1 vergeleken met 3,5% voor de dense retriever (een ~3x verbetering).
• Attention Ceiling: Over alle experimenten blijft de "attention ceiling" (het deel van de queries waarbij ten minste één head het gouden document als eerste rangschikt) nabij de 100%, wat bevestigt dat het retrieval-signaal intern aanwezig is maar verloren gaat tijdens de readout/generatie-fase door attention dilution.

Betekenis en Claims
Het artikel beweert de eerste systematische studie te bieden van in-context retrieval op de miljoen-token schaal, en stelt dit vast als een levensvatbaar alternatief voor klassieke dense retrieval onder specifieke condities.

• Levensvatbaarheid: Met architecturale aanpassingen (SSMax en routing) kunnen LM's de dense retrieval evenaren of overtreffen op standaard benchmarks en presteren ze significant beter op taken die vragen om niet-semantische gelijkenis.
• Efficiëntie: Het voorgestelde 0.6B model behaalt deze resultaten met 7x minder parameters dan gelijktijdige large-context modellen (MSA-4B) en zonder dat er hertraining nodig is wanneer het corpus verandert (in tegen tegenstelling tot generative retrieval).
• Fundamentele Bottleneck: Het werk identificeert attention dilution (softmax dispersie) als de primaire bottleneck voor schaalbare in-context retrieval. Hoewel eenvoudige interventies dit kunnen mitigeren, benadrukt het artikel dat het controleren van aandacht onder extreme contextgroei een fundamentele uitdaging blijft voor toekomstige retrieval-architecturen.

De auteurs blijven bescheiden en merken op dat hoewel hun methoden de functionele retrieval naar grotere corpora uitbreiden, ze het probleem van retrieval op extreme schaal niet volledig "oplossen", aangezien de prestaties afnemen naarmale $N$ toeneemt en er een resterende kloof tot dense retrieval blijft bestaan op sommige datasets.