• 问答
  • 技术
  • 实践
  • 资源
Hire-MLP: 华为诺亚提出分层重排 MLP,性能高达 83.4%

作者丨Sky
来源丨AIWalker
编辑丨极市平台

arXiv:https://arxiv.org/pdf/2108.13341.pdf

本文是华为诺亚\&北大\&悉尼大学在MLP方面的工作,针对MLP-Mixer存在的不灵活性与空间信息提取不够高效等问题,提出内区域重排与跨区域重排机制对MLP-Mixer进行改革。所提Hire-MLP在ImageNet数据集上取得了SOTA性能。比如Hire-MLP-L取得了83.4\%的top1精度,超过了其他Transformer与MLP架构,同时具有更佳的精度-吞吐量均衡。

Abstract

本文提出一种简单且有效的分层重排(Hierarchical Rearrangement)MLP架构Hire-MLP。已有MLP架构(如MLP-Mixer)对于不同图像尺寸的输入不够灵活,对于捕获空间信息不够高效。

Hire-MLP通过引入分层重排聚合全局与局部空域信息对已有MLP架构进行了改革,值得一提的是:Hire-MLP对于下游任务非常友好。具体来说,内区域(inner-region)重排捕获局部信息;为促进跨区域信息通信并捕获全局上下文信息,我们还提出了跨区域(cross-region)重排,它沿空间方向对所有tokens进行训练移位。所提Hire-MLP基于通道混叠MLP与重排操作而构建,因此具有高灵活性与高推理速度。

实验结果表明:Hire-MLP在ImageNet-1K基准数据集上取得了SOTA性能。比如,Hire-MLP取得了83.4\%的top1精度 ,超过了其他Transformer与MLP模型,具有更好地精度-吞吐量均衡。

Method

上图给出了本文所提Hire-MLP整体架构示意图,它通过堆叠多个Hire-MLP模块构成而成。类似MLP-Mixer,每个Hire-MLP包含两个子模块:Hire与Channel-MLP,分别用于进行空间信息与通道信息聚合。Hire-MLP可以描述如下:

$$
Y = \text{Hire-Module}(LN(X)) + X \
Z = \text{Channel-MLP}(LN(Y)) + Y
$$

相比MLP-Mixer,Hire-MLP的主要区别在于:采用Hire模块替代了MLP-Mixer中的token-mixing MLP模块。

Hierarchical Rearrangement Module

MLP-Mixer中的token-mixing MLP采用全连接层捕获跨位置信息,它采用全部token作为输入。由于全连接层的维度是固定的,因此MLP-Mixer无法处理变长序列,故下游迁移能力弱。此外,token-mixing操作仅捕获了厂区信息,而忽视了局部信息,而局部信息对于CV任务非常重要。

针对上述问题,我们提出了分层重排(Hierarchical Rearrangement, Hire)替换MLP-Mixer中的token-mixing MLP。每个Hire模块通过内区域重排聚合局部信息,通过跨区域重排聚合全局信息。所提Hire模块天然的可以处理变长序列。

Region Partition 我们首先将特征拆分为多个区域,在每个区域进行内区域重排。特征可以沿宽/高方向进行拆分。以高维度内区域重排为例,输入特征$X \in R^{H \times W \times C}$被拆分为g组:$X=\left[X{1}, X{2}, \cdots, X{g}\right]$。每个区域$X{i} \in R^{h \times W \times C}$包含h个高方向的token,注:$H=h \times g$。

Inner-Region Rearrangement 对于特征 $X{i}$ 来说,不同token间通过内区域重排操作均等通信。具体来说, 我们沿通道维度将区域内不同token进行拼接得到重排特征 $X{i}^{c} \in R^{W \times h C}$ ,见上图 $a{\circ}$ 然后重排特征将被送入到 MLP模块 $\mathcal{F}$ (它由两个两个bottleneck形式的线性投影层构成,即特征先降维再升维)进行信息混叠生成输出特 征 $X{i}^{o} \in R^{W \times h C}$ 。最后输出特征复原到原始形状得到 $X_{i}^{\prime} \in R^{h \times W \times C}$ 。宽方向的内区域重排与信息混叠操作类似,可参考下图b, 故不再赘述。

Cross-Region Rearrangement 通过内区域重排,输出token仅包含了局部区域的信息聚合,感受野与区域尺寸有关,感受野相对有限。我们又设计了一个跨区域重排,它通过移位操作进行跨区域信息聚合。

