无自注意力照样高效！RIFormer开启无需token mixer的Transformer结构新篇章

2023-06-13 04:06

©PaperWeekly 原创 · 作者 | 岳廷

研究方向 | 计算机视觉

引言

论文地址：

https://openaccess.thecvf.com/content/CVPR2023/papers/Wang_RIFormer_Keep_Your_Vision_Backbone_Effective_but_Removing_Token_Mixer_CVPR_2023_paper.pdf

问题：Vision Transformer 已取得长足进步，token mixer，其优秀的建模能力已在各种视觉任务中被广泛证明，典型的 token mixer 为自注意力机制，推理耗时长，计算代价大。直接去除会导致模型结构先验不完整，从而带来显著的准确性下降。本文探索如何去掉 token mixer，并以 poolformer 为基准，探索在保证精度的同时，直接去掉 token mixer 模块！

解决方案：

1. 基于重参数化思想，研究没有 token mixer 的模型架构。

2. 通过改进的学习范式以突破简单无 token mixer 骨干的限制，并总结实验经验为 5 条优化策略！

利用所提出的优化策略，在保证原有精度情况下，速度提升至少 17%！！！

主要思想

从强大的教师模型（带有 token mixer）进行定制化蒸馏知识到学生模型（无token mixer）。使用重参数化思想来扩大学生网络的建模能力但保持其高效性，如图 2。简单的仿射变换被引入到学生模型中，用来替换训练时的 token mixer。仿射变换的参数可以在推理过程中合并入 LayerNorm [2] 中，最终实现学生模型无需 token mixer 。

动机

3.1 token mixer 耗时分析

token mixer 效果好，运行效率低，模型定量分析中，单 token mixer 模块可以增加高达 1433.6ms 的延迟，占骨干网络的 46.3%。因此，非常有必要针对 token mixer 模块进行改进。但以 poolformer 进行实验，直接去掉会有较大性能丢失：

上图对比 poolformer 以及将 poolformer 中的池化操作替换成 identity 操作的结果，可以看到，直接替换会带来性能的效果显著下降。

五大改进思路

不使用 ground-truth label进行进行直接监督，直接使用纯 soft label 进行监督更有效。
重参数化仿射变换+原始蒸馏替换 token mixer，无法保持原有 token mixer 性能。
使用分块知识蒸馏（即 module imitation），可提升仿射变换操作建模能力。
大感受野老师模型有利于改进感受野有限的学生模型。
将老师模型（除 token mixer 外）的预训练权重加载到学生模型可改善收敛性和性能。

思路1：直接对无token mixer学生模型进行软蒸馏更有效。

大多数 KD 方法是为带有 token mixer 的模型设计。DeiT 中的一些实验结果表明，使用硬标签而不是软标签作为监督目标，可以显着改善 transformer 性能。但无 token mixer 结构没有 patch 聚合模块，其蒸馏方式应该不同。

为证明这点，表 2 比较了四种不同设置的模型性能。默认的教师是 GFNet-H-B [31]（5400 万参数）。使用真实标签进行硬蒸馏将准确性从 72.31% 提高到73.51%。这表明带有 token mixer 的教师对没有 token mixer 的学生具有积极影响。但如果只软蒸馏而不使用真实标签的组合表现最好，将网络性能提高到 74.05%。

总结 1：对没有 token mixer 的简化模型来说，使用真实标签进行监督学习不是最合适的方式。带有 token mixer 的教师可以帮助指导训练，但仍无法完全弥补去除 token mixer 所带来的性能差距，需要其他策略来解决。

思路2：直接使用重参数化仿射变换+一般蒸馏，无法避免性能下降。

采用结构重参数化的仿射变换+一般知识蒸馏方法，性能一样会下降，表 2 表 3 证明了这一点。

重参数化的仿射变换，具体而言，直接使用仿射变换操作替换 token mixer，仿射变换操作在训练时作为一个单独的层，在推理时融合到 layer norm 中。

