Pré-entraînement en contexte : Modélisation du langage au-delà des limites des documents

In-Context Pretraining: Language Modeling Beyond Document Boundaries

October 16, 2023

Auteurs: Weijia Shi, Sewon Min, Maria Lomeli, Chunting Zhou, Margaret Li, Victoria Lin, Noah A. Smith, Luke Zettlemoyer, Scott Yih, Mike Lewis

cs.AI

Résumé

Les grands modèles de langage (LMs) sont actuellement entraînés à prédire des tokens à partir de préfixes de documents, ce qui leur permet de réaliser directement des tâches de génération de texte long et de réponse à des prompts qui peuvent être réduites à la complétion de documents. Les pipelines de pré-entraînement existants forment les LMs en concaténant des ensembles aléatoires de documents courts pour créer des contextes d'entrée, mais les documents précédents ne fournissent aucun signal pour prédire le document suivant. Nous proposons plutôt le pré-entraînement en contexte (In-Context Pretraining), une nouvelle approche où les modèles de langage sont pré-entraînés sur une séquence de documents connexes, les encourageant ainsi explicitement à lire et à raisonner au-delà des limites des documents. Nous pouvons réaliser ce pré-entraînement en contexte simplement en modifiant l'ordre des documents pour que chaque contexte contienne des documents liés, et en appliquant directement les pipelines de pré-entraînement existants. Cependant, ce problème de tri des documents est complexe. Il existe des milliards de documents, et nous souhaitons que le tri maximise la similarité contextuelle pour chaque document sans répéter aucune donnée. Pour ce faire, nous introduisons des algorithmes approximatifs pour trouver des documents connexes grâce à une recherche efficace des plus proches voisins et pour construire des contextes d'entrée cohérents avec un algorithme de parcours de graphe. Nos expériences montrent que le pré-entraînement en contexte offre une approche simple et scalable pour améliorer significativement les performances des LMs : nous observons des améliorations notables dans les tâches nécessitant un raisonnement contextuel plus complexe, notamment l'apprentissage en contexte (+8 %), la compréhension de lecture (+15 %), la fidélité aux contextes précédents (+16 %), le raisonnement sur de longs contextes (+5 %) et l'augmentation par recherche d'information (+9 %).

English

Large language models (LMs) are currently trained to predict tokens given document prefixes, enabling them to directly perform long-form generation and prompting-style tasks which can be reduced to document completion. Existing pretraining pipelines train LMs by concatenating random sets of short documents to create input contexts but the prior documents provide no signal for predicting the next document. We instead present In-Context Pretraining, a new approach where language models are pretrained on a sequence of related documents, thereby explicitly encouraging them to read and reason across document boundaries. We can do In-Context Pretraining by simply changing the document ordering so that each context contains related documents, and directly applying existing pretraining pipelines. However, this document sorting problem is challenging. There are billions of documents and we would like the sort to maximize contextual similarity for every document without repeating any data. To do this, we introduce approximate algorithms for finding related documents with efficient nearest neighbor search and constructing coherent input contexts with a graph traversal algorithm. Our experiments show In-Context Pretraining offers a simple and scalable approach to significantly enhance LMs'performance: we see notable improvements in tasks that require more complex contextual reasoning, including in-context learning (+8%), reading comprehension (+15%), faithfulness to previous contexts (+16%), long-context reasoning (+5%), and retrieval augmentation (+9%).