Engenharia de Dados para Escalonamento de Modelos de Linguagem para Contextos de 128K

Data Engineering for Scaling Language Models to 128K Context

February 15, 2024

Autores: Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, Hao Peng

cs.AI

Resumo

Estudamos a receita de pré-treinamento contínuo para escalar o comprimento de contexto de modelos de linguagem para 128K, com foco na engenharia de dados. Nossa hipótese é que a modelagem de contexto longo, em particular a capacidade de utilizar informações em locais arbitrários de entrada, é uma habilidade que já é majoritariamente adquirida através de pré-treinamento em larga escala, e que essa capacidade pode ser prontamente estendida para contextos substancialmente mais longos do que os vistos durante o treinamento (por exemplo, de 4K para 128K) por meio de pré-treinamento contínuo leve em uma mistura de dados apropriada. Investigamos a quantidade e a qualidade dos dados para pré-treinamento contínuo: (1) para quantidade, mostramos que 500 milhões a 5 bilhões de tokens são suficientes para permitir que o modelo recupere informações em qualquer lugar dentro do contexto de 128K; (2) para qualidade, nossos resultados enfatizam igualmente o equilíbrio de domínio e a amostragem de comprimento. Concretamente, descobrimos que a amostragem ingênua de dados mais longos em certos domínios, como livros, uma prática comum de trabalhos existentes, resulta em desempenho subótimo, e que uma mistura equilibrada de domínios é importante. Demonstramos que o pré-treinamento contínuo do modelo completo em 1B-5B tokens de tais dados é uma estratégia eficaz e acessível para escalar o comprimento de contexto de modelos de linguagem para 128K. Nossa receita supera modelos de contexto longo de código aberto robustos e reduz a lacuna para modelos de fronteira, como o GPT-4 128K.

English

We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular the ability to utilize information at arbitrary input locations, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the quantity and quality of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize domain balance and length upsampling. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.