给RNN的记忆开一个三维的仓库

你有没有想过，AI模型在读一本很长的书时，到底是怎么"记住"前面讲过什么的？

这个问题比听起来复杂得多。2024年前后，整个AI圈子都在纠结一件事：如何让模型用更少的计算量处理更长的文本，同时还不丢失记忆。而这篇来自MIT和MIT-IBM联合实验室的论文，给出了一个挺巧妙的答案:把原本是"一张表格"的记忆，变成一个"立方体"。

**核心问题：记忆空间不够用了**

先说说背景。

现在主流的大语言模型大多基于Transformer架构，它靠一种叫注意力机制的东西来处理文本。

Transformer*：一种深度学习架构，通过"注意力机制"让模型在处理每个词时，能够同时关注输入序列中所有其他词的信息。

但Transformer有个麻烦事:它得把读过的每一个词都存下来，形成一个叫"键值缓存"的东西，文本越长，这个缓存就越大。处理6.4万个token（大约相当于一本中篇小说）的文本，缓存能占到12.6MB每千个token，算下来整本书读完缓存能撑到800MB开外。这既费内存，又拖慢速度。

于是有人想了个办法:用RNN（循环神经网络）来处理长文本，因为RNN只维护一个固定大小的"状态"，不管文本多长，状态大小都不变，读取速度是常数时间。

RNN*：循环神经网络，一种按顺序处理数据、用固定大小的"记忆"来概括历史信息的模型。

这里面最成功的分支叫"线性注意力"。它的核心思路是把RNN里原本只是一个向量的记忆状态，升级成一个矩阵。具体做法是，每读进一个词，就生成一个"键"和一个"值"，把它们做外积（outer product，相当于把两个向量交叉相乘生成一个矩阵），加到记忆矩阵里。读取的时候，用"查询"去和这个矩阵做乘法，就能把存过的信息捞出来。

线性注意力*：Linear Attention，一种用矩阵外积写入、矩阵乘法读取的记忆机制，能把Transformer那种需要存全部历史的缓存压缩成固定大小的状态。

这个设计很优雅，但也带来一个硬伤:矩阵的容量是有限的。如果键的维度是d，那这个记忆矩阵最多能无损存下d个互相垂直（专业说法叫"正交"）的键值对。超过这个数，新存进去的信息就会和老信息相互干扰，就像往一个只能装d格的抽屉里硬塞更多东西，东西会挤在一起，你拿的时候经常拿错。

论文里提到一个很关键的实验现象:现有的线性注意力模型在"长文本检索"任务上表现得明显不如Transformer，甚至没法精确复制超过一定长度的序列。这不是调参能解决的问题，是记忆容量这个天花板本身卡住了。

那怎么办？两条路:要么让记忆矩阵变大（比如加更多注意力头，或者把键值维度调大），要么找一种更聪明的方式扩容。前一条路很粗暴，直接导致参数量暴涨。研究者们想问的是:有没有办法只花很小的参数代价，就让记忆容量翻好几倍？

**三元线性注意力：给记忆加一个维度**

论文给出的答案叫"三元线性注意力"（Triadic Linear Attention）。

它的思路其实很直接:既然两个向量做外积能生成一个二维矩阵，那三个向量做外积，是不是就能生成一个三维的立方体？

具体来说，每个词除了生成原来的键k、查询q、值v之外，还额外生成一个"第二个键"k'和"第二个查询"q'，这两个的维度记作E。写入记忆的时候，把k、k'、v三个向量做三元外积，写进一个大小是d×E×d的三维张量（可以理解为一个立方体仓库）。读取的时候，用q和q'同时去"夹"这个立方体的两个键轴，把值取出来。

三元外积*：Triadic Outer Product，三个向量两两相乘后再组合，生成一个三维数组（张量）的运算，是从矩阵记忆升级到立方体记忆的数学基础。

这里有个关键点:这个"第二个键"k'完全是模型自己学出来的，不对应输入数据里任何具体的东西。它就像是给原来的每个键值对，额外打上了一个"标签"或者说"编号"，让模型可以把原本会互相冲突的信息，分别存进立方体的不同格子里，减少干扰。

论文算了笔账:记忆状态的容量从原来的d?，一下子涨到d?E。如果E设成8，容量直接翻8倍。但代价呢？只多了两个小小的线性投影层，用来生成k'和q'。论文里的实验数据显示，E从1涨到8，模型的非嵌入参数量只增加了1.2%。

