Triadic Linear Attention：三维递归记忆的中文研读笔记

论文：Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
作者：Oliver Sieberling、Bharat Runwal、David Jin、Ryan Chin、Rameswar Panda、Yoon Kim（MIT / MIT-IBM Computing Research Lab）
版本：arXiv:2609.36529v1，2026-09-29 提交；获取与研读日期：2026-10-09。
原文：arXiv · 本地 PDF · 英文汇报 PDF · 可编辑 TeX
作者提供的代码链接：TriadicLinearAttention。文中原图沿用论文图件，许可为 CC BY 4.0。以下页码均为从 1 开始的 PDF 页码，与论文印刷页码一致。

论文总览

线性注意力把历史压入固定大小的矩阵状态，推理成本因此不随历史长度增长，但记忆容量也受到限制。本文提出三元线性注意力（Triadic Linear Attention）：每个 token 额外产生一个短 key 和一个短 query，把原来的 $k\otimes v$ 写入改成 $k\otimes k'\otimes v$，让状态从 $d\times d$ 扩大到 $d\times E\times d$。检索权重由一次相似度变成两次相似度的乘积，联合地址空间扩大，而投影参数只增加少量。

论文的贡献有三个相互依赖的部分：三维记忆的参数化；与逐切片遗忘和联合 delta rule 的结合；利用乘积结构与 value 维分块实现可训练的 GPU 内核。400M GDN 的 $E=8$ 模型把状态从 6.3 MB 扩大到 50.3 MB，参数从 380.9M 增到 385.4M；PG19 的 16K–64K 区间困惑度从 14.15 降到 13.74，六任务召回均分从 26.2 升到 33.1。计算代价则在 1.3B、H100 的单块测试中测量：$E=8$ 的前向加反向耗时较普通 GDN 增加约 28%–30%。（Table 2, p. 7；Sec. 3.5, p. 8）

最有启发的结果是：状态容量可以成为独立于模型参数量的扩展方向。不过，长文 next-token 预测的改善和精确检索能力的改善幅度不同；Figure 3 中，纯 Triadic GDN 在 64K 的四项 NIAH 检索上仍明显落后于 Transformer。本文展示了容量、参数和计算之间更好的取舍，并没有消除固定状态的容量限制。

结构地图

一：动机——扩大可寻址空间，避免挤占其他模块

普通线性注意力的单头状态为 $S_t\in\mathbb R^{d\times d}$：

这里假设初始状态为零。$k_s$ 充当地址，$v_s$ 充当内容，$q_t$ 负责检索。如果查询与其他记录的 key 不正交，那些记录就会按相似度混入输出。这是叠加式关联记忆（associative memory）的串扰。$d$ 维地址空间至多包含 $d$ 个彼此正交的 key；更好的遗忘与替换策略可以管理干扰，却不会自动扩大地址空间。（Sec. 2.1, p. 2）

直接扩大状态有几条路径：增大每头维度、扩大 value、给一个 key 配多个 value head、增加 head 数。后几种方式也会扩大投影矩阵。在固定参数预算下，作者必须缩窄 MLP；部分 4 倍扩容对照把 MLP 宽度从 2816 缩到 640。因而本文研究的是一个预算分配问题：怎样用尽量少的慢权重参数，为每条序列提供更多快权重记忆？（Secs. 3.1–3.2, pp. 6–7）

二：方法——两个地址轴共同读写一个 value

2.1 符号与基本计算

无门控、无 delta 的基本形式是：

$\times_n$ 表示沿第 $n$ 个轴收缩。例如，输出的第 $j$ 个分量是 $\sum_{a,e}S_t[a,e,j]q_t[a]q'_t[e]$。因此，历史 token 需要同时在两套地址上与查询匹配，才会对输出产生大的贡献。第二 key 是模型学习的潜在地址，不需要数据自带第二种标签。（Eq. (2), p. 3）

这也可以展平成普通矩阵记忆：定义联合 key $\kappa_t=k_t\otimes k'_t\in\mathbb R^{dE}$，把状态视为 $\bar S_t\in\mathbb R^{dE\times d}$。于是三维表示等价于采用结构化特征映射的线性注意力，特征由两个短向量的外积产生。单次写入的联合地址有可分解约束，但多次写入累加后的状态不必保持秩一。

