Untie the Knots: An Efficient Data Augmentation Strategy for Long-Context Pre-Training in Language Models

TL;DR

This paper introduces 'Untie the Knots' (UtK), a data augmentation method that enhances large language models' ability to process long contexts efficiently by training them to untangle complex, knotted text structures.

Contribution

The paper presents a novel data augmentation strategy for long-context pre-training that improves model performance and training efficiency without changing data mixtures.

Findings

Achieves 75% and 84.5% accuracy on RULER at 128K context length for 7B and 72B models.

Significantly outperforms existing long-context strategies.

Models trained with UtK demonstrate better attention to relevant long-range information.

Abstract

Large language models (LLM) have prioritized expanding the context window from which models can incorporate more information. However, training models to handle long contexts presents significant challenges. These include the scarcity of high-quality natural long-context data, the potential for performance degradation on short-context tasks, and the reduced training efficiency associated with attention mechanisms. In this paper, we introduce Untie the Knots (\textbf{UtK}), a novel data augmentation strategy employed during the continue pre-training phase, designed to efficiently enable LLMs to gain long-context capabilities without the need to modify the existing data mixture. In particular, we chunk the documents, shuffle the chunks, and create a complex and knotted structure of long texts; LLMs are then trained to untie these knots and identify relevant segments within seemingly chaotic token sequences. This approach greatly improves the model's performance by accurately attending to relevant information in long context and the training efficiency is also largely increased. We conduct extensive experiments on models with 7B and 72B parameters, trained on 20 billion tokens, demonstrating that UtK achieves 75\% and 84.5\% accurracy on RULER at 128K context length, significantly outperforming other long context strategies. The trained models will open-source for further research.