Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

Guide de lecture

Problème et motivation

Les modèles de langue pré-entraînés emmagasinent d'importantes connaissances factuelles au sein de leurs paramètres (mémoire paramétrique). Cependant, ils présentent des inconvénients majeurs : ils peinent à étendre ou mettre à jour leurs connaissances sans réentraînement, ne peuvent pas facilement fournir la provenance ou l'interprétabilité de leurs sorties, et sont sujets aux hallucinations.

Bien que des modèles hybrides paramétriques/non paramétriques comme REALM et ORQA aient associé la recherche différentiable au masquage de modèle de langue, ils étaient limités au question-réponse extractif. Les architectures séquence-à-séquence (seq2seq) standards comme BART et T5 sont restées largement en boucle fermée (« closed-book »). Ce travail comble cette lacune en proposant la Retrieval-Augmented Generation (RAG), un cadre général d'affinage qui couple un extracteur neuronal dense pré-entraîné avec un générateur autorégressif seq2seq pour la génération et la classification en domaine ouvert.

Idées clés

• Mémoire hybride paramétrique et non paramétrique : La mémoire non paramétrique se compose d'un index vectoriel dense de 21 millions de passages Wikipedia de 100 mots, accessible via Maximum Inner Product Search (MIPS). La mémoire paramétrique est un générateur BART pré-entraîné.
• Affinage de bout en bout : L'encodeur de requêtes et le générateur sont entraînés conjointement en traitant les passages récupérés comme des variables latentes et en minimisant la log-vraisemblance marginale négative des tokens cibles sans supervision explicite de la recherche.
• Deux formulations de conditionnement :
  • RAG-Sequence : Se conditionne sur le même document extrait pour l'intégralité de la séquence générée.
  • RAG-Token : Permet à chaque token généré de prêter attention et de marginaliser sur différents passages récupérés.
• Connaissances du monde interchangeables à chaud : L'index externe peut être remplacé au moment de l'inférence (par exemple, remplacer un index Wikipedia de 2016 par un index de 2018) afin d'actualiser les connaissances du monde sans modifier les poids du modèle.

Méthode

Étant donné une séquence d'entrée $x$, le retriever $p_{\eta }(z|x)$ produit une distribution sur les documents $z$, et le générateur $p_{\theta }(y_i|x,z,y_{1:i-1})$ génère des tokens cibles conditionnés à la fois sur $x$ et $z$.

Le retriever s'appuie sur Dense Passage Retriever (DPR) : $$p_{\eta }(z|x)\propto \exp \left(\mathbf{d}(z{)}^{\top }\mathbf{q}(x)\right)$$ où $\mathbf{d}(z)={\text{BERT}}_d(z)$ et $\mathbf{q}(x)={\text{BERT}}_q(x)$. Les $k$ meilleurs candidats sont récupérés par MIPS avec FAISS. L'index de documents et l'encodeur de documents ${\text{BERT}}_d$ sont gelés, tandis que ${\text{BERT}}_q$ et les paramètres $\theta$ du générateur BART sont affinés.

RAG-Sequence

Le modèle traite le document récupéré comme une variable latente unique pour toute la séquence $y$ : $$p_{\text{RAG-Sequence}}(y|x)\approx \sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)\prod _i^N{p}_{\theta }(y_i|x,z,y_{1:i-1})$$ Le décodage peut être évalué à l'aide du Thorough Decoding (exécution d'une recherche par faisceau (beam search) par document et calcul des passes avant pour les hypothèses candidates à travers tous les documents) ou du Fast Decoding (approximation des probabilités manquantes du faisceau à 0).

RAG-Token

Le modèle marginalise sur les documents à chaque étape de token, ce qui permet un mélange de contenus issus de plusieurs sources au niveau du token : $$p_{\text{RAG-Token}}(y|x)\approx \prod _i^N\sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)p_{\theta }(y_i|x,z,y_{1:i-1})$$ Cela produit une probabilité de transition autorégressive standard qui s'intègre directement au beam search.

Résultats

• QA en domaine ouvert : RAG établit un nouvel état de l'art sur les benchmarks ouverts en utilisant la correspondance exacte (EM) :
  • Natural Questions : 44,5 (RAG-Seq) contre 41,5 (DPR) et 36,6 (T5-11B+SSM).
  • TriviaQA (split standard ouvert) : 56,8 (RAG-Seq) ; sur le split de test TQA-Wiki : 68,0 (RAG-Seq) contre 60,5 (T5-11B+SSM).
  • WebQuestions : 45,5 (RAG-Token) contre 41,1 (DPR).
  • CuratedTrec : 52,2 (RAG-Seq) contre 50,6 (DPR).
  • De manière notable, sur Natural Questions, RAG a obtenu 11,8 % d'EM lorsque la réponse correcte était totalement absente du contexte extrait (cas où les modèles extractifs obtiennent 0 %).
• Génération abstractive : Sur Open MS-MARCO, RAG-Sequence atteint 40,8 en Rouge-L et 44,2 en Bleu-1, surpassant BART (38,2 Rouge-L, 41,6 Bleu-1).
• Génération de questions Jeopardy : Des annotateurs humains ont jugé RAG-Token plus factuel que BART dans 42,7 % des cas, contre 7,1 % seulement à l'avantage de BART. Les évaluateurs ont également jugé RAG plus spécifique (37,4 % contre 16,8 %).
• Vérification de faits : Sur FEVER 3-classes, RAG atteint une exactitude d'étiquetage de 72,5 % (à 4,3 % des systèmes en pipeline supervisés par des preuves de référence) ; sur 2-classes, RAG atteint 89,5 % d'exactitude.
• Mise à jour des connaissances à chaud : Lors d'un test temporel portant sur 82 dirigeants entre 2016 et 2018, RAG a obtenu 70 % d'exactitude pour les dirigeants de 2016 avec l'index 2016 et 68 % pour les dirigeants de 2018 avec l'index 2018, s'effondrant drastiquement (12 % et 4 %) en cas de désaccord d'index.

Limites et questions ouvertes

• Effondrement de la recherche (Retrieval Collapse) : Pour les tâches nécessitant un conditionnement factuel explicite plus faible (comme la génération d'histoires ouvertes), le retriever peut s'effondrer en extrayant des documents statiques et invariants tandis que le générateur apprend à les ignorer.
• Indexation fixe des documents : L'encodeur de documents reste statique pendant l'affinage en raison du coût de calcul prohibitif d'une réindexation continue de 21 millions de vecteurs, ce qui limite l'adaptation du retriever.
• Ancrage à la source et biais : La mémoire non paramétrique reste vulnérable à la désinformation et aux biais présents dans le corpus de connaissances sous-jacent (Wikipedia).
• Pré-entraînement conjoint : Les composants d'extraction et de génération ont été pré-entraînés séparément avant l'affinage pour les tâches aval, ce qui laisse le pré-entraînement conjoint de zéro comme une piste ouverte.