FNet: Mixing Tokens with Fourier Transforms

This is a PyTorch implementation of the paper FNet: Mixing Tokens with Fourier Transforms.

This paper replaces the self-attention layer with two Fourier transforms to mix tokens. This is a $7\times$ more efficient than self-attention. The accuracy loss of using this over self-attention is about 92% for BERT on GLUE benchmark.

Mixing tokens with two Fourier transforms

We apply Fourier transform along the hidden dimension (embedding dimension) and then along the sequence dimension.

$$\mathcal{R}\left({\mathcal{F}}_{\text{seq}}\right({\mathcal{F}}_{\text{hidden}}(x)\left)\right)$$

where $x$ is the embedding input, $\mathcal{F}$ stands for the fourier transform and $\mathcal{R}$ stands for the real component in complex numbers.

This is very simple to implement on PyTorch - just 1 line of code. The paper suggests using a precomputed DFT matrix and doing matrix multiplication to get the Fourier transformation.

Here is the training code for using a FNet based model for classifying AG News.

FNet - Mix tokens

This module simply implements $$\mathcal{R}\left({\mathcal{F}}_{\text{seq}}\right({\mathcal{F}}_{\text{hidden}}(x)\left)\right)$$

The structure of this module is made similar to a standard attention module so that we can simply replace it.

The normal attention module can be fed with different token embeddings for $\text{query}$,$\text{key}$, and $\text{value}$ and a mask.

We follow the same function signature so that we can replace it directly.

For FNet mixing, $$x=\text{query}=\text{key}=\text{value}$$ and masking is not possible. Shape of query (and key and value ) is [seq_len, batch_size, d_model] .