Flexformer：具有可学习注意力核的灵活线性Transformer

# Flexformer：具有可学习注意力核的灵活线性Transformer
来源：https://arxiv.org/html/2606.27748  
Haoran Zhang, Feng Zhou  
中国人民大学应用统计学院 统计学院，北京，中国  
[email protected], [email protected]  

###### 摘要  
Transformer模型依赖注意力机制来捕捉长距离依赖关系，但受限于二次复杂度，难以扩展到长序列。基于核的线性注意力降低了这种复杂度，但通常依赖固定或弱可学习的核，限制了表达能力与性能。本文提出Flexformer，一种灵活的线性Transformer，以完全数据驱动的方式学习注意力核。Flexformer基于随机傅里叶特征线性注意力，并将频谱频率视为可训练参数，使模型能够学习广泛的注意力核家族。我们开发了平稳和非平稳两种变体，后者具有更强的表达能力。在语言建模和序列分类上的大量实验表明，Flexformer始终优于基线模型。此外，Flexformer可以从预训练的Transformer中高效蒸馏，恢复softmax注意力，并展现出跨领域的强核迁移能力，在长序列任务上同时实现高效率和竞争力。

## 1 引言  
Transformer架构（Vaswani et al., 2017）自提出以来在机器翻译任务中取得了显著成功。随后，Transformer的众多变体被广泛应用于人工智能的多个重要领域，包括自然语言处理（Devlin et al., 2019; Liu et al., 2019; Brown et al., 2020）、计算机视觉（Dosovitskiy et al., 2021; Liu et al., 2021）、时间序列分析（Zhou et al., 2021; Nie et al., 2023; Liu et al., 2024）以及音频处理（Schneider et al., 2019; Eilers and Jiang, 2023）。这些基于Transformer的模型中最关键的组件是点积注意力机制。它通过计算整个输入序列中成对token的相似度，使每个token能够动态聚合与自身最相关的信息，同时有效建模序列中的长距离依赖。与逐步处理token的循环神经网络（RNN; Sherstinsky, 2020）不同，Transformer注意力并行计算所有位置的表示，最大化计算硬件利用率。然而，这种有效的注意力机制存在效率瓶颈：它计算序列中所有位置对的注意力权重，导致时间和空间复杂度随序列长度呈二次增长。

为了解决这个问题，许多研究尝试优化注意力机制的计算。这些重新设计或优化的注意力计算方法应有效降低时间和空间复杂度，同时达到与softmax注意力相当甚至更优的性能。一些工作基于稀疏注意力（Kitaev et al., 2020; Roy et al., 2020; Zaheer et al., 2020），通过限制每个token的注意力范围仅关注关键子集来修改注意力模式。然而，这些注意力模式并未经过严格验证，通常基于经验得出的模式，难以保证在不同数据领域始终保持优越性能。另一类是核基线性注意力（Katharopoulos et al., 2020; Choromanski et al., 2021; Qin et al., 2022），它使用核来表示点积计算和softmax应用过程。核可以写成特征映射的内积。通过重新排列矩阵乘法的顺序，注意力的二次复杂度可以降低为线性复杂度。

图1展示了核基线性注意力和softmax注意力计算过程的差异。

（a）Softmax注意力

（b）核基线性注意力

