FNet：在Transformer编码器中使用傅里叶变换混合词元

1. 引言与概述

自Vaswani等人于2017年提出以来，Transformer架构已成为自然语言处理领域事实上的先进标准。其核心创新——自注意力机制——使得模型在处理每个词元时，能够动态权衡序列中所有词元的重要性。然而，这一机制带来了显著的计算成本，其复杂度随序列长度（$N$）呈二次方增长（$O(N^2)$），这限制了其在长文档或高吞吐量应用中的效率。

本文《FNet：使用傅里叶变换混合词元》提出了一种根本性的简化方案。作者探究了计算昂贵的自注意力子层是否可以被更简单、线性的词元混合机制完全取代。他们最令人惊讶的发现是：使用标准的、无参数的二维离散傅里叶变换，在GLUE基准测试上能达到BERT模型92-97%的准确率，同时在处理标准512词元序列时，GPU训练速度提升80%，TPU训练速度提升70%。

2. 方法与架构

2.1. 替换自注意力机制

其核心假设是：自注意力所执行的复杂、数据依赖的混合操作，或许可以通过固定的线性变换来近似或替代。作者首先尝试了参数化的线性混合层（稠密矩阵）。观察到有希望的结果后，他们探索了更快的、结构化的线性变换，最终选定了傅里叶变换。

2.2. 傅里叶变换子层

在FNet中，标准Transformer编码器块中的自注意力子层被一个二维傅里叶变换所取代。对于输入表示 $X \in \mathbb{R}^{N \times d}$（其中 $N$ 是序列长度，$d$ 是隐藏维度），混合操作如下：

$\text{FNet}(X) = \mathcal{F}_{\text{seq}}(\mathcal{F}_{\text{hidden}}(X))$

其中 $\mathcal{F}_{\text{hidden}}$ 沿隐藏维度（$d$）应用一维傅里叶变换，$\mathcal{F}_{\text{seq}}$ 沿序列维度（$N$）应用。仅保留变换结果的实部。至关重要的是，该子层没有可学习参数。

2.3. FNet模型架构

FNet编码器块保留了标准Transformer架构的其余部分：一个带有非线性激活函数（如GeLU）的前馈网络子层、残差连接以及层归一化。顺序为：傅里叶混合子层 → 残差连接与层归一化 → 前馈网络子层 → 残差连接与层归一化。

3. 技术细节与数学公式

对于长度为 $N$ 的序列 $x$，其一维离散傅里叶变换定义为：

$X_k = \sum_{n=0}^{N-1} x_n \cdot e^{-i 2\pi k n / N}$

对于应用于输入矩阵 $X$ 的二维变换，它被计算为两个连续的一维变换。使用快速傅里叶变换算法可将该操作的复杂度降至 $O(Nd \log N)$（针对序列维度变换），这显著优于标准自注意力在较大 $N$ 时的 $O(N^2 d)$ 复杂度。

关键见解在于：傅里叶变换在频域中对所有输入词元进行全局混合，这可能捕获与自注意力类似的全局依赖关系，但它是通过固定的数学基函数实现的，而非通过学习的、数据依赖的机制。

4. 实验结果与性能

4.1. GLUE基准测试结果

FNet模型（基础版和大型版）与对应的BERT模型进行了对比评估。结果令人瞩目：

FNet-Base 达到了BERT-Base平均GLUE分数的92.2%。
FNet-Large 达到了BERT-Large平均GLUE分数的97.3%。

这表明，精心调优的自注意力模型的大部分准确率，可以通过简单的傅里叶混合机制得以恢复。

4.2. 长序列竞技场（LRA）基准测试

在专为测试模型在长序列（1k至4k词元）上性能而设计的LRA基准测试中，FNet的准确率与最准确的“高效Transformer”模型相当。更重要的是，在GPU上，它在所有序列长度上的速度都显著快于最快的模型。

4.3. 速度与效率分析

性能提升显著：

训练速度： 在512序列长度下，GPU上比BERT快80%，TPU上快70%。
内存占用： 比标准Transformer更轻量，对于较小规模的模型尤其有益。
扩展性： FFT的 $O(N \log N)$ 复杂度扩展性，使得FNet在处理长序列时，即使在GPU上也比线性时间（$O(N)$）的注意力近似方法更具决定性优势，因为后者的方法通常具有较大的隐藏常数因子。

5. 分析框架与案例示例

案例：长文档文本分类
考虑诸如法律合同或科学文章分类的任务，这些文档通常超过2000个词元。标准的Transformer模型会因二次方的内存和计算成本而难以应对。一个“高效”的线性Transformer或许有帮助，但由于核化开销，在实践中可能速度较慢。

