Atlas: Spoločný tréning retrievera a readera prekonáva 540B-parametrové LLM modely s 11B parametrami

Atlas je pokračovaním práce Izacarda a Gravea na ich vlastnom článku o Fusion-in-Decoder, pričom rozširuje FiD na plne spoločne trénovaný systém, kde sú retriever a reader spoločne trénovaní od základov. Čítam to teraz, pretože to uzatvára architektonickú líniu od pôvodného článku o RAG cez FiD až po spoločne trénované vyhľadávanie – čo je presne ten rozhodovací priestor, v ktorom sa musí orientovať každý systém na odpovedanie na otázky (QA) nad účtovnou knihou.

O článku

„Atlas: Few-shot Learning with Retrieval Augmented Language Models“ (Izacard a kol., JMLR 2023) skúma, či modely rozšírené o vyhľadávanie dokážu konkurovať modelom LLM s obrovským počtom parametrov v úlohách s málo príkladmi (few-shot), ktoré sú náročné na vedomosti. Hlavným prínosom je starostlivo predtrénovaný systém rozšírený o vyhľadávanie, ktorý spoločne trénuje hustý retriever založený na modeli Contriever spolu s readerom Fusion-in-Decoder založenom na modeli T5. Kľúčovým poznatkom je, že spoločné predtrénovanie – nie architektúra – je to, čo poháňa výkonnosť vedomostí pri učení s málo príkladmi. Systém vyhľadá 20 najrelevantnejších dokumentov, každý nezávisle zakóduje v enkódere a potom ich spojí v rámci cross-attention v dekódere, čo je rovnaký dizajn FiD ako v článku autorov z roku 2021.

Kľúčové myšlienky

• Atlas-11B dosahuje presnosť 42,4 % v rámci Natural Questions len so 64 tréningovými príkladmi, čím prekonáva PaLM (540 miliárd parametrov) o približne 3 body pri použití 50-krát menej parametrov.
• V rámci TriviaQA (64-shot) dosahuje Atlas-11B úroveň 74,5 % na filtrovanej sade a 84,7 % na nefiltrovanom skrytom teste, čo ukazuje, že komponent vyhľadávania výrazne kompenzuje obmedzený dohľad nad úlohou.
• Hodnotia sa štyri ciele tréningu retrievera: Attention Distillation (ADist), EMDR2 (považovanie vyhľadaných dokumentov za latentné premenné), Perplexity Distillation (PDist) a LOOP (leave-one-out). Rozdiely vo výkonnosti medzi nimi sú malé; PDist sa prijíma kvôli efektívnosti výpočtov.
• Spoločné predtrénovanie na neoznačenom texte je najvýznamnejším faktorom: všetky konfigurácie predtrénovania rozšíreného o vyhľadávanie výrazne prekonávajú základnú líniu (baseline) využívajúcu iba jemné doladenie rozšírené o vyhľadávanie.
• Index dokumentov je možné aktualizovať po tréningu bez potreby opätovného trénovania modelu, čo je architektonicky dôležité pre dynamické znalostné bázy. Časovo nezhodné indexy citeľne znižujú výkon.
• V rámci MMLU (5-shot) dosahuje Atlas-11B úroveň 47,9 %, čím prekonáva uvádzaných 43,9 % pri GPT-3, a to napriek približne 16-krát menšiemu počtu parametrov.

Čo obstojí — a čo nie

Hlavné tvrdenie – že vyhľadávanie umožňuje výkonnosť vedomostí pri učení s málo príkladmi pri zlomku počtu parametrov – presvedčivo obstojí. Výsledok 42,4 % NQ so 64 príkladmi je pozoruhodný a porovnanie s PaLM je férové, pretože PaLM bol v tom čase najlepším benchmarkom v mierke.

