Atlas: Gezamenlijke Retriever-Reader Pre-Training Verslaat LLM's met 540 Miljard Parameters met Slechts 11 Miljard Parameters

Atlas is het vervolg van Izacard en Grave op hun eigen Fusion-in-Decoder-paper, waarbij FiD wordt uitgebreid naar een volledig gezamenlijk getraind systeem waarin de retriever en reader vanaf de basis samen worden getraind. Ik lees het nu omdat het de architecturale lijn sluit van de oorspronkelijke RAG-paper via FiD naar gezamenlijk getrainde retrieval—precies de beslissingsruimte waar elk QA-systeem voor grootboeken doorheen moet navigeren.

Het artikel

"Atlas: Few-shot Learning with Retrieval Augmented Language Models" (Izacard et al., JMLR 2023) stelt de vraag of retrieval-augmented modellen de prestaties van LLM's met enorme parameters kunnen evenaren op kennisintensieve few-shot taken. De belangrijkste bijdrage is een zorgvuldig voorgeprogrammeerd retrieval-augmented systeem dat een op Contriever gebaseerde dense retriever gezamenlijk traint met een op T5 gebaseerde Fusion-in-Decoder reader. Het belangrijkste inzicht is dat gezamenlijke pre-training—niet de architectuur—de drijvende kracht is achter few-shot kennisprestaties. Het systeem haalt de top-20 documenten op, codeert elk onafhankelijk in de encoder en fuseert ze vervolgens in de cross-attention van de decoder, hetzelfde FiD-ontwerp uit de paper van de auteurs uit 2021.

Belangrijkste ideeën

• Atlas-11B behaalt een nauwkeurigheid van 42,4% op Natural Questions met slechts 64 trainingsvoorbeelden, waarmee het PaLM (540 miljard parameters) met ongeveer 3 punten overtreft, terwijl het 50x minder parameters gebruikt.
• Op TriviaQA (64-shot) bereikt Atlas-11B 74,5% op de gefilterde set en 84,7% op de ongefilterde verborgen test, wat aantoont dat de retrieval-component sterk compenseert voor beperkte taaksupervisie.
• Vier trainingsdoelen voor de retriever worden geëvalueerd: Attention Distillation (ADist), EMDR2 (waarbij opgehaalde documenten als latente variabelen worden behandeld), Perplexity Distillation (PDist) en LOOP (leave-one-out). De prestatieverschillen tussen deze doelen zijn klein; PDist wordt gekozen vanwege de reken-efficiëntie.
• Gezamenlijke pre-training op ongelabelde tekst is de belangrijkste factor: alle retrieval-augmented pre-training configuraties presteren aanzienlijk beter dan de baseline die alleen retrieval-augmented fine-tuning gebruikt.
• De documentindex kan na de training worden bijgewerkt zonder het model opnieuw te trainen, wat architecturaal belangrijk is voor dynamische kennisbanken. Temporeel niet-overeenkomende indices verslechteren de prestaties merkbaar.
• Op MMLU (5-shot) bereikt Atlas-11B 47,9%, wat hoger is dan de gerapporteerde 43,9% van GPT-3, ondanks ongeveer 16x minder parameters.

Wat standhoudt — en wat niet

De belangrijkste claim—dat retrieval few-shot kennisprestaties mogelijk maakt met een fractie van het aantal parameters—houdt overtuigend stand. Het NQ-cijfer van 42,4% met 64 voorbeelden is een opvallend resultaat, en de vergelijking met PaLM is eerlijk omdat PaLM op dat moment de state-of-the-art schaalbenchmark was.