FNet应用： FNet模型可以高效处理这些长序列。傅里叶子层以 $O(N \log N)$ 的时间全局混合词元表示。随后的前馈网络层可以在这些混合表示上构建特征。在固定的延迟预算下，可以部署比同类Transformer更大的FNet模型，从而可能弥补在较短序列上观察到的微小准确率差距。

框架要点： FNet将归纳偏置从“数据驱动的关系加权”（注意力）转向“固定的全局频谱混合”。FNet的成功表明，对于许多NLP任务，全局组合信息的能力比特定的、学习到的组合方法更为关键。

6. 核心见解与批判性分析

核心见解： 皇帝的新衣可能比我们想象的要少。FNet的成功是对NLP正统观念的一次挑衅性挑战。它表明，自注意力这个常被认为是Transformer能力不可或缺来源的“圣牛”，可以被一个无参数的、拥有150年历史的数学运算所取代，且仅带来微小的性能损失，却获得巨大的效率提升。这表明Transformer能力的很大一部分源于其整体架构（残差连接、前馈网络、层归一化）及其全局信息流动的能力，而非注意力机制本身复杂、学习到的动态特性。

逻辑脉络： 本文的逻辑具有说服力。从昂贵的问题（二次方注意力）出发。假设更简单的混合可能有效。测试线性层（效果尚可）。意识到像FFT这样的结构化变换更快且扩展性极佳。测试它——令人惊讶的是，它几乎同样有效。从问题到迭代解决方案再到惊人发现的脉络清晰且科学合理。

优势与不足：
优势： 效率提升是毋庸置疑且具有实际意义的。本文在标准基准测试（GLUE， LRA）上进行了严格评估。其思想简洁优美，具有强烈的“我怎么没想到？”的吸引力。它为高效架构开辟了新的设计空间。
不足： 准确率差距虽小，但确实存在，对于追求最先进性能的应用可能很重要。本文没有深入分析傅里叶变换为何效果如此之好，或者丢失了哪些语言特性。人们怀疑其在需要非常精细的句法推理或复杂的多步推理（动态注意力至关重要）的任务上，性能可能达到瓶颈。其速度优势依赖于具有高度优化FFT内核的GPU/TPU，这是其速度声明的一个隐含依赖。

可操作的见解：
1. 对于从业者： 在吞吐量、延迟或成本是主要约束，且可以接受3-8%准确率下降的生产部署中，强烈考虑使用FNet。它是“足够好”的大规模文本处理的主要候选方案。
2. 对于研究者： 不要止步于傅里叶变换。本文为探索整个线性变换家族（小波变换、哈特利变换、离散余弦变换）和结构化矩阵作为注意力替代品开了绿灯。核心研究问题变为：“什么是足以满足语言理解需求的最小、最快的混合机制？”
3. 对于领域： 这项工作，连同视觉领域的MLP-Mixer等同期研究，标志着一场潜在的“回归基础”运动。在经历了多年架构复杂度的增加之后，我们可能正进入一个根本性简化的时代，质疑哪些组件是真正必不可少的。它提醒我们要定期挑战基本假设。

7. 未来应用与研究展望

混合模型： 将FNet层与稀疏或局部注意力层相结合，可以创建既高效又能在关键推理步骤中保持高精度的模型。
模态扩展： 将FNet的原理应用于多模态Transformer（视觉、音频）。通过傅里叶变换早期混合跨模态信号可能非常高效。
软硬件协同设计： 设计针对FFT操作优化的专用AI加速器，可以使FNet类架构在效率关键场景中更具优势。
理论理解： 深入分析傅里叶变换执行了哪些语言功能，以及前馈网络层如何补偿缺乏学习到的注意力机制，是未来工作的一个丰富领域。
长上下文建模： FNet是突破语言模型上下文长度限制的自然候选者，能够以可管理的计算量处理整本书籍或长对话。

8. 参考文献

Vaswani, A., et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems.
Devlin, J., et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL-HLT.
Tolstikhin, I., et al. (2021). MLP-Mixer: An all-MLP Architecture for Vision. Advances in Neural Information Processing Systems.
Tay, Y., et al. (2020). Efficient Transformers: A Survey. ACM Computing Surveys.
Wang, S., et al. (2020). Linformer: Self-Attention with Linear Complexity. arXiv preprint arXiv:2006.04768.
Katharopoulos, A., et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. International Conference on Machine Learning.
Google Research. FNet Official Code Repository. https://github.com/google-research/google-research/tree/master/f_net

目录