🤖 machine learning

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

Flexformer 是一种灵活的线性 Transformer，它通过以可训练频谱频率的方式以完全数据驱动的方式学习注意力核，增强了长序列的表达能力和可扩展性，在保持效率和迁移性的同时，始终优于基准模型。

原作者： Haoran Zhang, Feng Zhou

发布于 2026-06-29

📖 1 分钟阅读☕ 轻松阅读

原作者： Haoran Zhang, Feng Zhou

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图组织一场拥有数千名宾客的大型派对。在 AI 的世界里，这些宾客就是“Token”（文本或数据的片段），而目标是搞清楚谁需要与谁交流，从而理解整个故事。

问题所在：“握手”瓶颈

传统的 AI 模型（称为 Transformer）使用一种叫做 Softmax Attention 的方法。你可以把它想象成这样一条规则：每一位到场的宾客都必须与每一位其他宾客握手，以观察谁是相关的。

优点： 它非常精准。每个人都能对群体有一个完美的理解。
缺点： 如果有 1,000 名宾客，那就是 1,000,000 次握手。如果你有 10,000 名宾客，那就是 100,000,000 次握手。工作量呈二次方增长。这就像是在为一整座城市组织派对；计算机很快就会耗尽内存和时间。

旧的解决方法：“僵化的剧本”

为了解决这个问题，研究人员尝试了 Linear Attention（线性注意力机制）。他们不再让每个人都互相握手，而是使用了一种捷径。他们为每个人分配一个“标签”（特征映射），然后只比较这些标签。

代价： 大多数这类捷径都使用一个固定的剧本。它们假设存在一种唯一的正确方式来为人们贴标签（通常基于一个被称为“softmax kernel”的特定数学公式）。
缺陷： 仅仅因为某个剧本适用于某种类型的派对，并不意味着它适用于所有类型的派对。有时，这种“固定的标签”会遗漏重要的联系，导致 AI 变得不够聪明。

新的解决方案：Flexformer

论文作者提出了 Flexformer。你可以把 Flexformer 想象成一个智能、可定制的标签系统，它能在派对进行的过程中，学习出最合适的标签。

以下是它的工作原理，使用简单的类比：

1. “收音机调频”类比

想象 AI 连接人们的方式就像调节收音机。

旧的线性注意力机制： 收音机被固定在一个特定的频率上。它只能听到一个频道。
Flexformer： 收音机拥有一个可调旋钮。它不再被固定在一个频率上，Flexformer 将“频率”（决定人们如何连接的数学设置）视为可学习的旋钮。
它是如何学习的： AI 通过调高或调低这些旋钮，倾听数据，从而找到捕捉最重要的关系的完美“频道”。它不是在靠猜，而是精准地学习针对正在阅读的特定文本最有效的方法。

2. “灵活的橡皮筋”

论文创建了该系统的两个版本：

静态版本 (Flexformer_s)： 想象一根无论你怎么拉都会以同样方式伸展的橡皮筋。它很灵活，但伸展的规则是一致的。
非静态版本 (Flexformer_n)： 这就像一根超灵活、能变形的橡皮筋。它可以根据你在序列中的具体位置，进行伸展、扭转并改变其规则。论文声称这个版本更加强大，因为它能够适应“一致性”版本可能会错过的复杂模式。

为什么这很重要？

这篇论文证明了三件主要的事情：

快速且精简： 就像旧的线性方法一样，Flexformer 保持了较低的“握手”次数。它的扩展是线性的（1,000 名宾客 = 1,000 次握手，而不是 1,000,000 次）。这意味着它可以处理极长的文档（如整本书或长视频转录文本），而不会导致计算机崩溃。
更聪明： 因为它学习自己的“标签”而不是使用固定的剧本，它对数据的理解比旧的线性方法更好。在测试中（如阅读理解和预测句子中的下一个词），Flexformer 打败了所有其他的快速方法，甚至达到了或超过了那些缓慢、沉重的“金标准”模型的水平。
它可以“模仿”旧的方式： 如果你已经有一个缓慢但完美的 AI 模型，并想让它变快，你可以通过“蒸馏”（教导）的方法让 Flexformer 复制该模型的行为。Flexformer 可以学会表现得像那个缓慢、完美的模型一模一样，但运行速度却极快。此外，如果你在某个主题（如新闻文章）上教导它，它将比其他快速模型更好地将知识迁移到另一个主题（如科学论文）中。

总结

Flexformer 是一种让 AI 阅读长文本的新方法。它不再强迫 AI 使用一套僵化的、预先写好的规则书来连接想法，而是给了 AI 一组可调节的旋钮。AI 通过学习来旋转这些旋钮，找到连接想法的完美方式，从而产生一个既足够快以处理长文档，又足够聪明以理解复杂细节的系统。

技术摘要：Flexformer

问题陈述

Transformer 模型依赖点积注意力机制来捕捉长程依赖，但该机制相对于序列长度 $N$ 具有二次方的时间和空间复杂度 ( $O(N^2)$ )。这一瓶颈严重限制了其在长序列上的可扩展性。虽然基于核函数的线性注意力方法通过使用特征映射近似注意力机制，将复杂度降低到了线性水平 ( $O(N)$ )，但现有方法在效率与表达能力之间面临权衡：

