NVIDIA新方法让AI高效处理超长文本

作者:技术拾光时间:2026-08-01 14:55:19浏览:413

NVIDIA研究团队提出Jet-Long方法,通过动态双焦点RoPE将注意力视野分为本地和远程窗口,利用整数分组映射和纠偏旋转技巧,使已有大模型无需重新训练即可处理超长文本,在RULER等基准测试中准确率显著提升,计算速度几乎无损失。

这项来自NVIDIA研究团队的工作,以预印本形式在2026年7月8日上线,论文编号arXiv:2607.07740。需要深入了解技术细节的,可直接通过此编号查阅原文,代码也已开源,自行实验是获取直观理解的最佳途径。

试想一下,阅读一本厚厚的小说至第800页时,仍能清晰记得第1页主角的名字,以及第50页提到的那把钥匙的藏匿位置。人类大脑对长篇内容的整体感知能力,如同内置了一个无形的书架。然而,对于AI大语言模型而言,情况则颇为尴尬。它们在训练时仅接触过固定长度的文字,当实际输入超出此范围时,其表现如同近视者摘去眼镜,对远处的内容愈发模糊,直至完全失效。

NVIDIA团队提出的解决方案是一套名为Jet-Long的方法。该方法无需重新训练任何模型,即可使现有AI大模型处理远超其原始“视力范围”的超长文本,且计算速度几乎没有损失。这如同为旧眼镜更换一副自动变焦镜片——镜框无需更换,视野却得以数倍拓展。

一、AI的“视力问题”究竟从何而来

要理解Jet-Long解决了什么问题,首先需要明确AI大模型为何会“看不清”长文本。

现代大语言模型在处理文字时,会为每个字或词分配一个“位置编号”,类似于图书馆中为每本书编制的目录号,使模型能够知晓该词在句子中的具体位置。目前应用最广泛的位置编码方式为RoPE(旋转位置编码),其工作原理可类比于钟表的指针:每个位置对应一个特定角度,模型通过读取这些角度来判断两个词之间的距离关系。

问题在于,模型在训练时仅见过一定范围内的位置编号,例如最多至32768号。当实际使用中遇到第50000号、第100000号这样的位置时,对应的钟表指针便旋转至模型从未见过的角度,模型无法理解此类信息,注意力机制于是开始“乱猜”,输出质量急剧下降。这便是“位置越界”问题。

与此同时,还存在另一个难题。随着文本长度增加,模型需要同时关注的词越来越多,如同在一个嘈杂的大厅里试图听清一个人的说话声——背景噪声过大,注意力被稀释分散,真正重要的信息反而淹没在无关内容中。这被称为“注意力扩散”问题。

目前已有一些方法试图解决这两个问题,但它们均存在一个共同短板:需要在应用前预先设定一个固定的“调整系数”。调整过于激进,模型处理短文本时会出现偏差;调整过于保守,遇到超长文本又无法胜任。这如同为眼镜配制镜片,无论选择何种度数,总有一段距离无法看清。Jet-Long的思路则完全不同——它使镜片度数随观看距离自动调节,相当于一副智能变焦眼镜。

二、双焦点视野:同时看清近处和远处的秘密

Jet-Long的核心思想称为“动态双焦点RoPE”(Dynamic Bifocal RoPE)。双焦点这一概念源自眼镜学——双焦点眼镜包含两个区域,一个用于看近处,一个用于看远处,两者共存于同一副镜片之上。Jet-Long也将模型对文本的“注意力视野”划分为两个区域进行处理。

第一个区域称为“本地窗口”,大小固定为2048个词的范围。在此窗口内,模型完全按照原始训练时的模式工作,位置编码原封不动,如同使用眼镜的近视区看手边的书,清晰无误。第二个区域称为“远程窗口”,覆盖本地窗口之外的所有文本。此处的位置编码需要经过重新映射,以使那些“越界”的位置角度重新落回模型熟悉的范围之内。

