Transformer每层都开小会，KV缓存吃掉上百GB，UIC新架构砍掉104倍

大语言模型在回答问题时，内存里究竟悄悄存了些什么？这个问题很多人没想过。

答案或许出人意料：占内存大头的不是模型参数本身，而是一个叫“KV缓存”的东西，它可能悄悄吃掉上百GB。为了记住你此前说过的话，模型得为每一层、每一个词都留一份专属记忆副本。context越长、层数越多，这份缓存就像滚雪球一样膨胀。

伊利诺伊大学芝加哥分校的一篇论文提出了名为Extender的新架构，想从根子上改写这套记忆方式的规则。它没有在现有缓存上做压缩，而是直接改动了信息传递的通道设计。结果是：同样参数规模下，短文本任务不输传统Transformer，长文本任务反而更强；在他们测试的9.2亿参数模型上，持久化的注意力内存占用缩小了104倍。

104倍这个数字值得停下来想一想。它不是优化参数量化精度那种渐进式改良，而是架构层面重新设计带来的质变。

问题出在哪：每一层都要给后面所有层留一份完整的记忆

要弄懂Extender解决的是什么问题，得先看清楚标准Transformer里一个容易被忽略的细节。

标准Transformer*：眼下主流大语言模型采用的基础架构，由多层相同结构的模块堆叠而成，每层包含注意力机制和前馈网络两部分

标准Transformer里，所有信息都靠一条叫“残差流”的通道传递。可以把它想成一条传送带：每一层处理完数据，就把自己的“修改意见”叠加上去，传送带继续往下走，下一层接着改。这个设计很聪明，因为每一层都能直接从最终的损失函数那里拿到清晰的梯度信号，这是深度模型能稳定训练的关键原因之一。

但这条传送带有个致命的副作用：每层的修改都是全宽度的，叠加之后所有层的信息混在一起，分不出谁是谁。到了推理阶段，为了让注意力机制能正确回看之前的词，每一层都得单独存一份自己版本的键值对，也就是KV缓存，而且针对每一个token都要存。

键值对(KV缓存)*：这是注意力机制里的一张查找表，专门用来快速调取历史信息，Key 充当索引，Value 就是索引指向的内容；模型每生成一个新词，都得拿它跟前面所有词的 KV 逐一比对

这笔代价有多大？模型越宽、越深、处理的文本越长，缓存就线性增长。论文里提到，大模型的KV缓存不做任何优化的话，能轻松超过100GB。

好比一座大型图书馆，每层楼都要求自己单独留一份完整的馆藏目录副本，哪怕98%的内容和楼下那层的目录重复。不这么做会怎样？如果大家共用一份目录，万一楼下的目录被意外改动，楼上找资料就全乱了。标准架构里每一层的残差更新都是叠加式的，信息一旦混合就分不清谁贡献了什么，所以才不得已要求每层各自留档。问题的根源正在这里：叠加式的信息传递，逼着系统为每一层单独开一份完整备份。

Extender的解法：把“写日记”和“做广播”分开

Extender的核心想法说出来挺朴素：别让一条通道同时干两件不相关的事。

论文作者注意到，残差流其实承担着两种性质完全不同的工作。一种是预测下一个词所需的信息，这要求全宽度、精确叠加，因为预测任务需要细节到位。另一种是每一层想“告诉后面的层”的辅助发现，比如某个词和另一个词的距离、某种重复模式的计数器。这类信息往往体积很小，和下一个词预测本身没有直接关系，硬塞进残差流只会和其他写入操作打架。

于是Extender做了一件事：新开一条通道。

残差流h*：这条叠加式通道沿用自标准Transformer，承担的主要任务是预测下一个词，前馈网络(FFN)和注意力的查询投影(Q)都要从它这里读取信息

扩展流x*：这是 Extender 新加的拼接式通道，每层只在自己后面补上一小段新发现，体量比全宽度向量小得多，读取它的只有注意力的键值投影(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缓存（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”的计算代价，Extender没走这条路。它保守一些，也稳一些：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还能保持不错的分数。