FNet - A New Era In Deep Learning

 FNet Mixing Tokens with Fourier Transforms: A New Era in Deep Learning

Introduction

------------

In this blog post, we will discuss a groundbreaking paper titled "FNet: Mixing Tokens with Fourier Transforms" by James Lee Thorpe, Joshua Ainsley, Ilya Eckstein, and Santiago Antonion of Google Research. Although I apologize for being a bit late to the party, I believe this paper signifies a fascinating direction in machine learning, particularly in deep learning, sequence models, and image models. Specifically, we will explore the gradual shift away from attention mechanisms, such as those found in transformers.

Attention Mechanisms: The Long-Standing Challenge

-------------------------------------------------

For a long time, researchers have focused on transformers, which process a sequence input through attention layers composed of attention sublayers and feedforward layers. While the feedforward sublayers can efficiently parallelize and optimize computations, the attention sublayers have been a thorn in the side of many due to their high memory and computational requirements. The attention mechanism's job is to decide which information from the current layer's sequence goes to the next layer's sequence, acting as a routing problem with a complexity of O(n^2) for sequence length n. This complexity and memory requirement have prevented transformers from scaling to larger sequence lengths, limiting their application in computer vision, for example.

Linearizing the Attention Mechanism

-----------------------------------

Over the past couple of years, researchers have been chipping away at the attention mechanism's complexity, aiming to linearize it. This has led to the development of models like Linformer, Longformer, Reformer, Synthesizer, and Linear Transformer, all attempting to approximate the attention routing problem with more manageable complexities such as O(n) or O(n log n).

A New Era: Questioning the Need for Attention

---------------------------------------------

Even after the introduction of these linear or non-quadratic attention mechanisms, researchers have started questioning whether the attention layer is necessary at all. In recent times, several papers have emerged, attempting to actively remove the attention layer from sequence models. In this paper, we will explore one such approach, which replaces the attention layer with Fourier transforms.

Fourier Transforms: Mixing Tokens

---------------------------------

The paper presents a model that closely resembles a transformer, with an input split into word sequences, each receiving a word embedding, position embedding, and possibly a type embedding. The significant difference is that instead of the attention layer, the model employs a Fourier layer.

This Fourier layer applies a Fourier transform to the input x in both the hidden and sequence domains, followed by a one-dimensional Fourier transform in each dimension across the time domain. The fascinating aspect of this approach is that it does not involve any learned parameters, making it a purely linear transformation of the data. The only learned parameters in the entire setup are the normalization and feedforward parameters.

The Importance of Mixing Information

------------------------------------

The paper emphasizes the importance of mixing information between tokens rather than advocating for the Fourier transform as a superior technique. In language tasks, locality of information is not guaranteed, making it crucial to route information between elements of the sequence. The attention mechanism has been instrumental in facilitating these connections, but this paper suggests that the exact routing might not be as important as ensuring that information flows between tokens.

Experimental Results

--------------------

The authors compare the FNet model with other models, such as BERT, a linear encoder, and a feedforward-only model. In pre-training loss and masked language modeling, BERT outperforms the other models. However, in terms of speed, the FNet model demonstrates significant improvements, particularly on GPUs. The paper also evaluates the models on the GLUE benchmark, with BERT winning in most tasks but showing instability in some cases.

In the long-range arena, transformers still perform best, but the FNet model is not too far behind, using considerably less memory and compute while training way faster. Nevertheless, the authors acknowledge that the long-range arena results might be binary, with models either solving the tasks or not, without much nuance.

Conclusion

----------

In conclusion, this paper introduces a new approach that replaces the attention layer with Fourier transforms, emphasizing the importance of mixing information between tokens. While the Fourier transform is not superior to the attention mechanism, it offers a trade-off between speed, context size, and accuracy. As we continue to explore alternative transformations and mixing techniques, we may uncover even better methods to facilitate information flow between tokens.

Do check out the paper and its accompanying code to learn more about this exciting development in deep learning.

References

* FNet: Mixing Tokens with Fourier Transforms. James Lee Thorpe, Joshua Ainsley, Ilya Eckstein, and Santiago Antonion of Google Research. [Link to Paper](https://arxiv.org/abs/2105.03824)