Flexformer: Flexible Linear Transformer with Learnable Attention Kernel¶

讲者: Feng Zhou
会场: AI for Complex Systems: Methods and Applications
报告题目: Flexformer: Flexible Linear Transformer with Learnable Attention Kernel
链接: arXiv
来源: JCSDS 2026 · 返回会议总览

一、领域脉络与小综述¶

这个方向是什么¶

这个子方向是高效Transformer中的线性注意力机制。它要解决的根本问题是：标准Transformer的softmax注意力机制的计算复杂度和空间复杂度随序列长度N呈二次增长（O(N²)），这严重限制了其在长序列任务（如长文档、高分辨率图像、基因序列）上的可扩展性。当前，该方向已从早期的近似方法（如稀疏注意力、低秩近似）发展到更系统的核方法线性注意力，其核心是将注意力计算重写为核函数的内积形式，并通过改变矩阵乘法顺序将复杂度降至O(N)。然而，如何设计一个既保持线性复杂度、又具备足够表达力以匹配甚至超越softmax注意力的核函数，是该方向当前的核心瓶颈。

发展脉络（history）¶

奠基工作：Transformer与二次复杂度瓶颈
- Vaswani et al. (2017) [38]：提出Transformer架构，其核心是点积softmax注意力。该机制通过计算所有位置对的相似度来捕捉长程依赖，但复杂度为O(N²d)，成为长序列应用的瓶颈。这是整个方向的起点。
- Devlin et al. (2019) [9] & Brown et al. (2020) [5]：BERT和GPT-3等大规模预训练模型的成功，使得Transformer成为NLP领域的事实标准，也使得解决其二次复杂度问题变得愈发迫切。
主要进展：降低复杂度的两条路径
- 路径一：稀疏注意力。这类方法通过限制每个token的注意力范围来降低复杂度。
  - Kitaev et al. (2020) [15] (Reformer)：使用局部敏感哈希（LSH）将注意力复杂度降至O(N log N)。作者在本文中评价其“attention patterns are not rigorously validated and are often based on empirically derived patterns”。
  - Zaheer et al. (2020) [41] (BigBird)：提出一种结合局部、全局和随机注意力的稀疏模式，将复杂度降至O(N)。作者在本文中将其归为“sparse attention”一类。
  - Roy et al. (2020) [32] (Routing Transformer)：通过学习基于k-means的动态稀疏模式来分配注意力。作者同样将其归为“sparse attention”。
- 路径二：核方法线性注意力。这类方法将注意力核表示为特征映射的内积，通过改变计算顺序实现线性复杂度。
  - Katharopoulos et al. (2020) [13] (Linear Transformer)：首次提出将softmax替换为简单的核函数（如ELU+1），实现了线性注意力。作者评价其“efficient but limited in expressiveness”。
  - Choromanski et al. (2020) [7] (Performer) & Peng et al. (2021) [28] (RFA)：使用随机傅里叶特征（RFF）来获得softmax核的无偏估计。作者评价其“assuming softmax attention is optimal and fixing the kernel a priori”，即它们隐含地假设softmax是最优的，并将核函数固定下来，不具备可学习性。
  - Qin et al. (2022) [29] (Cosformer)：使用余弦函数作为核函数。作者评价其“efficient but limited in expressiveness”。
当前Frontier：可学习核函数
- Zhang et al. (2024) [42] (Hedgehog) & Meng et al. (2025) [22] (PolaFormer)：意识到固定核函数的局限性，开始引入可学习的特征映射。Hedgehog通过可学习的MLP和指数函数来模仿softmax的“低熵”和“单调性”特性。PolaFormer则显式建模正负查询-键交互，并使用可学习的幂函数来恢复“尖峰”特性。作者评价这些方法“empirically design kernels with desirable properties, but the expressiveness of their kernel families is not explicitly characterized”，即它们的设计是经验性的，其核函数族的表达能力没有被明确刻画。
本文的位置
- 本文（Flexformer）位于可学习核函数这一前沿。它区别于Hedgehog和PolaFormer的关键在于，其方法根植于谱表示理论（Bochner定理和Yaglom定理），通过将谱密度中的频率参数作为可训练参数，从而在理论上保证能够学习一个“substantially broader and more expressive family of kernels”。作者声称，Flexformer的核函数族严格包含softmax核的无偏估计，因此它既能恢复softmax，也有潜力学习到比softmax更好的注意力模式。

