Triadic Linear Attention：面向长上下文序列建模的三维循环状态

TL;DR Triadic Linear Attention 将线性注意力的矩阵状态扩展为三维张量状态，通过引入一个 E 维 second key 仅增加两个投影，换来 E 倍状态容量，并兼容遗忘与 delta rule，在 Gated DeltaNet 上显著提升长上下文建模。

问题

问题背景

近年来，RNN 和线性注意力因推理时恒定时间复杂度而重新受到关注。线性注意力 将普通 RNN 的向量隐藏状态扩展为矩阵状态，显著提升了记忆容量，而状态大小已被证明是长上下文建模性能的关键因素。

现有方法局限

线性注意力的矩阵状态通常形如 S_t = S_{t-1} + k_t v_t^T，其容量受限于 head_dim × head_dim。要扩大状态，常见做法是增大 head_dim 或堆叠多头，但这会带来参数量和计算量的线性甚至平方级增长，且写入操作仍是秩-1更新，记忆效率并未突破。已有工作尝试用门控或 delta rule 改进写入质量，但状态张量维度始终停留在二阶，缺乏在参数高效前提下直接提升状态阶数的方案。直接使用高阶张量又面临维度爆炸和训练并行性难题。

为什么难/重要

核心挑战在于如何在几乎不增加参数和计算开销的条件下，将状态从矩阵扩展为三阶张量。Triadic Linear Attention 通过引入第二个 key，利用三元外积将状态写成 k_t ⊗ k'_t ⊗ v_t，一个 E 维的 k'_t 即可让状态容量扩大 E 倍，仅需额外两个投影矩阵。同时该方法兼容数据依赖遗忘、delta rule 和 chunkwise 并行训练，突破了记忆容量与参数效率之间的权衡。在长上下文语言建模和召回任务中，固定内存预算下的有效容量直接决定模型能否利用远距离信息，因此该方向对实际工程部署有显著价值。

行业类比

类似 Transformer 推理时对 KV cache 的压缩：用更高阶张量在相同参数预算下换取更大记忆容量，如同用多维索引替代一维索引来提升检索召回率。

核心洞察

• 三元线性注意力通过引入第二个 key 的外积，将状态从矩阵扩展为三阶张量，在不增加大量参数的前提下实现状态容量 E 倍增长。传统线性注意力使用 key⊗value 外积形成矩阵状态，若要扩大状态尺寸（如增加 head dim 或 value dim）会导致参数和计算量平方级增长；而 Triadic Linear Attention 只需增加一个 E 维第二 key 投影，参数开销仅为两个线性层，却将状态规模从 D×D 提升到 D×E×D。这种参数高效的状态扩展策略为 RNN 类模型在资源受限下提升长上下文能力提供了新思路。
• 该工作将二阶联想记忆推广到三阶，读取时用两个 query 与状态张量进行双轴收缩，实质上是将查询分解为两个独立的检索模式，从而增强了对相似 key 的区分能力。与仅使用单个 query 的矩阵状态线性注意力相比，三阶张量状态能够编码 key1、key2 与 value 之间的三元关联，当输入序列中存在大量重复或相似 token 时，这种高阶关联可有效减少记忆单元之间的干扰，提升回忆精度。作者在 MQAR 等任务上验证了该机制对长程召回的一致改善。
• Triadic Linear Attention 不是孤立的设计，它与数据依赖遗忘、delta rule 以及 chunkwise-parallel 训练完全兼容，可无缝集成到 Gated DeltaNet 等现有高效 RNN 架构中。这一特点意味着实际部署时无需重写训练内核或牺牲线性注意力的线性复杂度优势，即可获得状态容量和回忆能力的提升。作者通过 continued pretraining 实验展示了对已有模型的低成本升级路径，突显了该方案在工业应用中的可落地性。

方法

输入与投影

序列 token 经线性变换得到 key1、key2、value、query1、query2，其中部分投影可与基线共享，仅新增两个投影即可引入第二 key 轴。

关键模块：三阶张量状态

• 将线性注意力的矩阵状态推广为 三阶张量状态：写入时计算三个向量的外积 key1 ⊗ key2 ⊗ value，累加到张量中；读取时用 query1 和 query2 分别收缩两个 key 轴，得到输出。
• 若第二 key 维度为 E，状态大小增加 E 倍，但额外参数仅来自两个投影层，参数效率高于直接扩大矩阵状态或增加 head 维度。
• 状态更新兼容 数据依赖遗忘 和 delta rule，可直接嵌入 Gated DeltaNet 或 scalar-gated linear attention，无需改变门控机制。

