[论文解读] Hire-MLP: Vision MLP via Hierarchical Rearrangement

Jianyuan Guo, Yehui Tang+6 more|arXiv (Cornell University)|Aug 30, 2021

Advanced Image and Video Retrieval Techniques参考文献 59被引用 6

一句话总结

Hire-MLP 提出了一种视觉 MLP 架构，通过分层重排机制捕捉局部和全局空间信息，实现灵活的输入分辨率和卓越的性能表现。通过结合区域内部和跨区域的 token 重排与通道混合 MLP，其在 ImageNet（83.8% 的 top-1 准确率）、COCO（51.7% 的框 AP）和 ADE20K（49.9% 的 mIoU）上均取得了最先进结果，优于以往的 MLP 和 Transformer 模型，且在准确率与吞吐量之间实现了更优的权衡。

ABSTRACT

Previous vision MLPs such as MLP-Mixer and ResMLP accept linearly flattened image patches as input, making them inflexible for different input sizes and hard to capture spatial information. Such approach withholds MLPs from getting comparable performance with their transformer-based counterparts and prevents them from becoming a general backbone for computer vision. This paper presents Hire-MLP, a simple yet competitive vision MLP architecture via extbf{Hi}erarchical extbf{re}arrangement, which contains two levels of rearrangements. Specifically, the inner-region rearrangement is proposed to capture local information inside a spatial region, and the cross-region rearrangement is proposed to enable information communication between different regions and capture global context by circularly shifting all tokens along spatial directions. Extensive experiments demonstrate the effectiveness of Hire-MLP as a versatile backbone for various vision tasks. In particular, Hire-MLP achieves competitive results on image classification, object detection and semantic segmentation tasks, e.g., 83.8% top-1 accuracy on ImageNet, 51.7% box AP and 44.8% mask AP on COCO val2017, and 49.9% mIoU on ADE20K, surpassing previous transformer-based and MLP-based models with better trade-off for accuracy and throughput. Code is available at https://github.com/ggjy/Hire-Wave-MLP.pytorch.

研究动机与目标

• 为解决现有视觉 MLP 模型（如 MLP-Mixer）因线性展平 patch 而对输入分辨率变化不灵活的问题。
• 使视觉 MLP 能够在不依赖自注意力机制的前提下，有效捕捉局部归纳偏置和全局上下文信息。
• 设计一种通用、高效且高精度的主干网络，适用于分类、检测和分割等多种计算机视觉任务。

提出的方法

• 提出一种两级分层重排机制：区域内部重排用于空间区域内特征的混合，跨区域重排通过循环移位实现区域间的通信。
• 在重排后应用通道混合 MLP 以建模通道间的特征交互，并通过恢复层保留空间结构。
• 在 Hire-MLP 模块中采用多分支架构：两个分支分别处理高度和宽度方向的重排与恢复，一个分支用于直接通道混合。
• 在通道混合头中采用瓶颈设计，包含两层全连接层，以在准确率与 FLOPs 之间取得平衡。
• 通过基于区域的重排解耦空间处理与 token 序列长度，实现灵活的输入分辨率支持。
• 借鉴 CNN 和 Swin Transformer 的金字塔特征层次结构，以支持多尺度下游任务。

实验结果

研究问题

• RQ1 视觉 MLP 架构是否能在不使用自注意力机制的前提下，有效捕捉局部与全局空间依赖关系？
• RQ2 如何使视觉 MLP 在保持性能的同时，对不同输入分辨率具有鲁棒性？
• RQ3 与标准的 MLP-Mixer 或循环全连接操作相比，token 的分层重排是否能提升特征表示能力？
• RQ4 纯 MLP 架构能否在多种视觉基准测试中实现与最先进 Transformer 模型相当的性能？
• RQ5 如何设计最优的重排与通道混合组件，以在准确率、FLOPs 和推理速度之间实现最佳平衡？

主要发现

• Hire-MLP 在 ImageNet-1K 上达到 83.8% 的 top-1 准确率，优于以往的 MLP 和基于 Transformer 的模型，且在准确率与吞吐量之间具有更优的权衡。
• 在 COCO val2017 上，基于 Hire-MLP 的 RetinaNet 达到 51.7% 的框 AP 和 44.8% 的掩码 AP，在相似 FLOPs 下优于 ResNet、PVT 和 CycleMLP。
• 在 ADE20K 的语义分割任务中，Hire-MLP 在多尺度测试下达到 49.9% 的 mIoU，比 Swin-T 高出 1.6 mIoU，且优于 PVT 和 CycleMLP。
• 通过移位的跨区域通信策略比 ShuffleNet 风格的打乱方式保留了更多相对位置信息，从而带来更好的性能表现。
• 通道混合头中采用两层瓶颈结构在准确率与计算成本之间提供了最佳权衡，超过两层后收益递减。
• 消融实验表明，性能提升源于分层重排机制本身，而非网络深度或层数的简单增加。