【新智元导读】来自英伟达和UIUC的华人团队提出一种高效训练方法，将LLM上下文长度从128K扩展至惊人的400万token SOTA纪录！基于Llama3.1-Instruct打造的UltraLong-8B模型，不仅在长上下文基准测试中表现卓越，还在标准任务中保持顶尖竞争力。

　　像是最新的Gemini 2.5 Pro在文本及代码上的顶尖性能，以及GPT-4o的原生生图能力都很好的证明了这点。

　　然而，很多实际应用场景，比如长文档和视频理解、上下文学习以及推理时扩展，都需要模型能够处理超长的token序列。

　　在这些场景中，模型的上下文窗口受限往往成为一大瓶颈，因为分布在长文档中的关键信息可能会被忽略。

　　这种方法可以从现有的指令微调模型出发，构建超长上下文的LLM，最高可将上下文长度推向400万token的极限！

　　研究人员利用上面方法训练的UltraLong-8B模型在长上下文任务上达到了顶尖水平，同时在标准任务上也保持了竞争力。

　　关键技术创新：研究者引入了特殊文档分隔符和基于YaRN的位置编码扩展技术，通过消融实验证明这些技术对长上下文建模至关重要。

　　高效的单步预训练策略：研究者发现，相比多步扩展方法，单步持续预训练在上下文扩展上更高效，在合成和真实世界长上下文基准测试中始终表现出色。

　　以Llama 3.1-8B-Instruct为基础，持续预训练阶段将模型的上下文窗口逐步扩展到目标长度（比如100万、200万、400万token）。随后，指令微调阶段优化模型的指令遵循能力和推理能力。

　　第一阶段通过持续预训练扩展模型的上下文窗口，采用特殊文档分隔符和基于YaRN的缩放技术来处理超长序列。第二阶段使用精心挑选的数据集进行指令微调，提升模型的指令遵循和推理能力

　　在第一阶段，研究者通过持续预训练将Llama-3.1-8B-Instruct的上下文窗口扩展到目标长度。

　　研究者对少于4000 token的短文档进行下采样，对超过8000 token的长文档进行上采样，最终形成一个包含10亿token的语料库。

　　这些文档被拼接成对应目标上下文长度的更长序列（比如100万、200万、400万token）。拼接时，他们使用特殊字符分隔不同文档，而不是用保留的开始和结束标记。

　　为了支持超长上下文，研究人员采用了基于YaRN的缩放方法，而不是之前工作中常用的NTK感知缩放策略。他们固定超参数α=1和β=4，并根据目标上下文长度计算缩放因子s。

　　当输入长度接近最大限制时，Llama-3.1模型的性能会下降。为解决这个问题，他们为RoPE嵌入采用了更大的缩放因子，从而更好地适应超长序列。

　　研究者针对三种上下文长度（100万、200万和400万token）构建了长上下文模型，并将RoPE缩放因子分别设置为128、256和512。

　　为了提升训练的可扩展性，他们使用了Megatron-LM框架。为了处理超长输入序列，采用了张量并行和上下文并行。

　　训练在256个NVIDIA H100 GPU上进行，1M、2M和4M模型的训练时间分别约为5小时、6小时和13小时。

　　在第二阶段，研究者通过监督微调（SFT）提升长上下文模型的指令遵循和推理能力，使用的是一些精心挑选的数据集。

　　为了进一步提升SFT数据集的质量，他们利用GPT-4o和4o-mini优化了这些数据集的回答内容。

　　值得注意的是，研究者的SFT数据集仅包含上述短上下文数据（少于8000 token的样本），没有加入合成长上下文指令数据。

　　最终，研究者构建了一个包含10万个样本的SFT数据集。对于每种目标上下文长度的模型，他们使用128的批大小和5×10⁻⁶的学习率。

　　研究者将他们的模型与基于Llama家族的最先进（SOTA）长上下文模型进行对比，以确保对训练方法公平且可控的评估。

