Pré-treinamento em Contexto: Modelagem de Linguagem Além das Fronteiras de Documentos

In-Context Pretraining: Language Modeling Beyond Document Boundaries

October 16, 2023

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

cs.AI

Resumo

Grandes modelos de linguagem (LMs) são atualmente treinados para prever tokens com base em prefixos de documentos, permitindo que realizem diretamente tarefas de geração de texto longo e de estilo de prompt que podem ser reduzidas à conclusão de documentos. Os pipelines de pré-treinamento existentes treinam LMs concatenando conjuntos aleatórios de documentos curtos para criar contextos de entrada, mas os documentos anteriores não fornecem sinal para prever o próximo documento. Em vez disso, apresentamos o Pré-Treinamento em Contexto, uma nova abordagem na qual os modelos de linguagem são pré-treinados em uma sequência de documentos relacionados, incentivando explicitamente que leiam e raciocinem além dos limites dos documentos. Podemos realizar o Pré-Treinamento em Contexto simplesmente alterando a ordenação dos documentos para que cada contexto contenha documentos relacionados e aplicando diretamente os pipelines de pré-treinamento existentes. No entanto, esse problema de ordenação de documentos é desafiador. Existem bilhões de documentos e gostaríamos que a ordenação maximizasse a similaridade contextual para cada documento sem repetir nenhum dado. Para isso, introduzimos algoritmos aproximados para encontrar documentos relacionados com busca eficiente de vizinhos mais próximos e construir contextos de entrada coerentes com um algoritmo de travessia de grafos. Nossos experimentos mostram que o Pré-Treinamento em Contexto oferece uma abordagem simples e escalável para melhorar significativamente o desempenho dos LMs: observamos melhorias notáveis em tarefas que exigem raciocínio contextual mais complexo, incluindo aprendizado em contexto (+8%), compreensão de leitura (+15%), fidelidade a contextos anteriores (+16%), raciocínio de contexto longo (+5%) e aumento de recuperação (+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%).