《FNet: Mixing Tokens with Fourier Transforms》曾获得2022年举办的计算机语言学会北美分会（NAACL）的最高效 NLP 论文奖，NAACL 委员会高度称赞该团队对大型语言模型效率所做出的贡献，并称这项创新还让模型能够处理更长的输入序列，让未来研究远程上下文成为可能。

论文链接：https://arxiv.org/abs/2105.03824

摘要

该论文发现，对于只含Encoder的Transformer（如Bert、GPT等模型），将自注意力层替换成完全不含参数的傅立叶变换操作（被称作FNet）就可以达到原有模型的92-97%的准确度。而与此同时带来了在GPU上80%的提速和在TPU上70%的提速。在长序列任务上，FNet可以在达到与现有长序列模型的准确度的同时，比这些模型在任意序列长度下的速度更快。

背景

当前Transformer的核心机制便是Attention机制，而现有的许多研究表示Attention矩阵能够获取到句子中各种不同的语法与语义关系。但该机制的一个问题在于其计算代价过高。由于Attention机制是要计算每个token之间的联系，因此该机制需要序列长度平方的时间复杂度和空间复杂度。因此本文作者尝试寻找更高效的网络结构来替代Attention层，同时还能保留Attention机制所带来的序列信息捕获能力。

前人的尝试中一个具有代表性的工作便是Mixed MLP，而放在NLP语境下便是对的输入（其中表示序列长度，代表隐藏层纬度），先在维做线性变换，而后在维做线性变换。这种做法可以先在隐藏层纬度将隐藏层的信息聚合起来，而后在序列的维度让不同的单词能够通过线性变换看到其他单词的信息。

而本文作者在尝试过这种方法并确认有一定效果后思考是否有更快速的信息聚合方式，而在尝试傅立叶变换之后，作者意外的发现仅采用傅立叶变换来聚合信息也能达到不错的效果。同时由于傅立叶变换目前已经有非常成熟的高效实现以及其不含参数的特点，将Attention机制完全替换成傅立叶变换的FNet可以在损失不到10%的精度的同时带来巨大的速度提升。同时作者还在文末提到，他们目前仅对不同的变换方式做了粗浅的调查，因此可能存在其他更优的信息聚合方式。

方法

1.离散傅里叶变换DFT

介绍本文如何使用傅里叶变换之前，需要简单阐释什么是傅里叶变换？用来干啥？学通信的同学可能对这个概念并不陌生，但计算机同学大部分看见这个概念时应该和笔者一样整个人扭曲成了问号。不急，接下来是一个不涉及高深数学的直观理解：

1.1频域&时域

也许勤于思考的你们曾想过这样一个问题：你眼中的世界是什么样的？我们看到的世界都是以时间贯穿，随时间改变的，如人的身高、汽车的轨迹、股票的走势等等。我们习惯以时间作为参照来观察动态世界，这种方式称为时域分析。但是如果抛开时间，用另一种方法来观察世界，这个世界居然也可以是永恒不变的，而这个静止的世界，就是频域。

引入一个非常“意合”的例子辅助大家理解。在你的认知中，一段音乐是什么呢？

左图这是我们对于音乐最直观且普遍的理解，一段随着时间变化的震动。而对于音乐小能手来说，一段音乐可能是右图这样的。他们可以分别被看作一段音乐在时域和频域的样子（映射）。所以对大家来说，频域这一概念并不陌生，只是我们从未意识到而已。

作为一段音乐的两种表现形式，时域中上下摆动变幻莫测的美妙旋律，在频域中却是一个个永恒不变的音符。所以

你眼中看似落叶纷飞变化无常的世界，实际上是躺在上帝怀中一份早已谱好的乐章。（宿命论狂喜？）

而傅里叶变换(Fourier Transformation)正是贯穿时域与频域的方法之一。接下来我们回归论文，关于背景知识更多精彩请自行移步https://zhuanlan.zhihu.com/p/19763358 (opens new window)。

1.2傅里叶变换

傅里叶变换的作用是将时域中非周期的连续信号转变为一个频域中非周期且连续的信号；

这句话听上去又让人云里雾里。如果拿钢琴谱举例，音乐是时域中非周期的连续信号，而琴谱显然是频域空间中非周期的离散信号，即一个个离散的音符，或者说频率（PS：当然变成周期的也可以，这个作曲人说了算）。而使用傅里叶变换，音乐会变成频域空间中的“连续谱”。

连续谱应该长啥样我们这里不做进一步探讨，但是把傅里叶变换拿到Transformer中，如果把连续的token序列（自然语言）看作时域中非周期的连续信号，傅里叶变换能够将其转换为频域空间的非周期连续信号（对于token的一种新的表示）。而使用傅里叶变换成熟的计算方法（快速傅里叶变换和矩阵乘）可以比自注意力机制更快的进行特征运算。

2.具体算法

针对一个序列{},∈[0,]，傅立叶变换可以写成如下方程：

也就是说，针对新序列中每一个位置，傅立叶变换均根据原有序列的全部信息生成一个新的表示。同时现有的傅立叶变换方法具有两个高效实现。一个高效实现便是广为人知的快速傅立叶变换，该实现方法具有的时间复杂度。