固定核（Fixed Kernels）： 如 Performer 和 RFA 等方法依赖随机傅里叶特征（RFF）来近似 softmax 核。然而，它们将核视为固定的且不可学习的，预设 softmax 是最优的，但这在所有场景下未必成立。
弱可学习核（Weakly Learnable Kernels）： 最近的研究（如 Hedgehog、Polaformer）引入了可学习映射并结合特定函数（指数、幂函数）来模拟 softmax 的低熵特性。然而，这些方法无法保证所学习的核族是否具有足够的表达能力以涵盖广泛的注意力模式，从而可能损害性能。

核心挑战在于构建一种灵活且可学习的注意力核，在保持线性复杂度的同时，提供足以超越固定核和弱可学习基准模型的表达能力。

方法论

作者提出了 Flexformer，这是一个利用随机傅里叶特征（RFF）以完全数据驱动方式学习注意力核的线性 Transformer 框架。

理论基础

该方法基于谱表示理论（Spectral Representation Theory）。

平稳核（Stationary Kernels）： 根据 Bochner 定理（定理 3.1），平稳核 $k(x-y)$ 对应于谱密度 $p(\omega)$ 。标准 RFF 通过从固定分布（例如高斯分布用于近似 softmax）中采样频率 $\omega$ 来进行近似。
非平稳核（Non-Stationary Kernels）： 根据 Yaglom 定理（定理 3.2），非平稳核可以通过联合谱密度进行表示。

Flexformer 架构

Flexformer 不再从固定分布中采样频率，而是将谱频率视为可训练参数，进行端到端的优化。

Flexformer $_s$ (平稳型)： 频率 $\{\omega_i\}$ 是可学习的参数。这使得模型能够直接从数据中学习最优的谱密度 $p(\omega)$ 。其特征映射构造为：
$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{2\sqrt{d}}\right) \phi_n(x)$
其中 $\phi_n(x)$ 由可学习频率的余弦项和正弦项组成。该公式严格包含了 softmax 核的无偏估计，保证了在需要时能够恢复 softmax 注意力。
Flexformer $_n$ (非平稳型)： 为了增加表达能力，作者利用 Yaglom 定理将框架扩展到非平稳核。这涉及学习频率对 $\{(\omega_{1i}, \omega_{2i})\}$ 以及一个缩放参数 $\tau$ （取代固定的 $2\sqrt{d}$ ）。其特征映射变为：
$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{\exp(\tau)}\right) \phi_n(x)$
在理论上，非平稳变体比平稳变体具有更强的表达能力。

复杂度

两种变体均保持了相对于序列长度 $N$ 的线性时间复杂度 $O(Ndd') $以及**线性空间复杂度**$ O(Nd + Nd' + dd') $，其中$ d $是隐藏层维度，$ d'$ 是特征映射维度。

核心贡献

可学习核框架： 提出了 Flexformer，它在随机傅里叶特征框架内利用可学习的谱频率，使模型能够以数据驱动的方式学习广泛的注意力核族。
理论表达能力： 证明了 Flexformer 核族严格包含 softmax 核（对于平稳变体），并且在非平稳变体中提供了更强的表达能力，解决了固定核或经验设计核的局限性。
蒸馏与迁移性： 展示了 Flexformer 可以有效地从预训练 Transformer 中进行蒸馏以恢复 softmax 注意力性能，并表现出强大的跨域核迁移能力。

实验结果

作者在语言建模、序列分类和蒸馏任务上对 Flexformer 进行了评估。

长程竞技场 (LRA)： Flexformer 在所有任务（ListOps, Text, Retrieval, Image, Pathfinder）中取得了最高的平均准确率。具体而言，Flexformer $_n$ 在文档检索、图像分类和 Pathfinder 上达到了最佳结果，相比于表现最好的现有线性注意力基准，其平均准确率相对提升了 4.4%。
自回归语言建模 (WikiText-103)：
- 在小模型（41M 参数）中，Flexformer 各变体优于所有其他线性注意力基准，并接近 vanilla Transformer 的性能。
- 在大模型（247M 参数）中，Flexformer $_n$ 超越了 vanilla Transformer，这表明可学习核能显著受益于增加的模型容量。
效率： 在文档检索任务上，与 vanilla Transformer 相比，Flexformer 实现了 2.6 倍的训练加速 和 84% 的显存占用减少，同时保持了时间与显存的线性缩放。
蒸馏 (GLUE 基准测试)： 在从预训练 RoBERTa 模型进行蒸馏时，Flexformer $_n$ 几乎完全恢复了 softmax 注意力的行为，优于非可学习线性方法（Performer, RFA），甚至在 CoLA 数据集上超过了标准的 softmax。
核迁移性： 在跨域实验中，由 Flexformer $_n$ 在一个 GLUE 数据集上学习到的核，比 Hedgehog 学习到的核能更有效地迁移到其他数据集，表现出在分布外任务上更小的性能下降。

意义与主张

本文声称 Flexformer 为长序列建模提供了一个有效的、可扩展的 softmax 注意力替代方案。其主要意义在于解决了线性注意力中效率与表达能力之间的矛盾：

它不再假设 softmax 是最优的，而是直接从数据中学习核。
它提供了一种基于谱表示理论的构建灵活核的有原则的方法，而非依赖于经验性的启发式方法。
它证明了线性注意力可以在保持计算效率的同时，匹配甚至超越二次方注意力的性能，特别是在规模化或蒸馏时。

局限性： 作者指出，Flexformer 假设学习到的注意力核是正定的。正定核是否在理论上是注意力机制的最优选择，仍是未来研究的一个开放性问题。