GFNet | MLP领域再发力，清华大学提出将FFT思想用于空间信息交互

AIWalker

发布于 2021-07-29 11:54:09

7230

发布于 2021-07-29 11:54:09

文章被收录于专栏：AIWalker

标题&作者团队

本文是清华大学周杰团队在MLP方面的探索，不用现有MLP方案在空域进行token间信息交换，提出了一种全局滤波器方案在频域进行token间信息交换。受益于FFT的log-linear计算复杂度，GFNet可以设计成分层架构形式，能够更高分辨率作为起点，比如

56\times 56

。虽然这篇文章的指标对比最新的VOLO、ViP等不算高，不过它为相关架构设计提供了一个非常不错的思路，值得学习。

Abstract

近期Transformer与MLP模型的发展证明了其具有以更少的归纳偏置取得更佳性能的潜力，这些模型往往基于从原始数据学习空间位置上相关性。然而，这些自注意力与MLP的计算复杂度会随图像尺寸迅速增长，这使得这些方法难以满足高分辨率特征需求。

在本文中，我们提出了全局滤波器网络(Global Filter Network, GFNet)，一种概念简单且计算高效的架构，它在频域以log-linear复杂度学习长距离空间依赖。所提架构通过如下三个关键操作替代ViT中的自注意力层：

2D 离散傅里叶变换；
频域特征与全局可学习滤波器的点乘操作；
2D逆傅里叶变换。

所提方案在ImageNet以以及下游任务上表现出了非常有力的精度-复杂度均衡。相比Transformer与CNN模型，所提方案在高效性、泛化性以及鲁棒性方面极具竞争力。

Method

在正式介绍之前，我们简单介绍一点关于傅里叶变换的基础知识。

Preliminaries

离散傅里叶变换(DFT)在数字信号处理领域起着非常重要的作用。我们以1D DFT为例进行介绍。给定长度为N的序列

x[n], 0\le n \le N-1

，1D DFT通过如下公式将其转换到频域：

注：

X[k]

表示序列

x[n]

的频谱。

此外，值得注意的是：DFT是一对一变换。因此，给定DFT

X[k]

，我们可以通过如下公式进行原始信号

x[n]

重建：

DFT在现代信号处理算法中得到了广泛应用，主要由以下两个原因：

DFT的输入与输出均为离散形势，可以通过计算机快速实现；
DFT存在高效算法，比如利用其对称与周期性的FFT。

Global Filter Networks

上图给出了本文所提方案整体架构示意图，它是一种类似ViT、DeiT的架构，即仅通过PatchEmbedding进行空间尺寸下降，然后通过多个核心模块进行处理，最后后接线性分类层进行分类。

所提方案的输入尺寸为

H \times W

并进行非重叠块拆分与线性投影得到维度D的词。GFNet的核心模块包含两部分：

全局滤波器层，它用于进行空间信息交换；
前馈网络，即MLP部分。

Global Filter Layer 我们提出了全局滤波器层作为自注意力层之外的选择：在频域进行空间信息交互。给定词

x \in R^{H \times W \times D}

，我们首先采用2DFFT变换将其变换到频域：

X = \mathcal{F}[x] \in R^{H\times W \times D}

然后，我们再采用可学习滤波器

K \in C^{H\times W \times D}

进行调制：

\tilde{X} = K \odot X

注：这里的滤波器

由于具有与X相同的维度而被称之为全局滤波器。最后，我们采用逆FFT将调制信号变换回空间并更新词：

x \leftarrow \mathcal{F}^{-1}[\tilde{X}]

上述核心部分的实现伪代码如下，就是这么的简单。

X = rfft2(x, dim=(1, 2))
X_tilde = X * K
x = irfft2(X_tilde, dim=(1, 2))

Architecture variants 考虑到自注意力与MLP的高计算复杂度问题，现有ViT、MLP采用快速降低分辨率的方式，即初始的PatchEmbedding尺寸非常大，比如

14\times 14

。然而，GFNet的的计算复杂度为log-linear，可以避免上述问题。因此，我们可以以更高分辨率(比如

56\times 56

)的特征作为起点，然后逐渐下采样。在这篇文章中，我们主要探索了两种形式的GFNet，即Transformer风格与CNN风格。

对于Transformer风格，类似DeiT与ResMLP-12，我们同样采用了12层模型并得到了三个尺寸的模型GFNet-Ti、GFNet-S以及GFNet-B(通过调整维度、深度等信息即可得到)；
对于类CNN分层风格，我们同样设计了三种复杂度的模型GFNet-H-Ti、GFNet-H-S、GFNet-H-B。相关信息见下表。