Extender重构Transformer内存路径：9.2亿参数模型KV缓存缩小104倍

大语言模型在回答问题时，内存里究竟存了些什么？这个问题可能很少有人细想。

答案或许出乎意料：占内存大头的不是模型参数，而是一种叫“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分，接近两倍。

论文推测，原因在于扩展流本身就提供了一条更稀疏、更干净的长程信息通道，不必与下一个词预测的信号相互争抢。于是当场景超出训练长度时，它的泛化能力就更强。

内存上的差距更夸张。模型参数量为9.2亿、上下文长度64k时，标准Transformer的持久KV缓存按bf16精度算大约要11.3GB，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，视觉上几乎能一一对应。设计思路本身也透着一种态度：不为了创新而创新，先把问题想清楚再动刀。

还有一点我觉得该单独说。论文提到，在"变量追踪"（VT）这项长文本任务上，Transformer反而赢了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仍能保持不错的分数。