　　本文研究者专注于Llama家族的模型，这样可以更清晰地展示他们扩展上下文长度训练方法的有效性，同时确保在标准任务上的性能依然具有竞争力。

　　RULER：这是一个专门评估长上下文语言模型的基准，通过生成不同序列长度的合成样本，覆盖四个任务类别。

　　LV-Eval：这是一个长上下文基准，包含最高256K token的五个长度级别，重点测试两种任务：单跳问答（single-hop QA）和多跳问答（multi-hop QA）。

　　InfiniteBench：这是一个长上下文基准，平均输入长度约200K token，最大长度超过2M token，包含合成任务和现实世界任务。

　　研究人员首先从「大海捞针」（Needle in a Haystack，NIAH）这一测试开始，然后再探讨长上下文和标准基准的评估。

　　研究人员通过NIAH密码检索测试，评估模型在长上下文检索方面的能力。在这个任务中，模型需要在一大段毫无意义的文本中，找到一个简单密码，比如一个随机的六位数字。

　　为了量化检索的准确性，他们测试了40种不同的输入序列长度。对于每种长度，密码会被随机插入到10个均匀分布的文档深度中。

　　结果如图2所示。对于本文的模型，测试了高达100万、200万和400万个token的输入长度；而对于基准模型，只测试了最高100万个token。

　　相比之下，如图2d到2f所示，研究者的超长（UltraLong）模型在所有输入长度和深度上都达到了100%的准确率，展现了强大的长上下文检索能力。

　　研究者在RULER、LV-Eval和InfiniteBench上的评估结果如表1所示vpn安卓版下载地址。加粗的数字表示性能超过了所有基准模型。

　　这些结果表明，研究者的训练方法有效扩展了语言模型的上下文窗口到超长输入，同时保持了原有输入长度的性能。

　　相比之下，基准模型中，Llama-3.1是为128K输入长度设计的，当输入超过128K token时，性能显著下降。ProLong是为512K上下文设计的，但即使它训练了更多token（410亿对比10亿），在512K长度上的表现也不如他们的模型。

　　Gradient是基准模型中支持最长上下文的（100万个token），但在LV-Eval和InfiniteBench上的表现较差，说明它的设计可能过于偏向人工任务，牺牲了现实任务的效果。

　　而本文的模型在人工（RULER）和混合（LV-Eval和InfiniteBench）基准测试中始终保持更高的分数，凸显了方法的高效性和可扩展性。

　　研究者还通过通用、数学和代码领域的标准基准测试评估了模型，以确保扩展上下文长度不会影响短上下文任务的性能。

　　特别值得一提的是，他们的模型在MMLU和MATH基准上表现出明显提升，同时在GSM8K和HumanEval等其他基准上的表现也极具竞争力。

　　相比之下，基准长上下文模型Gradient和ProLong在这些标准任务上的性能大幅下降，平均分数仅为37.36和40.81。

　　在这项工作中，研究人员提出了一种高效且系统化的训练方法，用于超长上下文语言模型，将上下文窗口扩展到100万、200万和400万个token，同时在标准基准测试中保持了竞争力。

　　这种结合了高效的持续预训练和指令微调，不仅提升了模型对长上下文的理解能力，还增强了其遵循指令的能力。

　　这一框架为可扩展的长上下文建模树立了新标杆，也为未来在实际应用中提升长上下文性能的研究铺平了道路。

　　伊利诺伊大学厄巴纳-香槟分校（UIUC）计算机科学博士研究生，导师是Bo Li教授。浙江大学计算机科学与技术专业学士学位，就读于CKC荣誉学院，导师是Shouling Ji教授和Siliang Tang教授。

　　专注于提升基础模型的安全性、可靠性和一致性，包括LLMs、多模态模型以及基于LLM的智能体。

　　加州大学欧文分校机器学习博士学位，热衷于构建用于文本、音频和多模态数据的尖端生成模型。此前，曾担任百度硅谷人工智能实验室（由吴恩达创立）的文本到语音团队负责人。