表 2 直接使用 identity 层替换 token mixer，并使用各种训练方案进行训练，包括使用原始 label，hard label 以及 soft label 进行各种组合进行训练。作为对比表 3 使用仿射变换层替换 token mixer，同样使用原始 label，hard label 以及 soft label 进行各种组合进行训练。从表 2 表 3 结果可以看到，区别不大，这说明，直接使用仿射变换+一般知识蒸馏，无法发挥仿射变换的性能，也无法实现替换 token mixer 的目的。

另外，师生模型的同构设计启发探索每层模块的知识转移的合适方法，也就是思路 3。

思路3：进行 block-wise 知识蒸馏，即模拟模块技巧（module imitation），有助于利用仿射算子的建模能力。更具体的，文章预训练了一个 PoolFormer-S12 网络作为蒸馏教师网络，并在学生网络中只使用简单的仿射变换来模拟 token mixer。如图三所示：

table1 中原始 poolformer 精度为 75.01，但是去掉 token mixer 后，精度下降到 72.31，但使用仿射变换替代 token mixer，通过 block wise 蒸馏，可以看到精度不但没有下降还上升 0.1 个百分点，而仿射变换在推理时合并到 LN 层，完全等价于 table1 中的 identity 操作。通过这种操作，实现在完全去除 token mixer 时没有精度下降甚至还有提升！

具体操作细节见上图，在 transformer 模块的输出，以及中间层，还有最终输出进行蒸馏。对应总的损失函数为：

其中第一项为前面提到的 soft 蒸馏，第二项为图 3 中对 transformer 模块的输出层进行蒸馏，第三项为对 transformer 模块的中间层进行监督，最后一项同样为对输出层进行监督，只不过使用关系矩阵（relation metrices）进行监督。

table 4 中列的参数，Feat 代表对上面损失函数 2，3 项进行监督的 epoch 数，Rel列为对损失函数最后一项进行监督的 epoch 数据，Layer 代表使用的中间层个数。

思路4：具有大感受野的教师有利于提高有限感受野的学生能力

表 5 比较不同教师架构下的学生性能。虽然 GFNet-H-B [31] 不属于表现最高的教师，但它仍然可以作为一个更好的选择，不管是否使用 module imitation 技巧。

结论 4。归因于教师和学生之间的感受野差异。偏差可以通过蒸馏从一个模型转移到另一个模型。具有大接受域的模型（如 GFNet 在频域中具有可学习的全局过滤器）可以为学生 RIFormer 提供更好的教学，后者感受野有限。

思路5：加载老师模型的预先训练权重（除了 token mixer）到学生模型中，可改进收敛和性能。

由于文章的目标仅是去除 token mixer，其余部分的参数仍保留且在以往过程中没有得到足够重视。文章观察到，初始化 RIFormer（除了仿射操作）的参数时使用对应的老师网络参数，性能得以从 75.13% 改善至 75.36%。这促使文章找到最终训练 RIFormer 的范式。

总括而言，到目前为止，文章已经完成了探索并发现了一个合适的范式来训练 RIFormer。它在宏观设计上与 MetaFormer [46] 大致相同，但不需要任何 token mixer。用所提出的优化方法配备 RIFormer，可以超过带有 token mixer 的复杂模型，对 ImageNet-1K 分类任务而言。

实验结果

5.1 实验设置

在 ImageNet-1K 数据集上，采用 [46] 中提出的训练方案，同时遵循第 4 节中的指南。数据增强包含 MixUp [50]、CutMix [48]、CutOut [52] 和 RandAugment [7]。文章的目的是去除 token mixer 的模型压缩工作，弥补去除 token mixer 导致的性能下降绝对是文章的首要任务，而非提出一个强有力的基线模型。因此使用 600 个延长的训练轮次。同时还对预训练模型进行 30 个 Epoch 的微调，输入分辨率为 384*384。

