From 128K to 4M: Efficient Training of Ultra-Long Context Large Language Models

📄 arXiv: 2504.06214v1 📥 PDF

作者: Chejian Xu, Wei Ping, Peng Xu, Zihan Liu, Boxin Wang, Mohammad Shoeybi, Bo Li, Bryan Catanzaro

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-04-08

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出高效训练方法，将LLM上下文长度扩展至4M tokens，并保持性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文模型 大型语言模型 持续预训练 指令调优 上下文扩展 高效训练 Transformer LLM

📋 核心要点

  1. 现有LLM在处理长文档、视频理解和复杂推理任务时，受限于上下文长度，性能面临挑战。
  2. 通过高效的持续预训练和指令调优，扩展LLM的上下文窗口，同时保持其指令遵循和推理能力。
  3. UltraLong-8B模型在长上下文基准测试中达到SOTA，并在标准基准测试中保持竞争力。

📝 摘要（中文）

本文提出了一种高效的训练方案，用于构建超长上下文的大型语言模型（LLM），从对齐的指令模型出发，将上下文长度从128K扩展到1M、2M和4M tokens。该方法利用高效的持续预训练策略来扩展上下文窗口，并采用有效的指令调优来保持指令遵循和推理能力。基于Llama3.1-Instruct构建的UltraLong-8B模型，通过该方案在各种长上下文基准测试中实现了最先进的性能。重要的是，使用该方法训练的模型在标准基准测试中保持了具有竞争力的性能，证明了长短上下文任务的平衡改进。此外，本文还深入分析了关键设计选择，强调了缩放策略和数据组成的影响。研究结果建立了一个强大的框架，用于在保持一般模型能力的同时，高效地扩展上下文长度。所有模型权重已开源。

🔬 方法详解

问题定义：现有大型语言模型在处理需要长程依赖的任务时，由于上下文窗口的限制，性能会显著下降。例如，在处理长篇文档、视频理解、复杂推理等任务时，模型无法有效地捕捉到全局信息。现有的方法要么计算复杂度高，要么难以保持模型在短上下文任务上的性能。

核心思路：本文的核心思路是通过高效的持续预训练策略来扩展模型的上下文窗口，并结合指令调优来保持模型的指令遵循和推理能力。通过这种方式，可以在不显著增加计算成本的前提下，提升模型处理长上下文任务的能力。

技术框架：整体框架包括两个主要阶段：1) 上下文窗口扩展：使用持续预训练策略，在包含长序列数据的语料库上对模型进行训练，逐步扩展模型的上下文窗口。2) 指令调优：使用指令数据对模型进行微调，以保持模型的指令遵循和推理能力。这两个阶段交替进行，以确保模型在长上下文任务和短上下文任务上都能保持良好的性能。

关键创新：最重要的技术创新点在于提出了一种高效的持续预训练策略，该策略能够在扩展上下文窗口的同时，避免模型性能的显著下降。此外，结合指令调优，可以有效地保持模型的指令遵循和推理能力。与现有方法相比，该方法在计算效率和模型性能之间取得了更好的平衡。

关键设计：在持续预训练阶段，采用了多种缩放策略，包括调整学习率、批量大小等。在指令调优阶段，使用了高质量的指令数据，并对损失函数进行了调整，以更好地优化模型的指令遵循能力。此外，还对数据组成进行了深入分析，发现合理的数据配比对于模型性能至关重要。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

UltraLong-8B模型在各种长上下文基准测试中实现了最先进的性能，证明了该方法的有效性。同时，该模型在标准基准测试中保持了具有竞争力的性能，表明该方法能够在扩展上下文窗口的同时，保持模型在短上下文任务上的性能。具体而言，该模型将上下文长度扩展至4M tokens，并在多个长文本benchmark上超越了现有模型。

🎯 应用场景

该研究成果可广泛应用于需要处理长序列数据的领域，如长文档理解、视频内容分析、代码生成、复杂推理等。通过扩展LLM的上下文窗口，可以提升模型在这些领域的性能，从而实现更智能、更高效的应用。例如，可以用于自动生成长篇报告、总结视频内容、辅助软件开发等。

📄 摘要（原文）

Long-context capabilities are essential for a wide range of applications, including document and video understanding, in-context learning, and inference-time scaling, all of which require models to process and reason over long sequences of text and multimodal data. In this work, we introduce a efficient training recipe for building ultra-long context LLMs from aligned instruct model, pushing the boundaries of context lengths from 128K to 1M, 2M, and 4M tokens. Our approach leverages efficient continued pretraining strategies to extend the context window and employs effective instruction tuning to maintain the instruction-following and reasoning abilities. Our UltraLong-8B, built on Llama3.1-Instruct with our recipe, achieves state-of-the-art performance across a diverse set of long-context benchmarks. Importantly, models trained with our approach maintain competitive performance on standard benchmarks, demonstrating balanced improvements for both long and short context tasks. We further provide an in-depth analysis of key design choices, highlighting the impacts of scaling strategies and data composition. Our findings establish a robust framework for efficiently scaling context lengths while preserving general model capabilities. We release all model weights at: https://ultralong.github.io/.