另一个高效实现便是直接将输入序列乘一个傅立叶变换矩阵即可完成傅立叶变换。其中傅立叶变换矩阵如下所示：

即便快速傅立叶变换方法具有的复杂度而此方法具有的复杂度，但矩阵乘操作对TPU来说更友好，且在相对较短的序列上具有更快的运算时间。

3.FNet结构

FNet的结构如下图所示：

针对网络的输入，网络的输出为：

其中代表傅立叶变换操作，代表取实数操作。这样做可以避免后续网络处理因傅立叶变换而带来的虚数。同时作者发现仅仅取傅立叶变换后的实数部分就已经能达到最优的效果。其中这个最优是和作者尝试的其他变换如阿达马（Hadamard）变换、哈特莱（Hartley）变换、离散余弦（Discrete Cosine）变换做比较的。

同时作者使用的Embedding层含词嵌入表示、绝对位置编码以及该单词的类型编码。同时作者也发现由于傅立叶变换本身就蕴含位置信息，因此不含位置编码的FNet也能有不错的性能表现。但是具体实验的时候为了和Bert等模型有相同模型设置，FNet依然使用了位置编码。

4.实现细节

作者发现在GPU上使用快速傅立叶变换能够在实验范围内的序列长度（512-8196个tokens）上取得更快的性能；而在TPU上，对于序列长度小于4096的序列，使用矩阵乘的方法更快。因此在GPU上的FNet实现完全采用快速傅立叶变换，而在TPU上的FNet仅在序列长度大于4096时才采用快速傅立叶变换，否则采用矩阵乘来实现。带来不同硬件上性能区别的原因在于TPU对于矩阵乘有更好的优化，且在GPU上具有更高效的快速傅立叶变换实现。

主要实验结果及分析

作者主要实验了以下模型并对比了实验结果，上表列举了这些模型可学习的参数量和前向传播时混合层(注意力子层及其各种替换)计算操作的数量级。FNet由于傅里叶子层没有可学习参数，和FF-only以及Random具有相同数量的可学习参数。而采用FFT与矩阵乘来优化傅里叶变换的方法分别具有和【其中n为序列长度，d为隐层维度】的时间复杂度：

• BERT-Base(传统的transformer encoder结构)

• FNet encoder(将自注意力子层替换为傅里叶子层)

• Linear encoder(将自注意力子层替换为两个可学习的稠密线性层，分别应用到隐层维度和序列维度)

• Random encoder(将自注意力子层替换为两个随机初始化的常数矩阵，分别应用到上述两个维度)

• FF-only encoder(完全删除自注意力子层，这种encoder不会混合token)

在GLUE基准上，作者通过实验发现在这些模型中，BERT-base具有最好的性能表现(在8个下游任务上取平均)，FNet和Linear模型表现十分相近，Random和FF-only在二分类任务上准确率接近50%，基本等于啥也没学到。

在训练稳定性方面，虽然使用稠密参数化的Linear encoder进行mix token会比无参数的FNet学习的更加灵活，但由于梯度爆炸等问题，Linear encoder更难训练。作者还在BERT-Large的基础上进行了试验，linear encoder的训练非常不稳定，导致其在GLUE基准上严重低于Base模型。相较于FNet，linear encoder还有以下缺陷：1、Linear encoder会占用更大的显存；2、在常规序列(长度为512)训练较慢，在长序列上拓展性更差；

在效率方面，base模型中FF-only速度较快，但性能很差。FNet训练速度明显快于BERT——在gpu上快80%，在tpu上快70%，FLOPS是BERT 的63%。而且FNet-Hybrid模型(最后两层encoder保留自注意力子层)达到了与FNet相近的加速效果和与BERT相近的性能。作者实验表明，傅里叶子层的正向和反向传递速度比自注意子层快了一个数量级(但FNet的整体训练速度会受到所有模型共享的前馈子层的限制)。

在LRA基准上，作者通过实验发现FNet和vanilla transformer（Tay等人（2021a）的LRA评估中第二准确模型）具有可比的效果，如表4a所示。效率方面，作者统计了在gpu(8个 V100芯片)上进行实验的训练速度和内存使用情况。对于不同的序列长度{512、1024、2048、4096、8192}执行实验。在GPU上，所有序列长度上FNet都比其他模型快得多，这可能因为GPU上FFT的实现比较高效。表4b还表明FNet会占用更少的内存。这是因为FNet傅里叶子层中没有可学习的参数，也可能部分由于FFT在处理较长的序列情况下效率较高。

结语

FNet(hybrid)在牺牲少量性能的基础上实现了令人吃惊的加速效果，主要归功于使用傅里叶变换在频域和时域(线性层可以看作对时域进行建模)中交替进行特征提取，并且傅里叶子层不需要可学习参数，相比于self attenion可以大大节省内存占用和计算量。没想到用傅里叶交换替换掉self attention能使其在transformer中发光发热，不禁让人觉得对于数学基础的掌握至关重要！

作者 | 刘新宇 张译 刘晓雯
单位 | 东北大学自然语言处理实验室