技术摘要：高阶线性注意力（HLA）

问题陈述

Transformer 架构中标准的缩放点积注意力机制，其计算和内存复杂度相对于序列长度 $n$ 为 $O(n^2)$。这种二次方成本是将自回归语言模型扩展到长上下文的主要瓶颈。虽然现有的替代方案（如线性注意力和状态空间模型 SSMs）提供了每 token $O(n)$ 或 $O(1)$ 的复杂度，但它们通常局限于一阶近似或基于核的方法。这些限制往往限制了模型的表达能力，使其无法捕捉标准注意力机制所能建模的 token 之间复杂的、高阶的交互。

方法论

本文介绍了高阶线性注意力（HLA），这是一种因果流式机制，通过紧凑的前缀充分统计量实现高阶交互。核心洞察在于，高阶注意力类算子可以分解为低阶矩（例如键的外积之和），从而在不实例化 $n \times n$ 注意力矩阵的情况下实现精确的因果流式处理。

1. 二阶 HLA

二阶变体通过为每个头维护一个捕获二阶矩的恒定大小状态，推广了线性注意力。

• 状态维护：HLA 不再维护单一的运行和，而是维护前缀摘要，包括键矩 $S^K_t = \sum_{i \le t} k_i k_i^\top$、查询 - 值累加器 $C^{QV}_t = \sum_{i \le t} q_i v_i^\top$ 以及查询质量 $m^Q_t = \sum_{i \le t} q_i$。
• 严格因果性：为了在不构建完整注意力矩阵的情况下强制执行严格的自回归因果性（防止未来 token 影响当前输出），作者推导出了扩展摘要。具体而言，他们引入了交叉摘要 $G_t$ 和 $h_t$，以修正朴素乘积 $S^K_t C^{QV}_t$ 中包含未来项的问题。
• 输出公式：时间 $t$ 处的未归一化输出计算为 $o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$。该公式允许每个 token 进行 $O(d^2 + d d_v)$ 的计算，且独立于序列长度。
• 归一化：还定义了一种比率归一化变体，通过将分子除以源自类似扩展摘要的掩码分母来进行归一化。

2. 非对称高阶线性注意力（AHLA）

本文提出了一种互补的非对称变体AHLA，它实现了算子 $Q(K^\top Q)(K^\top V)$。这种结构通过由中间键索引路由的单次传递对值进行加权，不同于标准 HLA 的对称 $AA^\top V$ 结构。它利用了一组独特的流式摘要（$P^{KV}_t, m^K_t, E_t, n_t$），并实现了相同的线性每 token 复杂度。

3. 三阶 HLA

该框架被扩展至三阶交互，定义为矩阵乘积 $AA^\top A$，其中 $A = QK^\top$。

• 流式代数：三阶掩码输出需要维护三个主要前缀矩（$S^K_t, S^Q_t, P^{KV}_t$）和六个额外的交叉摘要（$G^{(1)}_t, \dots, G^{(3)}_t$ 以及 $h^{(1)}_t, \dots, h^{(3)}_t$），以强制执行严格因果性。
• 复杂度：在线更新成本在特征维度上保持线性，尽管状态大小增加以容纳高阶矩。

4. 通过关联扫描进行并行训练

为了解决在 GPU 上纯训练循环模型的低效问题，作者开发了一种分块并行训练方案。

• 关联算子：他们定义了用于连接 token 片段的关联算子（具体为半直积）。这些算子考虑了片段之间产生的交叉项（例如 $S^K_B C^{QV}_A$），这些项在组合分块时会出现。
• 扫描等价性：通过对这些片段摘要应用独占 Blelloch 扫描，该方法产生的前向激活与串行递归的结果完全相同。这使得无需借助近似的随时间反向传播即可进行并行训练。
• 三阶扫描：对于三阶情况，扫描算子通过片段级线性映射（$M^{KQP}, M^{KQm}$）进行增强，以精确处理三阶校正状态的组合。

主要贡献

1. 二阶精确掩码流式处理：本文提供了扩展摘要的完整代数，产生了具有恒定每 token 成本的严格因果二阶 HLA。它包括建立掩码流式恒等式和在线更新的正式证明，区分了未归一化和比率归一化变体。
2. 匹配串行激活的关联扫描：作者定义了适用于无掩码和掩码设置（有和无指数衰减）的关联算子。他们证明了在这些算子上应用标准独占扫描所产生的前向激活与串行递归完全相同，从而实现了高效的并行训练。
3. 三阶扩展：该工作展示了完整的掩码三阶状态、在线更新、严格因果流式内核以及精确的分块并行算法。这包括推导精确组合三阶状态所需的片段映射。
4. 非对称变体（AHLA）：引入了一种新的非对称高阶注意力机制，提供了不同的归纳偏置，同时保持了相同的流式处理和复杂度特性。

结果与主张

本文主要关注算法结构和理论属性，而非特定数据集上的经验基准。这些主张基于数学推导：

• 复杂度：HLA 实现了每 token $O(d^2 + d d_v)$ 的时间和内存复杂度用于推理，其中 $d$ 是查询/键维度，$d_v$ 是值维度。这与序列长度 $n$ 无关。
• 精确性：分块并行训练方案被证明能精确复现串行递归的激活，避免了其他线性注意力训练方法中常见的近似误差。
• 因果性：提出的掩码公式严格强制执行自回归因果性，确保时间 $t$ 的输出仅依赖于 token $j \le t$。
• 可扩展性：通过避免 $n \times n$ 矩阵实例化并利用关联扫描，HLA 被定位为长上下文模型的可扩展构建模块。

意义

本文将 HLA 定位为“原则性、可扩展的构建模块”，弥合了类注意力、数据依赖混合与现代循环架构效率之间的差距。

• 表达能力与效率：与依赖固定特征映射的一阶线性注意力不同，HLA 学习数据依赖的度量（通过 $S^K_t$）并捕捉高阶交互（2 次或 3 次多项式），同时保留了 $O(1)$ 每 token 状态更新范式。
• 实现：这项工作强调，HLA 可以作为 Transformer 块中标准注意力子层的即插即用替代品，仅需替换内核，同时保留位置编码和掩码逻辑。
• 理论基础：通过为高阶项提供精确的流式恒等式和关联扫描算子，本文建立了一个严格的代数框架，用于将线性注意力扩展到一阶以上，同时不牺牲因果约束或训练效率。