[论文解读] FNet: Mixing Tokens with Fourier Transforms

James Lee-Thorp, Joshua Ainslie+2 more|arXiv (Cornell University)|May 9, 2021

Topic Modeling被引用 14

一句话总结

FNet 使用无参数的傅里叶变换替代 Transformer 编码器中的自注意力机制，以实现高效的序列建模。其在 GPU 上训练速度比 BERT 快 80%，在 TPU 上快 70%，且在长序列上具有更优的可扩展性，同时在小模型尺寸下表现强劲，GLUE 准确率可达 BERT 的 92–97%。

ABSTRACT

We show that Transformer encoder architectures can be sped up, with limited accuracy costs, by replacing the self-attention sublayers with simple linear transformations that "mix" input tokens. These linear mixers, along with standard nonlinearities in feed-forward layers, prove competent at modeling semantic relationships in several text classification tasks. Most surprisingly, we find that replacing the self-attention sublayer in a Transformer encoder with a standard, unparameterized Fourier Transform achieves 92-97% of the accuracy of BERT counterparts on the GLUE benchmark, but trains 80% faster on GPUs and 70% faster on TPUs at standard 512 input lengths. At longer input lengths, our FNet model is significantly faster: when compared to the "efficient" Transformers on the Long Range Arena benchmark, FNet matches the accuracy of the most accurate models, while outpacing the fastest models across all sequence lengths on GPUs (and across relatively shorter lengths on TPUs). Finally, FNet has a light memory footprint and is particularly efficient at smaller model sizes; for a fixed speed and accuracy budget, small FNet models outperform Transformer counterparts.

研究动机与目标

• 探究更简单的无参数线性变换是否能在不造成显著精度损失的前提下替代 Transformer 编码器中的自注意力机制。
• 评估傅里叶变换作为结构化、高效的标记混合机制替代注意力的有效性。
• 证明基于傅里叶的混合机制可实现更快的训练和推理速度，尤其是在长序列长度下。
• 探索将少量注意力层与傅里叶混合结合的混合模型，以实现更优的精度-速度权衡。
• 证明在实际应用中，替代性混合机制可能优于注意力的近似方法，尤其是在长距离任务中。

提出的方法

• FNet 在每个 Transformer 编码器块中，使用快速傅里叶变换（FFT）算法将自注意力子层替换为无参数的傅里叶变换。
• 傅里叶变换作用于序列维度，通过频域变换在位置之间混合标记。
• 每个块在傅里叶混合后应用层归一化，随后是使用 ReLU 和 Dropout 的标准前馈网络。
• 模型使用标准的位置嵌入和可学习的池化头用于分类任务。
• 对于混合模型，将少量自注意力层插入特定位置（例如顶层），其余部分使用傅里叶混合。
• 训练采用标准优化方法并配合学习率调度；推理在 GLUE 和 Long Range Arena 基准上进行评估。

实验结果

研究问题

• RQ1 无参数的傅里叶变换是否能有效替代 Transformer 编码器中的自注意力机制，同时保持具有竞争力的精度？
• RQ2 FNet 在标准和长序列上的训练与推理速度与 BERT 及高效 Transformer 变体相比如何？
• RQ3 在混合 FNet 模型中，注意力层的最佳位置和数量是多少，以实现精度与速度的最优平衡？
• RQ4 FNet 在不同输入长度下的性能和效率表现如何，尤其是在长距离任务中？
• RQ5 与参数化或稀疏注意力近似方法相比，使用 FFT 这类结构化线性变换替代注意力是否具有优势？

主要发现

• 在 GLUE 基准上，FNet 分别达到 BERT-Base 和 BERT-Large 性能的 92–97%，在输入长度为 512 时，GPU 上训练速度比 BERT 快 80%，TPU 上快 70%。
• 在 Long Range Arena 基准上，FNet 的精度与最准确的高效 Transformer 模型相当，且在所有序列长度下 GPU 上的推理速度均更优。
• 对于更长的序列，FNet 在 GPU 上比所有评估的高效 Transformer 模型更快，且在较短序列上 TPU 上的性能优于最快模型。
• 仅在顶层插入两个自注意力层的混合 FNet 模型，可达到 BERT 精度的 97–99%，同时训练速度仍快 40–70%。
• FNet 内存占用更小，在固定速度和精度预算下优于标准 Transformer 模型，尤其在小模型尺寸下表现更优。
• 傅里叶变换的 O(N log N) 复杂度使其在实际可扩展性上优于 O(N²) 的注意力机制，即使注意力机制被近似处理。