FFT-based Dynamic Token Mixer for Vision

摘要

配备多头自注意（MHSA）的模型在计算机性能方面取得了显著的性能。它们的计算复杂性与输入特征图中的像素平方成正比，导致处理缓慢，特别是在处理高分辨率图像时。新型的token Mixer 被提出作为MHSA的替代品，以规避这个问题：基于FFT的令牌混合器，在全局操作中类似于MHSA，但计算复杂度较低。然而，尽管它具有吸引人的特性，但基于FFT的令牌混合器尚未仔细检查其与快速发展的MetaFormer架构的兼容性。在这里，我们提出了一种名为动态滤波器、DFFormer和CDFFFormer的新型令牌混合器，使用动态滤波器来弥合上述差距的图像识别模型。CDFFormer实现了85.0%的Top-1精度，接近卷积和MHSA的混合架构。其他广泛的实验和分析，包括对象检测和语义分割，表明它们与最先进的架构具有竞争力；在处理高分辨率图像识别时，它们的吞吐量和内存效率是卷积和MHSA，与ConvFormer没有太大区别，远远优于CAFormer。我们的结果表明，动态过滤器是应该认真考虑的令牌混合器选项之一。

1. Introduction

变压器架构被推到了计算机视觉领域调查的最前沿。架构在各种视觉识别任务中定位中心，不仅包括图像分类[16，53，67，74，58，5]，还包括动作识别[37，3]，甚至点云理解[18，71，60]。视觉变压器（ViT）[16]及其变体引发了这次爆炸。ViT受到NLP中变压器的启发，并配备了多头自关注（MHSA）机制作为一个关键点。MHSA模块是低通过滤器[39]。因此，它们适合识别有关整个图像的信息。虽然MHSA模块取得了成功，但它一直面临挑战，特别是在由于全局注意力设计导致的二次计算复杂性方面。这个问题在ImageNet分类中并不令人痛苦，而是在语义分割等密集任务中，因为我们经常处理高分辨率输入图像。这个问题可以通过使用局部注意力设计[33、70、10、7]来解决，但令牌到令牌的交互是有限的，这意味着变形金刚的卖点之一被放错了。

GFNet[44]是ViT诱导的架构之一，具有两个迷人的特性：（1）全局滤波器，GFNet中的关键组件，在频域中乘以增加/减少特定频率的特征和复杂参数。它类似于大型内核计算[14,32]，但具有降低理论计算成本[44]的吸引力。此外，与大型内核卷积相比，GFNet的训练难度更小，而大型内核卷积则采用棘手的训练计划。此外，GFNet的计算复杂度是$O(HWC[lo{g}_2(HW))]+HWC$，它优于自注意$O(HW{C}^2+(HW{)}^{2}C)$，其中H是高度，W是宽度，C是通道。换句话说，输入图像的分辨率越高，GFNet的相对优势就越复杂。（2）全局滤波器预计将作为像MHSA一样的低通滤波器，因为全局滤波器是一种全局操作，有捕获低频形成空间。我们遵循[39]并比较了ViT和GFNet之间傅里叶变换的相对对数幅度，GFNet是一个类似于ViT的非分层架构，并显示在图2中。结果表明，GFNet在0附近保留了与ViT相同的低频信息。因此，我们需要解决MHSA的全局滤波器替代品的潜力，因为它在处理高分辨率输入图像时的计算优势及其与MHSA的相似性。

全局滤波器和MHSA在某些方面是不同的，尽管具有低通滤波器特性，就像它们在某些方面不同一样。例如，MHSA是由数据生成并应用到数据的数据依赖操作。相比之下，全局过滤器将数据和参数相乘，因此与数据的依赖性较低。在最近的作品中，基于MLP的模型（不太依赖数据）采用了依赖数据的操作[59，43]，实现了准确性的提高。从调查结果中可以推断出，将数据依赖的概念引入全局过滤器将提高准确性。

与此同时，人们应该注意，与经过充分研究的MHSA及其亲属的架构相比，全局过滤器并没有基于最新和最有效的架构。MetaFormer [65, 66] 关注变压器的整体架构。MetaFormer有一个由令牌混合器组成的主块，该混合器允许任意和通道混合器。它还使这个框架用于全局过滤器，因为它是最复杂的框架之一。如果全球过滤器被有效推进，作为MHSA模块的替代品，并且要实现竞争性准确性，则有必要消除这些差异。

我们建议DFFormer（基于FFT的网络）和CDFFormer（基于FFT的模块的混合架构）来填补上述空白。我们的架构继承MetaFormer和装备模块中，可以为图像中的每对功能通道动态生成全局过滤器，并考虑其内容。我们的实验研究表明，除MHSA外，所提出的方法在ImageNet-1K数据集的视觉模型中实现了最先进的性能。特别是，具有113M参数的CDFFormer-B36实现了85.0%的前1精度。此外，在处理高分辨率im-ages时，拟议的方法实现了比使用CNN和MHSA的基于MHSA的架构和架构更高的吞吐量。此外，提出的模型也足以进行密集预测。我们还发现，动态滤波器倾向于在分层Metaformer上学习低频而不是MHSA。

