FNet: 使用傅里叶变换混合标记

摘要

我们发现，通过用简单的线性变换来替代自注意力子层，从而“混合”输入标记，可以加速 Transformer 编码器架构，并且准确率损失有限。最令人惊讶的是，我们发现用标准的、无参数的傅里叶变换来替代 Transformer 编码器中的自注意力子层，在 GLUE 基准测试上可以达到 BERT 相关模型 92%-97% 的准确率，同时在标准的 512 输入长度下，在 GPU 上训练速度提高 80%，在 TPU 上训练速度提高 70%。在更长的输入长度下，我们的 FNet 模型速度更快：与 Long Range Arena（LRA）基准测试中的“高效 Transformer”相比，FNet 在准确率上可与最准确的模型匹敌，同时在所有序列长度上在 GPU 上比最快的模型更快（在 TPU 上则在相对较短的长度上更快）。此外，FNet 的内存占用较低，在较小的模型规模下尤其高效；在固定的速度和准确率预算下，小型 FNet 模型的表现优于 Transformer 相关模型。

1 引言

Transformer 架构（Vaswani 等，2017）在 NLP 领域迅速取得了主导地位。其核心是注意力机制——一种归纳偏差（inductive bias），它通过对每个输入标记赋予与其他所有标记的相关性权重来建立连接。许多论文都对 Transformer，尤其是其注意力子层进行了深入探讨，以更好地理解该架构，例如 Tenney 等（2019）、Vig 和 Belinkov（2019）、Clark 等（2019）、Voita 等（2019）。尽管这些研究可能在有效性上有所局限（Hewitt 和 Liang，2019），但总体上支持了这样一种直觉：Transformer 通过允许输入单元的高阶组合，能够灵活地捕捉多种句法和语义关系。

在本研究中，我们探讨了是否可以完全用更简单的标记混合机制替代 Transformer 编码器架构中相对复杂的自注意力层。首先，我们用两个参数化矩阵乘法替代注意力子层——一个用于混合序列维度，一个用于混合隐藏维度。在这一简单的线性混合方案取得良好结果后，我们进一步研究了更快的结构化线性变换的效果。令人惊讶的是，我们发现傅里叶变换（Fourier Transform），尽管完全没有参数，但其表现几乎与密集线性混合相当，并且能非常高效地扩展到长输入，特别是在 GPU 上（得益于其$O(N\log N)$复杂度的快速傅里叶变换 FFT 算法）。我们将这种新模型称为 FNet。

傅里叶变换此前曾用于近似或加速卷积神经网络（CNNs）（El-Bakry 和 Zhao，2004；Mathieu 等，2014；Highlander 和 Rodriguez，2015；Pratt 等，2017；Lin 等，2018；Chitsaz 等，2020；Goldberg 等，2020）、循环神经网络（RNNs）（Koplon 和 Sontag，1997；Zhang 和 Chan，2000；Zhang 等，2018）、Transformer（Choromanski 等，2020；Tamkin 等，2020）以及更一般的 MLP 层（Cheng 等，2015；Moczulski 等，2016；Sindhwani 等，2015）。然而，我们相信这是首次完全用傅里叶变换替代特定的神经网络子层。这种将傅里叶变换视为一种核心混合机制的方法，类似于 MLP-Mixer（Tolstikhin 等，2021）在计算机视觉任务中用 MLP 替代注意力的方式；但与 MLP-Mixer 不同的是，FNet 在空间维度上的混合没有可学习参数。

鉴于 FFT 具有良好的渐近复杂度，我们的工作也与“长序列”或“高效 Transformer”相关研究相联系，这些研究试图通过稀疏模式（Child 等，2019；Qiu 等，2020；Parmar 等，2018；Beltagy 等，2020；Ainslie 等，2020；Zaheer 等，2020；Wang 等，2020；Tay 等，2020a, 2020b；Kitaev 等，2020；Roy 等，2021；Vyas 等，2020；Liu 等，2018）或通过对注意力矩阵的线性化（Katharopoulos 等，2020；Choromanski 等，2021；Peng 等，2021）来优化注意力机制的可扩展性。然而，正如我们的实验所示，虽然其中一些研究能够实现$O(N)$复杂度的注意力计算，但其隐藏的常数项较大，使得它们在实际应用中比 FNet 的可扩展性更差。

我们的贡献包括：

