Инженерия данных для масштабирования языковых моделей до контекста в 128 тысяч токенов

Data Engineering for Scaling Language Models to 128K Context

February 15, 2024

Авторы: Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, Hao Peng

cs.AI

Аннотация

Мы исследуем подход к непрерывному предобучению для масштабирования длины контекста языковых моделей до 128K, уделяя особое внимание инженерии данных. Мы предполагаем, что моделирование длинного контекста, в частности способность использовать информацию из произвольных мест ввода, — это навык, который в основном уже приобретается в ходе крупномасштабного предобучения, и что этот навык можно легко расширить на контексты, значительно превышающие те, что встречались во время обучения (например, с 4K до 128K), с помощью легковесного непрерывного предобучения на подходящей смеси данных. Мы изучаем количество и качество данных для непрерывного предобучения: (1) по количеству мы показываем, что 500 миллионов — 5 миллиардов токенов достаточно, чтобы модель могла извлекать информацию из любого места в контексте длиной 128K; (2) по качеству наши результаты одинаково подчеркивают важность баланса доменов и апсэмплинга длины. Конкретно, мы обнаруживаем, что наивный апсэмплинг более длинных данных в определенных доменах, таких как книги, что является распространенной практикой в существующих работах, дает неоптимальные результаты, и что сбалансированная смесь доменов важна. Мы демонстрируем, что непрерывное предобучение полной модели на 1–5 миллиардах токенов таких данных является эффективной и доступной стратегией для масштабирования длины контекста языковых моделей до 128K. Наш подход превосходит сильные модели с длинным контекстом с открытым исходным кодом и сокращает разрыв с передовыми моделями, такими как 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.