2. Related Work

Vision Transformers and Metaformers

由于ViT[16]和DETR[6]的成功，NLP中提出的变压器[56]也已成为计算机视觉中的主导架构。不久之后，MLP-Mixer[52]证明了MLP也可以在Transformer中代理MHSA。MetaFormer [65]作为Transformer和MLP-Mixer的抽象类，被提议为宏架构。作者用池化测试了MetaFormer的假设。一般来说，与MHSA和MLP对应的抽象mod-ule被称为令牌混合器。它强调了MetaFormer的重要性，即新类MetaFormer，如使用RNN的音序器[50]和使用图形神经网络的Vision GNN [20]。此外，一项后续研究[66]表明，通过证明激活和与多种类型的令牌混合，可以开发出更高精度的模型。我们保留了[66]的宏观架构。此外，基于FFT的模块具有像MHSA这样的全球操作特征，以实现高分辨率图像的高效架构，而不会造成太大的准确性损失。

FFT-based Networks

近年来，有人提出了使用傅里叶变换的神经网络。FNet[28]，专为NLP设计，包含使用离散或快速傅里叶变换的模块，仅提取实部。一致地说，它们没有参数。Fast-FNet [45] 消除浪费并简化 FNet。相比之下，GFNet[44]是一个基于FFT的网络，专为视觉领域而设计的全局滤波。全局滤波器通过带有傅里叶滤波器的多层功能在频率空间中运行作为一个循环卷积。滤波器是一个参数，即所有样品都使用相同的滤波器。我们的动态滤波器可以使用MLP动态生成傅里叶滤波器。[17]是最相关的作品之一，并提出了AFNO。复杂权重的MLP运算不是元素乘积，而是用于表示自适应和不可分离的卷积。AFNO不是一个可分离的模块。因此，计算成本高于可分离模块。相反，我们提议的方法可以实现动态可分离的傅里叶滤波器，我们的模型在FLOP方面与全局滤波器没有太大区别。我们将在第4.4小节中提及这些吞吐量。并发工作[57]也试图生成动态滤波器，但结构是这样的，只有滤波器系数被改变。虽然滤波器的振幅和频率是动态的，但滤波器的主要属性无法更改（例如，高通滤波器不能更改为低通滤波器）。在同上，动态滤波器可以通过一阶耦合来表示高通和低通滤波器。

Dynamic Weights

目前，有一些对动态权重生成的相关研究。[25]实现滤波器参数动态变化，其中卷积的滤波器权重是模型参数。滤波器生成网络生成滤波器。[19]是一项当代研究。它提出的超网络是一种神经网络工作，用于创建另一个神经网络的权重。它们可以产生CNN和RNN的权重。[64，9]同时工作，两者都使会议动态滤波器。这些工作影响我们的动态滤波器：[64，9]预测标准卷积的实滤波器基的线性组合的实际系数，我们的工作预测了复杂参数滤波器基的线性组合的实数系数。这种方法可以从快速傅里叶变换的线性中作为正向传播的等效运算；因此，我们的工作与这些工作相似。然而，在反向传播中，是否可以学习拟议的动态滤波器并非易事。我们的架构也不需要这些研究中引用的培训困难，例如限制批次大小或调整软最大温度。在卷积之外，合成器[51]、DynaMixer[59]和AMixer[43]具有类似MLP-Mixer的令牌混合器，但具有动态的权重生成。特别是，AMixer包括与我们工作相关的线性组合。

3. Method

3.1. Preliminary: Global Filter

我们将回顾一个离散的傅里叶变换，然后引入一个全局滤波器[44]。我们讨论了2D 离散傅里叶变换（2D-DFT）。对于给定的二维信号$x(h,w)$，定义2D-DFT $\tilde{x}\left(h^{\prime },w^{\prime }\right)$如下：$$\tilde{x}\left(h^{\prime },w^{\prime }\right)=\sum _{h=0}^{H-1}\sum _{w=0}^{W-1}\frac{x(h,w)e^{-2\pi j\left(\frac{h{h}^{\prime }}{H}+\frac{w{w}^{\prime }}{W}\right)}}{\sqrt{HW}}\text{\hspace{1em}(1)}$$它的逆变换存在，被称为二维逆离散傅里叶变换（2D-IDFT）。一般来说，$\tilde{x}\left(h^{\prime },w^{\prime }\right)$对h′和w′是复数和周期性的。我们假设$x\left(h^{\prime },w^{\prime }\right)$是一个实数，那么与$\tilde{x}\left(h^{\prime },w^{\prime }\right)$关联的复矩阵$\tilde{X}\left(h^{\prime },w^{\prime }\right)$是埃尔米特的。$\tilde{x}\left(h^{\prime },w^{\prime }\right)$所属的空间被称为频域，可用于分析频率。此外，频域有一个重要的属性：频域中的乘法等同于原始域中的循环卷积，称为卷积定理。2D-DFT令人印象深刻，但具有$O(H^2{W}^2)$的复杂性。因此，提出了2D-FFT，并经常用于信号处理。它随着复杂性O（HW log2（HW））而得到改进。

