Linearized 2-Simplicial Attention

技术摘要：线性化 2-单纯形注意力 (Linearized 2-Simplicial Attention)

问题陈述
标准的 Softmax 注意力机制提供了基于内容的对过往 Token 的访问能力，但面临着二次计算复杂度 ($O(n^2)$) 和不断增长的 KV 缓存问题，限制了其可扩展性。虽然线性注意力机制通过将过去的信息压缩进固定大小的状态中，但往往在精确检索和选择性更新方面表现不佳。相反，2-单纯形注意力（2-simplicial attention）通过对查询与 Token 对（三元组）进行评分来引入高阶交互，这在理论上提升了处理三元依赖任务的能力。然而，精确的 2-单纯形注意力计算成本极高，需要 $O(n^3)$ 的计算量。近期试图使其工程化的尝试依赖于将两个 Token 轴都限制在固定窗口内，这牺牲了全局覆盖能力以换取效率。

方法论
作者提出了 LinSimp，一种线性化的 2-单纯形注意力机制，它实现了线性成本 ($O(n)$)，同时在其中一个轴上保持全局覆盖能力，并在另一个轴上保持局部精度。

1. 通过三线性重写实现线性化：
   核心创新在于将三线性评分 $\langle q_i, k_j, r_c \rangle$ 重写为复合查询与键之间的内积：
   $$\langle q_i, k_j, r_c \rangle = (q_i \odot r_c)^\top k_j$$
   这种重新表述使得对 Token 索引 $j$ 的求和可以被视为普通的指数点积核（exponential dot-product kernel），类似于标准的线性注意力。

2. 正向随机特征 (Positive Random Features)：
   为了近似指数核 $\exp((q_i \odot r_c)^\top k_j)$，作者采用了正向随机特征（具体为正交随机特征）。这使得对全局 Token 历史（$j$）的求和可以被压缩进一个固定大小的状态（$M_i$ 和 $a_i$）中，且该状态与序列长度无关。

   • 全局轴 ($j$)： 通过随机特征状态进行处理，从而实现对整个因果前缀的访问。
   • 局部轴 ($c$)： “锚点”（anchor）维度保持显式存在，但被限制在一个短滚动窗口内（例如 64 个 Token）。这创建了一种混合机制，每个近期的锚点都会生成一个唯一的复合查询并发送至全局状态。

3. 架构与实现：

   • 模型组成： 作者将 LinSimp 层集成到一个同时也利用 Kimi Delta Attention (KDA)（一种带门控 Delta 规则的快速权重模型）的模型架构中。他们构建了一个包含 18 层 KDA 和 6 层 LinSimp 的“无 Softmax”模型。
   • 自定义算子 (Custom Kernels)： 该方法使用自定义的 Triton 算子 实现。前向传播按块处理查询，并通过边界条件重建前缀状态，以避免实例化完整的 Token 对。反向传播采用两遍扫描设计，以在不保存大型中间张量 ($T \times m \times D_v$) 的情况下累积梯度。

核心贡献
论文概述了四项具体贡献：

1. 推导： 实现了 2-单纯形注意力的单模态核化（one-mode kernelization），将三线性评分转化为适用于线性注意力近似的形式。
2. 实现： 利用正交随机特征库处理全局模式，并利用短显式锚点窗口处理局部模式，实现了线性时间复杂度的实现。
3. 算子工程： 开发了自定义 Triton 前向和反向算子，包括针对 FP32 参考值的数值验证。
4. 实证比较： 在等 Token 数（iso-token）和等 FLOPs（iso-FLOP，即等计算量）设置下，针对标准注意力、精确窗口化 2-单纯形基准以及 KDA 混合模型进行了全面的评估。

结果
作者在 2k 和 16k 上下文长度的网页数据（FineWeb-Edu）和数学数据（FineMath-4+）上对模型进行了评估。

• 等 FLOPs 性能： 在匹配的计算预算下，在 3B-token 数学训练及 16k 上下文设置中，“无 Softmax” KDA + LinSimp 模型在对比架构中取得了最高的下游平均准确率。
• 长上下文 (16k)： 在 16k 上下文时，基于 LinSimp 的模型比 KDA 混合模型（使用标准 Softmax 注意力）的平均准确率提升了 0.0079，并将 LAMBADA 困惑度从 715.6 大幅降低至 602.6。
• 效率： 自定义算子实现了接近 Softmax 注意力的吞吐量（在 2k 上下文下，单个 LinSimp 层的时间为 Softmax 层的 1.03 倍），但尚未在原始速度上超越成熟的 Softmax 算子。
• 基准对比： 当计算量匹配时，精确窗口化 2-单纯形基准（Roy 等人，2025）显示的验证损失和困惑度高于线性化方法。

意义与主张
论文声称，LinSimp 成功弥合了高阶（三线性）交互的理论优势与线性时间推理的实际约束之间的鸿沟。通过将 2-单纯形交互分解为全局随机特征状态和局部显式窗口，作者证明了：

• 构建一个完全不含 Softmax 注意力层的高性能语言模型是可能的。
• 在等计算预算的长上下文场景下，该架构可以超越 KDA 混合模型（后者仍依赖标准注意力进行部分层计算）。
• 该方法提供了一种窗口化 2-单纯形注意力所缺乏的“全局覆盖能力”，同时保持了线性复杂度。

局限性与谦逊声明
作者明确指出，其实验属于初步研究，仅使用了单一随机种子，因此无法估计方差或置信区间。他们承认目前的反向传播算子尚未经过高度优化，且等 FLOPs 预算是一个分析性的计数度量，而非直接的能量或延迟测量。此外，该模型在超过测试参数量（约 330M）后的缩放行为尚不明确。