状态元素从 $d^2$ 增到 $d^2E$。主要新增投影参数约为 $2d_{\rm model}HE$，另有短卷积和门控的少量参数。把联合 key 直接投影到 $dE$ 维，则需要远大的投影矩阵。参数节省来自这种乘积分解，而实际状态仍是稠密存储，读写量随 $E$ 增加。

2.2 一个说明串扰如何减少的例子

下面是教学例子。设两条记录的第一 key 相同，都是 $e_1=(1,0)$，但第二 key 分别是 $f_1=(1,0)$ 和 $f_2=(0,1)$，value 分别为 $a,b$。普通矩阵记忆只能在地址 $e_1$ 下读出 $a+b$；三维记忆则保存

用 $(e_1,f_1)$ 查询，$b$ 项被 $f_1^\top f_2=0$ 消掉，恰好得到 $a$。一般地，取 $d$ 个正交第一 key 与 $E$ 个正交第二 key 的全部组合，可以形成 $dE$ 个正交联合地址。这是理想正交条件下的容量解释；真实模型学习的地址允许重叠。尤其主语言模型对第二 key 使用正值 softplus，正交例子描述的是机制极限。

当 $E=1$ 且 $k'_t=q'_t=1$ 时，公式严格退化成普通线性注意力。语言模型实现中第二 key/query 先经 softplus 再做 $\ell_2$ 归一化，因而一维情况下确实等于 1。（Sec. 2.2, p. 3；Appendix D, pp. 19–20）

2.3 遗忘与 delta rule：先衰减，再修正联合地址

逐切片遗忘先得到临时状态

Triadic GDN 随后读出当前联合 key 对应的已有内容，并写入残差：

这里的残差来自所有切片的联合读出，所以不能把它当作 $E$ 套互不相干的 DeltaNet。若记 $D_t$ 为展平后的对角衰减矩阵，原文 Eq. (5) 为：

归一化使 $\|\kappa_t\|=\|k_t\|\|k'_t\|=1$。由上式可直接推得，更新后在同一地址读到 $(1-\beta_t)\widehat v_t+\beta_tv_t$：$\beta_t$ 控制替换比例，取 1 才是完整替换。还可以把它理解为对内层损失 $\frac12\|\bar S^\top\kappa_t-v_t\|^2$ 做一次梯度更新，初值是衰减后的状态。这一推导解释了与快权重和 test-time training 的联系。（Eqs. (3)–(5), p. 4；上述替换关系为教学推导）

标量门控线性注意力 sGLA 则去掉擦除项，实际实验更新为

其写入强度 $\beta_t$ 仍然保留。主实验的 sGLA 是作者定义的受控基线，不能直接等同于完整的 Mamba-2 或 mLSTM。（Algorithm 2, p. 20）

2.4 实际语言模型模块

输入先经过 RMSNorm。$q,k,v$ 由线性投影、因果短卷积和 SiLU 生成；$q',k'$ 则由投影、短卷积和 softplus 生成。四个 query/key 向量分别做 $\ell_2$ 归一化，value 不做这一归一化。遗忘门的参数化为

状态更新并读出后，各头输出经过 RMSNorm 与 sigmoid 输出门，再拼接、投影、残差相加；之后是带残差的 SwiGLU MLP。第二 key/query 与原有投影、卷积合并执行，因此实现上不需要为每个新向量单独启动一整套 GEMM 和卷积。（Algorithm 2, p. 20；Appendix A, p. 18）

2.5 高效训练：把块内乘积核与跨块状态分开

对于一个长度为 $C$ 的 chunk，令 $Q,K,V\in\mathbb R^{C\times d}$，$Q',K'\in\mathbb R^{C\times E}$。先看无遗忘、无 delta 的情况。联合地址内积可分解为

因此，块内因果交互矩阵为

