基于领域分解的层次注意力

# 基于区域分解的分层注意力
来源：https://arxiv.org/html/2606.18525
[![[无标题图片]](https://arxiv.org/html/2606.18525v1/x1.png)Stephan Köhler](https://orcid.org/0000-0003-1015-8736) 和 [![[无标题图片]](https://arxiv.org/html/2606.18525v1/x2.png)Oliver Rheinbach](https://orcid.org/0000-0002-9310-8533) 数学与计算机科学学院，弗莱贝格工业大学，09596 弗莱贝格，德国

###### 摘要

我们提出了一种基于两层重叠 Schwarz 区域分解的分层注意力机制。该方法的动机源于观察到两层 Schwarz 区域分解方法将局部子域修正与一个用于传递全局、长程信息的粗层相结合。我们通过一个具有齐次 Dirichlet 边界条件的简单一维扩散问题，检验了该方法在有限维算子学习中的有效性。尽管该问题很基础，但它提供了一个受控的序列到序列设置，其中精确的非局部解算子是已知的。离散化后，学习解算子相当于逼近一个对称正定矩阵的逆。作为基线，我们使用一个全局的、无 softmax 的低秩注意力算子，形式为 QK^T。我们提出的构造用两层加性结构替换了这个稠密的全局分解：重叠子域上的局部低秩注意力块与一个粗注意力块相结合。得到的算子形式为：

M_θ^{-1} = Φ Q_0 K_0^T Φ^T + ∑_{i=1}^N R_i^T D_i^{1/2} Q_i K_i^T D_i^{1/2} R_i。

这里 R_i 限制到重叠子域，D_i 是单位分解权重，Φ 是一个粗插值（或延拓）矩阵。针对合成傅里叶右端项的数值实验表明，与全局低秩注意力基线相比，区域分解注意力算子能够更快地训练，并在使用显著更少参数的情况下提供更精确的逼近。

## 1 引言

注意力机制的引入 [1 (https://arxiv.org/html/2606.18525#bib.bib1)] 对现代大型语言模型的成功至关重要。与早期的循环网络不同，注意力能够在长序列上直接并行地对上下文依赖关系进行建模。形式上，注意力可以被理解为一类可学习的交互算子，其适用性源于在计算令牌交互时利用了低秩分解。

标准自注意力层从一个输入矩阵开始：

X ∈ R^{n × d_in}，

其中行对应 n 个令牌、网格点或自由度，d_in 列包含关联的特征。查询、键和值矩阵通过学习的线性映射获得：

Q = X W_Q, K = X W_K, V = X W_V, (1)
W_Q, W_K ∈ R^{d_in × d_k}, W_V ∈ R^{d_in × d_v}。

标准的带 softmax 的缩放点积注意力算子为：

Att_softmax(X) = softmax(Q K^T / √d_k) V，

其中 d_k 是 W_Q 和 W_K 的列数。算子 softmax 是逐行应用 softmax，即对于矩阵 G：

(softmax(G))_{ij} = exp(G_{ij}) / ∑_{ℓ=1}^n exp(G_{iℓ})。

矩阵 Q K^T ∈ R^{n × n} 包含 n 个令牌或我们情况下的网格点之间的成对交互。参数 d_k 对应于低秩分解 Q K^T 的秩 r，即：

r = rank(Q K^T) ≤ d_k。

在本工作中，我们将省略 softmax 算子。同时，我们也将省略由 d_k 缩放得到的因子 √d_k，因为 Q K^T 的对角线是长度为 d_k 的向量的标量积。引入此缩放是为了通过随机梯度下降或变体（如 Adam [2 (https://arxiv.org/html/2606.18525#bib.bib2)]）改善优化。

总之，我们得到线性注意力算子：

Att(X) = (Q K^T) V。

概括而言，在本工作中，我们考虑一种无 softmax 因而线性的注意力算子设置，并在有限维算子学习的意义下训练它来逼近一个已知的解算子。这使得我们可以专注于一个问题：来自数值区域分解 [3 (https://arxiv.org/html/2606.18525#bib.bib3), 4 (https://arxiv.org/html/2606.18525#bib.bib4), 5 (https://arxiv.org/html/2606.18525#bib.bib5)] 的方法能否用于设计更高效的注意力算子。

我们的模型问题是离散化的一维 Poisson 问题。离散化后，解映射是稀疏系统矩阵的逆。一个全局低秩注意力层可以表示该逆的一个秩为 d_k 的逼近。然而，椭圆型逆具有多尺度结构：局部耦合最重要，但也必须表示长程低频分量。例如，用于偏微分方程的两层 Schwarz 区域分解方法 [3 (https://arxiv.org/html/2606.18525#bib.bib3), 4 (https://arxiv.org/html/2606.18525#bib.bib4), 5 (https://arxiv.org/html/2606.18525#bib.bib5)] 正是利用了这种结构。

我们引入了一种受重叠 Schwarz 区域分解启发的两层注意力机制。细层由重叠子域上的局部低秩注意力块组成。粗层由一个作用于小型界面帽子粗基的注意力块组成。这种构造并非旨在替代经典求解器。相反，它是一个受控实验，旨在证明区域分解概念可以提供有用的结构，例如用于算子学习方法 [6 (https://arxiv.org/html/2606.18525#bib.bib6), 7 (https://arxiv.org/html/2606.18525#bib.bib7), 8 (https://arxiv.org/html/2606.18525#bib.bib8)] 或自然语言处理。现代重叠 Schwarz 方法在求解偏微分方程时是高度可并行扩展的，例如 [9 (https://arxiv.org/html/2606.18525#bib.bib9), 10 (https://arxiv.org/html/2606.18525#bib.bib10)]。

在自然语言处理 (NLP) 和大型语言模型 (LLM) 的背景下，我们在第 2 节 (https://arxiv.org/html/2606.18525#S2) 中提出的一维 Poisson 问题可以解释为一系列令牌，其中相邻令牌之间的局部交互最强。然而，较弱的远程依赖也必须表示。在我们的两层 Schwarz 注意力中，粗层表示较弱的远程交互，而强局部交互则由重叠子域捕获。

我们注意到，先前已经在不同的设置中研究过分层形式的注意力 [11 (https://arxiv.org/html/2606.18525#bib.bib11), 12 (https://arxiv.org/html/2606.18525#bib.bib12), 13 (https://arxiv.org/html/2606.18525#bib.bib13)]。这些方法与本文提出的方法不可直接比较，因为底层机制和预期应用都不同。在本工作中，层次结构由重叠区域分解诱导，并用于构造椭圆型解算子的无 softmax 注意力逼近。NLP 中的层次方法包括 [14 (https://arxiv.org/html/2606.18525#bib.bib14), 15 (https://arxiv.org/html/2606.18525#bib.bib15), 16 (https://arxiv.org/html/2606.18525#bib.bib16), 17 (https://arxiv.org/html/2606.18525#bib.bib17)]。

在完成本手稿后，我们发现我们在线性注意力算子学习方面的方法与工作 [18 (https://arxiv.org/html/2606.18525#bib.bib18)] 有密切联系，其中注意力被用于算子学习。

## 2 一维扩散模型问题

我们考虑一维 Poisson 问题：

- u''(x) = f(x), x ∈ (0,1) ⊂ R^1, u(0) = u(1) = 0。 (2)

令 n+2 为网格点数，则 n 是内部网格点数，并令：

x_j = j/(n+1), j = 1,...,n。 (3)

使用标准有限差分或有限元离散化得到：

A u = f, A = (1/h^2) * [2 -1; -1 2 -1; ...; -1 2 -1; -1 2], h = 1/(n+1)。 (4)

因此精确离散解算子为：

f ↦ u = A^{-1} f。 (5)

本文的学习任务是通过一个可训练的具有结构化的算子 M_θ^{-1} 来逼近 A^{-1}，其中 θ 表示参数。

对于我们的一维模型问题，输入矩阵 X 的一个典型选择是：

X = [x1 f1; x2 f2; ...; xn fn] ∈ R^{n×2}，

即 X 包含局部网格点和右端项作为特征。

然而，对我们来说，相比使用物理坐标作为输入特征，更方便的是用独热编码表示网格点，并通过右端项 f 增强这个表示，即我们的特征矩阵是：

X = [I_n, f] ∈ R^{n×(n+1)}。

因此 d_in = n+1，即前 n 列标识网格点，最后一列包含右端项 f 的值。

接下来，权重矩阵 W_V 被选为一个固定矩阵，它从 X 中提取右端项通道，即：

V = X W_V = f ∈ R^{n×1}, W_V = [0; 1]，

其中 W_V 中的 0 表示长度为 n 的零列向量。因此，在我们的上下文中，W_V 不包含可学习参数。

然后我们定义 W_Q 和 W_K 为：

W_Q = [Q; 0], W_K = [K; 0]，

其中矩阵 Q 和 K 包含 W_Q 和 W_K 中的可学习参数。这里 Q, K ∈ R^{n×d_k}，W_Q 和 W_K 中的 0 表示一个零行。

X 中的坐标独热编码然后从 W_Q 和 W_K 中提取 Q 和 K，即：

Q = X W_Q, K = X W_K。

作为我们定义的结果，在本文考虑的无 softmax 线性设置中，低秩分解 Q K^T 直接应用于右端项，即定义 u_θ = Att(X)，我们得到：

u_θ = (Q K^T) V = Q K^T f。 (6)

因此学习任务在于找到注意力因子使得：

Q K^T ≈ A^{-1}，

其中 A 是有限差分 Poisson 矩阵。

尽管我们的数值实验是针对一维椭圆模型问题进行的，但这种构造不应被视为局限于这个特定的 PDE 设置。相反，该模型问题提供了一个受控的设置，其中非局部序列到序列算子已知精确，因此可以用来研究由区域分解诱导的分层注意力的效果。这在一维序列模型中具有更广泛的相关性，因为在这种模型中，局部交互和长程耦合需要同时表示。

## 3 作为低秩逆的无 Softmax 注意力

我们的分层注意力的基线是应用于算子学习的全局低秩线性注意力，即：

M_{θ, global}^{-1} = Q K^T, Q, K ∈ R^{n × r_g}。 (7)

应用于右端项 f 时，预测解为：

u_θ = Q (K^T f); (8)

参见 (6)。算子的秩受限于秩参数 r_g。

## 4 通过两层重叠 Schwarz 区域分解实现的分层注意力

### 4.1 重叠子域和粗帽子函数

实现中使用的子域构造是对内部自由度的代数分解。因此，它与基于单元的有限元划分略有不同，在后者中，相邻子域在重叠添加到界面左右两侧之前共享一个界面节点。

令 n 个内部网格点的全局索引集为：

I = {1, ..., n}。 (9)

集合 I 被分成 N 个相等大小的连续不相交索引集。然后这些索引集通过 n_δ 个索引扩展到相邻索引集，形成重叠子域的索引集 I_i, i=1,...,N。因此，对于给定的 n_δ，如果相邻索引集大于重叠区域（如下面实验中情况那样），则相邻重叠子域共享 2n_δ 个公共网格点。

在这种构造中，重叠区域中没有中心网格点。因此，在离重叠中心最近的两个索引中，我们选择较大的作为相应粗界面帽子函数最大值的位置。

令 n_i = |I_i|，并令

R_i ∈ R^{n_i × n}

表示布尔限制矩阵，用于从 Rn 中的全局向量中提取 I_i 中的自由度。等价地，R_i u 是与重叠子域 I_i 相关的 u 的子向量，而 R_i^T 将局部向量以零扩展到全局索引集。

在所有数值实验中，我们使用重叠 n_δ = 2。因此，两个相邻重叠子域有 2n_δ = 4 个公共网格点。

### 4.2 重叠区域中的单位分解

在重叠区域中，多个子域可能覆盖同一个自由度。令 m_j 是包含网格点 j 的重叠子域的数量。在子域 i 上，定义对角权重矩阵：

D_i = diag(1/m_j : j ∈ I_i)。 (10)

在下面实验使用的分解中，重叠区域中的每个网格点由两个相邻子域覆盖，因此该处 m_j = 2，而在重叠区域外 m_j = 1。

那么这些权重构成了一个离散单位分解：

∑_{i=1}^N R_i^T D_i R_i = I。 (11)

在下面的实验中，我们使用对称重叠加权：

R_i^T D_i^{1/2} G_i D_i^{1/2} R_i。 (12)

### 4.3 局部注意力块

在每个 N 个重叠子域上，我们使用一个局部低秩注意力算子：

G_i = Q_i K_i^T, Q_i, K_i ∈ R^{n_i × r_ℓ}, (13)

其中 r_ℓ 表示局部注意力块的秩。

相应的细层算子为：

M_{θ, fine}^{-1} = ∑_{i=1}^N R_i^T D_i^{1/2} Q_i K_i^T D_i^{1/2} R_i。