• 我们证明了简单的线性变换（包括无参数的傅里叶变换）结合标准 MLP 前馈层，可以有效建模文本中的多种关系。这种简单的线性变换能够取得良好效果是令人惊讶的，这表明对于某些 NLP 任务而言，注意力机制可能并不是 Transformer 性能的决定性因素。

• 我们提出了一种新模型 FNet，该模型使用傅里叶变换作为混合机制。FNet 在速度、内存占用和准确率之间提供了极佳的折衷，在 GLUE 基准测试上分别达到 BERT-Base 和 BERT-Large 92% 和 97% 的准确率，同时在 GPU 上训练速度提高 80%，在 TPU 上提高 70%。

• 我们发现，仅包含两个自注意力子层的 FNet 混合模型在 GLUE 基准测试上的准确率可达 BERT 相关模型的 97%-99%，但运行速度仍然快 40%-70%。这表明，尽管注意力机制可以提高准确率，但并不一定需要在每一层都使用它。

• 我们证明了 FNet 对长输入具有极好的扩展性，并且在速度和准确率之间提供了比高效 Transformer 更优的折衷。在 LRA 基准测试上，FNet 取得了与最准确的高效 Transformer 架构相当的准确率，但在 GPU 上的训练和推理速度均快于所有评估的 Transformer 架构。在 TPU 上，FNet 在较短序列长度上更快；对于较长序列，唯一比 FNet 速度更快的高效 Transformer 在 LRA 基准测试上的准确率更低。因此，我们认为，与其寻求更高效的注意力近似方法，不如探索全新的混合机制可能更有价值。

  2 相关工作

  2.1 神经网络中的傅里叶变换

  傅里叶分析在神经网络通用逼近性质的研究中占据重要地位；例如，可参见（Cybenko, 1989; Barron, 1993）。在实际应用方面，离散傅里叶变换（DFT），尤其是快速傅里叶变换（FFT），已被用于解决信号处理问题，例如将神经网络拟合到心电图信号的FFT（Minami等，1999；Gothwal等，2011；Mironovova和Bíla，2015）和振动信号（Zhang等，2013），或用于演化偏微分方程的解（Li等，2021）。由于频域中的普通乘法对应于时域中的卷积，FFT已被部署在卷积神经网络中以加速计算，在循环神经网络中加速训练并减少梯度爆炸和消失问题，并通常用于近似密集的线性层以降低计算复杂度；参见第1节中引用的参考文献。DFT也在一些Transformer工作中被间接使用。Performer（Choromanski等，2020）通过利用随机傅里叶特征来近似softmax核的高斯表示，从而线性化Transformer的自注意力机制。在我们的工作中，我们没有近似注意力，而是用傅里叶变换替换注意力，将其作为一种替代的隐藏表示混合机制。Tamkin等（2020）使用频谱滤波器生成分层特征，表明滤波后的嵌入在不同任务（词级、句级或文档级）中表现良好，具体取决于过滤的频率尺度。与FNet不同，他们分离了傅里叶频率，而不是使用变换来组合特征。最后，通过个人交流，我们了解到一项并行的未发表工作（Backurs等，2021），其中描述了一种基于FFT的神经网络模型，与FNet非常相似。

  2.2 通过注意力建模语义关系

  注意力模型在几乎所有自然语言处理（NLP）任务甚至一些图像任务（Dosovitskiy等，2021）中取得了最先进的结果。这一成功通常归因于注意力的灵活性和容量。尽管一些工作（Ramsauer等，2021）致力于更深入地理解注意力，但普遍的直觉是，注意力模型的成功源于不同层中依赖于token的注意力模式；例如，参见（Tenney等，2019）。然而，很自然地会问：我们真的需要注意力的灵活性及其相关成本吗？

  Tay等（2020a）在其Synthesizer模型（与我们的“线性”基线相关）中通过实验研究了注意力机制中点积操作的重要性。他们发现，学习到的依赖于token的注意力权重具有高度表达能力，但对于实现准确的NLP模型并不一定是关键。You等（2020）将Transformer编码器和解码器中的注意力权重替换为无参数化的高斯分布，表明只要保留可学习的交叉注意力权重，性能下降极小。类似地，Raganato等（2020）发现，当将编码器中每个注意力层除一个注意力头外的所有注意力头替换为固定的、不可学习的位置模式时，准确性几乎没有下降。最后，Tolstikhin等（2021）提出了MLP-Mixer，其中注意力被多层感知机（MLP）取代，在图像分类任务中性能下降有限。

  2.3 高效与长序列模型

  标准的注意力机制（Vaswani等，2017）在序列长度上具有二次时间和内存瓶颈。这限制了其在涉及长距离依赖任务中的适用性。大多数提高注意力效率的努力都基于稀疏化注意力矩阵。Tay等（2020c）综述了许多近期的高效注意力工作；另见第1节中的引用。一些“高效Transformer”实现了$\mathcal{O}(N\sqrt{N})$甚至O(N)的理论复杂度。然而，这种表示法隐藏的常数可能很大。例如，在Longformer（Beltagy等，2020）、ETC（Ainslie等，2020）和BigBird（Zaheer等，2020）等模型中，注意力作为输入长度的函数是O(N)，但在“全局token”数量上是二次的；后者必须足够大以确保良好的性能。

  长距离竞技场基准测试（Tay等，2021a）尝试在一系列需要长距离依赖的任务中比较多种高效Transformer，发现Performer（Choromanski等，2021）、Linear Transformer（Katharopoulos等，2020）、Linformer（Wang等，2020）和Image Transformer（局部注意力）（Parmar等，2018）在TPU上速度最快且每个设备的峰值内存使用量最低。然而，在本文中，我们完全用不同的混合机制（即傅里叶变换）替换了自注意力，其提供了：（1）性能，（2）减少的模型大小（无可学习参数），以及（3）简洁性。

  最后，我们注意到，为了研究不同的token混合机制，我们将普通的BERT模型（Devlin等，2019）与普通的FNet进行了比较，忽略了最近的Transformer优化，我们认为这些优化与本文工作正交；例如，参见（Narang等，2021；Kim和Hassan，2020；Shleifer和Rush，2020）。

