Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
Oliver Sieberling, Bharat Runwal, David Jin, Ryan Chin, Rameswar Panda, Yoon Kim
cs.LG, cs.CL
2026-09-29
把 key、第二 key、value 的三重外积写进 3D 张量状态,E=8 只加 1.2% 参数,长上下文与 recall 全面胜过同状态量基线。
线性注意力把 softmax 换成外积读写,整个历史被压进一个固定大小的矩阵状态,推理耗时不再随上下文增长。代价是容量:d×d 的状态最多分开 d 个相互正交的 key,已有结论是固定状态的模型连稍长一点的序列都复制不动。Gated DeltaNet、Mamba-2 这一代靠遗忘门和 delta 规则把状态管得更好,底子仍是那个 d² 矩阵,recall 密集任务上明显输给 Transformer。
直接扩状态要付参数税。加 head、加宽 value 都会放大 sequence mixer 的全部投影,论文里 4× 状态的两个基线变体,参数多到要把 MLP 宽度从 2816 砍到 640 才能对齐预算。问题于是变成:能否几乎不花参数把状态撑大。
出发点是一个观察:线性注意力的矩阵状态,就是经典 RNN 向量状态被 key⊗value 外积抬了一阶,这个构造可以继续往上推。
工程上两步保住训练效率:联合 key 的 Kronecker 结构让 chunk 内注意力代价停在 C²(d+E),而非 C²·dE;状态沿 value 轴切成 32 列一块,常驻各 thread block 的寄存器,任何 SM 都不必装下完整状态(E=8 时单 head 状态 512 KiB,已是 Hopper SM 寄存器文件的两倍)。
消融还固定了一个细节:第二 key 与 query 必须过非负激活(softplus 或 sigmoid),带符号的版本会让同一 token 以相反符号读写不同 slice,聚合时相互抵消。
| 400M 模型 | WikiText ↓ | PG19 16k–64k ↓ | Recall ↑ |
| Transformer | 11.11 | 13.92 | 41.6 |
| GDN 基线 | 11.25 | 14.15 | 26.2 |
| 4× 状态:更大 head(d=512) | 11.25 | 14.24 | 28.5 |
| 4× 状态:grouped values | 11.35 | 14.36 | 27.3 |
| 4× 状态:Triadic(E=4) | 10.93 | 13.79 | 31.1 |
同参数量、同状态量下,triadic 在所有 PG19 区间与 recall 上都最优;4× 状态时三种常规做法相比基线不升反降。合成任务 MQAR 上,E=16 存下约 16 倍键值关联,非嵌入参数只多 1.08 倍;放到语言模型,E=8 只多 1.2% 参数。长上下文上,E=8 的 triadic GDN 在 64k 全程的 next-token 预测都好于 Transformer,而 vanilla GDN 约 10k 处就被反超;1.3B 规模 recall 从 36.1 升到 44.4。训好的 GDN 还能 upcycle:长上下文扩展阶段把 E 从 1 扩到 8,拿回从头训练一半到四分之三的收益(1.3B:42.0 对 44.4)。对学到门的分析显示,八个 slice 的半衰期从 0.2 token 铺到 1080 token,而 GDN 单门中位数只有 5.1 token,新轴确实被用来装不同时间尺度的记忆。速度上,1.3B 单卡 H100,E=8 比 vanilla GDN 慢 28–30%,4k 处比 Transformer 慢 3%,64k 处快 5.1 倍。
对做线性注意力与混合架构(Kimi Linear、Qwen3 这类 GDN 加 attention 的分层设计)的团队,这给了一个几乎零参数代价的状态扩容旋钮,内核开源。混合实验直接回答了一个实际问题:同样预算下,把 GDN 层状态扩 4 倍比把 attention 层 KV head 翻倍更划算,困惑度、zero-shot、大海捞针全面更好(56.3 对 53.0),64k 时显存还省近一半。upcycling 给出一条训练路径:短上下文预训练用小状态,长上下文扩展时再撑大。作者顺带提出一个更远的解释:TTT(test-time training)变体的部分收益,可能只是来自更大的状态容量。规模停在 1.3B、65B token,属于架构研究,但方向清楚、代码可复现。
作者自认三条:E=8 训练慢约 30%;部分 recall 任务仍输 Transformer(400M 上 41.6 对 33.1,差距不小);只在 GDN 与 sGLA 两个 mixer 上验证过。另有两点论文没有覆盖。其一,只报了训练耗时,没报解码吞吐:400M 模型状态从 6.3 MB 涨到 50.3 MB,每步状态读写带宽同比变大,decode 场景的实际代价未量化。其二,语言建模收益随状态增大递减,recall 收益能否带到更大规模,1.3B 且单 seed 的证据还撑不住。