构造它的成本从直接展平地址时的 $O(C^2dE)$ 降为 $O(C^2(d+E))$。设 $q'_e,k'_e$ 为第二 query/key 矩阵的第 $e$ 列，跨块状态的读写为

块内 token 可以并行，跨 chunk 才递归。省掉的是块内地址点积的乘法扩张；状态读写的 $E$ 倍成本仍然存在。（Sec. 2.5, pp. 4–5，Eqs. (6)–(8)）

完整门控 delta 形式用每片累计衰减 $\gamma_e^r=\prod_{j=1}^r\alpha_e^j$，定义因果矩阵 $(\Gamma_e)_{rs}=\gamma_e^r/\gamma_e^s$（$r\ge s$；其他位置为零），以及

所有切片共享同一个 $C\times C$ 的三角系统：

这里 $\operatorname{tril}(\cdot,-1)$ 只保留严格下三角部分，表示当前 token 的写入依赖此前写入造成的修正。由它求得块内经 delta 修正的 value：

向量上的乘除均按元素进行。门控沿第二 key 轴变化后，不能只预计算一份标准 DeltaNet 的 $TK$；作者选择直接把共享 $T$ 作用到残差上。（Appendix A, p. 17，Eqs. (9)–(11)）

2.6 GPU 实现与数值稳定性

在 $d=128,E=8$ 时，一个头的 FP32 状态有 $128\times8\times128\times4=512$ KiB，超过单个 Hopper SM 的寄存器容量。value 各列的计算互不依赖，作者按 32 个 value 通道切分，单块只需保存 128 KiB 的状态，同时保留这部分 value 的全部 $E$ 个切片。块数不足时使用 16 列。这样无需把整个头的状态放到一个 SM 上。（Sec. 2.5, p. 5；Appendix A, p. 18）

衰减比值直接计算为 $\exp(\log\gamma_e^r-\log\gamma_e^s)$。因果范围内指数非正，避免把一对可能溢出的正指数因子相乘。作者展示，先分解成两个矩阵再对指数做截断的便宜方案，在极强衰减时会破坏近对角项；所构造的 FP64 模拟中，块内总衰减为 $e^{-200}$ 和 $e^{-600}$ 时，相对输出误差分别达 0.05 和 1.5。这是选择逐元素形成 $R,R'$ 的具体原因。（Appendix A, p. 17）

CuTe 内核使用 FP32 三角求解，并在训练前向中保存 BF16 的块边界状态供反向使用。因此，“固定状态”描述的是逐 token 推理时的历史记忆；训练仍需保留随 chunk 数增长的中间信息。表格中的 2 字节状态计量与内核 FP32 寄存器工作集也属于不同口径。（Appendix A, p. 18）

三：实验——分别检验容量、语言建模和资源分配

3.1 独立容量实验：MQAR

多查询关联召回（Multi-Query Associative Recall, MQAR）先提供 $N$ 个随机 key–value 对，再以随机顺序查询全部 key。每个 key–value 对被打包为一个 token，因此总序列长为 $2N$；查询位置只读、不更新状态。key 与 value 分别从两个互不重叠的 8K 词表中无放回采样。

模型只有两层、4 个头，每头第一 key/value 维度为 16，模型宽度 128，MLP 宽度 384。序列混合器去掉卷积、遗忘、delta rule 和额外激活，仅保留基本外积机制及归一化；模型仍有 SwiGLU MLP。扫描 $N=32,64,\ldots,4096$ 和 $E=1,2,4,8,16$，每个组合训练 25 个种子。每次训练 50K steps，每步约 250K tokens，报告 3K 条留出序列上的平均正确率。（Sec. 2.3, p. 3；Appendix C, p. 19）

3.2 语言模型与训练协议

预训练使用 FineWeb-Edu；扩展阶段混合 FineWeb-Edu、PG19 和科学 PDF 数据。优化器是 AdamW，两阶段峰值学习率分别为 $3\times10^{-4}$ 和 $10^{-4}$，weight decay 为 0.1，采用线性 warmup 后余弦衰减至峰值的 10%。所有语言任务结果来自长上下文扩展后的最终 checkpoint。（Sec. 3.1, pp. 5–6）