3 模型

3.1 离散傅里叶变换

傅里叶变换将一个函数分解为其组成的频率。给定一个序列 { x n } \{ x _ { n } \} {xn​}，其中$n\in [0,N-1]$，离散傅里叶变换（DFT）由以下公式定义：
$$X_k=\sum _{n=0}^{N-1}{x}_n{e}^{-\frac{2\pi i}{N}nk},0\le k\le N-1.\text{\hspace{1em}(1)}$$
对于每个$k$，DFT 生成一个新的表示$X_k$，它是所有原始输入标记$x_n$与所谓的“旋转因子”之和。计算 DFT 主要有两种方法：快速傅里叶变换（FFT）和矩阵乘法。标准的 FFT 算法是 Cooley-Tukey 算法（Cooley 和 Tukey，1965；Frigo 和 Johnson，2005），它递归地将长度为$N=N_1{N}_2$的序列的 DFT 表示为$N_1$个大小为$N_2$的较小 DFT，从而将计算时间减少到$O(N\log N)$。

另一种方法是直接将 DFT 矩阵应用于输入序列。DFT 矩阵 ( W ) 是一个范德蒙矩阵，其元素为单位根的幂次，并包含一个归一化因子：

$$W_{nk}=\left(e^{-\frac{2\pi i}{N}nk}/\sqrt{N}\right),\text{\hspace{1em}(2)}$$

其中，n,$k=0,\dots ,N-1$。此矩阵乘法是一个O(N²)操作，其渐近复杂度高于快速傅里叶变换（FFT），但在TPU上处理相对较短的序列时，速度更快。

3.2 FNet 架构

FNet 是一种无需注意力机制的 Transformer 架构，其中每一层由一个傅里叶混合子层和一个前馈子层组成。该架构如图 1 所示。本质上，我们将每个 Transformer 编码器层的自注意力子层替换为一个傅里叶子层，该子层对其（序列长度，隐藏维度）的嵌入输入应用二维离散傅里叶变换（2D DFT）——一个一维 DFT 沿序列维度${\mathcal{F}}_{\mathrm{seq}}$，另一个一维 DFT 沿隐藏维度${\mathcal{F}}_{\mathrm{h}}{:}^3$

$$y=\Re \left({\mathcal{F}}_{\mathrm{seq}}\left({\mathcal{F}}_{\mathrm{h}}(x)\right)\right).\text{\hspace{1em}(3)}$$