高效实现

• 分离联合 key 投影：将 key2 的生成与 key1 解耦，避免显式构造大外积，降低计算与内存开销。
• Tiling the state：对张量状态分块存储，支持 chunkwise 并行训练，数值稳定性沿用线性注意力的标准技巧。

输出

经张量收缩得到的输出向量送入后续门控和前馈层，输出 token 预测或表示。

差异点：相比通过增加 head 维度或低秩分解来扩大矩阵状态的方法，triadic linear attention 引入额外 key 轴形成三阶交互，在相同状态大小下参数更少、长上下文回忆能力更强。

实验

实验设计围绕 Gated DeltaNet 和 scalar-gated linear attention 两个基线模型，在 长上下文语言建模 和 回忆任务（含 MQAR 合成基准）上评估。主要对比包括状态扩展（scaling the state）与状态匹配比较（state-matched comparison），并考察从预训练模型升格（upcycling）和混合模型变体。训练效率与消融实验分析第二关键维度和激活函数。

关键发现：triadic linear attention 通过引入三维递归状态，在参数高效的前提下大幅增加状态容量，显著提升长上下文建模和回忆能力，优于其他扩大状态的方法。这一提升在 Gated DeltaNet 和 scalar-gated linear attention 上均得到验证。

与基线对比：相比仅直接扩大状态尺寸的方案，triadic 构造以少量新增投影获得指数级状态增长，避免参数开销线性膨胀；且与数据依赖遗忘、delta rule 及 chunkwise 并行训练兼容，兼具性能与效率优势。

行业影响

落地场景

Triadic Linear Attention 将线性注意力的矩阵状态扩展为三阶张量，以极少的额外参数（仅增加两个投影）提升记忆容量，适合需要长上下文且低延迟推理的产品：

• 电商客服对话：多轮会话中保留用户偏好、历史订单等长程信息，减少重复提问。
• 金融/法律文档分析：流式处理超长合同、年报，精准提取跨段落条款关联。

商业价值

• 降本：保持单步 O(1) 更新和常数量推理，可在 CPU/边缘设备部署，大幅降低推理成本。
• 体验提升：长上下文 recall 准确率更高，直接改善问答、摘要和推荐质量，提升用户留存。

集成接口

• 可无缝替换现有 Gated DeltaNet / scalar-gated linear attention 层，无需改动整体架构。
• 支持 chunkwise-parallel training 和 upcycling，可从已预训练模型继续训练，节省算力。
• 提供 CUDA 内核，便于接入 PyTorch/JAX 等主流训练框架。

参考实现：GitHub

局限

• 论文在实验部分主要评估了语言建模和合成回忆任务，对于长文档问答、多模态输入或代码生成等实际长上下文场景缺乏验证，因此三阶状态的泛化能力尚不明确。此外，作者在 Discussion 中承认，三阶张量状态虽然参数高效，但引入的额外计算和内存开销可能限制其在资源受限环境下的部署，且对更长上下文的 scaling 行为需要进一步分析。这些局限意味着该方法目前更适合研究探索，距离生产环境的大规模应用还有距离。
• 从方法本身看，三阶线性注意力在读取时需要两个 query 分别收缩两个 key 轴，这增加了每个 token 的计算量；同时三阶张量状态的存储和更新对硬件缓存不友好，可能导致显存带宽成为瓶颈。尽管论文提出了 chunkwise 并行训练和 tiling 策略，但在实际 GPU 上，三阶状态相比矩阵状态（如 DeltaNet）或更紧凑的状态空间模型（Mamba）可能带来更高的延迟。此外，第二 key 的维度 E 作为关键超参数，对性能敏感，需要仔细调参，增加了使用成本。
• 与同类工作对比，现有的线性注意力变体如 Gated Linear Attention、RetNet、RWKV 等都采用二阶矩阵状态，本文的三阶扩展虽然自然，但目前仅在 Gated DeltaNet 和 scalar-gated linear attention 上验证有效，对于其他门控机制或非门控线性注意力是否同样受益未知。此外，与直接增加矩阵状态维度（如增大 head dimension 或使用多 head）相比，三阶状态的优势需要更多消融实验支撑，且该方法尚未与近来兴起的混合架构（如 Mamba-Transformer 混合）进行系统比较，因此其在更大规模模型上的竞争力仍有待观察。