Atlas: El pre-entrenament conjunt de recuperador i lector supera els LLM de 540 mil milions de paràmetres amb 11 mil milions de paràmetres

Atlas és la continuació d'Izacard i Grave al seu propi article sobre Fusion-in-Decoder, ampliant FiD cap a un sistema totalment entrenat de forma conjunta on el recuperador i el lector s'entrenen conjuntament des del principi. Ho estic llegint ara perquè tanca el llinatge arquitectònic des de l'article original de RAG a través de FiD i cap a la recuperació entrenada conjuntament: exactament l'espai de decisió que qualsevol sistema de preguntes i respostes per a llibres majors ha de navegar.

L'article

"Atlas: Few-shot Learning with Retrieval Augmented Language Models" (Izacard et al., JMLR 2023) es pregunta si els models augmentats per recuperació poden igualar els LLM amb un nombre massiu de paràmetres en tasques de pocs exemples intensives en coneixement. La contribució principal és un sistema augmentat per recuperació acuradament pre-entrenat que entrena conjuntament un recuperador dens basat en Contriever al costat d'un lector Fusion-in-Decoder basat en T5. La idea clau és que el pre-entrenament conjunt —no l'arquitectura— és el que impulsa el rendiment del coneixement amb pocs exemples. El sistema recupera els 20 documents principals, codifica cadascun de forma independent en el codificador i després els fusiona en l'atenció creuada del descodificador, el mateix disseny FiD de l'article dels autors de 2021.

Idees clau

• Atlas-11B aconsegueix una precisió del 42,4% en Natural Questions amb només 64 exemples d'entrenament, superant PaLM (540B paràmetres) per aproximadament 3 punts mentre utilitza 50 vegades menys paràmetres.
• A TriviaQA (64 exemples), Atlas-11B arriba al 74,5% en el conjunt filtrat i al 84,7% en la prova oculta no filtrada, demostrant que el component de recuperació compensa amb escreix la supervisió limitada de la tasca.
• S'avaluen quatre objectius d'entrenament del recuperador: Destil·lació d'Atenció (ADist), EMDR2 (tractant els documents recuperats com a variables latents), Destil·lació de Perplexitat (PDist) i LOOP (deixar-ne un fora). Les diferències de rendiment entre ells són petites; s'adopta PDist per eficiència de càlcul.
• El pre-entrenament conjunt en text no etiquetat és el factor més important: totes les configuracions de pre-entrenament augmentades per recuperació superen amb escreix la base de referència que només realitza ajustament fi.
• L'índex de documents es pot actualitzar després de l'entrenament sense tornar a entrenar el model, la qual cosa és arquitectònicament important per a bases de coneixement dinàmiques. Els índexs desajustats temporalment degraden el rendiment de manera notable.
• A MMLU (5 exemples), Atlas-11B arriba al 47,9%, superant el 43,9% informat de GPT-3, malgrat tenir aproximadament 16 vegades menys paràmetres.

Què es manté i què no

L'afirmació principal —que la recuperació permet un rendiment del coneixement amb pocs exemples amb una fracció del nombre de paràmetres— es manté de manera convincent. La xifra del 42,4% en NQ amb 64 exemples és un resultat sorprenent, i la comparació amb PaLM és justa perquè PaLM era el referent d'escala més avançat en aquell moment.

Però tinc tres reserves. En primer lloc, la precisió de la recuperació no és excel·lent fins i tot després de l'entrenament conjunt: anàlisis independents mostren que Contriever omet almenys una afirmació de referència en aproximadament el 85% dels casos i aconsegueix una precisió de recuperació de preguntes i respostes al voltant del 47%. L'entrenament conjunt millora la recuperació respecte a les bases de referència no entrenades conjuntament, però el lector està fent un treball enorme per compensar una recuperació imperfecta; les xifres principals de pocs exemples reflecteixen el sostre del sistema, no la qualitat del component de recuperació. En segon lloc, el cost d'infraestructura és real: actualitzar els índexs de documents durant el pre-entrenament afegeix aproximadament un 30% de sobrecàrrega computacional, i l'índex complet de Wikipedia+CommonCrawl requereix 587 GB en fp16. Això és gestionable per a un entorn de recerca, però és una restricció operativa real per al desplegament en producció. En tercer lloc, la fuga de dades s'accepta però no es resol: el 2,8% de les preguntes de MMLU apareixen textualment al corpus CCNet utilitzat per al pre-entrenament, inflant els resultats de MMLU en un marge desconegut.

També hi ha una limitació arquitectònica més subtil que l'article no aborda totalment: FiD codifica cada passatge recuperat de manera independent abans de la fusió, cosa que ajuda al paral·lelisme però significa que el codificador no té atenció entre passatges. Les cadenes de raonament multietapa llargues que necessiten connectar informació entre diversos passatges han de fer tota aquesta feina en el descodificador —i amb 20 passatges recuperats, l'atenció creuada del descodificador suporta una càrrega pesada.

Per què això és important per a la IA financera

Per a les preguntes i respostes dels llibres majors de Beancount, la contribució més rellevant d'Atlas és la demostració empírica que l'entrenament conjunt de recuperador i lector val la pena en configuracions de pocs exemples —i la seva comptabilitat honesta de quan no és així. Un agent de Beancount que consulta l'historial de transaccions de diversos anys s'enfronta precisament al problema de l'índex dinàmic: noves entrades arriben diàriament, i un índex que té un mes d'antiguitat produeix respostes incorrectes. Atlas demostra que l'índex es pot canviar en calent sense tornar a entrenar, cosa que és arquitectònicament encoratjadora.

No obstant això, les xifres de precisió de recuperació són alliçonadores. Si Contriever omet l'entrada del llibre major rellevant en el 53% dels intents de recuperació, fins i tot després de l'entrenament conjunt en text general, un agent del domini financer que operi sobre llibres majors de Beancount —amb els seus noms de matèries primeres (commodities) específics del domini, jerarquies de comptes i directives de bean— necessitarà un entrenament del recuperador adaptat al domini o una recuperació augmentada per mètodes de consulta estructurats (coincidència exacta de comptes, filtratge per dates). La recuperació a l'estil RAG per si sola, fins i tot entrenada conjuntament, no serà suficient per a operacions de llibres majors d'alta precisió.

La comparació amb PaLM també clarifica el compromís arquitectònic: la recuperació permet comprimir el coneixement en menys paràmetres, reduint el cost d'inferència. Per a un producte com Beancount.io on el cost d'inferència importa a escala, la filosofia de disseny d'Atlas és atractiva. Però el cost de l'índex de 587 GB trasllada la càrrega a la infraestructura d'emmagatzematge i recuperació —un tipus diferent de restricció operativa que no apareix en les xifres dels referents.

Què llegir a continuació

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — el marc anterior de pre-entrenament conjunt de recuperador i lector que Atlas amplia; essencial per entendre què millora realment Atlas i què deixa sense canvis.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — aconsegueix un rendiment competitiu amb Atlas utilitzant l'ajustament d'instruccions en lloc del pre-entrenament conjunt des de zero; suggereix que la bretxa entre l'entrenament conjunt i l'independent es pot tancar sense el cost d'infraestructura.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — l'enfocament de DeepMind per a la recuperació durant el pre-entrenament a una escala diferent; completa la visió dels enfocaments de pre-entrenament augmentats per recuperació abans de prendre decisions arquitectòniques per a les preguntes i respostes dels llibres majors.