其次，我们为特征$X\in {\mathbb{R}}^{C\times H\times W}$定义全局滤波器。全局过滤器G定义以下内容：$$\mathcal{G}(\mathbf{X})={\mathcal{F}}^{-1}(\mathbf{K}\odot \mathcal{F}(\mathbf{X}))\text{\hspace{1em}(2)}$$ $K\in {\mathbb{C}}^{C\times H\times \frac{W}{2}}$是一个可学习的滤波器，F是2D-FFT，其中冗余组件被减少（即rfft2），因为F(x)是埃尔米特的。请注意，该操作等价于基于卷积定理的滤波器${\mathcal{F}}^{-1}(X)$的循环卷积。

众所周知，全局过滤器具有一些属性。（1）理论复杂度为$O(HWC\lceil lo{g}_2(HW)\rceil +HWC)$。当输入是高分辨率时，它比变压器和MLP更有利。（2）由于插值过滤器，全局过滤器可以轻松地放大输入分辨率。有关更多详细信息，请参阅[44]。

3.2. Dynamic Filter

本小节讨论了一种动态滤波器，其中神经网络动态决定了足够的全局滤波器。图4a显示了动态滤波器，具有全局滤波器，其尺寸为N。线性的耦合全局滤波器用于每个通道，而不是可学习的全局滤波器。系数由MLP M统治。我们将整个动态滤波器D指定如下：$$\mathcal{D}(\mathbf{X})={\mathcal{F}}^{-1}\left({\mathcal{K}}_{\mathcal{M}}(\mathbf{X})\odot \mathcal{F}\circ \mathcal{A}(\mathbf{X})\right)\text{\hspace{1em}(3)}$$其中$K_M$表示决定动态滤波器的函数。A是连续的实映射，包括点式卷积和单位映射。

Generating Dynamic Filter
记全局滤波基为$\mathbb{K}$使得 K = { K 1 , . . . , K N } \mathbb{K}=\{K_1,...,K_N\} K={K1​,...,KN​}，$K_1,...,K_N\in C^{H\times \lceil W^2\rceil }$。滤波器$K_M(X)\in C^{C^{\prime }\times H\times \lceil W^2\rceil }$与M相关联进行加权，并由以下定义：
$$\begin{gathered} {\mathcal{K}}_{\mathcal{M}}(\mathbf{X}{)}_{c,:,:}:=\sum _{i=1}^N\left(\frac{e^{s_{(c-1)N+i}}}{{\sum }_{n=1}^N{e}^{s_{(c-1)N+n}}}\right){\mathcal{K}}_i \\ where{\left(s_1,\dots ,s_{N{C}^{\prime }}\right)}^{\top }=\mathcal{M}\left(\frac{{\sum }_{h,w}{\mathbf{X}}_{:,h,w}}{HW}\right)\text{\hspace{1em}(5)} \end{gathered}$$
在本文中，我们使用N = 4，这是K的维度
MLP for weighting
我们描述了MLP M，用于在整个特定计算公式中进行加权：$$M(X)=W_{2}StarReLU(W_{1}LN(X))\text{\hspace{1em}(6)}$$StarReLU(·)是[66]提出的激活函数，$W_1\in {\mathbb{R}}^{C\times int({\rho }^C)}$，$W_2\in {\mathbb{R}}^{int({\rho }^C)\times N{C}^{\prime }}$表示MLP层的矩阵，ρ是中间维度与MLP输入维度的比率。我们选择ρ = 0.25，但其他情况请参阅第4.4节。

3.3. DFFormer and CDFFormer

我们构建符合MetaFormer的DFFormer和CDFFormer。DFFormer和CDFFormer主要由MetaFormer块组成，例如DFFormer块和ConvFormer块。DFFormer和CDFFormer块符合MetaFormer块。$$\mathcal{T}(\mathbf{X})=\mathbf{X}+{\text{ Conv }}_{pw2}\circ \mathcal{L}\circ \text{ StarReLUoConv }{v}_{pw1}\circ \mathrm{LN}(\mathbf{X})\text{\hspace{1em}(7)}$$Convpw1(·)和Convpw2(·)是逐点卷积，因此输出通道的数量分别为(C ′ =)2C和C，$L(\cdot )$是任何维度不变函数。如果$L(\cdot )$是可分离卷积，则在[66]处将T命名为ConvFormer块。如果$L(\cdot )=F^{-1}({\mathcal{K}}_M(X)\odot F(\cdot ))$，我们定义T是DFFormer块。图4b和4c显示了示意图，以了解这些块的结构。

总体框架也遵循了MetaFormer [65, 66]。换句话说，我们使用图5中的四阶段模型。我们为DF-Former（主要由DFFormer块）和CDFFormer（DFFormer块和ConvFormer块的混合模型）准备了四种尺寸的模型。每个模型都配备了一个MLP分类器，以Squared ReLU[47]作为激活。去尾模型结构如表1所示。

4. Experiments

实验自己做去吧，有代码