三元线性注意力：面向长上下文序列建模的三维循环状态 Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
把线性注意力的矩阵状态升级为三维张量状态，以极小参数代价将记忆容量放大E倍
前置知识
线性注意力与矩阵状态
线性注意力（Katharopoulos et al., 2020）把 RNN 的向量隐状态升级为矩阵状态 $S_t$：每步用键值外积做秩一更新 $S_t = S_{t-1} + k_t v_t^{\top}$，读取时算 $o_t = S_t^{\top} q_t$。历史被压进固定大小状态，推理每步恒定代价，训练可用分块并行；等价于 Schmidhuber 的快速权重编程。
三元线性注意力的全部改动都发生在这个矩阵状态上——理解其写入与读取方式，才能理解三维张量状态如何对它做严格泛化。
张量积绑定与高阶联想记忆
Smolensky（1990）的张量积表示用外积把填充物绑定到角色：$k \otimes v$ 展开后是 $d^2$ 维绑定向量，多个绑定相加成叠加态，用匹配角色的向量收缩即可解绑。$n$ 个向量的外积给出 $n$ 阶张量记忆，角色正交时可精确存储 $d^{n-1}$ 个关联（Baldi & Venkatesh, 1987）。
三元构造的理论依据正是『把绑定本身再绑定一次』：线性注意力是二阶（dyadic），加入第二键即升为三阶（triadic），容量从 $d$ 跃升到 $d^2$。
Delta 规则与 Gated DeltaNet
Delta 规则在写入新关联前，先按比例擦除状态中该键已存的内容：$S_t = (I - \beta_t k_t k_t^{\top}) S_{t-1} + \beta_t k_t v_t^{\top}$，$\beta_t \in (0,1)$ 为数据依赖写强度。Gated DeltaNet 在此之上加数据依赖遗忘门 $\alpha_t$，让更新从纯累加变为带擦除的插值式更新。
论文必须证明三元状态同样兼容遗忘门与 delta 规则，这是该构造从玩具想法走向现代可用架构的关键一步。
分块并行（chunkwise）训练
把序列切成 $C=64$ 的块：块内交互用带因果掩码的注意力一次并行算完（$O = Q S^{[i]} + \mathrm{tril}(QK^{\top})V$），块间只循环传递边界状态 $S^{[i+1]} = S^{[i]} + K^{\top}V$。既保住 RNN 的常数状态，又吃满 GPU tensor core 的并行度。
三维状态若按 $d \cdot E$ 维联合键直接分块，计算与状态都贵 $E$ 倍且放不进片上内存；论文靠克罗内克积分解与状态平铺解决，是全文工程上最微妙处。
MQAR 多查询联想召回
MQAR（Arora et al., 2023）先给 $N$ 个键值对，再随机顺序重列这 $N$ 个键，要求模型逐一预测对应值。它不依赖门控或非线性，纯考联想记忆容量：$N$ 超过状态容量后精度骤降，精度-规模曲线就是容量的直接量度。
论文用 MQAR 曲线随 $E$ 右移来验证『$E$ 倍状态 = $E$ 倍容量』这一核心主张，是读懂 Figure 1 的前提。
研究动机
现代线性注意力把全部历史压进一个固定大小的 $d \times d$ 矩阵状态，而状态大小硬性决定能精确分离的键数：$d$ 维键最多容纳 $d$ 个相互正交的键，超出后干扰不可避免。后果非常直观——固定大小的状态连超过一定长度的复制任务都无法完成（Jelassi et al., 2024），召回密集任务上线性注意力明显落后：本文 Table 1 中，380M 参数下 Gated DeltaNet 在六个召回任务上的平均分只有 26.2，同规模 Transformer 为 41.6；PG19 16k–64k 段困惑度 GDN 为 14.15，逊于 Transformer 的 13.92，且交叉点出现在约 10k token 处。想扩状态却很贵：加头要放大全部投影，加值维度或每键挂多个值头会显著增大值投影与输出投影。论文实验显示，为在 4 倍状态下保持参数匹配，部分基线被迫把 MLP 宽度从 2816 压到 640，结果 PG19 各段困惑度反而全面退化。
本文的目标是本文的目标是找到一种参数高效的方式扩大线性注意力的状态容量：状态扩大 $E$ 倍时几乎只增加状态本身，而参数量几乎不变；同时不牺牲现代线性注意力的三大支柱——数据依赖遗忘、delta 规则和分块并行训练；并且配套高效 GPU 内核，让容量放大不以可用性为代价。最终在 400M 与 1.3B 参数的语言建模上验证：更大的状态能否带来长上下文困惑度与召回能力的实质提升，并在同等推理内存下超越其他扩容方案、尽可能逼近甚至局部超过 Transformer。
与已有工作不同的是，本文的独特切入是把扩容问题翻译成表示理论的语言：线性注意力的矩阵状态本质上是 Smolensky 张量积表示中的二阶（dyadic）绑定，而绑定本身是一个 $d^2$ 维向量，可以继续绑定到第二个角色，从而得到三阶张量状态。最接近的前作 Schlag & Schmidhuber（2018）把三阶状态用于图遍历的组合推理，其更新含写/移动/回链三个图专用操作；本文则剥掉所有领域特化部件，用纯外积写入加双查询读取来放大现代线性注意力层的状态容量，并配上让单个流式多处理器永不持有完整三维状态的 GPU 内核，使这一构造能以约 1.2% 的参数开销训练到 1.3B 规模。
核心方法
直觉上，三元线性注意力只是让每个 token 在产出 $q_t, k_t, v_t \in \mathbb{R}^d$ 之外，再用两个小线性投影额外产出第二键 $k'_t$ 和第二查询 $q'_t \in \mathbb{R}^E$；写入时把三个向量做三元外积存进 $d \times E \times d$ 的三阶张量 $S_t = S_{t-1} + k_t \otimes k'_t \otimes v_t$，读取时用两个查询分别收缩两条键轴 $o_t = S_t \times_1 q_t \times_2 q'_t$。$E$ 维第二键让状态从 $d^2$ 变成 $d^2 E$ 个条目——$E=8$ 时每头状态放大 8 倍，参数只多两个投影。技术路线分三步兼容现代组件：给每个第二键切片配独立遗忘门实现数据依赖遗忘；把 delta 规则推广为按两键联合读出旧值再插值擦除；利用联合键 $\kappa_t = k_t \otimes k'_t$ 的克罗内克结构把块内注意力分解为 $d$ 维键的内积与 $C \times C$ 的第二键矩阵，并沿值轴平铺状态，使任何流式多处理器都不必持有完整的三维状态。
核心创新在于指出：以往扩大状态的所有手段（更大头、更宽值、更多头、分组值）本质上都在放大投影矩阵，状态/参数比不会改善，参数预算一大还会挤占 MLP；而三元构造让状态/参数比随扩容渐近上升——因为按照张量积表示，绑定本身（$k_t \otimes v_t$，$d^2$ 维）就是向量，可以再绑定一次角色，第三阶只需要 $k'$ 与 $q'$ 两个新投影。与 Schlag & Schmidhuber（2018）的本质区别在目的与用法：他们用三阶状态编码图结构、按节点-关系对解绑检索；本文的第二键不对应数据中任何实体，模型自由学习如何在两条键轴间分配存储，把三阶状态当作通用联想记忆扩容器。取 $E=1$ 且 $k'_t = q'_t = 1$ 时严格退化为普通线性注意力，因此它是一个真正的严格泛化，而非另起炉灶的新架构。
方法步骤详情
第一步，投影：一个 GEMM 同时产出 $q,k,v \in \mathbb{R}^d$ 与 $q',k' \in \mathbb{R}^E$，一个因果深度卷积统一处理，$q,k,v$ 用 SiLU、$k',q'$ 用 softplus，随后各自 $l_2$ 归一化。第二步，遗忘：每切片 $S_t[:,e,:]$ 配独立标量门 $\alpha_{t,e}$，先做 $S_t = S_{t-1} \times_2 \mathrm{diag}(\alpha_t)$。第三步，写入：GDN 版执行 $S_t \leftarrow S_t + \beta_t k_t \otimes k'_t \otimes v_t - S_t \times_1 k_t \times_2 k'_t$（联合两键读出旧值再插值擦除），sGLA 版只累加。第四步，读取：$o_t = S_t \times_1 q_t \times_2 q'_t$，经输出门控与 RMSNorm 后拼接各头。第五步，分块并行：块长 $C=64$，联合键内积因子化为 $(q_r^{\top}k_s)(q_r'^{\top}k'_s)$，衰减折进 $R = \sum_e k'_e k_e'^{\top} \odot \Gamma_e$，全部切片共享同一 UT 变换 $T$。第六步，内核：状态沿值轴按 32 列分块给线程块、$E$ 个切片驻留寄存器；衰减因子按累计对数门之差的指数直接计算，无溢出、无需 clamp。
技术新颖性
技术新颖性有四层。其一，概念层面：把 dyadic 到 triadic 的升级看作线性注意力当年把经典 RNN 向量状态升级为矩阵状态的同一操作的重演，容量按 $d^{n-1}$ 律从 $d$ 升到 $d^2$，使『扩状态』第一次有了参数效率的阶梯，而非靠堆投影。其二，算法层面：三维状态与遗忘门、delta 规则、分块形式全部兼容——联合键视角说明它等价于键维 $d \cdot E$、每切片标量衰减的 Gated DeltaNet，但内积因子化让实际计算远便宜于 $d \cdot E$ 维键。其三，系统层面：状态沿值轴平铺（每线程块 32 列、$E$ 个切片驻留寄存器）绕开片上内存限制；用 CuTe DSL 写的 warp 专业化内核保证训练时间不超过 GDN 的 1.3 倍（$E=8$）。其四，应用层面：提出 upcycling——把预训练 $E=1$ 模型在长上下文续训阶段扩展为 $E=8$，类比 Transformer 的 KV 缓存随上下文增长，恢复从头训练一半以上的收益。
实验结果
实验链条完整且相互印证。(1) 容量测试（Figure 1）：$E$ 每翻倍 MQAR 精度曲线右移约一倍 $N$，$E{=}16$ 存储量约 16 倍于普通线性注意力，非嵌入参数仅增 1.08 倍。(2) 语言建模（Figure 2）：GDN 在约 10k token 处被 Transformer 反超，$E{=}8$ 在 64k 处仍更优，且此时 Transformer KV 缓存早已超过循环状态；召回增益集中在需长文档复制的 FDA 与 SWDE。(3) 状态匹配对比（Table 1）：4× 状态下 Triadic GDN ($E{=}4$) 全面最佳——Wiki 10.93、PG19 三段 14.80/14.08/13.79、Recall 31.1；替代基线 2× 时几乎无提升，4× 时 PG19 全段退化。(4) Upcycling（Table 2）：400M 召回 26.2→29.7（从头 33.1），1.3B 36.1→42.0（从头 44.4）。(5) 混合架构（Table 3）：Triadic ($E{=}4$) 困惑度与 NIAH（56.3）最佳，64k 状态 220MB 不到 GQA-4（407MB）六成。(6) 效率（Figure 4）：$E{=}8$ 比 GDN 慢 28–30%，但 64k 处比 Transformer 快 5.1 倍。(7) 消融与分析：$d_k{=}128, E{=}8$ 最优、非负激活更优（Table 4）；$\beta$ 均值 0.31→0.50、切片半衰期 0.20–1080 token（Figure 5）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MQAR 联想召回（d=16 最小模型，N 取 32–4096） | 可精确存储的关联数 N（精度曲线拐点） | E=16 时容量约为基线的 16 倍 | 普通线性注意力（E=1） | 容量 ×16，非嵌入参数仅 ×1.08 |
| WikiText-2 困惑度（400M，4× 状态） | 困惑度（越低越好） | 10.93（Triadic GDN, E=4） | GDN base 11.25；四种扩容基线均不低于 11.18 | -0.32，且在 2× 与 4× 两档均全面领先 |
| PG19 16k–64k 困惑度（400M，4× 状态） | 困惑度 | 13.79 | GDN base 14.15；Transformer 13.92 | 较 GDN -0.36，长上下文段反超 Transformer |
| 召回套件平均（DROP/NQ/TriviaQA/FDA/SQuAD/SWDE） | 平均准确率 | 400M：31.1（E=4）；1.3B：44.4（E=8） | GDN base：400M 26.2；1.3B 36.1；Transformer 400M 41.6 | 400M +4.9；1.3B +8.3，但仍未达 Transformer 水平 |
| 3:1 GDN/GQA-8 混合模型 RULER NIAH 平均（400M） | 八任务平均准确率 | 56.3（Triadic E=4，状态 220MB@64k） | 基线混合模型 51.5；加大 KV 缓存（GQA-4）53.0（407MB@64k） | +4.8，且长上下文内存不到一半（约 4.6k token 后反超） |
| 单块训练耗时（1.3B，batch 4，单张 H100，2k–64k） | 前向+反向时间相对 Transformer | E=8 在 64k 处快 5.1 倍（4k 处仅慢 3%） | Transformer（FlexAttention） | 长上下文大幅加速，代价是相对 vanilla GDN 慢 28–30% |
局限与改进
作者明确承认三点局限：其一，更大的状态拖慢训练——即使有定制内核，$E=4$ 约慢 15%、$E=8$ 约慢 30%；其二，部分召回密集任务仍落后 Transformer，尽管循环状态在长上下文中远小于 KV 缓存——Table 1 中 400M Triadic 召回 31.1 对 Transformer 41.6，差距仍可观；其三，构造只应用于 GDN 与 sGLA 两种混合器，仅覆盖线性注意力文献的一小角。我自己的观察补充如下：收益随 $E$ 增大递减（Figure 2 中 $E$ 从 4 到 8 的提升小于从 2 到 4），提示三阶状态也有自己的容量天花板；upcycling 只能恢复一半到四分之三收益，说明第二键切片的分工在预训练早期就已成型，后期难以完全补齐；非负激活显著重要，但作者只给出『正负切片贡献在读出时相消』的猜想性解释，缺乏理论刻画；此外所有语言建模实验都在不超过 1.3B 参数、20B/65B token 的规模进行，更大规模与 MoE 场景未验证，$E$ 与 $d$ 的最优配比是否随规模漂移亦属未知。
独立分析的弱点
第一，训练开销：$E=8$ 慢 28–30% 主要来自状态搬运与 $E$ 份矩阵乘，可尝试 BF16/FP8 状态精度、切片间更深的流水线化，或让不同切片共享部分计算。第二，召回绝对水平：31.1 对 41.6 说明三阶稠密状态仍不足以完全替代精确检索，实践中应与少量 softmax 层或稀疏记忆路由（如 Mixture-of-Memories）组合——论文的混合实验（Table 3）已证明这种组合的收益可以叠加。第三，第二键语义不明：模型如何在两键轴间分配存储完全隐式，建议增加探针实验或辅助正则（例如鼓励不同切片对应不同时间尺度，Figure 5 右已观察到这种自发分工），以提升可解释性与可控性。第四，密集扩容的推理内存代价：$d^2 E$ 状态在推理时也要常驻，若与乘积键寻址结合做稀疏化，可能以更低内存达到同等容量。第五，激活函数选择完全经验化：sigmoid 与 softplus 的召回差异（33.3 对 32.2）暗示门控形状尚有优化空间，可在 $E$ 轴引入逐元素（而非逐切片）门控做进一步消融。
未来方向
作者提出的方向：把三阶构造推广到矩阵状态文献的其余进展，如恒等加低秩转移的推广形式与更灵活的参数化，以及测试时训练（TTT）路线——并指出部分 TTT 变体的收益可能正是来自其更大的状态，三元构造提供了一个参数高效的替代；额外键轴还可能催生全新变体；最终目标是让纯线性注意力模型站上前沿。基于本文成果可自然延伸：一是课程式状态增长——先以 $E=1$ 在短上下文预训练、长上下文阶段 upcycle 到更大 $E$，Table 2 已验证可行，值得系统化为标准训练配方；二是与稀疏记忆（product-key 寻址、memory slots）结合，用三元状态做稠密底座、稀疏路由做超大容量补充；三是探索非均匀因子化——消融显示 $d_k$ 与 $E$ 越对称性能越差，说明最优比例值得自动搜索；四是把切片半衰期谱（0.20–1080 token）作为结构先验，显式设计多时间尺度门控；五是在 MoE 与 7B 以上规模验证状态/参数比优势是否进一步放大。
复现评估
复现条件较好。代码已开源（github.com/OliverSieberling/TriadicLinearAttention）；附录 A 给出含遗忘与 delta 规则的完整分块公式、数值稳定性推导与内核细节（CuTe DSL、warp 专业化、TMA 双缓冲、固定归约顺序、梯度位级可复现），并在 1.3B 训练形状下对照 FP64 参考验证全部七个输入的输出与梯度（相对 ℓ2 误差低于 1%）；附录 C/D 给出 MQAR 与语言建模的全部超参（24 层、峰值学习率 3e-4 与 1e-4、批量 0.5M/1M token、$C=64$、$E \in \{1,2,4,8\}$ 等）。数据全部公开：Fineweb-Edu、PG19、WikiText-2、lm-eval 十任务、RULER 与召回套件。算力未披露总 GPU 时，但 400M/20B 与 1.3B/65B token 属学术实验室可负担规模；主要门槛在 Hopper GPU 与 CUTLASS/CuTe 工程能力。MQAR 容量实验（3.5M 参数）单卡可复现。总体：实验复现中等，内核复现中高难度。
论文图表