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

Large language models (LLM) have focused on expanding the context window in order to incorporate more information effectively.However, training models to handle long contexts poses significant challenges.These include the scarcity of high-quality natural long-context data, the potential of performance degradation on short-context tasks, and the reduced training efficiency associated with attention mechanisms.In this paper, we introduce Untie the Knots (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 of modifying the existing data mixture.In particular, we chunk the documents, shuffle the chunks, and create a knotted structure of long texts; LLMs are then trained to untie these knots and identify relevant segments within seemingly chaotic token sequences.This approach substantially enhances the model's performance by accurately attending to relevant information in long contexts, while also greatly improving the training efficiency.We conduct extensive experiments on models with 7B and 72B parameters, trained on 20 billion tokens, demonstrating that UtK achieves 75% and 84.5% accuracy on RULER at 128K context length, significantly outperforming other long-context strategies.The trained models and data processing code are open-sourced for further research.https://github.com/rgtjf/Untie-the-Knots