Flexformer: Flexible Linear Transformer with Learnable Attention Kernel¶

会议: NeurIPS 2026
arXiv: 2606.27748
OpenReview: https://openreview.net/forum?id=D8frDvzgAf
领域: 学习理论
关键词: 线性注意力, 随机傅里叶特征, 可学习核函数, 谱表示理论, 注意力蒸馏

一句话总结¶

Flexformer 基于谱表示理论将随机傅里叶特征中的谱频率与尺度参数转化为端到端可学习参数，构建包含 Softmax 核且表达力更强的平稳与非平稳注意力核族，在保持 $O(N)$ 线性时空复杂度的同时显著超越现有线性注意力基线。

研究背景与动机¶

标准 Transformer 依赖 Softmax 点积自注意力机制动态汇聚全局上下文，但计算所有位置对之间的相似度导致时间与显存复杂度随序列长度 $N$ 呈二次方增长（$O(N^2 d)$ 与 $O(N^2 + Nd)$），成为处理长上下文、极长文档与像素级序列的主要效率瓶颈。为了打破二次复杂度，基于核方法的线性注意力（如 Linear Transformer、Cosformer）通过将核函数分解为特征映射内积 $k(\mathbf{x}, \mathbf{y}) = \phi(\mathbf{x})^\top \phi(\mathbf{y})$，利用矩阵乘法结合律优先计算键值外积汇总，将计算开销降至线性的 $O(N d d')$。早期基于随机傅里叶特征（RFF）的近似方法（如 RFA、Performer）默认 Softmax 核是最优选择，将谱采样分布固化为固定的高斯密度，导致其在 ListOps、文档检索等长序列结构任务上精度大幅崩塌。

现有改进方案如 Hedgehog 和 Polaformer 尝试引入可学习映射或借助幂函数、指数函数模拟 Softmax 的低熵聚集特性，但其特征映射设计主要基于经验启发式规则，无法从理论上明确保证核函数族的完备性与表达能力上限，且跨域泛化能力脆弱。这就引出了线性注意力设计的核心矛盾：既要摆脱固定 Softmax 核次优性的桎梏并赋予模型充分拟合复杂模式的自由度，又要提供严格的数学理论依托以确保所学核函数族覆盖甚至超越标准注意力。

本文的切入角度是回到调和分析的谱表示理论——既然连续正定核与谱测度存在一一对应关系，那么学习核函数在本质上等价于直接从数据中学习其谱密度。核心 idea：Flexformer 突破随机傅里叶特征依赖固定先验采样的范式，直接将傅里叶谱频率与范数缩放因子转化为端到端可学习参数，依据 Bochner 与 Yaglom 定理推导出平稳核与非平稳核映射，在理论上严格包含标准 Softmax 核的同时赋予线性注意力充分的数据自适应表达力。

方法详解¶

整体框架¶

Flexformer 的核心目标是构建一个不仅计算复杂度为线性，且核函数形式完全由数据驱动优化的注意力机制。方法由三大核心支柱构成：基于范数分解与 Bochner 定理的平稳核构建、基于 Yaglom 定理的双频交叉乘积非平稳核扩展，以及面向预训练模型的无梯度结构注意力权重蒸馏机制。输入序列的查询向量与键向量首先经受可学习尺度的范数加权，随后映射进由可学习谱频率张成的有限维傅里叶特征空间，通过重排乘法顺序在 $O(N)$ 复杂度下完成序列聚合。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 X 映射为 Q, K, V"] --> B["谱频率与缩放尺度参数化<br/>初始化可训练谱参数与温度 τ"]
    B --> C["平稳核特征映射 Flexformer_s<br/>基于 Bochner 定理单频三角映射"]
    B --> D["非平稳核特征映射 Flexformer_n<br/>基于 Yaglom 定理双频交叉调制"]
    C --> E["全局低秩聚合与线性注意力计算<br/>利用结合律线性聚合上下文状态"]
    D --> E
    E --> F["双路径训练与部署<br/>从头预训练或预训练权重注意力蒸馏"]

关键设计¶

1. 谱频率与缩放尺度的端到端参数化：打破固定高斯谱采样的表达局限 传统 RFA 采用随机傅里叶特征近似高斯核时，必须预先从标准正态分布采样固定的正交或独立频率向量，模型只能被动逼近带有固定方差的 Softmax 相似度，无法针对特定下游语料调整注意力锐度与聚集拓扑。Flexformer 借鉴高斯过程的稀疏谱学习思想，将 $n$ 组谱频率向量直接定义为可学习参数矩阵。以平稳核变体 $\text{Flexformer}_{\text{s}}$ 为例，基于 Bochner 定理，任何有界连续正定平稳核均可由非负测度的逆傅里叶变换表示。模型将谱频率 $\{\omega_i\}_{i=1}^n \in \mathbb{R}^d$ 设为可训练权重，直接通过反向传播自适应优化谱测度分布。此外，针对标准 Softmax 核展开中依赖向量模长的因子，Flexformer 将原本固定的尺度参数 $2\sqrt{d}$ 替换为对数尺度可学习参数 $\exp(\tau)$，使特征映射不仅能自适应调节空间基底方向，还能动态拉伸高范数特征与低范数特征之间的对比度。

2. 基于 Yaglom 定理的非平稳核特征映射：突破平移不变性约束 平稳核要求相似度仅取决于差值向量 $\mathbf{x} - \mathbf{y}$，这一平移不变假设在序列建模中并不总是成立，限制了对查询与键之间非对称或绝对位置特征的建模。为此，$\text{Flexformer}_{\text{n}}$ 引入现代调和分析中的 Yaglom 理论，将正定核扩展至非平稳双频域 $\mathbb{R}^d \times \mathbb{R}^d$。为确保联合谱密度 $p(\boldsymbol{\omega}_1, \boldsymbol{\omega}_2)$ 的半正定性，采用任意联合概率密度的对称化分解，并结合三角和差化积技巧推导出紧凑实数特征映射 $\tilde{\phi}_n(\mathbf{x}) \in \mathbb{R}^{2n}$： $$ \tilde{\phi}n(\mathbf{x}) = \exp\left(\frac{|\mathbf{x}|^2}{\exp(\tau)}\right) \frac{\sigma}{2\sqrt{n}} \left[ \cos({\omega') \cos({\omega'}}^\top \mathbf{x{21}}^\top \mathbf{x}), \dots, \sin({\omega') \right]^\top $$ 其中旋转辅助频率通过成对可学习谱向量的均值与半差构建：}}^\top \mathbf{x}) \cos({\omega'_{2n}}^\top \mathbf{x$\omega'_{1i} = \frac{\omega_{1i} + \omega_{2i}}{2}$，$\omega'_{2i} = \frac{\omega_{1i} - \omega_{2i}}{2}$。由于高斯先验只是该参数空间的一个特例，非平稳特征映射不仅严格包含了标准 Softmax 的无偏估计，还具备远超平稳核的函数逼近自由度，每个注意力头仅增加 $2d^2$ 个参数即可实现完全自适应的核函数学习。

3. 注意力权重蒸馏策略：兼顾无损替换与跨域泛化迁移 针对已有大规模预训练 Transformer 无法轻易从头重训的痛点，Flexformer 设计了轻量级的注意力权重交叉熵蒸馏机制。在蒸馏阶段完全冻结 Transformer 主干的所有投影与前馈层权重，仅将原版的 Softmax 注意力计算替换为 Flexformer 的可学习特征映射，专门优化谱频率和尺度参数。目标函数直接最小化归一化核权重与教师模型 Softmax 概率分布之间的逐位置 KL 散度： $$ \mathcal{L}{awd} = -\frac{1}{N}\sum}^{N}\sum_{j=1}^{N}\frac{\exp(\frac{\mathbf{Q{i}^{\top}\mathbf{K}}}{\sqrt{d}})}{\sum_{l=1}^{N}\exp(\frac{\mathbf{Q{i}^{\top}\mathbf{K}}}{\sqrt{d}})}\log\frac{\tilde{\phi{n}(\mathbf{Q}})^{\top}\tilde{\phi{n}(\mathbf{K}})}{\sum_{l=1}^{N}\tilde{\phi{n}(\mathbf{Q}})^{\top}\tilde{\phi{n}(\mathbf{K} $$ 该蒸馏过程不仅收敛速度显著快于经验设计的 Hedgehog（在 8,000 步即可收敛稳定），而且所习得的谱参数在保留原模型语义表征的同时具备出色的域外迁移能力，单一源域蒸馏出的核参数可直接迁移至异构下游任务并维持高精度。})

损失函数 / 训练策略¶

从头预训练目标：在自回归语言建模（WikiText-103）中使用标准自回归交叉熵损失，序列块大小为 512，训练 100,000 步，采用 Adam 优化器配合学习率 $5 \times 10^{-4}$，warmup 4,000 步，梯度范数裁剪阈值为 0.1。
蒸馏训练目标：采用上述 $\mathcal{L}_{awd}$ 进行 20,000 步梯度下降，学习率设为 $1 \times 10^{-3}$，批大小为 32。
复杂度与参数预算：设置谱采样数 $n = d$，输出特征维度 $d' = 2d$。在多头设定下，各头参数保持独立，单层额外参数仅占约 1%（Stationary 为每头 $d^2$，Nonstationary 为每头 $2d^2$），计算与显存占用严格保持为 $O(N)$。

实验关键数据¶

主实验¶

论文在极具挑战的长序列基准 Long Range Arena (LRA) 上系统评估了从头训练分类性能（序列长度 1K–4K），并对比了标准 Transformer 与现有线性注意力模型。

模型	ListOps (2K)	Text (4K)	Retrieval (4K)	Image (1K)	Pathfinder (1K)	平均准确率 (%)
Transformer (Softmax)	37.23	65.66	64.49	41.89	73.46	56.55
Reformer	36.62	63.52	57.60	40.06	73.61	54.28
Longformer	36.46	63.09	59.24	41.23	72.44	54.49
Linear Transformer	17.42	62.64	58.92	43.40	75.11	51.49
Performer (固定RFF)	16.97	63.67	55.17	41.23	75.38	50.48
RFA (固定RFF)	17.20	64.89	58.76	41.00	74.51	51.27
Cosformer	37.44	65.16	64.04	43.56	72.32	56.50
Hedgehog (经验学习)	37.02	64.10	68.67	41.64	75.89	57.47
Polaformer (经验学习)	36.93	69.65	65.39	40.95	73.26	57.23
$\text{Flexformer}_{\text{s}}$ (Ours)	37.09	67.38	71.55	43.39	76.15	59.11
$\text{Flexformer}_{\text{n}}$ (Ours)	37.18	65.72	73.94	45.42	77.66	59.99

消融与语言建模实验¶

在 WikiText-103 语言建模任务上评估不同模型规模下的困惑度（PPL，越低越好），验证增加可学习核容量对语言建模能力的收益。

模型	小模型 (41M) 验证集 PPL	小模型 (41M) 测试集 PPL	大模型 (247M) 验证集 PPL	大模型 (247M) 测试集 PPL	相对参数增量 (%)
Vanilla Transformer	30.0	31.3	22.6	23.8	0.00%
Linear Transformer	34.1	35.6	27.9	29.2	0.00%
RFA	33.8	35.3	26.7	28.0	0.00%
Cosformer	32.8	34.2	24.9	26.2	0.00%
Hedgehog	31.7	33.2	23.0	24.2	+0.43%
Polaformer	31.5	32.9	23.1	24.3	+6.74%
$\text{Flexformer}_{\text{s}}$	31.1	32.7	22.2	23.4	+0.43%
$\text{Flexformer}_{\text{n}}$	30.3	31.6	21.2	22.3	+0.86%

关键发现¶

超越 Softmax 的大模型扩展表现：在小规模模型（41M）下，$\text{Flexformer}_{\text{n}}$ 测试集 PPL（31.6）已极其贴近 Vanilla Transformer（31.3）；而在 247M 大模型配置下，$\text{Flexformer}_{\text{n}}$ 实现了 22.3 的测试集 PPL，显著反超原生 Softmax（23.8）达 1.5 点，表明可学习非平稳核能有效吸纳高容量参数带来的拟合增益。
推理与显存的高效线性扩展：在 $N=4000$ 的双塔长文档检索任务中，Flexformer 相较于标准 Softmax 实现了 2.6 倍的端到端训练加速，并大幅降低显存消耗达 84%，前向传播显存与耗时随长度增加严格呈斜率平缓的线性增长。
谱参数量与采样维度敏感性：实验显示当采样频率数 $n$ 从较小值增至 $d$ 时，模型拟合能力持续提升；在自然语言大模型任务上继续增加 $n$ 能带来稳健泛化，而在 LRA 简单像素级任务上 $n=d$ 已达到最佳性能功耗平衡点。

亮点与洞察¶

将随机傅里叶特征从固定采样转化为可学习参数：以往研究将傅里叶特征单纯视为逼近固定高斯核的蒙特卡洛积分工具，本文从调和分析对偶理论出发，将学习正定核转化为直接学习离散谱测度，打通了核函数族自由学习与高效线性计算之间的阻隔。
Yaglom 理论赋能非平稳双频交互：巧妙借助和差化积将非平稳核积分转化为双正弦余弦乘积形式，打破了平移不变性限制，仅需为每个注意力头引入微量二次级参数（$<1\%$ 总模型参数），却赋予模型捕捉非对称依赖的能力。
注意力蒸馏的高跨域复用价值：实验表明在单一下游语料（如 QNLI）上蒸馏的注意力核参数，迁移至其他 7 个完全异构的 GLUE 任务时，平均准确率降幅显著低于基于经验映射的 Hedgehog，为大模型向线性注意力架构的轻量化无损转换提供了可行范式。

局限与展望¶

依赖正定核假定的理论盲区：论文依赖 Bochner 与 Yaglom 定理推导特征映射的前提是核函数必须是正定的。尽管实验表明正定性能有效抑制退化注意力，但注意力机制在理论上是否必须严格满足半正定性仍缺乏根本性数学证明。
模态与长上下文覆盖局限：目前的评测主要局限在 NLP 任务与 LRA 合成测试集，尚未在包含百万级上下文长度的长文本理解、高分辨率视频扩散模型或生物大分子长序列等极端异构长距离场景上充分验证。
硬件感知的高性能算子实现空间：当前实现主要依赖标准 PyTorch 算子，若结合类似 FlashLinearAttention 或 Triton 编写针对双频乘积映射与因果前缀和融合的高性能定制 GPU Kernel，吞吐量有望进一步提升。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙利用 Bochner 与 Yaglom 定理将傅里叶特征谱转为可微可训参数，立意高远且数学推导严密自洽。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 LRA 分类、WikiText 自回归语言建模、RoBERTa 权重蒸馏以及跨域迁移实验，多角度验证且具备误差线。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨，方法推导清晰流畅，理论定理与实验验证层层递进。
价值: ⭐⭐⭐⭐☆ 为长序列大模型摆脱二次计算瓶颈提供了一种兼具理论完备性、高效线性推理与良好蒸馏替换能力的优秀方案。