5.2 主要结果

表 6 展示了 RIFormer 在 ImageNet 分类上的结果。文章主要关注吞吐量指标，因为首要考量是满足边缘设备的延迟要求。如预期所示，比其他类型的骨干拥有明显的速度优势，因为 RIFormer 其构建块不包含任何 token mixer。

令人惊奇的是，使用文章的训练方法，RIFormer 成功地去除了所有 token mixer 而不影响性能。例如，RIFormer-M36 可以在分辨率为 224224 的条件下每秒处理 1185 张图像以上，top-1 准确率达到 82.6%。与此相比，最新基线 PoolFormer-M36 [46] 带有 Pooling token mixer，每秒可以处理 1009 张相同尺寸的图像但更差的 82.1% 的准确率。

文章还与一个高效的骨干 GFNet [31] 进行了比较。它通过全局滤波器进行 token 混合，总的计算复杂度为 O(NlogN)。GFNet-H-B 的吞吐量为 939，top-1 准确率为 82.9%，而文章的 RIFormer-M48 可以达到类似的 82.8% 的准确率与吞吐量 897 持平。与此同时，推理时 RIFormer 主要由 11 conv following LN 组成，没有复杂的 2D FFT 或注意力，更适合硬件专用化。

值得注意的是，没有 token mixer，RIFormer 甚至不能在其构建块中进行基本的 token 混合操作。然而，ImageNet 的实验表明，使用所提出的训练范式，RIFormer 仍然表现出令人期待的结果。文章认为是优化策略起着关键作用。

消融实验

module imitation 的有效性，作为额外仿射算子学习适当权重的重要方法，模组模仿是基于蒸馏的。因此，文章将其与隐藏状态特征蒸馏方法（带有关系）进行比较。采用第 4.2 节的范式，软蒸馏而不考虑交叉熵损失，文章得到表 7 中的结果。使用特征蒸馏，准确率比模组模仿低 0.46%，说明模组模仿正向影响额外权重的优化。

不同加速策略的比较。接下来，文章讨论拆除 token 是否优于其他稀疏化策略。基于 PoolFormer [46] 基线，文章首先构建了一个更薄的 PoolFormer-S9 和 PoolFormer-XS12，分别通过减少深度到 9 和保持其宽度（即嵌入维度）大约为原来的 5/6，以获得与文章的 RIFormer-S12 相当的推理速度。文章也跟随第 4.2 节的软蒸馏范式。表 8 显示结果。直接减少深度或宽度无法比文章无需延迟高昂的 token mixer 更好。

总结来说，本文方法带来了显著加速，达到或超过其他类型的加速策略。但是，文章的方法需要一个合适的训练范式和模组模仿技巧来弥补由于去除 token mixer 而造成的性能损失。

不同老师模型的泛化性。为了验证所提出的训练范式是一种通用的压缩技术，本文采用 [47] 中提出的架构修改作为学生模型，同时将老师模型更换为其他 4 种 MetaFormer基线 [47]，其 token mixer 分别为随机矩阵、池化、可分离深度卷积和自注意力。表 9 显示，文章的方法在不同深度设置和老师模型情况下都有积极效果。

文章的方法可以很好地适应不同的老师模型。加载权重、模组模仿技巧和优化策略一起，所提出的训练范式可以泛用于移除 token mixer 压缩模型。

总结与局限性

作为一篇论文，本文调查了视觉骨干构建块中的 token mixer，并针对其相当大的延迟成本提出了移除的方法，在边缘设备工程应用中有较大的实用性。为了保持剩余结构的有效性，文章彻底重新审视了训练范式。文章观察到，适当的优化策略可以有效地帮助一个没有 token mixer 的模型从另一个模型中学习有用的知识，提高其性能并弥补不完整结构造成的差异。但是，本文的局限性在于没有讨论更多的视觉任务，包括检测、去模糊等。

更多阅读