We leverage the geometric structure of images in order to build an efficient token mixer in comparison to self-attention. Inspired by global convolution, we borrow Fourier Neural Operators (FNO) from PDEs for mixing tokens and propose principled architectural modifications to adapt FNO for images. Specifically, we impose a block diagonal structure on the weights, adaptive weight sharing, and sparsify the frequency with soft-thresholding and shrinkage. We call the proposed mixer Adaptive Fourier Neural Operator (AFNO) and it incurs quasi-linear complexity in sequence length. Our experiments indicate favorable accuracy-efficiency trade-off for few-shot segmentation, and competitive high-resolution segmentation compared with state-of-the-art. There are still important avenues to explore for the future work such as exploring alterantives for the DFT such as the Wavelet transform to better capture locality as in [1]}.