Ik heb echter drie kanttekeningen. Ten eerste is de retrieval-nauwkeurigheid zelfs na gezamenlijke training niet geweldig: onafhankelijke analyses tonen aan dat Contriever in ongeveer 85% van de gevallen ten minste één 'gold statement' mist, en een QA-retrieval-nauwkeurigheid van ongeveer 47% behaalt. Gezamenlijke training verbetert de retrieval ten opzichte van niet-gezamenlijk getrainde baselines, maar de reader verricht enorm veel werk om de imperfecte retrieval te compenseren—de spectaculaire few-shot cijfers weerspiegelen het plafond van het systeem, niet de kwaliteit van de retrieval-component. Ten tweede zijn de infrastructuurkosten reëel: het verversen van documentindices tijdens de pre-training voegt ongeveer 30% computationele overhead toe, en de volledige Wikipedia+CommonCrawl-index vereist 587 GB in fp16. Dat is beheersbaar voor een onderzoeksomgeving, maar is een reële operationele beperking voor productie-implementatie. Ten derde wordt datalekken erkend maar niet opgelost: 2,8% van de MMLU-vragen komt letterlijk voor in het CCNet-corpus dat voor pre-training is gebruikt, wat de MMLU-resultaten met een onbekende marge opblaast.

Er is ook een subtielere architecturale beperking waar de paper niet volledig op ingaat: FiD codeert elk opgehaald tekstfragment onafhankelijk vóór de fusie, wat helpt bij parallellisatie, maar betekent dat de encoder geen cross-passage attention heeft. Lange multi-hop redeneerketens die informatie tussen fragmenten moeten verbinden, moeten al dat werk in de decoder doen—en bij 20 opgehaalde fragmenten draagt de decoder cross-attention een zware last.

Waarom dit belangrijk is voor financiële AI

Voor QA op Beancount-grootboeken is de meest relevante bijdrage van Atlas de empirische demonstratie dat gezamenlijke retriever-reader-training loont in few-shot scenario's—en de eerlijke weergave van wanneer dat niet het geval is. Een Beancount-agent die een transactiegeschiedenis van meerdere jaren bevraagt, wordt geconfronteerd met precies het probleem van de dynamische index: dagelijks komen er nieuwe boekingen bij, en een index die een maand oud is, produceert foute antwoorden. Atlas laat zien dat de index 'hot-swapped' kan worden zonder hertraining, wat architecturaal bemoedigend is.

De cijfers voor retrieval-nauwkeurigheid zijn echter ontnuchterend. Als Contriever de relevante boeking in het grootboek mist in 53% van de retrieval-pogingen, zelfs na gezamenlijke training op algemene tekst, dan zal een agent in het financiële domein die werkt met Beancount-grootboeken—met hun domeinspecifieke commodity-namen, rekeninghiërarchieën en bean-directives—behoefte hebben aan ofwel domein-adaptieve retriever-training, ofwel retrieval die wordt aangevuld met gestructureerde query-methoden (exacte rekening-matching, datumfiltering). RAG-stijl retrieval alleen, zelfs gezamenlijk getraind, zal niet voldoende zijn voor precisie-operaties op grootboeken.

De vergelijking met PaLM verduidelijkt ook de architecturale afweging: retrieval stelt je in staat om kennis te comprimeren in minder parameters, wat de inferentiekosten verlaagt. Voor een product als Beancount.io waar inferentiekosten op schaal van belang zijn, is de ontwerpfilosofie van Atlas aantrekkelijk. Maar de indexkosten van 587 GB verschuiven de last naar opslag- en retrieval-infrastructuur—een ander soort operationele beperking die niet direct zichtbaar is in de benchmarkcijfers.

Wat nu te lezen

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — het eerdere raamwerk voor gezamenlijke retriever-reader pre-training dat Atlas uitbreidt; essentieel om te begrijpen wat Atlas daadwerkelijk verbetert en wat ongewijzigd blijft.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — behaalt concurrerende prestaties met Atlas door middel van instruction-tuning in plaats van gezamenlijke pre-training vanaf nul; suggereert dat de kloof tussen gezamenlijke en onafhankelijke training gedicht kan worden zonder de infrastructuurkosten.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — de aanpak van DeepMind voor retrieval tijdens pre-training op een andere schaal; maakt het beeld van retrieval-augmented pre-training benaderingen compleet voordat er architecturale keuzes worden gemaakt voor grootboek-QA.