如公式（3）所示，我们仅保留结果的实部；因此，我们无需修改（非线性）前馈子层或输出层以处理复数。我们发现，当仅在傅里叶子层结束时提取总变换的实部时，即应用${\mathcal{F}}_{\mathrm{seq}}$和${\mathcal{F}}_{\mathrm{h}}$之后，FNet 取得了最佳结果。我们还尝试了哈达玛变换、哈特利变换和离散余弦变换。在这三种变换中，哈特利变换是最强的替代方案，其精度与公式（3）相当；详见附录 A.3。

傅里叶变换最简单的解释是作为一种特别有效的 token 混合机制，它为前馈子层提供了对所有 token 的充分访问。由于傅里叶变换的对偶性，我们还可以将每个交替的编码器块视为应用交替的傅里叶变换和逆傅里叶变换，将输入在“时间”域和频域之间来回转换。由于在频域中与前馈子层系数相乘等效于在时域中进行卷积（使用一组相关的系数），因此 FNet 可以被视为在乘法和卷积之间交替。我们使用了与 Devlin 等人（2019）相同的嵌入层；即，我们结合了词嵌入、token 的绝对位置嵌入和句子的类型嵌入。由于公式（1）中的傅里叶变换编码了位置信息（参见 n, k 索引），FNet 在没有位置嵌入的情况下表现同样出色。尽管如此，我们仍包含位置嵌入，以便与 BERT 进行更清晰的比较。

3.3 实现

根据实验，我们发现在 GPU 上：对于我们所考虑的所有序列长度（512 − 8192 个 token），FFT 比矩阵乘法更快；而在 TPU 上：对于相对较短的序列（≤ 4096 个 token），缓存 DFT 矩阵然后通过矩阵乘法计算 DFT 比使用 FFT 更快；对于更长的序列，FFT 更快。因此，我们的 GPU FNet 实现始终使用 FFT，而我们的 TPU 实现对于长度不超过 4096 的序列使用矩阵乘法计算 2D DFT，对于更长的序列使用 FFT。GPU 与 TPU 的差异可能主要源于两个因素：（1）TPU 对矩阵乘法的优化程度比 GPU 更高，（2）GPU 提供了比 TPU 更高效的 FFT 实现。我们推测，随着 TPU 上 FFT 实现的改进，FNet 在 TPU 上的性能只会变得更好。我们的模型使用 JAX，特别是 Flax 框架。核心模型代码见附录 A.7，完整源代码可在网上获取。

4 结果

4.1 迁移学习

我们在常见的迁移学习设置中比较了 FNet 和 Transformer 架构。为了更全面地了解，我们比较了多种模型（参见表 1 中“Base”配置的参数数量）：

• BERT-Base：一种 Transformer 编码器模型。
• FNet 编码器：我们将每个自注意力子层替换为傅里叶子层。
• 线性编码器：我们将每个自注意力子层替换为两个可学习的密集线性子层，一个应用于隐藏维度，另一个应用于序列维度。
• 随机编码器：我们将每个自注意力子层替换为两个常数随机矩阵，一个应用于隐藏维度，另一个应用于序列维度。
• 仅前馈（FF-only）编码器：我们完全移除了自注意力子层，因此该模型没有 token 混合机制。

尽管线性基线模型非常简单，但其准确性和速度却出人意料地高。我们的线性模型类似于 MLP-Mixer（Tolstikhin 等人，2021）（用于视觉任务）以及 Random Synthesizer（Tay 等人，2020a），但通过移除多头和 softmax 投影进一步简化了后者模型，最终在混合子层中仅包含两次矩阵乘法。

可以合理预期，使用密集参数化混合层的线性编码器会比使用无参数混合层的 FNet 学习得更灵活。然而，尽管线性 Base 模型在 GLUE 上略微优于 FNet-Base（0.3 分），但它相对于 FNet 有几个效率劣势：内存占用更大（见表 4b），在常规 512 序列长度上训练速度较慢（见表 3），并且在长序列长度上扩展性显著较差（见表 4b-4c）。我们还发现，由于梯度爆炸问题，线性 Large 模型更难训练（见表 2 中的“Large”分数）。