子线索聚类¶

稀疏注意力：Reformer, BigBird, Routing Transformer, Longformer。核心思路是设计固定的或可学习的稀疏模式，只计算部分token对的注意力。优点是直观，但缺点是模式设计依赖经验，且难以保证在所有数据上一致有效。
固定核线性注意力：Linear Transformer, Performer, RFA, Cosformer。核心思路是使用一个固定的、非学习的核函数来近似softmax或替代它。优点是计算高效，但缺点是表达力受限，且隐含假设softmax是最优的。
可学习核线性注意力：Hedgehog, PolaFormer, Flexformer (本文)。核心思路是引入可学习的参数来构造特征映射，使注意力模式能适应数据。这是当前最活跃的方向，Flexformer的独特之处在于其理论根基（谱表示）和由此带来的更广的核函数族。

这个方向在追问的核心问题¶

表达力 vs. 效率的权衡：如何设计一个特征映射，使其在保持O(N)复杂度的同时，其核函数族足够大，能够逼近甚至超越softmax注意力的性能？
可学习性的理论保证：可学习核函数的参数空间是否足够大？其对应的核函数族是否包含softmax？能否从理论上刻画其表达能力？
跨领域泛化：在一个任务/领域上学到的注意力核，能否有效地迁移到另一个任务/领域？
正定性的必要性：当前所有核方法（包括softmax）都假设核函数是正定的。这是否是注意力机制的最优选择？是否存在非正定但更有效的核函数？

⚠️ 作者的framing¶

作者把缺口frame成什么：作者将现有工作的缺口定位为“核函数表达力不足且缺乏理论刻画”。具体来说，作者认为：
- 固定核方法（Performer, RFA）隐含假设softmax最优，限制了潜力。
- 可学习核方法（Hedgehog, PolaFormer）是经验性设计，其核函数族的表达能力没有被明确刻画，不能保证足够灵活。
- 因此，本文的“显然的下一步”是：基于谱表示理论，将频率作为可学习参数，从而在理论上保证能学习一个更广泛、更具表达力的核函数族。
哪些竞争路线被他淡化或回避了：
- 稀疏注意力：作者在引言中一笔带过，称其“not rigorously validated and are often based on empirically derived patterns”，将其定位为一种缺乏理论保证的替代方案。这淡化了稀疏注意力在某些任务上的成功和其理论分析（如BigBird的通用逼近性）。
- 状态空间模型（SSM）：本文完全没有提及近年来在长序列建模上表现优异的状态空间模型（如Mamba）。这是一个明显的回避。SSM是线性注意力之外的另一条重要技术路线，其成功对“线性注意力是否必要”构成了挑战。
什么明显该被引/该存在、却没出现在intro里？
- 状态空间模型（SSM）：如Mamba (Gu & Dao, 2023) 等。这些模型在长序列任务上取得了与Transformer竞争甚至更好的性能，且同样具有线性复杂度。不讨论SSM使得本文的动机（“我们需要线性注意力”）显得不够完整。
- 更近期的线性注意力工作：论文发表于2026年，但引用的最晚近的可学习核工作是2025年的PolaFormer。可能存在2025-2026年间发表的、同样基于谱方法或更先进方法的线性注意力工作未被引用。

张力¶

未见明显对立引用。所有被引工作都认同“降低注意力复杂度”是核心目标，分歧在于实现路径（稀疏 vs. 核方法）和核函数的设计哲学（固定 vs. 可学习）。本文与Hedgehog/PolaFormer的竞争关系是直接的，但作者将其定位为“更优的理论框架”，而非根本性的对立。

二、最核心、最简单的例子 / 数学问题¶

第一步：把符号、模型、可观测数据交代清楚¶