这就好比你家里原来有一个储物柜，东西放多了经常翻不着。现在你没有换新柜子，而是在柜子里加装了几层可以左右滑动的抽屉隔板。每层隔板几乎不占地方，但你能存放的东西类别一下子多了好几倍，取的时候也更准，不容易拿错。如果不做这个隔板设计，你要么把柜子本身做大（也就是加更多头、加大键值维度），那样柜子（模型参数）整体体积会暴涨；要么继续忍受东西混在一起、经常拿错的问题（检索准确率下降）。

论文做了个很直观的小实验来验证这个想法，叫"多查询联想召回"（MQAR）：给模型一堆键值对，然后乱序问它每个键对应的值是什么，看模型答对多少。实验用的是一个极简的两层模型，没有加任何forgetting或delta rule之类的花哨机制。

结果很清楚:E每翻一倍，模型能够准确记住的键值对数量上限也大致翻一倍。比如E=16时，模型能记住的键值对数量差不多是E=1时的16倍，而参数量只涨了8%。这个实验等于是在"隔离变量"的条件下，证明了这套设计确实有效，不是靠别的机制蒙混过关的。

**和现代机制的兼容：遗忘门与Delta Rule**

光有容量还不够，现代的线性注意力模型还有两个重要的升级机制，论文需要证明三元结构能和它们配合使用。

第一个叫"数据依赖的遗忘"（data-dependent forgetting），通俗讲就是让模型自己学会该忘记什么、该记住什么，而不是一股脑把所有历史信息都同等对待地往前叠加。这就像你读书做笔记，不会把每一页的每一句话都抄下来，而是会判断哪些是重点、哪些可以放一放。

遗忘门*：Forget Gate，一个取值在0到1之间的"开关"，控制记忆状态在每一步更新前要保留多少之前的信息。

三元结构的巧妙之处在于，它给立方体记忆的每一个"切片"（沿着第二键轴切出来的二维矩阵）都配了一个独立的遗忘门。也就是说，立方体的不同隔层可以有不同的"记性"，有的隔层记性短，很快就忘，有的隔层记性长，能保留很久的信息。论文后面的分析显示，这种设计是真实起作用的：在训练好的模型里，同一个记忆头的8个切片，半衰期从0.2个token到1080个token不等，形成了一个很宽的时间尺度谱系，而普通的GDN只有单一一个遗忘门，半衰期固定在5.1个token左右。

第二个叫Delta Rule，这是一种更精细的写入机制。

Delta Rule*：写入新信息前，先把该键位置原来存的值读出来、减掉，再写入新值，避免新旧信息简单叠加导致内容混乱，类似"先擦除再书写"而不是"直接往上盖"。

论文把这个机制也扩展到了三维状态上：写入新的键值对之前，先用两个键一起去读一遍立方体里当前存的内容，把这部分"擦掉"，再写入新内容。有意思的是，这套三元Delta Rule经过数学变形后，可以等价成一个普通的Gated DeltaNet，只不过键的维度变成了d乘以E这么大。这意味着理论上这套机制是站得住脚的，不是生造出来的新东西。

**训练效率：如何不让立方体拖垮GPU**

这里出现了一个很现实的工程问题。

如果真的把立方体当成一个d·E维的普通矩阵来训练，那训练时每一步矩阵运算的开销会乘以E倍，而且这个立方体状态本身也会大到GPU片上内存（on-chip memory）根本放不下。

论文在这里用了两个技巧来解决这个问题。

第一个技巧是利用"外积可以拆开算"这个数学性质。两个联合键的内积，可以拆成两部分分别计算再相乘，而不需要先把两个键拼成一个大向量再算内积。这样原本要花C?(d·E)次运算的部分，降到了C?(d+E)，省了不少。

第二个技巧是把立方体状态按"值"这个维度切成一块一块的，分给GPU上不同的计算单元（线程块）去并行处理，每个计算单元只需要把自己负责的那一小块立方体放在寄存器里，不需要一次性把整个立方体都塞进去。

这就像一个大型仓库的管理问题。如果你非要把整个仓库的货架图纸都打印出来，钉在每一个搬运工的工作台上，台子根本放不下，效率也低。更聪明的做法是把仓库划分成若干个区域，每个搬运工只负责自己那片区域的货架信息，不用关心别的区域存了什么。论文里提到，E=8时整个状态占512KB，是GPU计算单元寄存器空间的两倍，但切分成小块之后，每一块只占128KB，正好放得下。

