基于FFT的动态令牌混合器：突破视觉Transformer的性能瓶颈

摘要

本文聚焦于视觉Transformer中计算复杂度高的问题，介绍基于FFT的动态令牌混合器。提出动态滤波器及DFFormer、CDFFormer模型，经实验验证其在图像分类和下游任务中性能出色，处理高分辨率图像时吞吐量和内存效率优势显著。研究还分析模型特性，为视觉模型发展提供新方向。
关键词：视觉Transformer；动态滤波器；FFT；图像识别；MetaFormer

一、引言

Transformer架构在计算机视觉领域掀起了研究热潮，从图像分类到动作识别，再到点云理解，它都展现出强大的能力。其中，视觉Transformer（ViT）及其变体的成功，很大程度上得益于多头自注意力（MHSA）机制。然而，MHSA也面临着严峻的挑战，其计算复杂度与输入特征图中像素数量的平方成正比，这使得在处理高分辨率图像时，计算成本急剧增加，处理速度大幅下降。

在一些对实时性要求较高的视觉任务中，如自动驾驶场景下的目标检测，高分辨率图像能提供更丰富的细节信息，但MHSA的高计算复杂度成为了阻碍。为了解决这一问题，研究人员提出了多种方法，其中基于FFT的令牌混合器成为了关注的焦点。它具有与MHSA类似的全局操作能力，同时计算复杂度更低，有望成为MHSA的有效替代方案。本文将详细介绍基于FFT的动态令牌混合器在视觉领域的应用，包括其原理、模型构建以及实验验证等方面。

二、相关工作

2.1 Vision Transformers和Metaformers

Transformer最初在自然语言处理领域大放异彩，随着ViT和DETR的出现，它迅速在计算机视觉领域占据重要地位。MLP-Mixer的诞生，证明了MLP可以替代Transformer中的MHSA。MetaFormer作为Transformer和MLP-Mixer的抽象类，其核心由任意令牌混合器和通道混合器组成。这一架构的提出，引发了一系列新的MetaFormer类别出现，如使用RNN的Sequencer、使用图神经网络的Vision GNN以及使用保留机制的RMT等。这些研究不断拓展着Transformer在视觉领域的应用边界 。

2.2 FFT - 基于的网络

近年来，基于傅里叶变换的神经网络逐渐兴起。FNet专为NLP设计，通过离散或快速傅里叶变换提取实部，且没有参数。Fast - FNet则对FNet进行了优化，去除冗余。GFNet是为视觉设计的基于FFT的网络，它的全局滤波器通过在频域中与特征相乘来操作，等效于循环卷积，计算复杂度为$O(HWC[{\log }_2(HW)]+HWC)$，在处理高分辨率图像时具有优势。而本文提出的动态滤波器，与传统的全局滤波器不同，它可以使用MLP动态生成傅里叶滤波器，进一步提升了模型的性能 。

2.3 动态权重

在动态权重生成方面，已有许多有价值的研究。Jia等人实现了卷积滤波器参数的动态变化，由滤波器生成网络生成滤波器。Ha、Dai和Le提出的超网络可以生成CNN和RNN的权重。Yang等人和Chen等人也分别使卷积滤波器动态化。这些工作为本文动态滤波器的设计提供了重要的参考 。

三、方法

3.1 全局滤波器

在介绍动态滤波器之前，先回顾离散傅里叶变换。对于二维信号$x(h,w)$，其二维离散傅里叶变换（2D - DFT）$\tilde{x}(h^{\prime },w^{\prime })$定义为：
$$\tilde{x}(h^{\prime },w^{\prime })=\sum _{h=0}^{H-1}\sum _{w=0}^{W-1}\frac{x(h,w)e^{-2\pi j(\frac{h{h}^{\prime }}{H}+\frac{w{w}^{\prime }}{W})}}{\sqrt{HW}}$$
其中$H,W\in \mathbb{N}$， h , h ′ ∈ { z ∈ Z ∣ 0 < z < H } h, h' \in \{z \in \mathbb{Z} | 0 < z < H\} h,h′∈{z∈Z∣0<z<H}， w , w ′ ∈ { z ∈ Z ∣ 0 < z < W } w, w' \in \{z \in \mathbb{Z} | 0 < z < W\} w,w′∈{z∈Z∣0<z<W}。其逆变换为二维逆离散傅里叶变换（2D - IDFT）。2D - DFT复杂度较高，为$O(H^2{W}^2)$，而2D - FFT的复杂度优化为$O(HW{\log }_2(HW))$，在实际中应用更广泛。

对于特征$X\in {\mathbb{R}}^{C\times H\times W}$，全局滤波器$G$定义为：
$$\mathcal{G}(X)={\mathcal{F}}^{-1}(K\odot \mathcal{F}(X))$$
其中$\odot$是逐元素乘积，$K\in {\mathbb{C}}^{C\times H\times \lceil \frac{W}{2}\rceil }$是可学习的滤波器，$\mathcal{F}$是二维快速傅里叶变换（rfft2）。根据卷积定理，该操作等效于滤波器$F^{-1}(X)$的循环卷积 。