这种重新映射的方式非常直观。假设模型的训练窗口为32768个词,而当前文本有65536个词,约为训练窗口的两倍。Jet-Long将远程区域的每两个相邻位置“合并”成一个,让它们共用同一个位置编号——如同将两层楼高的书架压缩至一层,虽然书籍排列得更为紧密,但每本书的编号仍在模型认识的范围内。这个合并的比例称为“分组大小”G,其计算方式为将当前文本长度除以训练窗口长度并向上取整,确保合并后的位置永远不会超出训练范围。

最为关键的一点是,这个分组大小G会随着当前输入文本的长度实时变化。文本较短时,G等于1,相当于不做任何合并,模型行为与原始状态完全一致;文本越长,G越大,合并的力度越强,始终维持恰到好处的压缩,既不过度也不不足。这便是“动态”的涵义——它并非一个预先设定的固定参数,而是根据实际需要即时计算的活参数。

使用整数分组而非连续插值的原因也颇具趣味。模型在训练时接触过的位置角度是有限的、离散的一组数值,就像一把只有固定刻度的尺子。如果使用连续插值,产生的角度可能落在刻度之间,模型依然会感到陌生;而整数分组确保每个远程位置都映射到某个模型在训练时确实见过的整数位置,如同将书签插回已有的刻度上,模型能够准确识别。

三、生成时不重算:聪明的“纠偏旋转”技巧

双焦点视野的构想虽然巧妙,但落地时面临一个工程难题:模型在生成文字时,会将之前计算过的键值(KV缓存)存储起来,以便后续调用,避免重复计算。然而,如果分组大小G随着文本增长而变化,那么之前用旧G计算出的缓存岂非全部作废?每次G变化,就需要清空整个缓存并重新计算,效率将会极其低下。

Jet-Long利用一个巧妙的数学技巧解决了这个问题。它规定缓存中永远只存储原始的、未经任何映射的位置信息,如同档案馆中永远保存原始档案一样。当需要以远程窗口的方式读取某条缓存记录时,并不修改档案本身,而是在读取的瞬间即时计算一个“纠偏量”——具体而言,就是目标位置与原始位置之间的差值。

RoPE位置编码有一个美妙的数学性质:两个位置角度相加等于两个旋转矩阵相乘。这意味着,要从位置A的编码变为位置B的编码,仅需额外乘以一个对应差值(B-A)的旋转矩阵即可,无需从头重算。Jet-Long正是利用这一性质,在每次注意力计算时,对当前的查询向量施加一个纠偏旋转Δq,对从缓存中取出的键向量施加一个纠偏旋转Δk,使两者瞬间转换为远程窗口所需的映射后版本,而缓存本身始终保持原样。此操作在寄存器层面完成,耗时极短,无需修改任何已存储的数据,也不会因G的变化而致使之前的计算白费。

四、三次合一:预填充阶段的数学魔法

在模型处理一段长文本的开头阶段(称为“预填充”),需要同时计算所有位置对之间的注意力关系。问题在于,某些位置对应使用本地窗口的原始RoPE处理,而另一些则应使用远程窗口的映射RoPE处理。如果严格区分,就需要对每一对位置判断其距离是否超过本地窗口大小2048,这相当于在注意力矩阵上打一个巨大的布尔掩码——对于128K长度的文本,此矩阵包含128000×128000个格子,仅存储它就会将显存占满。

Jet-Long运用了一个来自数学的容斥原理来解决此问题,完全无需存储那个巨大的掩码矩阵。具体做法是将整个计算拆分为三个标准的FlashAttention(一种高效注意力计算方法)调用:第一次,在完整序列上使用远程映射的位置编码进行全局因果注意力,记为调用B;第二次,仅在本地2048窗口内进行滑动窗口注意力,但依然使用原始位置编码,记为调用A;第三次,同样在本地2048窗口内进行滑动窗口注意力,但使用远程映射的位置编码,记为调用C。

