Hire-MLP: 华为诺亚提出分层重排MLP,性能高达83.4%

2021-09-17 14:06:06 浏览数 (1)

arXiv: 2108.13341: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 = [X_1, X_2, cdots, X_g]

。每个区域

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 hC}

,见上图a。然后重排特征将被送入到MLP模块

mathcal{F}

(它由两个两个bottleneck形式的线性投影层构成,即特征先降维再升维)进行信息混叠生成输出特征

X_i^o in R^{W times hC}

。最后输出特征复原到原始形状得到

X_i^{'} 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 人点赞