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

TL;DR

Untie the Knots (UtK) introduces a data-augmentation strategy for continue pre-training to bolster long-context capabilities in LLMs without changing the data mixture. It relies on a Tangling phase (chunking, knotting, tying), a Backtracing task, and an Untying phase to train models to locate relevant segments across tangled sequences, with an additional longer-than-claimed training length to further improve results. Evaluated on 7B and 72B Qwen2 models with 128K contexts, UtK surpasses ABF, upsampling, DCA, and other baselines on long-context benchmarks like RULER and LV-Eval, while maintaining short-context performance and improving training efficiency. The approach yields practical gains, with open-sourced Qwen2-UtK models at 128K, and demonstrates a scalable, efficient pathway to robust long-context understanding in LLMs.

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.

Figures (8)

• Figure 1: Comparison of various long-context strategies based on the Qwen2-base (7B) model on the RULER benchmark. UtK more effectively maintains performance at the 128K context length.
• Figure 2: Illustration of the UtK Pre-training process. In the Tangling phase, documents are split into chunks, which are then randomly tied together. Knot Tokens are inserted at the split points to guide the model in locating the partitions during the Untying phase. The Chunk IDs of each chunk are appended to the last chunk of the document to help the model learn to correctly backtrace the original document structure.
• Figure 3: The top panel shows the UtK-augmented expected conditional information for the same four documents, while the bottom panel displays the changes in the histogram of relative positional embedding distances from the original to the UtK-augmented.
• Figure 4: Distribution of document lengths categorized by token counts. The ratios represent the number of tokens within each document length category proportional to the total number of tokens.
• Figure 5: Performance with varying numbers of chunks on the RULER 128K benchmark.