这三次调用背后的逻辑是:调用B覆盖了所有位置对,包括本地和远程,但本地部分使用的是远程映射(不正确);调用C专门覆盖了本地部分,也使用的是远程映射(同样不正确,但与调用B的本地部分完全一致);两者相减,本地部分的“远程映射”贡献恰好抵消,剩下的就是纯粹的远程部分贡献。再加上调用A提供的本地部分真实贡献,三者按照特定权重组合,即可得到完全正确的双焦点注意力结果。

权重的计算使用了FlashAttention内置输出的LogSumExp统计量,确保数值稳定,不会因大数减小数而出现精度损失。整个合并过程被融合进一个统一的CUDA内核(CuTe kernel)中,在硬件层面直接执行,避免了三次独立调用之间的数据传输开销。

五、实验结果:真正的成绩单

NVIDIA团队在三个不同规模的Qwen3基础模型上测试了Jet-Long,这三个模型分别拥有约17亿、40亿和80亿个参数,原始训练窗口均为32768个词,测试时将上下文扩展至最长131072个词(约128K),相当于训练窗口的四倍。

在RULER基准测试上,这是一套包含13种不同任务的综合测试,考察模型在长文本中的信息检索能力。与现有最优的对比方法相比,Jet-Long在17亿参数规模上领先4.79个百分点,在40亿规模上领先2.18个百分点,在80亿规模上领先2.03个百分点。尤其值得关注的是在65536词长度下的分任务表现:在“多键检索”这类需要同时记住多个散落在文本各处的关键信息的任务中,Jet-Long以61.80的得分遥遥领先于次优方法的20.00分(17亿规模),在80亿规模的三键检索任务中同样以70.80对39.80大幅领先。这表明Jet-Long在处理“大海捞针”类型的长程依赖时具有显著优势。

在HELMET-RAG基准测试上,这是一套模拟真实问答应用场景的测试,涵盖了自然问题、常识问答、多跳推理等任务,被研究者认为是预测模型实际应用表现最可靠的指标之一。Jet-Long在40亿和80亿规模上均排名第一,仅在17亿规模上以0.73个百分点的微小差距落后于Self-Extend方法。

在PG-19困惑度测试上,困惑度是衡量语言模型对文本预测能力的指标,越低越好。裸模型在超过32K的文本上困惑度急剧飙升,以17亿模型为例,在128K长度时困惑度高达71.00,而Jet-Long处理同样长度时困惑度仅为11.41,是所有方法中最低的,即预测最为准确。

在32K以内的文本上,Jet-Long的表现与原始模型完全一致——数学上即为同一计算,没有任何额外损耗。这一点非常重要,意味着使用Jet-Long不会使原本就能处理的短文本变差,是一种纯粹的能力扩充,而非能力置换。

六、速度表现:快到几乎感觉不到

准确率高固然重要,但如果速度过慢,实际应用中也难以推广。NVIDIA团队专门在H100显卡上对推理效率进行了系统测试,对比了原始FlashAttention-2、最新的FlashAttention-4(目前仅支持H100特定功能),以及Jet-Long的两个版本:未融合版本(三次FlashAttention分开调用)和融合CuTe内核版本。

测试结果非常清晰。在32K以内,四种配置的速度几乎完全相同,误差在1%以内,Jet-Long不引入任何额外开销。超过32K后,未融合版本的预填充速度下降至FlashAttention-2的89%至93%,而文字生成速度则暴跌至14%至34%,完全无法接受。融合版本则截然不同:预填充速度不仅不下降,反而超过FlashAttention-2,达到1.28至1.39倍的速度(以8B模型128K长度为例),非常接近仅限H100的FlashAttention-4的1.53倍水平,而Jet-Long不依赖任何H100专属特性;文字生成速度始终保持在FlashAttention-2的96%至100%之间,额外开销不超过4%。

此性能表现意味着,Jet-Long的准确率提升几乎是“免费”获得的,无需为此付出明显的计算代价。

七、举一反三:从纯注意力到混合架构

Jet-Long不仅在标准的全注意力Transformer上有效,研究团队还将其应用至Jet-Nemotron这一混合架构模型。混合架构是指在普通的注意力层之间插入了线性注意力层,是近年来兼顾效率与性能的一种新型设计。