Mám však tri výhrady. Po prvé, presnosť vyhľadávania nie je skvelá ani po spoločnom tréningu: nezávislé analýzy ukazujú, že Contriever vynechá aspoň jedno kľúčové tvrdenie (gold statement) v približne 85 % prípadov a dosahuje približne 47 % presnosť vyhľadávania QA. Spoločný tréning zlepšuje vyhľadávanie oproti nespoločným baseline modelom, ale reader vynakladá obrovské úsilie na kompenzáciu nedokonalého vyhľadávania – tie hlavné few-shot čísla odrážajú strop systému, nie kvalitu komponentu vyhľadávania. Po druhé, náklady na infraštruktúru sú reálne: aktualizácia indexov dokumentov počas predtrénovania pridáva približne 30 % výpočtovej réžie a kompletný index Wikipedia + CommonCrawl vyžaduje 587 GB v fp16. To je zvládnuteľné vo výskumnom prostredí, ale predstavuje to skutočné prevádzkové obmedzenie pre nasadenie v produkcii. Po tretie, únik dát (data leakage) je priznaný, ale nevyriešený: 2,8 % otázok z MMLU sa doslovne objavuje v korpuse CCNet použitom na predtrénovanie, čo zvyšuje výsledky MMLU o neznámu mieru.

Existuje aj jemnejšie architektonické obmedzenie, ktorým sa článok plne nezaoberá: FiD kóduje každú vyhľadanú pasáž nezávisle pred fúziou, čo pomáha paralelizácii, ale znamená to, že enkóder nemá cross-passage attention. Dlhé viacúrovňové (multi-hop) argumentačné reťazce, ktoré potrebujú prepojiť informácie naprieč pasážami, musia všetku túto prácu vykonať v dekódere – a pri 20 vyhľadaných pasážach nesie cross-attention dekodéra veľkú záťaž.

Prečo je to dôležité pre AI vo financiách

Pre QA nad účtovnou knihou Beancount je najrelevantnejším prínosom Atlasu empirický dôkaz, že spoločný tréning retrievera a readera sa vypláca v nastaveniach s málo príkladmi – a jeho čestné zhodnotenie toho, kedy to tak nie je. Agent Beancount dopytujúci sa na viacročnú históriu transakcií čelí presne problému dynamického indexu: nové záznamy pribúdajú denne a index, ktorý je mesiac starý, poskytuje nesprávne odpovede. Atlas ukazuje, že index je možné vymeniť za chodu (hot-swap) bez opätovného trénovania, čo je architektonicky povzbudivé.

Čísla presnosti vyhľadávania sú však vytriezvením. Ak Contriever vynechá relevantný záznam v účtovnej knihe v 53 % pokusov o vyhľadávanie aj po spoločnom tréningu na všeobecnom texte, finančný agent pracujúci nad účtovnými knihami Beancount – s ich špecifickými názvami komodít, hierarchiami účtov a direktívami bean – bude potrebovať buď doménovo adaptívny tréning retrievera, alebo vyhľadávanie rozšírené o štruktúrované metódy dopytovania (presná zhoda účtov, filtrovanie podľa dátumu). Samotné vyhľadávanie v štýle RAG, dokonca aj spoločne trénované, nebude stačiť na vysoko presné operácie v účtovnej knihe.

Porovnanie s PaLM tiež objasňuje architektonický kompromis: vyhľadávanie vám umožňuje komprimovať znalosti do menšieho počtu parametrov, čím sa znižujú náklady na inferenciu. Pre produkt ako Beancount.io, kde na nákladoch na inferenciu záleží pri škálovaní, je filozofia dizajnu Atlasu príťažlivá. Náklady na index s veľkosťou 587 GB však presúvajú záťaž na infraštruktúru úložiska a vyhľadávania – iný druh prevádzkového obmedzenia, ktorý sa v číslach benchmarkov neobjavuje.

Čo si prečítať ďalej

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu a kol., arXiv:2002.08909, ICML 2020) – skorší rámec spoločného predtrénovania retrievera a readera, ktorý Atlas rozširuje; nevyhnutné pre pochopenie toho, čo Atlas skutočne zlepšuje a čo ponecháva nezmenené.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin a kol., arXiv:2310.01352, ICLR 2024) – dosahuje konkurencieschopný výkon s Atlasom pomocou ladenia podľa inštrukcií namiesto spoločného predtrénovania od základu; naznačuje, že medzeru medzi spoločným a nezávislým tréningom je možné uzavrieť bez nákladov na infraštruktúru.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud a kol., arXiv:2112.04426, ICML 2022) – prístup DeepMind k vyhľadávaniu počas predtrénovania v inom rozsahu; dopĺňa obraz prístupov k predtrénovaniu rozšírenému o vyhľadávanie pred prijatím architektonických rozhodnutí pre QA nad účtovnou knihou.