Atlas : le pré-entraînement conjoint de l'extracteur et du lecteur surpasse les LLM à 540 milliards de paramètres avec seulement 11 milliards

Atlas est la suite des travaux d'Izacard et Grave après leur article sur Fusion-in-Decoder, étendant FiD en un système entièrement entraîné conjointement où l'extracteur et le lecteur sont co-entraînés depuis la base. Je le lis actuellement car il boucle la lignée architecturale allant du papier RAG original à la récupération entraînée conjointement en passant par FiD — exactement l'espace de décision qu'un système de questions-réponses (QA) pour grands livres doit explorer.

L'article

« Atlas : Few-shot Learning with Retrieval Augmented Language Models » (Izacard et al., JMLR 2023) se demande si les modèles augmentés par récupération peuvent égaler les LLM à paramètres massifs sur des tâches à forte intensité de connaissances avec peu d'exemples (few-shot). La contribution centrale est un système augmenté par récupération soigneusement pré-entraîné qui entraîne conjointement un extracteur dense basé sur Contriever avec un lecteur Fusion-in-Decoder basé sur T5. L'idée clé est que c'est le pré-entraînement conjoint — et non l'architecture — qui stimule la performance des connaissances en few-shot. Le système récupère les 20 meilleurs documents, encode chacun indépendamment dans l'encodeur, puis les fusionne dans l'attention croisée du décodeur, reprenant la conception FiD de l'article des auteurs de 2021.

Idées clés

• Atlas-11B atteint une précision de 42,4 % sur Natural Questions avec seulement 64 exemples d'entraînement, surpassant PaLM (540 milliards de paramètres) d'environ 3 points tout en utilisant 50 fois moins de paramètres.
• Sur TriviaQA (64-shot), Atlas-11B atteint 74,5 % sur l'ensemble filtré et 84,7 % sur le test caché non filtré, montrant que la composante de récupération compense fortement la supervision limitée de la tâche.
• Quatre objectifs d'entraînement de l'extracteur sont évalués : Attention Distillation (ADist), EMDR2 (traitant les documents récupérés comme des variables latentes), Perplexity Distillation (PDist) et LOOP (leave-one-out). Les différences de performance entre eux sont minimes ; PDist est adopté pour son efficacité de calcul.
• Le pré-entraînement conjoint sur du texte non étiqueté est le facteur le plus important : toutes les configurations de pré-entraînement augmentées par récupération surpassent largement la base de référence de réglage fin (fine-tuning) uniquement augmentée par récupération.
• L'index des documents peut être mis à jour après l'entraînement sans ré-entraîner le modèle, ce qui est architecturalement important pour les bases de connaissances dynamiques. Des index temporellement décalés dégradent notablement les performances.
• Sur MMLU (5-shot), Atlas-11B atteint 47,9 %, dépassant les 43,9 % rapportés par GPT-3, malgré environ 16 fois moins de paramètres.

Ce qui tient la route — et ce qui ne la tient pas

L'affirmation principale — selon laquelle la récupération permet une performance de connaissance en few-shot avec une fraction du nombre de paramètres — tient la route de manière convaincante. Le chiffre de 42,4 % sur NQ avec 64 exemples est un résultat frappant, et la comparaison avec PaLM est juste car PaLM était la référence de l'état de l'art à l'époque.

Cependant, j'ai trois réserves. Premièrement, la précision de la récupération n'est pas excellente même après l'entraînement conjoint : des analyses indépendantes montrent que Contriever manque au moins un énoncé de référence dans environ 85 % des cas, et atteint environ 47 % de précision de récupération QA. L'entraînement conjoint améliore la récupération par rapport aux bases de référence non entraînées conjointement, mais le lecteur fait un travail énorme pour compenser une récupération imparfaite — les chiffres phares du few-shot reflètent le plafond du système, pas la qualité de la composante de récupération. Deuxièmement, le coût d'infrastructure est réel : rafraîchir les index de documents pendant le pré-entraînement ajoute environ 30 % de surcharge de calcul, et l'index complet Wikipedia+CommonCrawl nécessite 587 Go en fp16. C'est gérable pour un cadre de recherche mais constitue une véritable contrainte opérationnelle pour un déploiement en production. Troisièmement, la fuite de données est reconnue mais non résolue : 2,8 % des questions MMLU apparaissent textuellement dans le corpus CCNet utilisé pour le pré-entraînement, gonflant les résultats MMLU d'une marge inconnue.

Il existe également une limitation architecturale plus subtile que l'article n'aborde pas pleinement : FiD encode chaque passage récupéré indépendamment avant la fusion, ce qui aide au parallélisme mais signifie que l'encodeur n'a pas d'attention croisée entre les passages. Les longues chaînes de raisonnement à sauts multiples (multi-hop) qui doivent connecter des informations à travers les passages doivent effectuer tout ce travail dans le décodeur — et avec 20 passages récupérés, l'attention croisée du décodeur supporte une lourde charge.

Pourquoi cela est important pour l'IA financière

Pour le QA sur les grands livres Beancount, la contribution la plus pertinente d'Atlas est la démonstration empirique que l'entraînement conjoint extracteur-lecteur porte ses fruits dans des contextes de few-shot — ainsi que son compte rendu honnête des moments où ce n'est pas le cas. Un agent Beancount interrogeant un historique de transactions sur plusieurs années est confronté exactement au problème de l'index dynamique : de nouvelles écritures arrivent quotidiennement, et un index vieux d'un mois produit des réponses erronées. Atlas montre que l'index peut être remplacé à chaud sans ré-entraînement, ce qui est encourageant sur le plan architectural.

Les chiffres de précision de récupération sont toutefois décevants. Si Contriever manque l'écriture comptable pertinente dans 53 % des tentatives de récupération même après un entraînement conjoint sur du texte général, un agent du domaine financier opérant sur des journaux Beancount — avec leurs noms de commodités spécifiques au domaine, leurs hiérarchies de comptes et leurs directives Beancount — aura besoin soit d'un entraînement de l'extracteur adapté au domaine, soit d'une récupération augmentée par des méthodes de requête structurées (correspondance exacte de compte, filtrage par date). La récupération de style RAG seule, même entraînée conjointement, ne suffira pas pour des opérations de grand livre de haute précision.

La comparaison avec PaLM clarifie également le compromis architectural : la récupération vous permet de compresser les connaissances dans moins de paramètres, abaissant le coût d'inférence. Pour un produit comme Beancount.io où le coût d'inférence compte à grande échelle, la philosophie de conception d'Atlas est séduisante. Mais le coût de l'index de 587 Go déplace la charge sur l'infrastructure de stockage et de récupération — un type différent de contrainte opérationnelle qui n'apparaît pas dans les chiffres de référence.

Lectures recommandées

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — le cadre de pré-entraînement conjoint extracteur-lecteur antérieur qu'Atlas étend ; essentiel pour comprendre ce qu'Atlas améliore réellement et ce qu'il laisse inchangé.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — atteint des performances compétitives avec Atlas en utilisant le réglage par instructions plutôt que le pré-entraînement conjoint à partir de zéro ; suggère que l'écart entre l'entraînement conjoint et indépendant peut être réduit sans le coût d'infrastructure.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — l'approche de DeepMind pour la récupération pendant le pré-entraînement à une échelle différente ; complète le panorama des approches de pré-entraînement augmentées par récupération avant de faire des choix architecturaux pour le QA de grands livres.