对照 Transformer 使用 GQA-8（每 8 个 query head 共享一个 KV head）、QK-norm 和 RoPE；RoPE base 从预训练的 10K 提高到扩展时的 2M。混合模型另用 NoPE，不能把两个设置混为一谈。各模型统一使用 32K 词表的 Llama-2 tokenizer。（Appendix D, pp. 19–20；Sec. 3.4, p. 8）

400M 的主对照与混合结果是三次训练种子的均值；Table 2 明确 400M 为三种子均值、1.3B 为单种子结果。Table 4 的激活消融行与 Table 2 的主实验行数值略有不同，应在各自表内比较，不把两表当成同一次运行。

3.3 指标口径

Zero-shot 其余七项为 LAMBADA、PIQA、ARC-Easy、WinoGrande、SciQ、BoolQ、COPA。Figure 3 只展示四项代表性 NIAH：单 needle、多 key、多 query、同 key 多 value。Recall 均分、NIAH 均分与逐位置预测损失回答的是不同问题。（Appendix D, p. 20）

四：结果——更大记忆在哪里有效，代价是什么

4.1 简化 MQAR 支持容量随第二 key 维度扩大

原图 Figure 1，PDF p. 3。横轴是 key–value 对数，按倍增取点；纵轴是正确率；从深红到黄色表示 $E$ 从 1 增到 16。随着 $E$ 增大，准确率下降的位置明显右移。作者概括为每次把 $E$ 翻倍，可记住的关联数量大致翻倍；$E=16$ 的非 embedding 参数仅约为基线的 1.08 倍。

这个实验的价值是把复杂门控和 delta rule 去掉后，仍能看到容量提升。曲线同时显示饱和边界：$E=16$ 到 2K、4K 关联时正确率也很低。

4.2 扩大状态改善长文预测和平均召回

原图 Figure 2 左下，PDF p. 5。纵轴是模型损失减去 Transformer 损失，负值更好；横轴是上下文位置。普通 GDN 在短上下文有优势，但约 10K 后开始落后；$E=8$ 到 64K 仍低于 Transformer。竖向虚线表示 Transformer 的 KV cache 超过相应递归状态的长度，不是准确率曲线的交叉点。

以下为精确表格数据。PG19 一列均指 16K–64K token 位置区间，Recall 单位为百分数；状态按全模型、每元素 2 字节计量。

来源：Table 1, p. 6；Table 2, p. 7。400M 的 $E=8$ 相对基线增加约 1.18% 参数，Recall 增加 6.9 个百分点，PG19 PPL 相对下降约 2.90%。1.3B 对应参数增加约 1.33%，Recall 增加 8.3 个百分点，PG19 PPL 相对下降约 2.55%。这些比例由表内数值计算。400M 的 PG19 收益随 $E$ 增长逐步收窄：从 $E=4$ 到 8 只再下降 0.05 PPL。

原图 Figure 2 右下，PDF p. 5。横轴列出六任务与均值；纵轴是正确率，彩色柱对应不同 $E$，空心柱为 Transformer。FDA 与 SWDE 的提升最突出，作者把它们与长文复制需求联系起来。并非每个任务都随 $E$ 单调改善，例如图中的 DROP、NQ 和 SQuAD 就存在波动。平均召回改善也不代表赶上 Transformer；400M 的 $E=8$ 均分为 33.1，Transformer 为 41.6。

4.3 等状态、近似等参数：乘积地址优于直接加宽

400M GDN 的 4 倍状态比较如下。四行状态均为 25.2 MB；参数在 380.6M–383.0M 之间。

来源：Table 1, p. 6。Triadic 在这些 PPL 与 Recall 指标上均领先；2 倍状态比较也呈现同一趋势。4 倍“更多 heads”没有进入表中，因为仅序列混合器的投影就会超过参数预算。

这组结果同时反映地址参数化和全模型预算分配：Triadic 保留了完整 MLP，而宽 value / 多 value 对照必须大幅压缩 MLP。因此，结论是“在给定参数预算下怎样扩状态更有效”。更大 head 的对照不需要同样增大投影宽度，也仍逊于 Triadic，说明优势并不能全部归因于 MLP 被削减。

