[论文解读] Rethinking Token-Mixing MLP for MLP-based Vision Backbone

Tan Yu, Li Xu+3 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network Applications参考文献 32被引用 9

一句话总结

本文提出了一种循环矩阵通道特定（CCS）的 token 混合 MLP，作为纯 MLP 视觉主干网络中标准 token 混合 MLP 的空间不变且通道特定的替代方案。通过利用循环权重结构和分组通道特定操作，CCS 将参数量减少了高达 90%，同时在 MLP-Mixer 中将 ImageNet1K 的 top-1 准确率从 77.2% 提升至 79.8%，在 ResMLP 中从 79.8% 提升至 80.6%，且参数量和 FLOPs 更少。

ABSTRACT

In the past decade, we have witnessed rapid progress in the machine vision backbone. By introducing the inductive bias from the image processing, convolution neural network (CNN) has achieved excellent performance in numerous computer vision tasks and has been established as \emph{de facto} backbone. In recent years, inspired by the great success achieved by Transformer in NLP tasks, vision Transformer models emerge. Using much less inductive bias, they have achieved promising performance in computer vision tasks compared with their CNN counterparts. More recently, researchers investigate using the pure-MLP architecture to build the vision backbone to further reduce the inductive bias, achieving good performance. The pure-MLP backbone is built upon channel-mixing MLPs to fuse the channels and token-mixing MLPs for communications between patches. In this paper, we re-think the design of the token-mixing MLP. We discover that token-mixing MLPs in existing MLP-based backbones are spatial-specific, and thus it is sensitive to spatial translation. Meanwhile, the channel-agnostic property of the existing token-mixing MLPs limits their capability in mixing tokens. To overcome those limitations, we propose an improved structure termed as Circulant Channel-Specific (CCS) token-mixing MLP, which is spatial-invariant and channel-specific. It takes fewer parameters but achieves higher classification accuracy on ImageNet1K benchmark.

研究动机与目标

• 解决现有纯 MLP 视觉主干网络中 token 混合 MLP 的空间敏感性和通道无关性局限。
• 在提升特征表示能力的同时，增强模型对空间平移的鲁棒性。
• 在不损失性能的前提下，减少 token 混合 MLP 层的参数数量。
• 与原始 MLP 基础模型相比，在更少参数和 FLOPs 下实现更高的 ImageNet1K 准确率。
• 证明循环结构和分组通道特定操作在类似 Vision Transformer 架构中的有效性。

提出的方法

• 提出一种循环矩阵通道特定（CCS）的 token 混合 MLP，用循环权重矩阵替代标准全连接层，以实现空间不变性。
• 通过将通道划分为若干组并为每组应用独立的循环矩阵，实现通道特定的混合，从而增强表示能力。
• 利用快速傅里叶变换（FFT）高效计算循环矩阵的矩阵-向量乘法，降低计算成本。
• 引入分组设计，使每组通道独立进行 token 混合，实现更丰富的模式编码。
• 用所提出的 CCS 模块替换 MLP-Mixer 和 ResMLP 中的标准 token 混合 MLP，以评估性能与效率。
• 引入超参数 G（分组数量）以控制模型容量与参数量之间的权衡。

实验结果

研究问题

• RQ1 标准 token 混合 MLP 的空间特定设计在多大程度上影响模型对空间平移的鲁棒性？
• RQ2 现有 token 混合 MLP 的通道无关性在多大程度上限制了视觉主干网络中的特征表示能力？
• RQ3 循环权重结构能否在减少参数量的同时实现空间不变性？
• RQ4 通道特定混合是否能提升纯 MLP 视觉架构的性能？
• RQ5 所提出的 CCS 模块能否在更少参数和 FLOPs 下实现比标准 token 混合 MLP 更高的准确率？

主要发现

• 在 MLP-Mixer-B/16 中，将原始 token 混合 MLP 替换为 CCS 后，ImageNet1K 的 top-1 准确率从 77.2% 提升至 79.8%，参数量仅 57M。
• 在 ResMLP-36 中，CCS 模块将 top-1 准确率从 79.8% 提升至 80.6%，同时将参数量从 44M 减少至 43M。
• 当 G=8 组时，模型在准确率与参数效率之间达到最佳平衡，top-1 准确率为 80.6%，参数量为 43M。
• 当 G 超过 8 后，准确率趋于饱和，但参数量从 43M 增加至 46M，表明收益递减。
• CCS 在更少参数和 FLOPs 下，性能与 S²-MLP-wide 和 S²-MLP-deep 相当或更优。
• 循环结构通过 FFT 实现高效计算，但当 patch 数量较小时（如 N=196），优势不那么明显。