![]()
你有没有想过,一个大语言模型在回答你问题的时候,内存里到底在偷偷存着什么?
答案可能会让你吃惊:不是模型参数本身占用最大,而是一个叫"KV缓存"的东西,可能悄悄吃掉上百GB内存。这个数字放在今天看也很惊人,一个模型为了记住你之前说过的话,要为每一层、每一个词都存一份专属的记忆副本。context越长,层数越多,这个缓存就滚雪球一样越大。
这篇来自伊利诺伊大学芝加哥分校的论文,提出了一个叫Extender的新架构,试图从根本上改写这套记忆方式的游戏规则。它不是在现有缓存上做压缩,而是直接改变了信息传递的通道设计。结果是什么?在同样参数规模下,短文本任务表现不输传统Transformer,长文本任务反而更强,而持久化的注意力内存占用,在他们测试的9.2亿参数模型上,缩小了104倍。
这个数字值得你停下来想一想。104倍不是优化参数量化精度那种渐进式的改良,这是架构层面的重新设计带来的质变。
问题出在哪:每一层都要给后面所有层留一份完整的记忆
要理解Extender在解决什么问题,得先搞懂标准Transformer里一个你可能从没注意过的细节。
标准Transformer*:目前主流大语言模型采用的基础架构,由多层相同结构的模块堆叠而成,每层包含注意力机制和前馈网络两部分
在标准Transformer里,所有信息的传递都靠一条叫"残差流"的通道。你可以把它想象成一条传送带,每一层处理完数据后,把自己的"修改意见"叠加到传送带上,传送带继续往下走,下一层接着改。这个设计很聪明,因为它让每一层都能直接从最终的损失函数那里拿到清晰的梯度信号,这是深度模型能稳定训练的关键原因之一。
但这个传送带有一个致命的副作用:因为每一层的修改都是全宽度的,叠加之后,所有层信息混在一起,没法分离。等到推理阶段,为了让注意力机制(attention)能够正确地回看之前的词,每一层都必须单独存一份自己版本的键值对(KV缓存),而且是针对每一个token都要存。
键值对(KV缓存)*:注意力机制中用来快速检索历史信息的查找表,Key相当于索引,Value相当于对应的内容,模型生成每个新词时都要和之前所有词的KV做比对
这个代价有多大?随着模型变宽、变深、处理的文本变长,缓存大小是线性增长的。论文里提到,大模型的KV缓存不加任何优化的话,能轻松超过100GB。
这就好比一个大型图书馆,每一层楼都要求自己单独保留一份完整的馆藏目录副本,哪怕98%的内容和楼下一楼的目录重复。如果不这样做会怎样?如果大家共用一份目录,万一楼下的目录被意外修改,楼上找资料就全乱了,因为标准架构里每一层的残差更新都是叠加式的,信息一旦混合就无法区分谁贡献了什么,所以才不得已要求每层各自留档。这正是问题的根源:叠加式的信息传递,逼着系统为每一层单独开一份完整的备份。
Extender的解法:把"写日记"和"做广播"分开
Extender的核心想法说出来其实挺朴素:别让一条通道同时干两件不相关的事。
论文作者观察到,残差流其实承担了两种性质完全不同的工作。第一种是预测下一个词需要用到的信息,这个需要全宽度、精确叠加,因为预测任务要求细节到位。第二种是每一层想要"告诉后面的层"一些辅助性的发现,比如某个词和另一个词的距离、某种重复模式的计数器。这类信息往往体积很小,和下一个词预测本身没有直接关系,硬塞进残差流里,只会和其他写入操作打架。
于是Extender做了一件事:新开一条通道。
残差流h*:延续自标准Transformer的叠加式通道,主要负责下一个词的预测,前馈网络(FFN)和注意力的查询投影(Q)都会读取它
扩展流x*:Extender新增的拼接式通道,每一层只往后面追加(append)一小段自己的发现,体积远小于全宽度向量,只有注意力的键值投影(KV)会读取它
这个"追加而非叠加"的设计是整个方案最关键的一步。叠加意味着信息混合,一旦混合就无法还原;追加意味着历史完整保留,前面所有层写的东西,原样待在那里,像写日记一样,今天写的不会覆盖昨天写的。
这就带来一个美妙的性质:当一个token走完全部L层之后,它积累出来的扩展流,论文里叫x*,天然就包含了每一层想给注意力看的全部信息,而且全部层可以共用这一份扩展流缓存,不需要像标准Transformer那样各留一份。
打个比方,这就像一个团队共用一本公共笔记本,而不是每个人单独记一份会议记录。如果不这样设计会怎样?传统做法里,13个人开会,13个人各自记一份几乎重复的笔记,存储空间全被这种冗余吃掉了。改成公共笔记本之后,新来的信息往后翻页追加就行,之前翻过的页码永远不会被篡改,所有人读同一本本子,省下的纸张(内存)自然就省了。而这个"永远不被篡改"的性质,恰恰来自于追加式写入而非叠加式写入的设计选择。
具体数字上,标准Transformer需要的KV缓存内存和$2Ld_{model}$成正比(两个全宽度向量乘以层数),Extender只需要$sum|epsilon_ell|$,也就是每层那一小段扩展的总和。论文里用的扩展宽度是32,相比动辄上千的模型宽度,这个比例差距就是内存节省的来源。
查询(Q)依然要看全局:注意力到底在看什么
光拆分两条通道还不够,还要搞清楚注意力机制里的查询(Query)、键(Key)、值(Value)三个角色分别该读哪条通道。
论文给出的方案是:注意力的键值投影只读扩展流x,但查询投影要同时看扩展流和残差流,具体做法是把二者各自做归一化(RMSNorm)后相加。这个设计不是随便拍脑袋决定的,论文做了消融实验来验证。
RMSNorm*:一种归一化技术,用来稳定数值范围,防止某些特征的数值过大压过其他特征,比传统的LayerNorm计算更简单
实验结果很直接:如果查询只看残差流h,或者只看扩展流x,长文本任务的表现都明显下滑。具体数字上,在8k长度的RULER测试里,默认设计能拿到13.5分,只看h掉到12.4分,只看x更惨,只有8.0分,几乎少了一大半。
这说明一件事:查询需要同时感知"下一个词要预测什么"和"过去每一层发现了什么"这两种信息,缺了任何一边,模型对长距离依赖的捕捉能力都会打折扣。这就像一个翻译在工作时,既要盯着眼前这句话的语法结构(残差流h),又要记得前文提到的人名地名专有信息(扩展流x),只盯一头,翻译出来的句子不是语法不通,就是张冠李戴把人名搞混。如果只给翻译看语法而藏起专有名词表,他翻得再流畅也会把人名译错,这正是只用h做查询时长文本任务崩盘的原因。
临时KV缓存:省内存不是没有代价的
这里有个问题:如果KV缓存被扩展流取代了,那注意力计算的时候,键值投影到底怎么算出来?
答案是实时计算,但不是每次都从头算。注意力的K和V本质上是扩展流x乘以对应的权重矩阵$W_k$和$W_v$算出来的。Extender的做法是维护一个临时的、"用完就能扔"的KV缓存(ephemeral kv-cache)。
临时KV缓存*:在一次对话或一次推理过程中临时生成的KV数据,计算完当前这轮交互就可以释放,不需要像传统KV缓存那样长期占用内存
这个设计的巧妙之处在于,它把"持久占用的内存"和"临时计算的开销"做了切割。在一轮对话内部,临时缓存该怎么用还是怎么用,内存占用和标准Transformer一样,没有减少。但一轮对话结束,这个临时缓存就可以释放,下次需要时,从那份小小的、永久保存的扩展流里重新投影出来就行。
这就像图书馆的影印服务。你要查一份资料,图书馆不会把全套资料的复印件常年摆在你桌上占地方,而是你来查的时候,现场从保存完好的原始档案里影印一份给你用,用完归还。原始档案(扩展流)体积小、常年保存,影印件(临时KV缓存)用的时候才生成,用完就扔。如果反过来,让每个来访者的影印件永远占着桌子不收,图书馆早就被堆满了,这正是标准Transformer持久KV缓存臃肿的原因所在。
当然,这不是没有代价的,重新投影需要额外的计算开销。论文提到,这个开销可以在多个输入输出token上分摊,所以实际影响不大,而且如果搭配分组查询注意力(GQA)这种降维技术,重新投影的速度还能进一步提升。
分组查询注意力(GQA)*:一种让多个查询头共享同一组键值头的技术,能把KV缓存压缩4到8倍,是目前大模型里很常用的内存优化手段
这个设计到底行不行:实验数据说话
说了这么多架构设计,到底效果如何?论文在三个参数规模上做了对比:1.98亿、4.36(4.38)亿、9.2(9.24)亿参数,分别对应短文本测试集DCLM CORE和长文本测试集RULER。
DCLM CORE*:一套包含22个子任务的短文本语言理解基准测试,涵盖常识推理、阅读理解等任务
RULER*:专门用来评测模型在超长文本(从1千到十几万token)场景下的信息检索与推理能力的基准测试
在短文本CORE任务上,Extender和标准Transformer打了个平手,差异基本在误差范围内。但在长文本RULER任务上,差距就拉开了。
| 训练阶段 | 模型 | 1k | 8k | 16k | 32k | 64k |
| ClimbMix预训练18B | Extender | 59.0 | **19.9** | **8.4** | **3.8** | **1.9** |
| ClimbMix预训练18B | Transformer | 59.4 | 0.0 | – | – | – |
| ProLong微调+5B | Extender | **68.7** | **49.9** | **51.7** | **43.8** | **32.7** |
| ProLong微调+5B | Transformer | 63.2 | 40.5 | 39.4 | 35.1 | 17.0 |
这组数字很能说明问题。在只用2k长度文本预训练之后,标准Transformer在8k长度上的RULER分数直接归零,等于完全失效,而Extender还能保持19.9分。经过长文本微调之后,差距依然存在,在64k长度上,Extender是32.7分,Transformer只有17.0分,几乎是接近两倍的差距。
这背后的原因,论文推测是因为Extender的扩展流本质上提供了一种更稀疏、更干净的长程信息通道,不需要和下一个词预测的信号互相竞争,所以在超出训练长度的场景下,泛化能力更强。
内存层面的对比更夸张。对于9.2亿参数的模型,在64k上下文长度下,标准Transformer的持久KV缓存需要大约11.3GB(bf16精度),Extender只需要109MB,整整104倍的差距。这个比例还会随着模型变宽进一步扩大,因为Extender每层只增加固定的$d_epsilon$(本文设为32)宽度,而标准Transformer是和整个$d_{model}$成比例的。
论文还做了知识蒸馏的实验,用一个9.24亿参数的大模型(无论是Extender还是Transformer架构)去教一个1.99亿参数的小Extender,结果显示两种老师教出来的学生表现接近,说明Extender这种架构本身具备良好的可蒸馏性,不会因为换了一种信息传递方式就变得难以训练。
另外一个值得提的细节是分组查询注意力(GQA)和Extender搭配得很好。把13个头压缩到GQA4/5分组之后,参数量从9.24亿降到8.13亿,但长文本任务的表现几乎没有损失,有些长度上甚至还略有提升。这说明Extender架构本身的内存优化和传统的降维压缩技术并不冲突,可以叠加使用。
这个方案目前的局限
论文作者自己也坦率地提到了限制:目前测试的模型规模最大只到9.24亿参数,训练的token预算基本锁定在参数量的20倍左右,受限于计算资源,还没能验证在更大规模、更大token预算下这个结论是否依然成立。
另外有一个细节值得琐碎地提一下。论文在设计扩展宽度$d_epsilon$和每一层能读取哪部分历史信息的"滑动窗口"时,发现窗口设计对结果影响很大。比如键(K)的窗口始终用"最近的那一段"效果最好,但值(V)就不一定,因为很多任务需要模型原封不动地复述最初的输入信息,如果V的窗口只截取最近的内容,反而会把最早的输入信息挤掉。最后论文选择的方案是K用最近窗口,V用随机采样窗口,算是一种折中。
这个细节其实揶揄了一件事:看似简单的"滑动窗口"背后,藏着不少需要靠实验摸索才能确定的工程取舍,并不是一拍脑袋就能定下来的架构参数。
写在后面
读到这篇论文最让我意外的一点,是Extender并没有去挑战注意力机制本身,它没有像线性注意力(linear attention)或者状态空间模型(Mamba)那样试图绕开"每个token都要回看所有历史token"这件事的计算代价。它选的是一条更保守但也更稳妥的路:保留softmax注意力的全部能力,只改变注意力"看"的是哪条通道里的数据。
这种克制挺有意思的。很多新架构论文一上来就想颠覆点什么,Extender选择的是在既有框架里做一次外科手术式的改动,拆开一根线,改成两根线。这种改动的好处是容易验证、容易对比,论文里那张左右对照的架构图,一个标准Transformer,一个Extender,视觉上几乎能一一对应,这种设计思路本身也是一种态度:不是为了创新而创新,是先把问题想清楚再动刀。
还有一个细节我觉得值得单独说一下。论文提到,Transformer在"变量追踪"(VT)这一项长文本任务上反而赢了Extender,而且赢得很明显。作者没有深挖原因,只说这可能暗示两种架构在不同任务上有不同的"归纳偏好"。这种诚实挺难得的,很多论文会选择性忽略对自己不利的数据点,这篇论文把它摆出来,还承认"我们没有继续调查"。这种留白反而让整篇论文显得更可信。
如果再往前想一步,这种"拆分预测通道和记忆通道"的思路,会不会在其他领域也有类似的影子?人类大脑里,工作记忆和长期记忆走的也不是同一套神经回路,这未必是巧合,也可能只是我瞎联想。但至少这篇论文提醒了我一件事:很多时候系统效率低下,不是因为算法本身不行,而是因为我们把两件不相关的事硬塞进了同一条通道里。
Q&A
Q1:Extender架构是什么?
A:Extender是一种改进版Transformer架构,它把信息传递拆成两条通道,一条残差流负责预测下一个词,一条扩展流专门给注意力机制存历史信息,从而大幅降低模型推理时的内存占用。
Q2:Extender相比标准Transformer内存能省多少?
A:论文测试的9.24亿参数模型里,在64k长文本场景下,Extender的持久注意力内存占用比标准Transformer小104倍,而且这个优势会随模型变宽进一步扩大。
Q3:Extender在实际任务表现上比Transformer强吗?
A:在短文本任务上两者表现基本持平,但在长文本任务RULER上,Extender明显胜出,尤其是在超出训练长度的场景下,标准Transformer表现会大幅下滑甚至归零,Extender依然保持不错的分数。





京公网安备 11011402013531号