400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

编辑:犀牛

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

大语言模型(LLM)在文本和多模态任务上已经展现出惊艳的表现。

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

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

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

为了解决这些问题,来自英伟达和UIUC的研究者提出了一种高效的训练方法。

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

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

论文地址:https://arxiv.org/pdf/2504.06214

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

主要贡献:

  • 高效且可扩展的训练方法

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

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

  • 全面的实验验证:研究者在多个基准测试上进行了广泛实验,包括RULER、LV-Eval、InfiniteBench、MMLU、MMLU-Pro、MATH、GSM-8K和HumanEval,证明UltraLong-8B模型在长上下文和标准任务上均优于现有基线。

实验方法

如图1所示,本文方法主要分为两个阶段:持续预训练指令微调

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

这两个阶段结合,让模型既能高效处理超长输入,又能在长短上下文任务中表现出色。

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

第一阶段通过持续预训练扩展模型的上下文窗口,采用特殊文档分隔符和基于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。

每个模型在10亿token的语料上训练一个epoch,学习率为3×10⁻⁵。

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

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

指令微调

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

他们整合并优化了多个开源SFT数据集,覆盖三个关键领域:通用领域、数学和代码。

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

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

他们发现,仅依靠短上下文数据就足以取得优异效果,这与之前研究的观察一致。

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

训练依然基于Megatron-LM 框架,在256个NVIDIA H100 GPU上进行,张量并行度设为tp=8。每次训练大约需要 30 分钟。

基线模型与评估基准

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

  • Llama-3.1 (Llama-3.1-8B-Instruct):这是他们的基础模型,支持128K的上下文窗口。

  • ProLong (Llama-3-8B-ProLong-512k-Instruct):基于Llama-3构建的长上下文模型,拥有512K的上下文窗口。

  • Gradient (Llama-3-8B-Instruct-Gradient-1048k):另一个基于Llama的长上下文模型,支持高达1M的上下文窗口。

本文研究者专注于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。

如图2a到2c所示,在基准模型中,只有Llama-3-8B-Instruct和Gradient-1048k通过了NIAH测试,而Llama-3.1-8B-Instruct和Llama-3-8B-ProLong-512k-Instruct即使在它们声称的上下文长度内也出现了错误。

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

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

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

总体来说,他们的三个模型在大多数情况下都取得了最高分。

在RULER基准测试中,UltraLong模型在512K和100万个token的输入长度上表现最佳。在LV-Eval中,他们的模型在128K和256K token长度内的平均F1分数最高。

此外,他们在InfiniteBench上也取得了最佳表现。

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

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

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

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

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

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

如表2所示,他们的模型性能与基础模型Llama-3.1-8B-Instruct相当甚至更高,平均分数分别为62.47、61.06和60.95,而Llama-3.1-8B-Instruct为61.45。

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

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

这些结果表明,研究者的方法不仅有效扩展了上下文窗口,还保持甚至提升了模型的通用任务能力。

而Llama-3-8B-Instruct-Gradient-1048k和Llama-3-8B-ProLong-512k-Instruct的显着性能下降,表明它们的超长上下文方法可能存在局限性。

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

结论

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

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

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

作者介绍

Chejian Xu

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

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

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

Wei Ping

400万token新SOTA!英伟达UIUC联手:兼顾长短上下文顶尖性能

NVIDIA应用深度学习研究团队的资深研究科学家,专注于大型语言模型和生成模型的研究。

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

参考资料:

https://arxiv.org/abs/2504.06214

内容来源于网络。发布者:科技网btna,转转请注明出处:https://www.btna.cn/6824.html

(0)
科技网btna的头像科技网btna
上一篇 2025年5月1日
下一篇 2025年5月1日

相关推荐

  • 从游戏少女到AI女神!她用物理驯服AI,让飓风预测快1000倍

    编辑:英智 【导读】25年前,10岁的Rose Yu收到一台改变命运的电脑,从玩游戏到成为AI领域的领军人物,她如何用物理学与深度学习破解交通、飓风和核聚变的难题?跟随她的脚步,揭秘一位天才科学家的开挂人生! 10岁那年,Rose Yu收到一份改变命运的生日礼物——一台电脑。 在25年前,电脑在中国可是稀罕物。这份来自叔叔的礼物,从此开启了她开挂般的人生! …

    2025年5月2日
    4400
  • Transformer原作打脸DeepSeek观点?一句Wait就能引发反思,RL都不用

    编辑:KingHZ 【导读】Transformer作者Ashish Vaswani团队重磅LLM研究!简单指令:「Wait,」就能有效激发LLM显式反思,表现堪比直接告知模型存在错误。 只要预训练,LLM就能涌现自我反思、自我纠正! Transformer首席作者Ashish Vaswani带队的新研究引来万人围观! 强化学习已证明能有效激发大语言模型的反思…

    2025年4月23日
    5900
  • 新·摩尔定律诞生:AI智能体能力每4个月翻一番!智能爆炸在即

    编辑:英智 犀牛 【导读】AI编程智能体的能力正在飞速增长,最新研究揭示了这一「新摩尔定律」,如果AI智能体的任务时长继续以每4个月翻倍的速度增长,到2027年它们可能完成长达167小时的月级任务! 编程智能体能完成的任务量,正以爆炸式的速度增长! 而且,这一趋势还在加速。 2024到2025年,它们的任务处理能力每4个月翻一倍;而在2019到2025年,这…

    2025年4月30日
    5500
  • AI「自我复制」能力曝光!RepliBench警示:大模型正在学会伪造身份

    编辑:定慧 【导读】科幻中AI自我复制失控场景,正成为现实世界严肃的研究课题。英国AISI推出RepliBench基准,分解并评估AI自主复制所需的四大核心能力。测试显示,当前AI尚不具备完全自主复制能力,但在获取资源等子任务上已展现显着进展。 想象一种失控的、具有自我复制能力的数字生命体,需要多久占领地球? 从《2001太空漫游》中的HAL 9000到《终…

    2025年5月5日
    5900
  • 诺奖得主、AlphaGo之父:AI远不如人类科学家,根本提不出有价值的猜想

    编辑:桃子 【导读】去年10月,DeepMind联合创始人Demis Hassabis因开发AlphaFold以解决蛋白质复杂结构预测问题,而获得了诺贝尔化学奖;2016年,他和他的团队开发的AlphaGo战胜了人类围棋世界冠军李世石,自此声名鹊起;作为一个天才少年,他更是在13岁时就晋升为了国际象棋「大师」。 一直以来,Demis Hassabis都被认为…

    2025年4月30日
    5600

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

联系我们

400-800-8888

在线咨询: QQ交谈

邮件:admin@example.com

工作时间:周一至周五,9:30-18:30,节假日休息

关注微信