符号：
- N：输入序列的长度（token数量）。
- d：每个token的嵌入维度（hidden dimension）。
- X ∈ ℝ^{N×d}：输入序列矩阵，每一行X_i是一个token的嵌入向量。
- Q, K, V ∈ ℝ^{N×d}：查询（Query）、键（Key）、值（Value）矩阵，由X通过线性变换得到。
- sim(·,·): 相似度函数，用于计算查询和键之间的注意力权重。在softmax注意力中，sim(Q_i, K_j) = exp(Q_i^T K_j / √d)。
- k(x, y): 注意力核函数，是sim(·,·)的泛化。在核方法线性注意力中，k(x, y) = ϕ(x)^T ϕ(y)。
- ϕ: ℝ^d → ℝ^{d'}: 特征映射，将d维的查询/键向量映射到d'维的特征空间。
- d': 特征映射的维度。在Flexformer中，d' = 2n，其中n是频率采样数。
- ω ∈ ℝ^d: 频率向量，是谱密度p(ω)的采样点。在Flexformer中，ω是可学习的参数。
- n: 频率采样数（Monte Carlo样本数）。
- τ: 可学习的缩放参数，用于非平稳变体中的范数项。
- σ²: 谱密度的总测度（总质量），通常被视为一个常数或可学习参数。
模型：
- 本文研究的是Transformer模型中的自注意力层。给定输入序列X，自注意力层的输出O由下式计算： O_i = ( Σ_{j=1}^N sim(Q_i, K_j) V_j ) / ( Σ_{j=1}^N sim(Q_i, K_j) )
- 在核方法线性注意力中，sim(Q_i, K_j)被替换为一个核函数k(Q_i, K_j) = ϕ(Q_i)^T ϕ(K_j)。通过矩阵乘法的结合律，可以避免显式计算N×N的注意力矩阵，从而将复杂度从O(N²d)降至O(N d d')。
- Flexformer的模型：它定义了一个特定的、可学习的特征映射ϕ。这个ϕ基于随机傅里叶特征，但其核心参数——频率{ω_i}——不再是采样自一个固定的分布（如高斯分布），而是被当作可训练的参数，通过反向传播从数据中学习。
可观测数据：
- 可观测：输入序列X（或其经过嵌入层后的表示）。这是研究者能直接获得的。
- 想要但观测不到：最优的注意力核函数k*(·,·)。我们不知道对于给定任务，什么样的相似度函数是最好的。softmax只是一个先验选择。Flexformer的目标就是通过数据驱动的方式，学习到一个能逼近这个未知最优核的核函数。

第二步：讲最小内核¶

本文的核心思路是：通过将随机傅里叶特征中的频率参数从“采样”变为“学习”，从而将固定核函数扩展为一个可学习的、更广泛的核函数族。

最简特例：一维输入 (d=1) 下的平稳核学习

为了看清本质，我们考虑一个极端简化的场景： - 输入维度d=1。那么查询Q_i和键K_j都是标量。 - 我们只考虑平稳核，即k(x, y) = k(x-y)。这对应于Flexformer的平稳变体Flexformer_s。 - 我们暂时忽略softmax核中的范数项exp(||x||² / (2√d))，只关注其核心的高斯核部分：k_G(x, y) = exp(-(x-y)² / (2√d))。

在这个特例下，发生了什么？

传统方法（RFA/Performer）：
- 根据Bochner定理，高斯核k_G(x-y)的谱密度p(ω)是一个均值为0、方差为1/√d的高斯分布。
- 传统方法从这个固定的高斯分布中采样n个频率{ω_i}。
- 然后构造特征映射ϕ_n(x) = (σ/√n) [cos(ω_1 x), ..., cos(ω_n x), sin(ω_1 x), ..., sin(ω_n x)]^T。
- 这样，ϕ_n(x)^T ϕ_n(y)就是k_G(x-y)的一个无偏估计。
- 关键：这个核函数是固定的。无论数据是什么，它都只能逼近这个特定的高斯核。
Flexformer的方法：
- Flexformer不从任何分布中采样{ω_i}。
- 它将{ω_i}初始化为一些随机值（例如，从标准正态分布中采样），然后将其声明为可训练的参数。
- 在训练过程中，通过反向传播和梯度下降，这些{ω_i}会根据任务目标（如语言建模的交叉熵损失）被优化。
- 优化后的{ω_i}不再服从一个简单的高斯分布。它们可能聚集在某些特定频率附近，从而学习到一个非高斯、数据自适应的谱密度p*(ω)。
- 根据Bochner定理，这个学习到的谱密度p*(ω)对应着一个新的、数据驱动的平稳核k*(x-y)。这个核函数可能比高斯核更适合当前任务。

这个最小内核说明了什么？

核心数学困难：如何在不显式参数化整个核函数（这在高维下不可行）的情况下，学习一个灵活的核函数？
本文的关键想法：利用Bochner定理（平稳核）和Yaglom定理（非平稳核），将“学习核函数”这个困难问题，转化为“学习其谱密度”的问题。而谱密度可以通过一组可学习的频率参数{ω_i}来参数化。这相当于在频域空间中进行学习，而频域空间的维度（由n控制）可以独立于输入维度d，从而提供了灵活性。
结论：这篇论文在数学上干了一件什么事？它将核方法线性注意力中的固定核函数，替换为一个由可学习频率参数定义的、基于谱表示的核函数族。这个核函数族理论上可以逼近任何满足条件的平稳/非平稳核，从而极大地提升了线性注意力的表达能力。

三、这篇论文做了什么¶

三句话¶

研究了什么问题：如何设计一个既保持线性计算复杂度、又具备足够表达力以匹配甚至超越标准softmax注意力的可学习注意力核函数。
核心工具/方法：基于Bochner定理和Yaglom定理，将注意力核的学习转化为其谱密度的学习，具体地，将随机傅里叶特征中的频率参数{ω_i}作为可训练参数，通过端到端优化从数据中学习注意力核。提出了平稳（Flexformer_s）和非平稳（Flexformer_n）两种变体。
主要结论：在长序列分类（LRA）、自回归语言建模（WikiText-103）、softmax注意力恢复（GLUE）和跨域核迁移等实验中，Flexformer一致性地优于所有现有的线性注意力基线（包括RFA, Performer, Hedgehog, PolaFormer等），并在大模型配置下超越了标准softmax Transformer。非平稳变体Flexformer_n通常表现最佳。

关键设定与假设¶

设定：标准的Transformer自注意力层。输入序列X ∈ ℝ^{N×d}。目标是设计一个核函数k(x, y)来替代sim(Q_i, K_j) = exp(Q_i^T K_j / √d)，使得注意力计算可以写成k(Q_i, K_j) = ϕ(Q_i)^T ϕ(K_j)的形式，从而利用矩阵乘法的结合律实现线性复杂度。
假设1：核函数的正定性。这是应用Bochner定理和Yaglom定理的前提。作者在Limitations中明确指出，该假设“remains theoretically unexplored”，即正定核是否是注意力机制的最优选择尚无理论保证。这是一个重要的潜在弱点。
假设2：softmax核的分解形式。作者将softmax核分解为exp(||x||²/(2√d)) * exp(-||x-y||²/(2√d)) * exp(||y||²/(2√d))。其中，蓝色项（范数项）被保留为固定形式（或在非平稳变体中用一个可学习参数τ替换其分母），红色项（高斯核）被替换为可学习的核。这个分解是合理的，但它隐含地假设了注意力核的“位置信息”和“交互信息”是可分离的。
假设3：谱密度的对称性。在推导平稳核的RFF近似时，作者假设谱密度p(ω)是对称的，即p(ω) = p(-ω)，这使得虚部期望为零。这是一个标准假设，用于保证核函数是实值的。
相比已有文献的放宽/强化：
- 放宽：相比RFA和Performer，Flexformer放宽了“核函数必须是softmax”的假设，允许学习任意满足正定条件的核。
- 强化：相比Hedgehog和PolaFormer，Flexformer强化了核函数族的理论边界。作者声称其核函数族“strictly contains the unbiased estimate of the softmax kernel”，并且基于谱表示理论，其表达能力在理论上更清晰。

主要结果¶

长序列分类（LRA）：Flexformer_n在5个任务上取得了平均59.99%的准确率，显著优于所有基线。相比最好的现有线性注意力方法Hedgehog（57.47%），相对提升约4.4%。在Document Retrieval任务上，Flexformer_n达到73.94%，远超第二名Hedgehog的68.67%。同时，Flexformer_n保持了与Linear Transformer相近的线性复杂度，相比标准Transformer实现了2.6倍训练加速和84%的内存节省。
自回归语言建模（WikiText-103）：
- 小模型（41M）：Flexformer_n在测试集上达到31.6 PPL，优于所有线性注意力方法（Hedgehog: 33.2, PolaFormer: 32.9），并接近标准Transformer（31.3）。
- 大模型（247M）：Flexformer_n在测试集上达到22.3 PPL，超越了标准Transformer（23.8 PPL）。这是一个关键结果，表明可学习核在大模型下能利用额外参数获得超越softmax的性能。作者将此归因于“learnable kernels benefit more from increased model capacity”。
恢复Softmax注意力（GLUE）：通过注意力蒸馏，Flexformer_n在8个GLUE任务上的平均得分（83.7）非常接近标准RoBERTa（85.5），并显著优于不可学习的Performer（63.3）和RFA（64.9）。在CoLA任务上，Flexformer_n（63.3）甚至超过了RoBERTa（61.9），表明其学到了比softmax更好的模式。
跨域核迁移：Flexformer_n在几乎所有跨域迁移场景下都优于Hedgehog，尤其是在CoLA任务上，Hedgehog的性能下降严重（从62.5降至55.6），而Flexformer_n的下降较小（从63.3降至60.3），展示了更强的泛化能力。

证明路线与技术技巧¶

本文是方法/应用型论文，没有传统意义上的“定理证明”。其核心“证明”是理论上的保证，即通过Bochner/Yaglom定理保证其核函数族的表达能力。技术路线如下：

理论奠基：引用Bochner定理（定理3.1）和Yaglom定理（定理3.2），建立核函数与其谱密度之间的一一对应关系。这为“学习核函数等价于学习谱密度”提供了理论基础。
参数化谱密度：将谱密度p(ω)（或g(ω₁, ω₂)）通过一组可学习的频率{ω_i}（或{(ω₁i, ω₂i)}）进行参数化。这是关键跳跃点：从“从固定分布采样”到“将采样点作为可学习参数”。这个技巧借鉴了高斯过程文献中的“Sparse Spectrum Gaussian Process Regression” [17]。
构造可学习特征映射：利用Monte Carlo近似，将谱密度的积分形式转化为求和形式，从而构造出可学习的特征映射ϕ_n(x)。对于平稳变体，ϕ_n(x)由cos(ω_i^T x)和sin(ω_i^T x)组成；对于非平稳变体，ϕ_n(x)由cos(ω'_{1i}^T x)cos(ω'_{2i}^T x)和sin(ω'_{1i}^T x)cos(ω'_{2i}^T x)组成，其中ω'_{1i}, ω'_{2i}是(ω₁i, ω₂i)的线性变换。
端到端学习：将构造好的ϕ_n(x)嵌入到线性注意力框架中，整个模型（包括Transformer的其他部分和频率参数{ω_i}）通过反向传播进行端到端训练。频率参数的梯度可以通过自动微分框架（如PyTorch）轻松计算。
非平稳扩展：通过Yaglom定理，将方法从平稳核扩展到非平稳核。非平稳变体通过引入两组频率(ω₁i, ω₂i)和额外的可学习缩放参数τ，获得了“strictly greater expressiveness”。

技术技巧点名： - 随机傅里叶特征（Random Fourier Features）：用于近似平稳核。这是整个方法的基础。 - 谱表示理论（Bochner/Yaglom定理）：提供了将核学习问题转化为谱学习问题的理论依据。 - 可学习频率参数：将采样点视为参数，是本文的核心创新点，借鉴自稀疏谱高斯过程。 - 注意力蒸馏损失：使用交叉熵损失来监督Flexformer的注意力权重分布去匹配预训练Transformer的softmax注意力分布，从而实现模型转换。

真实例子与应用¶

本文有丰富的真实数据实验，是典型的应用/方法型论文。

数据/场景：
- 长序列分类：LRA基准，包含ListOps（层次结构推理）、IMDb（文本分类）、AAN（文档检索）、CIFAR-10（像素序列分类）、Pathfinder（空间推理）。序列长度1K-4K。
- 语言建模：WikiText-103，包含超过1亿词的大规模语言建模基准。
- 注意力恢复与迁移：GLUE基准，包含8个不同的自然语言理解任务（语法可接受性、情感分析、释义检测、语义相似度、问题匹配、自然语言推理等）。
方法应用：
- 从零训练：在LRA和WikiText-103上，直接将Flexformer作为Transformer的注意力模块进行训练。
- 蒸馏：在GLUE上，先冻结预训练的RoBERTa模型，用Flexformer替换其注意力层，然后通过最小化注意力权重分布的KL散度（蒸馏损失）来训练Flexformer的核参数。之后，再对整个模型（包括核参数和RoBERTa的其他参数）进行微调。
结果：
- LRA：Flexformer_n在5个任务中的3个上取得最佳，平均准确率最高。
- WikiText-103：大模型下，Flexformer_n的PPL（22.3）低于标准Transformer（23.8），证明了其超越softmax的潜力。
- GLUE：蒸馏后的Flexformer_n几乎完全恢复了RoBERTa的性能（平均83.7 vs 85.5），并在CoLA上超越。
- 跨域迁移：在一个任务上蒸馏的核，迁移到其他任务后，Flexformer_n的性能下降远小于Hedgehog。
例子想说明什么：
- LRA实验：验证Flexformer在长序列任务上的有效性和效率。
- WikiText-103实验：验证Flexformer在大规模、复杂任务上的表达力，特别是其在大模型下超越softmax的能力。
- GLUE蒸馏实验：验证Flexformer的实用价值——可以作为一个“即插即用”的模块，将已有的高性能Transformer模型高效地转换为线性版本，而性能损失很小。
- 跨域迁移实验：验证Flexformer学到的核函数具有更好的泛化性和鲁棒性，不是简单地过拟合到特定任务。

🔎 结论是否比证明窄¶

是。作者在结论和Limitations中承认了这一点。
具体语句：
- “Flexformer assumes the learned attention kernel is positive definite, but whether positive definite kernels are truly the best choice for attention remains theoretically unexplored.” —— 这是对核心假设的质疑。论文的所有理论保证（Bochner/Yaglom定理）都建立在正定性之上，但作者无法证明正定性对于注意力机制是最优的。因此，论文的结论（“Flexformer can learn a broad family of kernels”）实际上被限制在了“正定核”这个子集内。
- “Our experiments are primarily conducted on tasks and datasets related to natural language processing.” —— 论文的实证结论主要局限于NLP领域。虽然LRA提供了其他模态的测试，但作者承认这“remains insufficient to fully reflect the diverse range of attention patterns required across different modalities and domains”。因此，关于Flexformer“通用有效性”的结论，其证据基础是相对狭窄的。

四、开放问题¶

正定核假设的理论验证：本文的核心假设是注意力核应为正定。但作者在Limitations中承认，这是否最优“remains theoretically unexplored”。一个开放问题是：能否从注意力机制的功能（如信息聚合、特征选择）出发，证明或证伪正定核的必要性？ 如果非正定核在某些场景下更优，Flexformer的框架能否扩展以学习它们？（扎根于论文第6节Limitations第一段）
非平稳核的理论保证：本文的非平稳变体Flexformer_n基于Yaglom定理，但为了进行Monte Carlo采样，作者对谱密度p(ω₁, ω₂)做了一个对称化假设（p(ω₁, ω₂) = 1/4 (g(ω₁, ω₁) + g(ω₂, ω₂) + g(ω₁, ω₂) + g(ω₂, ω₁))）。这个假设对最终学到的核函数族有何影响？能否在不做此假设的情况下，设计出更高效或更具表达力的非平稳核学习方案？（扎根于论文第3.2节，公式(6)之前的推导）
跨模态泛化的系统性研究：论文的实验主要集中在NLP和LRA（包含图像、文本等）。一个开放问题是：Flexformer学到的核函数在不同模态（如时间序列、图数据、蛋白质序列）之间是否具有可迁移性？ 其性能与针对特定模态设计的注意力机制（如Swin Transformer的局部注意力、Informer的ProbSparse注意力）相比如何？（扎根于论文第6节Limitations第二段）
计算-统计权衡：Flexformer通过增加可学习参数（频率）来提升表达力。这引入了一个新的计算-统计权衡：增加频率数n（从而增加参数和计算量）带来的性能提升，是否存在一个理论上的饱和点？ 这个饱和点与序列长度N、隐藏维度d以及任务复杂度有何关系？能否从信息论或统计学习理论的角度给出一个指导性的n的选择准则？（扎根于论文附录D.1的超参数敏感性分析，图3-5）

Maintained by 陈星宇 · Homepage · Source on GitHub