3.2 动态滤波器

动态滤波器是本文的核心创新点之一，神经网络根据输入动态确定合适的全局滤波器。其定义为：
$$\mathcal{D}(X)={\mathcal{F}}^{-1}({\mathcal{K}}_{\mathcal{M}}(X)\odot \mathcal{F}\circ \mathcal{A}(X))$$
其中${\mathcal{K}}_{\mathcal{M}}$表示确定动态滤波器的函数，$\mathcal{A}$是连续实映射，包括逐点卷积和恒等映射。

全局滤波器基$K$表示为 K = { K 1 , … , K N } \mathbb{K}=\{K_1, \ldots, K_N\} K={K1​,…,KN​}，$K_1,\dots ,K_N\in {\mathbb{C}}^{H\times \lceil \frac{W}{2}\rceil }$。与$M$相关的滤波器$K_M(X)\in {\mathbb{C}}^{C^{\prime }\times H\times \lceil \frac{W}{2}\rceil }$通过加权定义：
$${\mathcal{K}}_{\mathcal{M}}(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$$
其中$(s_1,\dots ,s_{N{C}^{\prime }}{)}^{\top }=\mathcal{M}\left(\frac{{\sum }_{h,w}{X}_{i,h,w}}{HW}\right)$。本文中取$N=4$以避免过度计算 。

用于加权的MLP $M$的计算公式为：
$$\mathcal{M}(X)=W_2\text{StarReLU}(W_1\text{LN}(X))$$
其中$\text{LN}(\cdot )$是层归一化，$\text{StarReLU}(\cdot )$是激活函数，$W_1\in {\mathbb{R}}^{C\times int(\rho C)}$，$W_2\in {\mathbb{R}}^{int(\rho C)\times N{C}^{\prime }}$，$\rho$是MLP中间维度与输入维度的比率，本文取$\rho =0.25$ 。

动态滤波器组件中，“任意”模块都允许进行连续的实数映射。动态滤波器Transformer模块（DFFormer Block）是一种配备了动态滤波器的元Transformer模块（MetaFormer block）。

3.3 DFFormer和CDFFormer

DFFormer和CDFFormer基于MetaFormer架构构建。MetaFormer块的定义为：
$$\mathcal{T}(X)=X+{\text{Conv}}_{pw2}\circ \mathcal{L}\circ \text{StarReLU}\circ {\text{Conv}}_{pw1}\circ \text{LN}(X)$$
当$\mathcal{L}(\cdot )$是可分离卷积时，$\mathcal{T}$是ConvFormer块；当$\mathcal{L}(\cdot )={\mathcal{F}}^{-1}(K_M(X)\odot \mathcal{F}(\cdot ))$时，$\mathcal{T}$是DFFormer块。

DFFormer主要由DFFormer块组成，CDFFormer是DFFormer块和ConvFormer块的混合模型。它们都采用四阶段模型，并配备以平方ReLU为激活函数的MLP分类器 。

动态滤波器Transformer（DFFormer）和改进型动态滤波器Transformer（CDFFormer）的模型设置。

“DF”和“CF”分别表示动态滤波器Transformer模块（DFFormer block）和改进型滤波器Transformer模块（CFFormer block）。“K”、“S”、“P”、“C”和“L”分别表示内核大小、步长、填充、通道数以及模块数量。每个元组的索引与各个阶段的顺序相对应。

四、实验

4.1 实验设置

4.1.1 图像分类

在ImageNet - 1K数据集上进行实验，该数据集包含1000个类别，1,281,167张训练图像和50,000张验证图像。训练策略依据Touvron等人的方法，使用多种数据增强技术，如MixUp、CutMix、随机擦除和RandAugment，同时采用随机深度和标签平滑进行正则化。使用AdamW优化器训练300个epoch，批量大小为1024，基础学习率为$\frac{批量大小}{512}\times 5\times 10^{-4}$ ，采用线性热身和余弦退火衰减学习率，权重衰减为0.05 。

4.1.2 语义分割

在ADE20K数据集上进行语义分割实验，采用语义FPN作为基础框架。训练图像调整大小并裁剪为$512\times 512$，测试图像短边调整为512像素。使用AdamW优化器，初始学习率为$2\times 10^{-4}$，多项式调度器，幂为0.9，批量大小为32，训练40K次迭代 。

4.1.3 对象检测

在COCO基准上评估模型在对象检测任务中的性能，采用RetinaNet作为检测框架。使用在ImageNet - 1K上预训练的模型作为骨干网络，训练12个epoch，批量大小为16，AdamW优化器，初始学习率为$10^{-4}$ 。

4.2 实验结果

4.2.1 图像分类结果

将DFFormer和CDFFormer与多种模型进行比较，包括基于CNN的模型、基于注意力的模型、基于保留机制的模型、基于MLP的模型、基于FFT的模型以及混合模型。结果表明，在不使用注意力或保留机制的模型中，DFFormer和CDFFormer的Top - 1准确率表现出色。DFFormer在传统基于FFT的模型中性能领先，比其他同类模型高出0.5%以上，且优于基于MLP的模型。CDFFormer结合卷积，性价比更高，最大的模型CDFFormer - B36性能优于DFFormer - B36 。

模型	类型	参数(M)	FLOP(G)	吞吐量(img/s)	Top - 1(%)
ConvNeXt - T	C	29	4.5	1471	82.1
ConvFormer - S18	C	27	3.9	756	83.0
CSWin - T	A	23	4.3	340	82.7
…	…	…	…	…	…
DFFormer - S18	F	30	3.8	535	83.2
CAFormer - S18	CA	26	4.1	741	83.6
CDFFormer - S18	CF	30	3.9	567	83.1
…	…	…	…	…	…

4.2.2 语义分割结果

基于DFFormer和CDFFormer的模型在ADE20K数据集上的语义分割任务中表现优异，优于基于其他模型（如PoolFormer）的方法。DFFormer - S36的mIoU比PoolFormer - S36高5.5个点，CDFFormer - M36达到了48.6的mIoU 。

骨干网络	参数(M)	mIoU(%)
ResNet - 50	28.5	36.7
PVT - Small	28.2	39.8
PoolFormer - S24	23.2	40.3
…	…	…
DFFormer - S18	31.7	45.1
CDFFormer - S18	31.4	44.9
…	…	…

4.2.3 对象检测结果

在COCO对象检测任务中，以DFFormer和CDFFormer为骨干网络的RetinaNets表现出色，优于可比的ResNet和PoolFormer骨干网络。例如，DFFormer - S36的平均精度（AP）比PoolFormer - S36高5.8个点 。

骨干网络	参数(M)	AP(%)	AP50	AP75(%)	AP S(%)	AP M(%)	AP L(%)
ResNet - 50	37.7	36.3	55.3	38.6	19.3	40.0	48.8
PoolFormer - S24	31.1	38.9	59.7	41.3	23.3	42.1	51.8
…	…	…	…	…	…	…	…
DFFormer - S18	38.1	43.6	64.5	46.6	27.5	47.3	58.1
CDFFormer - S18	37.4	43.4	64.7	46.3	26.3	47.1	57.3
…	…	…	…	…	…	…	…

4.3 消融研究

4.3.1 滤波器

改变动态滤波器的超参数，如将动态滤波器基的维度$N$和中间维度$\rho C$减半，结果显示吞吐量和参数数量几乎不变，准确率分别下降0.2%和0.1%。将DFformer - S18的动态滤波器替换为全局滤波器得到GFFormer - S18，实验表明DFformer - S18比GFFormer - S18性能高0.3%。用AFNO替换动态滤波器后，基于AFNO的模型准确率下降幅度更大，且动态滤波器的吞吐量更好 。

消融	变体	参数(M)	FLOPs(G)	吞吐量(img/sec)	Top - 1 Acc.(%)
基线	DFFormer - S18	30	3.8	535	83.2
滤波器	N = 4 → 2	29	3.8	534	83.0
	ρ = 0.25 → 0.125	28	3.8	532	83.1
	GFFormer - S18	30	3.8	575	82.9
	DF → AFNO	30	3.8	389	82.6

4.3.2 激活函数

DFFormer和CDFFormer使用StarReLU作为激活函数，用GELU和ReLU分别替换StarReLU后，DFFormer - S18比替换后的模型性能高出0.5%以上，证明StarReLU对动态滤波器的有效性。同时，实验还发现即使去除StarReLU的影响，动态滤波器仍对模型性能有显著提升 。

4.4 分析

4.4.1 高分辨率下的优势

随着分辨率的提高，基于MHSA的CAFormer吞吐量显著下降，而DFFormer和CDFFormer的吞吐量与ConvFormer接近。在峰值内存方面，CAFormer随分辨率增加而增加，DFFormer和CDFFormer与ConvFormer相当。这表明DFFormer和CDFFormer在处理高分辨率图像时，在速度和内存效率上具有优势，适合如语义分割这类对分辨率要求高的任务 。

4.4.2 表示相似性

通过小批量线性CKA度量模型间的相似性，发现GFFormer - S18和DFFormer - S18在第3阶段之前非常相似，ConvFormer - S18和DFFormer - S18也有类似情况。而CAFormer - S18和CDFFormer - S18在第2阶段之前相似，第3阶段之后差异较大。对第3阶段的傅里叶分析表明，CAFormer - S18表现像高通滤波器，CDFFormer - S18在除0附近频率外衰减更多，说明MHSA和基于FFT的令牌混合器在学习表示上存在差异 。

4.4.3 动态滤波器基分析

对DFFormer - S18的动态滤波器基进行可视化，发现同一层中滤波器的冗余度降低，且存在高通、低通和带通滤波器，这展示了动态滤波器在频率处理上的多样性和有效性 。

五、结论

本文深入研究了视觉模型中全局滤波器和MHSA的异同，提出了动态滤波器以及基于此的DFFormer和CDFFormer模型。实验证明，这些模型在图像分类和下游任务中性能优异，在处理高分辨率图像时，吞吐量和内存效率方面具有显著优势。