我们采用了与原始 BERT（Devlin 等人，2019）相同的固定“Base”和“Large”模型及训练配置，不同之处在于我们在更大的 C4 数据集（Raffel 等人，2020）上进行预训练，并使用 32000 的 SentencePiece 词汇模型（Kudo 和 Richardson，2018）（完整预训练细节见附录 A.1）。在 GLUE 基准（Wang 等人，2018）上进行微调时，我们发现相同基础学习率的不同 BERT 运行可能会产生略微不同的结果。因此，对于 Base（Large）模型，我们分别对每个基础学习率进行了 3（6）次试验，并报告了所有实验中的最佳结果。这反映了我们观察到的 BERT-Large 比 BERT-Base 更不稳定的现象，正如 Devlin 等人（2019）所指出的那样。

我们在表 2 中报告了 GLUE 验证集上最佳基础学习率（无早停）的结果。对于 Base 模型，结果与预训练指标一致（见附录 A.1）：BERT 表现最佳。FNet 和线性模型的表现均比 BERT 低 7.5% − 8%。从表 3 中可以看出，尽管 FNet 的准确性较低，但其训练速度显著快于 BERT——在 GPU 上快 80%，在 TPU 上快 70%——并且仅执行了 BERT 63% 的浮点运算量。单独测量时，傅里叶子层的前向和反向传播速度比自注意力子层快一个数量级（见附录 A.4），但 FNet 的整体训练速度受到所有模型共享的前馈子层的拖累。

回到表 2：仅前馈模型的表现严重落后于其他所有模型：正如预期的那样，token 混合对模型的表达能力至关重要。例如，在二分类任务（QNLI、SST-2、RTE）上 50% 的准确率表明模型未能学会任务。随机模型的低准确率表明，并非任何混合机制都有效，而是需要一种结构化的混合机制。我们还提供了混合 FNet 注意力模型的指标。在混合模型中，我们将 FNet 的最后两个傅里叶子层替换为自注意力子层——其他配置也是可能的，但我们通常发现替换最后几层效果最佳；详见附录 A.5。通过仅添加两个自注意力子层，混合 FNet 模型分别达到了其对应 BERT 模型 97% 和 99% 的准确率，且速度下降有限（见表 3）。

有趣的是，对于 Large 模型，BERT 和 FNet 之间的差距缩小到仅 3%；这可能是由于 FNet-Large 在训练过程中比 BERT-Large 更稳定。线性 Large 模型由于训练不稳定性，在 GLUE 基准上的表现严重低于其 Base 版本。我们通常发现，线性模型和 BERT 比混合子层中无参数的模型（即 FNet、随机和仅前馈模型）更不稳定。

GPU（8 个 V100 芯片）预训练的速度与 MLM 准确率曲线如图 2 所示（TPU 结果见附录 A.2）。TPU 和 GPU 模型均按照 Devlin 等人（2019）的方法训练了 100 万步。受 Turc 等人（2019）中考虑的模型启发，我们评估了多种模型规模；详见附录 A.1 中的表 6。我们发现较小的模型架构受益于较大的学习率，因此我们使用 10^−3 和 $10^{-4}$ 对所有模型选择最佳结果。GPU（图 2）和 TPU（附录 A.2 中的图 3）结果显示了相同的趋势。对于较大、较慢的模型，BERT 和 FNet-Hybrid 定义了帕累托速度-准确率效率前沿。对于较小、较快的模型，FNet 和线性模型定义了效率前沿。

4.2 长距离竞技场（LRA）基准测试

在 Tay 等人（2021a）评估的高效 Transformer 模型中，他们的结果表明：（1）标准 Transformer 以微弱优势成为第二准确的模型，（2）Performer（Choromanski 等人，2021）是最快的模型。我们使用 Tay 等人（2021a）的代码库并在相同的硬件（4 × 4 TPU v3 芯片）上运行，将 FNet 的准确性与这两种模型进行了基准测试；结果如表 4a 所示。为了确保公平比较，我们还报告了我们自己对标准 Transformer 的实验结果（详见附录 A.6）。

表 4a 表明，总体而言，标准 Transformer 和 FNet 获得了相当的结果。鉴于标准 Transformer 是 Tay 等人（2021a）评估的第二准确的模型，且表 4a 中平均准确率得分的相对差异较小，我们的结果表明 FNet 在 LRA 上与最准确的高效 Transformer 模型具有竞争力。

