In-Context-Pretraining: Sprachmodellierung über Dokumentgrenzen hinaus

In-Context Pretraining: Language Modeling Beyond Document Boundaries

October 16, 2023

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LMs) werden derzeit darauf trainiert, Token basierend auf Dokumentpräfixen vorherzusagen, wodurch sie in der Lage sind, direkt langfristige Generierungsaufgaben und Prompting-Aufgaben zu bewältigen, die auf die Vervollständigung von Dokumenten reduziert werden können. Bestehende Vorverarbeitungspipelines trainieren LMs, indem sie zufällige Sätze kurzer Dokumente verketten, um Eingabekontexte zu erstellen, wobei die vorherigen Dokumente jedoch kein Signal für die Vorhersage des nächsten Dokuments liefern. Stattdessen präsentieren wir In-Context Pretraining, einen neuen Ansatz, bei dem Sprachmodelle auf einer Sequenz verwandter Dokumente vortrainiert werden, wodurch sie explizit dazu angeregt werden, über Dokumentgrenzen hinweg zu lesen und zu schlussfolgern. Wir können In-Context Pretraining durchführen, indem wir einfach die Dokumentreihenfolge so ändern, dass jeder Kontext verwandte Dokumente enthält, und bestehende Vorverarbeitungspipelines direkt anwenden. Dieses Dokumentensortierungsproblem ist jedoch herausfordernd. Es gibt Milliarden von Dokumenten, und wir möchten, dass die Sortierung die kontextuelle Ähnlichkeit für jedes Dokument maximiert, ohne Daten zu wiederholen. Um dies zu erreichen, führen wir approximative Algorithmen ein, um verwandte Dokumente durch effiziente Nearest-Neighbor-Suche zu finden und kohärente Eingabekontexte mit einem Graphtraversal-Algorithmus zu konstruieren. Unsere Experimente zeigen, dass In-Context Pretraining einen einfachen und skalierbaren Ansatz bietet, um die Leistung von LMs signifikant zu verbessern: Wir bemerkenswerte Verbesserungen bei Aufgaben, die komplexeres kontextuelles Denken erfordern, einschließlich In-Context-Lernen (+8%), Leseverständnis (+15%), Treue zu vorherigen Kontexten (+16%), Langzeitkontext-Schlussfolgerung (+5%) und Retrieval-Augmentierung (+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%).