图1：Softmax注意力和核基线性注意力的对比。Softmax注意力需要O(N²d)时间复杂度和O(N²+Nd)空间复杂度，而核基线性注意力实现了O(Ndd')时间复杂度和O(Nd+Nd'+dd')空间复杂度，其中N, d, d'分别表示输入序列长度、隐藏维度和特征映射维度。当N≫d且N≫d'时，线性注意力在时间和空间上均实现显著节省。

本文聚焦于核基线性注意力，其核心在于构建有效的特征映射。尽管许多方法旨在恢复softmax注意力，但尚无证据表明softmax在所有情况下始终最优。因此，更有效的方法是使用直接从数据中学习的可学习核。然而，如何构造足够灵活且保证强性能的可学习核仍是一个开放问题。Hedgehog（Zhang et al., 2024）和Polaformer（Meng et al., 2025）通过将可学习映射与尖峰函数（如指数函数和幂函数）结合，经验性地利用了softmax注意力的低熵特性。但这些方法无法保证学习的核具备足够表达能力以涵盖多种核类型，这可能会损害其性能。

本文提出了一种具有可学习核的灵活线性Transformer（Flexformer）。Flexformer通过使特征映射完全可学习，扩展了基于随机傅里叶特征的核基线性注意力。与现有可学习核方法相比，Flexformer能够建模更广泛、更具表现力的核族。具体地，谱表示理论指出，核可以通过其谱测度表示为逆傅里叶变换，并可通过蒙特卡洛采样近似。在Flexformer中，我们不从固定的谱密度中采样频率，而是将频率视为可学习参数，直接从数据中优化，从而以数据驱动的方式学习注意力核。我们提出了Flexformer的平稳和非平稳两个版本，其中非平稳形式具有严格更强的表达能力。理论上，我们可以证明Flexformer的核族包含softmax核。这保证了Flexformer至少能够学习到softmax注意力，并可能学习到优于softmax的注意力模式。Flexformer可以从头训练，也可以通过从预训练Transformer中蒸馏（Zhang et al., 2024）来训练。

我们在多个基准上进行了大量实验，包括语言建模和分类任务，展示了Flexformer相对于现有基线的优越性能，以及其恢复softmax核和学习核迁移的强大能力。作为线性注意力的变体，它在处理极长序列时显著提升了效率并降低了内存消耗。具体来说，我们的贡献总结如下：

- 我们提出了Flexformer，利用通过随机傅里叶特征构建的可学习注意力核，在保持序列长度线性和空间复杂度的同时实现了高表达能力。
- 我们通过在语言建模和序列分类任务上的从头训练评估了Flexformer。结果表明，Flexformer始终优于其他线性注意力基线，突显了其强大的有效性和效率。
- Flexformer可以通过注意力权重蒸馏忠实近似softmax注意力，实现高效且保持性能的线性注意力替代，并在跨领域设置中展现出强核迁移能力。

## 2 预备知识  
本节介绍自注意力机制和核基线性注意力的背景知识。

### 2.1 自注意力机制  
给定输入序列X = [X₁, …, X_N]ᵀ ∈ ℝ^{N×d}，其中N和d分别表示输入长度和嵌入维度。每个token X_i通过线性映射投影生成查询、键和值张量，形式化表示为：Q_i = W_Q X_i, K_i = W_K X_i, V_i = W_V X_i, i=1,…,N。对于每个查询张量，计算其与所有键张量的相似度度量。这些相似度得分经归一化形成注意力权重，用于对相应的值张量进行加权求和，从而生成上下文感知的表示。形式上，第i个token的自注意力模块输出可以表示为：

O_i = (∑_{j=1}^{N} sim(Q_i, K_j) V_j) / (∑_{j=1}^{N} sim(Q_i, K_j)),   (1)

其中sim(·,·): ℝ^d × ℝ^d → ℝ₊ 表示相似度度量函数。在标准softmax注意力的Transformer（Vaswani et al., 2017）中，sim(Q_i, K_j)定义为 sim(Q_i, K_j) = exp(Q_iᵀ K_j / √d)。

### 2.2 核基线性注意力  
如果直接按照公式(1)的计算顺序，先计算所有位置对查询和键的相似度得到N×N的注意力权重矩阵，再与值相乘，将导致O(N²d)的时间复杂度和O(N²+Nd)的空间复杂度，随着N增大带来显著的计算和内存开销。为解决上述问题，我们假设一个取值在ℝ₊的注意力核k(x,y)作为相似度函数，并令φ: ℝ^d → ℝ^{d'}（d'为核诱导特征空间的维度）表示相应的特征映射，使得：

k(x, y) = φ(x)ᵀ φ(y)。   (2)

则公式(1)中自注意力的计算可以表示为：

O_iᵀ = [∑_{j=1}^{N} φ(Q_i)ᵀ φ(K_j) V_jᵀ] / [∑_{j=1}^{N} φ(Q_i)ᵀ φ(K_j)] = [φ(Q_i)ᵀ ∑_{j=1}^{N} φ(K_j) V_jᵀ] / [φ(Q_i)ᵀ ∑_{j=1}^{N} φ(K_j)]。   (3)

这种方法避免了显式计算任意查询和键之间的成对相似度。相反，我们首先预计算∑_{j=1}^{N} φ(K_i) V_jᵀ 和 ∑_{j=1}^{N} φ(K_j)，这些计算对所有位置查询共享，然后与每个位置查询Q_i相乘。因此，该方法的时间复杂度为O(N d d')，空间复杂度为O(Nd + Nd' + dd')，实现了序列长度的线性缩放。当N >> d, d'时，与标准softmax注意力相比，表现出显著的效率提升和内存节省。

## 3 方法论  
本节我们首先回顾如何使用随机傅里叶特征近似softmax核。然后我们将该公式扩展到完全可学习的核，从而提出Flexformer。

### 3.1 使用随机傅里叶特征的Softmax核  
在softmax注意力中，相似度度量可以表示为以下核。我们之后将其称为softmax核：

k_{SM}(x, y) = exp(xᵀ y / √d)。

通过泰勒展开，我们可以推导出满足公式(2)的特征映射φ。不幸的是，φ的维度是无限的（Sieber et al., 2024）。因此，我们需要寻找有限维的φ，使得softmax核能被有效近似。

Softmax核可以按如下方式分解：

k_{SM}(x, y) = exp( (||x||² + ||y||² - ||x - y||²) / (2√d) ) = exp(||x||²/(2√d)) · exp(-||x - y||²/(2√d)) · exp(||y||²/(2√d))。   (4)

蓝色项仅依赖于向量范数，使得所有位置都有可能强烈关注关键位置处具有大范数的某些键，对于保持softmax的尖峰特性至关重要。红色项仅关注x和y之间的特定相似度，本质上是一个带宽σ² = √d的高斯核。