sGLA 的结论与 GDN 相近：$E=4$ 的 PG19 长区间 PPL 为 13.97，其他 4 倍状态对照为 14.52–14.59；Recall 为 30.7，对照为 28.3–29.0。一个细节例外是 2 倍状态下 grouped values 的 sGLA Recall 为 28.5，Triadic 为 28.4；因此不宜把结果写成所有模型、所有指标全面领先。

4.4 NIAH：扩容延迟检索崩溃，64K 仍有明显缺口

原图 Figure 3，PDF p. 7。四个面板依次为单 needle、多 key、多 query、多 value；横轴为上下文长度，纵轴为正确率，虚线是 Transformer。增大 $E$ 能把准确率下降的区间推后，尤其单 needle 在中等长度处获益明显。然而到 64K，四项任务上的纯 Triadic 曲线仍处于低位，Transformer 保持明显优势。

这一结果解释了为何“64K 书籍困惑度优于 Transformer”和“64K 精确检索弱于 Transformer”能够同时成立：语言预测平均损失可从压缩后的统计与语义信息获益，NIAH 则要求找回具体记录。后者仍明显受到有限地址容量与干扰影响。

4.5 Upcycling：扩容可以推迟到长上下文训练阶段

作者从已预训练的 $E=1$ 模型出发，把遗忘门复制到 8 个切片，随机初始化新增第二 key/query 投影及卷积，再进行长上下文扩展。此过程需要继续训练来适配新地址；论文没有把它描述成无需训练的等价变换。（Sec. 3.3, pp. 7–8）

来源：Table 2, p. 7。PG19 仍为 16K–64K 区间；扩容后与从头 Triadic 均为 $E=8$。例如 400M GDN 的 PG19 改善恢复率为 $(14.15-13.86)/(14.15-13.74)\approx71\%$，Recall 改善恢复率约为 51%；1.3B GDN 对应约 77% 与 71%。这支持在主要短上下文预训练结束后再扩大状态的路线，效果仍低于全程使用三维状态的对应模型。

4.6 混合模型：扩线性状态比加倍 KV 缓存更划算

作者在 400M、3:1 GDN/softmax 混合架构上比较三种方案。GQA-8 每 8 个 query head 共享一个 KV head，改成 GQA-4 会把 KV head 数量加倍；Triadic 方案则保留 GQA-8，并把 GDN 状态扩大到 $E=4$。

来源：Table 3, p. 8；三种子均值。此表 PG19 区间合并为 4K–64K，与前面的 16K–64K 不同。Triadic 的 Wiki PPL 10.63 和零样本均分 53.8 也最好。它在 64K 时比更大 KV 方案少用约 45.9% 的状态内存，同时 NIAH 高 3.3 个百分点；但 Recall 低 0.8 个百分点，作者称该差距在种子波动范围内。

在约 4.6K 之后，更大 KV 方案才比 Triadic 方案占用更多记忆；4K 时 Triadic 的表内状态还略大。混合模型保留 softmax 层，因此总状态仍随序列长度增长。这里给出的是递归状态与 KV cache 的计量，不包含参数、优化器及全部训练激活。

4.7 实测训练成本：相对 GDN 变慢，相对长序列 Transformer 变快

原图 Figure 4，PDF p. 8。横轴为序列长度，纵轴为每条序列耗时（ms），插图放大 2K–8K 区间。测试对象是 1.3B 模型的一个 block 的前向加反向，batch size 为 4，硬件为单张 H100。Triadic 使用 CuTe 内核，GDN 使用 FlashQLA 的 TileLang 内核，Transformer 使用 FlexAttention。

64K 时，$E=8$ 的该项测试比 Transformer 快约 5.1 倍。这是给定硬件、内核和单块前反向协议下的结果；不能直接当作完整训练吞吐、端到端服务吞吐或逐 token 解码加速。论文也没有通过这个图验证其他 GPU 上相同的速度比。

4.8 消融：总地址维度相同，如何分配仍重要