测试结果表明,Jet-Long在混合架构上同样有效。在Jet-Nemotron-2B和4B模型上,32K以内的行为与原始模型完全一致,超过32K后原始模型的RULER准确率在128K时分别跌至8.54和5.65,而加入Jet-Long后分别回升至33.78和33.14,提升幅度超过25个百分点。平均七个长度的综合得分,Jet-Long分别提升了10.01和11.31个百分点。这表明Jet-Long的双焦点机制具有良好的架构通用性,并非仅为某一种特定模型结构量身定制的解决方案。

八、调参有多难:其实基本不用调

Jet-Long只有一个需要人工设置的参数:本地窗口的大小w0,默认设置为2048。研究团队专门测试了将此值改为不同数字时的影响,从512一直测试至8192。

结果显示,在512到4096这个较大范围内,所有设置下的RULER准确率彼此相差不超过2个百分点,差异非常微小。只有当w0被设为0时(本地窗口完全消失,模型只能看远处,看不清近处),准确率才会崩溃至接近零,因为缺乏清晰的近距离参照,远程信息也就失去了意义。另一端,当w0设置至8192这样大时,本地窗口占据了整体上下文的较大比例,对远程信息的处理资源相对减少,性能有轻微下滑,约为1.4到2.1个百分点的差距。

研究团队还测试了另一个设计选择:在远程窗口中,是使用整数位置分组(Jet-Long的默认做法),还是使用连续频率插值(类似YaRN的做法)。在65536词长度下,整数分组平均领先连续插值6.99个百分点;在131072词长度下,整数分组依然领先4.30个百分点。在某些特定任务(如两键检索和第二类问答)上,超长文本时连续插值表现稍好,说明两种方法各有其适用边界,但在128K及以内的范围内,整数分组是更稳健的默认选择。

总而言之,Jet-Long提供了一个既实用又有趣的答案:AI模型看不清长文本,并非因为它“笨”,而是因为它用于定位文字位置的那套方法,在超出训练范围后便开始“迷路”。Jet-Long的做法是给它配上一副自动变焦眼镜——近处使用原装镜片,清晰如初;远处自动调整焦距,将看不清的位置压缩映射至熟悉的范围,随后利用一套精妙的数学技巧使两个视野天衣无缝地融合在一起,并且整个过程快至几乎察觉不到任何额外开销。

这对普通用户意味着什么?当使用AI助手处理一份几万字的合同、一个几十万行的代码库,或者连续多轮的复杂对话记录时,Jet-Long这类技术将使AI的理解能力保持在更高水准,而非随着文本越来越长而越来越敷衍。尽管当前这项工作主要针对研究者和模型部署者,但其影响最终将传导至每一个依赖AI处理长文本的普通用户。

希望进一步研究技术细节的读者,可通过arXiv编号2607.07740找到完整论文,NVIDIA团队也已公开代码仓库,供感兴趣者自行实验。

Q&A

Q1:Jet-Long在短文本上会不会比原始模型差?

A:不会。Jet-Long在设计上保证了当文本长度不超过模型原始训练窗口(例如32768词)时,分组大小G自动等于1,远程映射退化为恒等映射,整个计算与原始模型完全等价,数学上没有任何差别,性能也完全一致。

Q2:Jet-Long需要重新训练模型吗?

A:不需要。Jet-Long是一种“零样本”方法,完全在推理阶段运行,不修改模型的任何参数权重。它仅改变注意力计算时的位置编码方式,以及注意力的合并逻辑,可直接叠加在任何已有的开源模型检查点上使用,无需任何微调或继续训练。

Q3:Jet-Long和Self-Extend、DCA这些方法最大的区别是什么?

A:核心区别在于动态性。Self-Extend和DCA均需提前固定一个分组大小或缩放系数,在所有长度下使用同一个参数,因此短文本和长文本之间存在难以兼顾的取舍。Jet-Long的分组大小G会根据当前输入文本的实际长度实时计算,始终是“刚好够用”的最小压缩,既不过度损失短文本精度,也不会在极长文本时压缩不足。