上图给出了高方向跨区域重排示意图。它通过沿特定方向给定步长s循环移位token实现,此时局部覆盖的token会发生变化。为得到全能局感受野,跨区域重排会插入在内区域重排之前,见上面的Figure1。类似的,跨区域重排同样需要复原,见上图。宽方向的跨区域重排与复原见下图,不再赘述。

Hire Module 考虑到输入特征同时包含高、宽维度,信息通讯演两个方向通过两个分支实现。受启发于ResNet、ViP,我们还引入了一个不带空域通信的分支。也就是说,输入特征首先被送入到三个分支分别处理,然后将三分支处理结果相加作为输出,可参见下图。

Overall Architecture

下表给出了不同复杂度Hire-MLP的配置信息,它采用金字塔类型架构。整个网络包含四个阶段,输出分辨率从 $H / 4 \times W / 4$ 到 $H / 32 \times W / 32$ 。金字塔架构包含不同层级的空间特征,便于向下游任务迁移。

Experiments

上图给出了所提Hire-MLP与其他方案的性能对比,从中可以看到:

  • Hire-MLP-Smalle取得了81.8\%的精度,而计算量仅为4.2G Flops,优于其他MLP方案。相比AS-MLP、CycleMLP,所提Hire-MLP性能更佳,同时不需要复杂的移位操作或者全连接层变种操作。
  • Hire-MLP-B与Hire-MLP-L分别取得了83.1\%与83.4\%的精度,而计算量分别为81.G与13.5G。
  • 相比DeiT、Swin以及PVT,所提方案具有更快的推理速度;
  • 相比RegNetY,所提方案具有更高的精度,同时具有相似的模型大小和复杂度。
  • Happy语:相比AS-MLP,Hire-MLP好像并没有什么优势,性能相当,速度反而AS-MLP更快

上表对比了区域划分中区域尺寸的影响,从中可以看到:浅层需要更大的区域尺寸 ;但过大会导致性能下降。我们认为:Bottleneck结构中可能存在信息损失。

上表对比了跨区域重排中的步长s的影响,从中可以看到:不进行移位对性能会有影响,说明了移位的重要性

上表对比了不同padding方式的影响,从中可以看到:Circular Padding更适合于所提Hire模块

上表对比了Hire模块中不同成分的影响,从中可以看到:用于捕获局部信息的内区域重排最重要 ;跨区域重排次之。

上表比价了不同跨区域信息通讯的方案,可以看到:移位方式取得了比置换更优的性能,这说明:移位操作可以更好的保持相对位置信息。

吐槽时刻

首先,看上表。从Hire-MLP与AS-MLP的对比来看,Hire-MLP并未看到明显优势。精度基本相当,但AS-MLP的推理速度明显更快啊 ,这个论文并没有提到。原文是这样描述的:

When compared to recently proposed AS-MLP and CycleMLP, our Hire-MLP can obtain slightly better performances without any complicated shift oprations or variants of fully connected layer.

上图是AS-MLP的移位操作,再回过头来看看Hire-MLP的移位,可参考前面的图示。都是类似的移位操作,AS-MLP咋就变成了“complicated shift oprations”? 不懂!

我们再来看一下ViP与Hire-MLP的差异,见上图。事实上,ViP中也有Hire-MLP中的内区域重拍类似的操作,称之为H-C Permute。但从这一点来看,Hire-MLP中的内区域重排与ViP并无本质区别 。但论文并未对两者的关联进行对比说明,全文只是提到了这样一句:

Inspired by the shortcut in ResNet and ViP, an extra branch withoutspatial communication is alse added ...

至于说指标对比,ViP中对应的模块只用了一个Linear,而Hire模块则用了两个Lienar,难说表中0.2-0.3\%的性能提升有多少是来自额外的一个Linear。从ViP与Hire-MLP的对比来看,能看到的改进只有跨区域重排 这一块了。

除此之外,论文前面一直在强调MLP-Mixer架构无法向下游任务迁移,Hire-MLP适合向下游任务迁移,但实验部分并没有向下游任务迁移的实验。这个....

最后吐槽一下:华为的开源速度真的有点慢啊,7月份的CMT尚未开源,Hire-MLP又待何期?AI的方法每周都会出一堆,不开源很难让人follow。

以上仅为笔者学习Hire-MLP过程中的一点疑惑,不一定正确,仅供参考。

  • 0
  • 0
  • 90
收藏
暂无评论