在效率方面，表 4b 提供了我们在 GPU（8 个 V100 芯片）上的实验训练速度和内存使用统计数据；TPU 结果见附录 A.2。我们对序列长度进行了扫描，范围包括 {512, 1024, 2048, 4096, 8192}。在 GPU 上，由于 GPU 上高效的 FFT 实现，FNet 在所有序列长度上都比其他模型快得多。表 4b 还表明，FNet 的内存占用更轻（这在 GPU 和 TPU 上均成立；详见附录 A.2 的扩展结果）。这部分是因为 FNet 的混合子层中没有可学习参数，同时也归功于 FFT 的效率，尤其是在较长的序列长度上。最后，表 4c 显示训练速度的提升通常也会转化为推理速度的提升（详细的 TPU 结果见附录 A.2）。

5 结论

在本研究中，我们探索了用于类 Transformer 编码器架构的简化 token 混合模块，并做出了多项贡献。首先，我们证明了简单的线性混合变换与前馈层中的非线性结合，能够有效地建模文本中的多样化语义关系。其次，我们提出了 FNet，这是一种类 Transformer 模型，其中自注意力子层被无参数的傅里叶变换取代。FNet 在 GLUE 基准测试中分别达到了 BERT-Base 和 BERT-Large 模型 92% 和 97% 的准确率，但在 GPU/TPU 上的训练速度快了 70% − 80%。第三，由于其良好的扩展性，FNet 在长距离竞技场（LRA）基准测试中与“高效 Transformer”模型相比极具竞争力，在匹配最准确模型的同时，速度更快且内存占用更少。

我们的研究凸显了线性单元作为注意力机制替代方案在文本分类任务中的潜力。我们发现，由于快速傅里叶变换（FFT）的速度，傅里叶变换是一种特别高效且有效的混合机制。然而，我们仅对其他线性变换进行了初步调查（另见附录 A.3），更多快速替代方案值得进一步探索。

鉴于较小 FNet 模型相对于 Transformer 的速度和准确性优势，我们推测 FNet 将作为一种轻量级、蒸馏后的学生模型在资源受限的环境中（如生产服务或边缘设备）有效部署。随着对巨型模型的兴趣增加（Raffel 等人，2020；Brown 等人，2020；Lepikhin 等人，2021），对此类轻量级服务模型的需求预计只会增长。在这方面，一个自然的探索方向是从更大的 Transformer 教师模型中对小型 FNet 模型进行知识蒸馏，例如 Sanh 等人（2019）；Jiao 等人（2020）；Turc 等人（2019）。

另一个值得进一步研究的方面是混合 FNet-注意力模型。我们发现，仅在 FNet 中添加少量自注意力子层，可以以速度换取准确性。具体而言，将最后两个傅里叶子层替换为自注意力子层，可以在有限的速度损失下提供 BERT 97% − 99% 的准确率。

在本研究中，我们主要关注编码器。FNet 解码器可以通过对 Vandermonde 矩阵进行“因果”掩码来设计，但需要在更低层次的实现中为 FFT 引入因果掩码。如何将傅里叶混合应用于编码器-解码器的交叉注意力仍是一个开放问题，因为有证据表明交叉注意力可能对性能至关重要（You 等人，2020）。我们专注于不需要生成的任务，因此将 FNet 解码器和编码器-解码器设置留给未来的工作；尽管我们确实提到，FNet 编码器可以作为 Transformer 的直接替代品，正如其他研究已成功展示的那样；例如，（Zaheer 等人，2020；Guo 等人，2021）。

A 附录

A.1 预训练细节

我们采用了与原始 BERT（Devlin 等人，2019）相同的固定“Base”和“Large”模型及学习配置。我们在更大的 C4 数据集（Raffel 等人，2020）上进行训练，并使用在 C4 的 1 亿句子子集上训练的 32000 SentencePiece 词汇模型（Kudo 和 Richardson，2018）。我们的 TPU 实验使用与 Devlin 等人（2019）相同的批量大小 256，并在 4 × 4 TPU v3 芯片上运行。我们的 GPU 实验使用较小的批量大小 64，并在 8 个 V100 芯片上运行。由于训练配置直接取自 Devlin 等人（2019），可能略微偏向 BERT 注意力模型。

