Interdomain Attention: 超越令牌级键值记忆

# Interdomain Attention：超越令牌级键值记忆
来源：https://arxiv.org/html/2605.24330
Harrison Bo Hua Zhu、Riccardo El Hassanin、Zhuo Sun、Wenlong Chen、Samir Bhatt、Yingzhen Li

###### 摘要

Transformer 和深度状态空间模型 (SSM) 位于一个基本设计选择的两端：注意力通过基于内容的匹配，以平方成本在每个查询上通过不断增长的键值 (KV) 缓存进行路由，而深度 SSM 则将上下文压缩为一个固定大小的循环状态，该状态不直接由查询-键匹配访问。我们提出了 **Interdomain Attention**，它通过核方法将 SSM 集成到注意力模块中：注意力核通过一个有限特征图进行近似，得到的键特征和值被投影到由单个 SSM 循环维护的一组共享基函数上，每个查询通过自身的特征图关注这些压缩系数，从而在固定大小状态上恢复受查询约束的注意力。可扩展的层是该推导的一种学习了松弛版本，我们通过消融实验验证其组件。在 FineWeb-Edu 上进行的一项 125M–1.3B 自回归语言建模研究中，在匹配的循环状态预算下，Interdomain Attention 在每个尺度上都优于 SSM 令牌混合器，在 1.3B 尺度上，其在验证困惑度和八任务常识套件上超过了相同配方的 softmax 基线，并将固定状态核心的长度平坦行为继承到训练上下文的 3.5 倍。消融实验表明，受查询约束的投影是增益的主要来源。

注意力机制，状态空间模型，线性注意力，长上下文序列建模

## 1 引言