实测下来，E=8时训练速度比普通GDN慢28%到30%，E=4慢14%到15%，E=2慢9%到11%。听起来是有代价的，但换算成实际吞吐量，在1.3B参数规模、6.4万token长度下，三元GDN（E=8）比Transformer还要快5.1倍。换句话说，这套机制虽然比"没有扩容"的线性注意力慢一些，但比起依然要处理巨大缓存的Transformer，依然是划算的买卖。

**实验结果：扩容之后到底有多大提升**

论文在400M和1.3B两个参数规模上做了系统的对比实验，应用到两种基础机制上：Gated DeltaNet（GDN）和一种叫sGLA的标量门控线性注意力。

先看长文本语言建模效果。论文用了一个很直观的指标:模型在不同上下文长度下，相对Transformer的损失差距。结果显示，普通GDN在大约1万token之后就开始落后于Transformer，但三元GDN（E=2）把这个临界点明显往后推，E=8的版本甚至在6.4万token长度下依然比Transformer表现更好。

在六个检索密集型任务（包括DROP、NQ、TriviaQA、FDA、SQuAD、SWDE）上，三元GDN的提升更明显，尤其在FDA和SWDE这类需要从长文档中精确复制信息的任务上，提升幅度相当可观。RULER基准测试里的"大海捞针"（needle-in-a-haystack）任务也显示出同样的趋势：状态越大，模型能在越长的上下文里准确找到藏起来的信息。

更有说服力的是和其他扩容方法的对比。论文设计了一组"状态匹配"实验，把三元注意力和四种传统扩容方式（加大头维度、加宽值维度、每个键配多个值头、增加头数量）放在同样的状态大小下比较。

下面这张表摘录了部分关键数据（参数量单位百万，困惑度越低越好，召回率越高越好）：

| 模型 | 状态大小(MB) | 参数量 | PG19困惑度(16k-64k) | 召回率 |

|---|---|---|---|---|

| GDN基线 | 6.3 | 380.9 | 14.15 | 26.2 |

| 加大头维度(2倍) | 12.6 | 380.7 | 14.17 | 26.6 |

| 多值头(2个/键) | 12.6 | 378.2 | 14.12 | 27.1 |

| 加宽值维度 | 12.6 | 377.9 | 14.14 | 27.4 |

| 增加头数(16个) | 12.6 | 381.6 | 14.28 | 27.5 |

| **三元注意力(E=2)** | 12.6 | 381.9 | **13.95** | 28.4 |

| 加大头维度(4倍) | 25.2 | 380.6 | 14.24 | 28.5 |

| **三元注意力(E=4)** | 25.2 | 383.0 | **13.79** | **31.1** |

可以看到，在同样的状态大小下，三元注意力的困惑度始终是最低的，召回率也是最高的。而且4倍扩容时，有两种传统方法（加宽值维度、多值头）因为参数量暴涨得太厉害，不得不砍掉MLP层的宽度（从2816砍到640）来凑参数预算，结果性能反而比基线还差。三元注意力因为只加了两个小投影层，完全不需要动MLP，保住了模型其他部分的表达能力。

**旧模型能不能"升级"成三元结构**

论文还做了一个很实用的实验：能不能把一个已经训练好的普通线性注意力模型，后期改造成三元结构，而不用从头重新训练？

做法是把E=1的模型扩展到E=8，把原来的遗忘门复制到每个新切片上，第二键和第二查询的投影层从头初始化，然后在长上下文扩展阶段继续训练。

结果显示，这种"升级"（论文里叫upcycling）确实有效，升级后的模型比原来的基线模型有明显提升，大约能拿到从头训练三元模型所获提升的一半到四分之三。这个结果挺有实际意义的：如果你已经训练好一个模型，不用把钱全部打水漂重新训练，可以用较低的成本给它"扩容记忆"。

这里还隐含了一个有意思的推测，论文提到:短文本数据需要的状态其实很小，但短文本数据占训练语料的大头，所以也许未来可以先用小状态训练模型，等到后期专门做长上下文训练时，再把状态容量撑大，这样整体训练成本能省不少。

**混合架构：该扩哪一部分**

现实中很多模型并不是纯粹用线性注意力或纯Transformer，而是两者混合。比如论文提到Kimi Linear和Qwen3这类工业界的模型，都采用了线性注意力层和普通softmax注意力层交替堆叠的设计。

这就引出一个问题：如果要在混合架构里扩容，到底该扩大线性注意力层的状态，还是扩大Transformer那部分的键值缓存？