表 5 总结了不同模型的预训练指标；预训练速度见正文中的表 3。尽管线性模型和 FNet 的准确性指标较弱，但它们在 GPU 上的训练速度比 BERT 快近 80%，在 TPU 上快 70%（见表 3）。我们还发现，混合层中没有可学习参数的三种模型（即 FNet、随机模型和仅前馈模型）在训练过程中最为稳定。

BERT 在 MLM 预训练任务上的较高准确性并不仅仅是由于它比其他模型拥有更多参数。事实上，表 5 显示 BERT-Base 实际上比参数数量是其两倍以上的 FNet-Large 更准确。BERT 可能更具表现力，因为其混合（注意力）权重既是任务特定的，也是 token 依赖的，由 token-token（查询-键）点积决定；另见 Tay 等人（2020a）。而 FNet 的混合权重既不是任务特定的，也不是 token 依赖的。

最后，表 6 显示了用于构建图 2（正文）和图 3（附录 A.2）的模型规模。

A.2 TPU 结果

在本节中，我们报告了 FNet 在 TPU 上的效率结果；正文主要关注 GPU。图 3 显示了在 TPU（4 × 4 v3 芯片）上训练时的速度与 MLM 预训练准确率曲线。与 GPU 类似，FNet 和线性模型定义了较小、较快模型的帕累托效率前沿，而 BERT 定义了较大、较慢模型的前沿。

表 7 显示了在 TPU（4 × 4 v3 芯片）上长距离竞技场文本分类任务的效率结果。对于序列长度 ≤ 2048 和 512，线性模型和 FNet 的训练速度分别快于所有高效 Transformer 模型。对于更长的序列，FNet 比 Performer 慢，并且根据 Tay 等人（2021a）的结果，可能也比其他通过线性化注意力实现的高效 Transformer 模型（如局部注意力（Parmar 等人，2018）、Linformer（Wang 等人，2020）和线性 Transformer（Katharopoulos 等人，2020）慢。然而，值得注意的是，表 4a 表明 FNet 比上述所有模型更准确。此外，我们预计随着 TPU 上 FFT 实现的改进，GPU 上的速度优势将转移到 TPU 上。

A.3 我们尝试的其他配置

我们尝试了许多额外的想法来改进 FNet。

傅里叶变换算法：在 GPU 上，FFT 是我们实验的所有序列长度（512 − 8192）中计算 DFT 最快的算法。在 TPU 上，对于相对较短的序列长度（最多 4096），直接使用矩阵乘法计算 DFT 更快（见表 7）。TPU 上矩阵乘法和 FFT 之间的效率边界会因矩阵乘法的 XLA 精度而变化。我们发现，尽管在计算 DFT 时需要（较慢的）最高 XLA 精度才能非常准确地重现 FFT，但（较快的）默认 XLA 精度足以促进模型的准确收敛。

修改傅里叶变换计算：为了保持整个 FNet 架构的简单性，傅里叶子层接受实数输入并返回实数输出。FNet 中的标准傅里叶子层在计算 2D DFT 后仅提取实部。我们发现，如果在整个计算过程中仅使用 DFT 的实部，FNet 在训练期间的准确性和稳定性会降低。简单地提取绝对值（而不是实部）也会导致模型的准确性显著下降。由于前馈子层混合了隐藏维度，我们尝试在傅里叶子层中仅沿 token 维度应用 1D DFT（即傅里叶子层中没有隐藏维度混合）。这带来了一些训练速度的提升，但损害了准确性。仅进行 1D（仅 token 混合）DFT 的模型仍然显著优于（无 token 混合的）仅前馈模型，表明 token 混合是傅里叶子层中最重要的机制。

其他变换：我们尝试了三种傅里叶变换的自然替代方案：

• 离散余弦变换（DCT）：DCT 与 DFT 密切相关，但将实数输入转换为实数输出。然而，我们发现 DCT 模型的表现不如 FNet（准确性下降约 4%）。
• 哈达玛变换：尽管哈达玛变换比 DFT 稍快，但其结果的准确性较低（准确性下降约 2%）。
• 哈特利变换：哈特利变换将实数输入转换为实数输出，可以用傅里叶变换描述为 H = ℜ { F } − ℑ { F } { \mathcal { H } } = \Re \left\{ { \mathcal { F } } \right\} - \Im \left\{ { \mathcal { F } } \right\} H=ℜ{F}−ℑ{F}。我们发现哈特利变换在 GLUE 上的表现与傅里叶变换相当（76.7 vs. 76.7）。