Softmax 注意力和状态空间模型 (SSM) 位于序列模型基本设计选择的两端。注意力 (Vaswani et al., 2017 (https://arxiv.org/html/2605.24330#bib.bib1)) 维护一个每个令牌的键值缓存，并让每个查询通过基于内容的匹配进行路由，这提供了清晰的召回，但代价是 O(N_q N) 的工作量和 O(N) 的 KV 状态，其中 N_q 是查询长度，N 是键值长度。硬件感知的核 (Dao et al., 2022 (https://arxiv.org/html/2605.24330#bib.bib54); Dao, 2024 (https://arxiv.org/html/2605.24330#bib.bib2))、KV 缓存 (Shazeer, 2019 (https://arxiv.org/html/2605.24330#bib.bib38)) 和分布式分片 (Shoeybi et al., 2019 (https://arxiv.org/html/2605.24330#bib.bib59)) 缓解了但并未消除这种扩展问题。深度 SSM，建立在 HiPPO (Gu et al., 2020 (https://arxiv.org/html/2605.24330#bib.bib3)) 之上，并由 S4 (Gu et al., 2022b (https://arxiv.org/html/2605.24330#bib.bib4))、S4D (Gu et al., 2022a (https://arxiv.org/html/2605.24330#bib.bib5)) 和 Mamba (Gu and Dao, 2024 (https://arxiv.org/html/2605.24330#bib.bib6); Dao and Gu, 2024 (https://arxiv.org/html/2605.24330#bib.bib7)) 实现，采取了相反的立场：它们将整个上下文压缩成一个固定大小的循环状态，该状态每步以 O(1) 时间更新。这种压缩的代价是固定状态不直接由查询-键匹配访问。

混合架构将注意力的细粒度、基于内容的检索与 SSM 的高效长程压缩交织在一起 (Lenz et al., 2025 (https://arxiv.org/html/2605.24330#bib.bib55); Glorioso et al., 2024 (https://arxiv.org/html/2605.24330#bib.bib66); Ren et al., 2025 (https://arxiv.org/html/2605.24330#bib.bib56); De et al., 2024 (https://arxiv.org/html/2605.24330#bib.bib67); Fu et al., 2023 (https://arxiv.org/html/2605.24330#bib.bib32); Brixi et al., 2026 (https://arxiv.org/html/2605.24330#bib.bib9))，并且越来越多的线性和次二次注意力直接降低了注意力的成本 (Katharopoulos et al., 2020 (https://arxiv.org/html/2605.24330#bib.bib10); Peng et al., 2021 (https://arxiv.org/html/2605.24330#bib.bib11); Poli et al., 2023 (https://arxiv.org/html/2605.24330#bib.bib8); Peng et al., 2023 (https://arxiv.org/html/2605.24330#bib.bib57); Sun et al., 2023 (https://arxiv.org/html/2605.24330#bib.bib58); Yang et al., 2024a (https://arxiv.org/html/2605.24330#bib.bib27), b (https://arxiv.org/html/2605.24330#bib.bib28))。

我们采取不同的路线：不是将循环层堆叠在注意力旁边，而是询问 SSM 是否可以在保持固定大小状态效率的同时，仍然允许每个查询关注压缩的历史。我们将循环核心固定为 S4D，并将每个令牌的循环状态固定为预算，研究是什么缩小了 S4D 令牌混合器与受查询约束的混合器之间的差距。我们用 **Interdomain Attention** 来回答这个问题，这是一种令牌混合器，其中键和值通过单个复数 S4D 循环映射到共享的 SSM 基。在每个位置，查询关注这些压缩系数，而不是关注每个过去的令牌。该构造的动机是通过一个有限特征图来表示注意力核，并将键特征投影到 HiPPO 基函数上，从而得到一个与序列长度无关的固定大小状态。可扩展的实现并非核推导的字面实现：它使用一个学习的 SiLU/l_2 特征图、输入归一化和一个无分母的读出。因此，我们将该推导作为设计动机，并通过消融实验从经验上评估所得层。

我们的贡献是：
- • 从注意力的核回归视角和 HiPPO 风格基投影出发，构建了一个受查询约束的固定状态令牌混合器，并在理想推导与可扩展实现之间划清了界限（第3.1节 (https://arxiv.org/html/2605.24330#S3.SS1) 和 3.3节 (https://arxiv.org/html/2605.24330#S3.SS3)）。
- • 在 125M 参数下的机制分解，分离了双键/值输入和受查询约束投影的贡献，并将投影确定为主导轴（附录B.1 (https://arxiv.org/html/2605.24330#A2.SS1)）。
- • 在 FineWeb-Edu 上进行的 125M–1.3B 等状态语言建模研究，其中 Interdomain Attention 在每个尺度上都优于 S4D 令牌混合器，在 1.3B 尺度上，其在验证困惑度和八任务常识套件上超过了相同配方的 softmax 基线，并将固定状态核心的长度平坦行为保持到训练上下文的 3.5 倍（图2 (https://arxiv.org/html/2605.24330#S4.F2)，附录B.4 (https://arxiv.org/html/2605.24330#A2.SS4)，表2 (https://arxiv.org/html/2605.24330#S4.T2)）。

## 2 背景

在本节中，我们简要回顾注意力机制和状态空间模型。

#### 注意力作为核回归。
标准点积注意力 (Vaswani et al., 2017 (https://arxiv.org/html/2605.24330#bib.bib1)) 将输入令牌 x_n ∈ R^d 映射为查询、键和值，通过 q_n = W_q x_n ∈ R^d, k_n = W_k x_n ∈ R^d, v_n = W_v x_n ∈ R^d，并计算第 i 个令牌的输出为：
o_i = (∑_{n=1}^N K(q_i, k_n) v_n) / (∑_{n'=1}^N K(q_i, k_n'))，
其中对于 softmax 注意力，K(q,k) = exp(q^T k / √d)。这是 Nadaraya-Watson 核回归估计量 (Nadaraya, 1964 (https://arxiv.org/html/2605.24330#bib.bib13); Watson, 1964 (https://arxiv.org/html/2605.24330#bib.bib14)) 的一个实现，这一联系在若干关于核注意力的工作中已被指出 (Tsai et al., 2019 (https://arxiv.org/html/2605.24330#bib.bib22); Katharopoulos et al., 2020 (https://arxiv.org/html/2605.24330#bib.bib10); Choromanski et al., 2021 (https://arxiv.org/html/2605.24330#bib.bib15))。这一观点使得核 K 的选择成为一个设计自由度：将 softmax 核替换为允许有限或学习特征表示的核，可以实现次二次计算和内存。然后注意力在测试时对值 v_n 执行*非参数回归*，其中 K 决定了对上下文的加权，这与更广泛的测试时回归 (Wang et al., 2025 (https://arxiv.org/html/2605.24330#bib.bib17)) 或记忆化（例如 Titans; Behrouz et al. 2025 (https://arxiv.org/html/2605.24330#bib.bib20)）相联系。

#### 状态空间模型和 HiPPO。
线性状态空间模型将输入信号 z(t) ∈ R 映射到隐状态 u(t) ∈ R^M，通过：
u̇(t) = A(t) u(t) + B(t) z(t)。
HiPPO (Gu et al., 2020 (https://arxiv.org/html/2605.24330#bib.bib3)) 给出了 A(t) ∈ R^{M×M} 和 B(t) ∈ R^M 的初始化，使得状态 u(t) 维护输入历史在 M 个时变正交基函数 {φ_m^{(t)}}_{m=1}^M 上的最优投影。这构成了诸如 S4 (Gu et al., 2022b (https://arxiv.org/html/2605.24330#bib.bib4))、S5 (Smith et al., 2023 (https://arxiv.org/html/2605.24330#bib.bib50))、S4D (Gu et al., 2022a (https://arxiv.org/html/2605.24330#bib.bib5)) 和 Mamba (Gu and Dao, 2024 (https://arxiv.org/html/2605.24330#bib.bib6); Dao and Gu, 2024 (https://arxiv.org/html/2605.24330#bib.bib7)) 等深度 SSM 架构的基础。

## 3 Interdomain Attention

(a) 标准注意力
Q K V A = softmax(Q K^T / √d) N_q × N 分数 值聚合 A (N_q×N) · V (N×d) 输出

(b) Interdomain 注意力
Q K V ξ_ω ξ_ω SSM 循环 核查询 F_q 循环内状态 (U, Γ, η) 状态空间读出 F_q U^T / (F_q U^T η) (N_q×M) · Γ (M×d) 输出

图 1: (a) 标准注意力从 Q 和 K 计算 N_q×N 个分数，然后乘以 V 以产生输出。(b) Interdomain 注意力将查询和键映射到一个共享的特征空间，以产生核查询矩阵 F_q (N_q×R)。键和值通过 SSM 循环压缩成 M 个循环内状态：U (M×R，键特征投影)、Γ (M×d，值投影) 和 η (M×1，归一化常数)。状态空间读出通过 N_q×M 的互协方差 F_q U^T 计算输出。

核回归观点提出了一种使注意力循环的直接方法：将键和值总结为一组固定大小的基系数，然后让每个查询关注这些系数。HiPPO 风格的 SSM 为在线维护此类系数提供了一种自然机制。基于 HiPPO-SVGP 和 SSM 之间循环内核对核计算的联系 (Chen et al., 2025 (https://arxiv.org/html/2605.24330#bib.bib16))，我们在下面推导这个构造，然后描述实验中使用的学习层。图1 (https://arxiv.org/html/2605.24330#S3.F1) 展示了该架构。

### 3.1 核注意力的特征图视角

假设注意力核允许一个有限特征表示：
K(q,k) ≈ ξ(q)^T ξ(k), ξ(·) ∈ R^R。 (3)
将此表示代入方程1 (https://arxiv.org/html/2605.24330#S2.E1) 得到：
o_i ≈ Ô_i = (∑_{n=1}^N ξ(q_i)^T ξ(k_n) v_n) / (∑_{n'=1}^N ξ(q_i)^T ξ(k_n'))。 (4)

注意力的核和特征图视图已在先前的分析和高效注意力变体中使用过 (Tsai et al., 2019 (https://arxiv.org/html/2605.24330#bib.bib22); Katharopoulos et al., 2020 (https://arxiv.org/html/2605.24330#bib.bib10); Choromanski et al., 2021 (https://arxiv.org/html/2605.24330#bib.bib15))。随机傅里叶特征 (Rahimi and Recht, 2007 (https://arxiv.org/html/2605.24330#bib.bib12)) 是平稳核的一个标准实例：根据 Bochner 定理，
K(x, x') = E_{p(ω)}[ξ_ω(x)^T ξ_ω(x')]， (5)
ξ_ω(x) = [cos(ω^T x), sin(ω^T x)]^T， (6)