论文用一个3:1的GDN和分组查询注意力（GQA）混合模型做了对比，一边是把GQA的键值缓存翻倍，另一边是把GDN层换成三元结构、状态扩大4倍。结果显示，三元混合模型在各项困惑度指标上表现最好，零样本任务和NIAH任务平均分也是最高的，只在召回率上略逊于扩大缓存的版本（而且差距在误差范围内）。更关键的是，在6.4万token长度下，扩大缓存的方案需要的内存几乎是三元方案的两倍。

这说明，至少在这个实验设置里，花同样的"内存预算"，扩大线性注意力的状态比扩大Transformer缓存更划算。

**消融实验：细节决定成败**

论文还做了几个小而重要的消融实验。

一个是比较第一个键的维度和第二个键维度E的分配方式。固定两者乘积为1024不变，分别试了dk=128/E=8、dk=64/E=16、dk=32/E=32几种配置，发现两个维度越接近，性能反而略微下降。这提示如果想做更极端的因子分解，设计上需要更谨慎。

另一个是关于第二键和第二查询要用什么激活函数。实验比较了softplus、sigmoid、SiLU和不加激活函数四种方案，发现只能取非负值的激活函数（softplus、sigmoid）表现明显更好。论文给出的解释是，如果激活函数允许负值，模型可能会让不同切片的贡献相互抵消，反而削弱了扩容带来的好处。这个细节很容易被忽略，但却是让整个设计真正发挥作用的关键一环。

**局限性：没有免费的午餐**

论文也坦诚地列了几条局限。

训练速度确实会变慢，E=4慢15%左右，E=8慢30%左右，即使用了专门优化的GPU内核依然如此。部分检索密集型任务上，三元注意力虽然比普通线性注意力好很多，但距离Transformer依然有差距，只有状态足够大才能追平。另外，论文只把这套机制应用在了GDN和sGLA这两种具体模型上，线性注意力领域还有很多其他变体（比如RWKV系列、Mamba系列），这些设计有没有可能从三元结构里受益，论文没有涉及。

**写在后面**

读完这篇论文，有个地方让我反复琢磨了很久：论文提到的"第二个键不对应输入数据里任何具体的东西"，这件事其实挺反直觉的。

一般我们理解AI模型的输入输出，总习惯认为每个变量都应该有明确的现实含义，比如键对应词、值对应词的语义。但这个第二个键完全是模型凭空学出来的一个"内部索引"，它的作用纯粹是让记忆系统更好地给信息分类归档，不承载任何语义。这让我想到图书馆的编目系统，书架上的书本身有内容，但索书号这个东西完全是人为设计出来的分类工具，它不是书的内容，却决定了你能不能快速找到那本书。模型学会了给每条信息自己贴一个"隐形索书号"，这个发现本身挺有意思的。

另外一个让我停下来想的地方，是论文里提到的那个遗忘门半衰期分布的实验。三元GDN里8个切片的半衰期从0.2个token到1080个token，跨越了四个数量级。这说明模型不是简单地"扩大了容量"，而是自发学会了用不同的时间尺度去组织记忆，有的切片专门记最近几个词的细节，有的切片专门记很久以前的大致脉络。这种分工不是人为设计的，是训练过程自己涌现出来的结构。这让人忍不住想，如果继续增加切片数量，这种时间尺度的分工会不会出现更细腻的层次，模型会不会学出类似"短期记忆、中期记忆、长期记忆"这样接近人类认知结构的东西？

这个问题，论文没有回答，但留下来挺值得继续想的。

Q&A

Q1：三元线性注意力是什么？

A：三元线性注意力是一种改进的循环神经网络记忆机制，通过把键、第二个键、值三个向量做外积，写入一个三维张量状态，用两个查询同时读取，从而用很小的参数代价大幅提升模型能记住的信息量。

Q2：三元线性注意力和普通线性注意力相比有什么优势？

A：普通线性注意力的记忆矩阵容量有限，长文本处理容易丢信息或检索出错。三元线性注意力通过增加第二个键维度E，能让记忆容量提升E倍，而参数量只增加很少，在长文本语言建模和检索任务上效果明显更好。

Q3：三元线性注意力会不会让模型训练变得很慢？

A：确实会有一定开销，E=8时训练速度比普通GDN模型慢28%到30%左右，但由于避免了Transformer那种随文本长度线性增长的缓存负担，在长文本场景下三元模型整体速度反而比Transformer快数倍。