向傅里叶子层引入可学习参数：我们尝试向傅里叶子层引入可学习参数，但这些尝试要么有害，要么无关紧要，通常会略微减慢模型速度。对于每个傅里叶子层中的（序列长度，隐藏维度）输入，我们尝试了两种引入可学习参数的方法：（1）与（序列长度，隐藏维度）矩阵进行逐元素乘法，（2）与（序列长度，序列长度）和（隐藏维度，隐藏维度）矩阵进行常规矩阵乘法。我们在各种配置中尝试了这些方法：在 DFT 之前和/或之后，以及与逆 DFT 结合（例如，转换到频域，应用逐元素乘法，转换回时域），但大多数设置会降低准确性并减少训练稳定性，而少数设置不会改变准确性但会导致速度小幅下降。在另一组实验中，为了为模型提供更多灵活性，我们向 2D DFT 矩阵添加了（复数）可学习权重。该模型稳定，但未带来任何准确性提升，表明 DFT 在某种意义上是局部最优的。

FNet 块修改：标准的 FNet 编码器块结构遵循 Transformer 的结构：傅里叶子层后接前馈子层，每个子层后有残差连接和层归一化；见图 1。我们尝试了几种修改此结构的方法，基于在乘法之间进出频域的直觉。例如，将傅里叶、前馈、傅里叶（或逆傅里叶）子层夹在一起，并仅对最终结果应用残差连接和层归一化，产生了一种更接近卷积的结构。然而，这些设置会降低准确性，并导致训练期间模型更不稳定。在这种分层中添加额外的前馈子层，或将前馈子层替换为更简单的密集子层，也没有帮助。

A.4 混合层速度

表 8 总结了不同混合层的推理和训练速度。对于每个 Base 和 Large 配置，我们移除了所有其他子层和变换，然后计算每批输入示例的速度。FNet 的训练速度特别快，因为没有参数需要更新。在 TPU 上，线性模型的推理速度比 FNet 快，因为它执行的是实数矩阵乘法，而 FNet 执行的是复数矩阵乘法；见公式（2）。

尽管傅里叶混合子层本身的前向和反向传播速度显著快于自注意力子层，但 FNet 整体比 BERT 快 70-80%，因为整体训练和推理速度受到所有模型共享的前馈子层的瓶颈限制。

A.5 FNet-Hybrid 消融实验

表 9 显示了在 FNet-Hybrid 模型中改变注意力子层数量和注意力布局的效果。对于“BOTTOM”布局，所有注意力子层都放置在前几个编码器层中，替换傅里叶混合子层。对于“TOP”布局，注意力子层放置在最后的编码器层中；对于“MIDDLE”布局，它们放置在中间层；对于“MIXED”布局，它们分布在模型中。从表 9 中，我们可以得出两个结论：

1. 更多的注意力子层会提高准确性，但以速度为代价，最终收益递减；
2. 将注意力层放置在模型的顶部会带来最佳的准确性结果。鉴于我们对速度的关注，我们选择在论文正文中重点研究 2 个注意力层的“TOP”配置变体。

A.6 关于长距离竞技场超参数设置的说明

关于长距离竞技场的设置，Tay 等人（2021a）未描述几个超参数，并且论文中描述的配置与代码库之间存在一些不匹配之处。在可能的情况下，我们优先采用论文中描述的配置，只有两个例外。首先，对于 CIFAR10（图像）任务，我们在 [1; 2; 3; 4] 范围内对层数进行了扫描。我们发现 1 层对所有模型效果最好；而 Tay 等人（2021a）建议 3 层效果最佳。其次，对于 Pathfinder 任务，我们发现基础学习率为 0.001（代码库中给出的值）比 Tay 等人（2021a）中提到的 0.01 值对所有模型效果更好。我们还对嵌入维度和批量大小进行了非常小的扫描，这些在 Tay 等人（2021a）中未列出。

我们还注意到，在比较我们的运行结果与 Tay 等人（2021a）的结果时，应谨慎对待，因为我们发现某些任务（尤其是文本和检索任务）的结果在不同运行之间可能会有较大差异，特别是对于 Transformer；我们报告的是最佳结果。

A.7 FNet 代码

论文名称：
FNet: Mixing Tokens with Fourier Transforms
论文地址：
https://arxiv.org/pdf/2105.03824