Table 4 上半部分固定联合 key 维度 $d_kE=1024$、value 维度和头数，每个头只有一个共享标量遗忘门。投影缩小后增加 MLP 宽度以匹配参数预算。

来源：Table 4, p. 9。相同元素个数并不足以决定性能。更均衡的因子维度可以节省投影参数，却使困惑度略差；这一对照还包含 MLP 宽度的配套变化。共享标量门的 128×8 行在表内表现很好，但论文没有把这些行组织为带不确定性估计的“共享门与逐切片门”专项比较，不能据此断言逐切片遗忘必然更好或更差。

Table 4 下半部分比较第二 key/query 的激活函数：

非负激活在 PPL 上较好。作者的解释是：允许负值时，同一 token 在不同切片上的读写贡献可能异号，在聚合时互相抵消。Recall 没有完全遵循同一排序，所以这个解释应与其支持最直接的语言建模结果联系起来。

4.9 机制分析：更多状态如何被使用

原图 Figure 5，PDF p. 18；相关文字延续至 p. 19。三个面板分别检验：

1. 同一预测目标能否利用更多远程上下文。 固定每本书最后 4K tokens 为预测目标，改变此前读入的上下文。左图以 4K 上下文下的损失为基准，越负越好。普通 GDN 约在 16K 开始饱和，Triadic 到 32K 仍持续获益。固定目标避免把“位置不同、目标不同”误当成上下文利用。
2. 写入强度如何变化。 中图横轴为 $\beta$，纵轴为分布密度；均值从 0.31 增到 0.50，分布也更宽。更大状态对应更强、更多样的更新。
3. 切片是否分工。 右图在每个 head 内按遗忘半衰期排序，纵轴使用对数尺度。八个排序位置的中位半衰期从 0.20 到 1080 tokens，普通 GDN 的单门中位半衰期为 5.1 tokens。小于一个 token 表示单步衰减已经超过一半。

这些观察支持“不同切片形成不同时间尺度”的解释。右图度量的是遗忘门导致的衰减，完整状态还受 delta 擦写影响，因此半衰期不能直接等同于具体事实的可召回时长。（Appendix B, pp. 18–19；最后一点为基于更新公式的分析）

五：结论与研究启发

本文把两个短地址向量组合成更大的联合地址空间，用很少的新增参数扩大稠密递归记忆。它的实证优势同时依赖可训练的状态管理与利用结构的内核：两组相似度相乘降低块内计算，value 维分块使大状态适配 GPU，联合 delta 更新与遗忘门负责管理记忆内容。

与已有工作的关系也很明确。高阶关联记忆和三阶张量积表示并非本文首次提出；Sec. 5 特别指出 Schlag 与 Schmidhuber 的 2018 年工作已使用三阶状态和图遍历操作。本文把相近构造接入现代 GDN/sGLA、多头架构与可扩展内核，并在 400M、1.3B 语言模型上验证，创新重点在这一组合及其成本收益。（Sec. 5, pp. 9–10）

对后续架构研究，有三点值得保留：

• 比较记忆机制时应同时看参数、状态容量与实际计算。仅匹配模型参数，可能把更大的快权重容量带来的收益误认为更新规则本身的收益；本文关于较大 TTT 状态的讨论也是这一层面的假设，并未直接实验分解各 TTT 方法的贡献。
• 状态大小可以随训练阶段改变。Upcycling 支持短上下文阶段使用小状态、长上下文阶段再扩容，但新增地址需要适配，达到全程三维预训练的效果仍有差距。
• 长文本语言建模和精确记录检索需要分别评估。Triadic 的容量收益很清楚，纯模型在最长 NIAH 上的缺口也很清楚；混合模型的结果提示，增大递归记忆与保留部分精确 KV 存储可以互补。

实际取舍是明确的：$E=8$ 约增加 1.2%–1.3% 参数，却把递归状态扩大 8 倍，并使所测单块训练耗时增加约 30%。这些代价换来了更好的语言建模和召回，但内核测试范围、最长检索表现，以及等参数对照中的 MLP 预算调整，都会影响如何把结果用于新的模型设计。