Vision MLP之Hire-MLP: Vision MLP via Hierarchical Rearrangement

875 阅读6分钟

Hire-MLP: Vision MLP via Hierarchical Rearrangement

原始文档:www.yuque.com/lart/papers…

image.png

这篇文章非常好读,没有复杂的单词和句式,从头看到尾非常顺畅。非常喜欢这样的写作风格。

从摘要了解文章

This paper presents Hire-MLP, a simple yet competitive vision MLP architecture via hierarchical rearrangement.

问题:

  1. 是什么?针对空间MLP的简化处理策略。
  2. 为什么?分而治之,简化处理。
  3. 如何做?将原有的无差别全局空间MLP替换为先局部处理后再对各个区域进行关联。

Previous vision MLPs like MLP-Mixer are not flexible for various image sizes and are inefficient to capture spatial information by flattening the tokens.

本文的目的:

  1. 去除曾经的MLP方法对于输入数据尺寸的依赖
  2. 更有效的捕获空间信息的方式

Hire-MLP innovates the existing MLP-based models by proposing the idea of hierarchical rearrangement to aggregate the local and global spatial information while being versatile for downstream tasks.

局部和全局空间信息的集成:局部可能依赖卷积操作;在MLP的语境下,全局空间信息如果不使用空间MLP那又该如何处理呢?难不成是池化操作?——从后面看,这里实际上更类似于SwinTransformer的处理方式,先局部范围内的密集连接,再对各个区域进行关联。 对于下游任务(分割检测等)更加通用,这说明本文的方法会采用多尺度结构。 那这里关于特征的下采样又是如何处理的?池化?跨步卷积?还是patch融合?——从论文中看,使用的是跨步卷积。

Specifically, the inner-region rearrangement is designed to capture local information inside a spatial region. Moreover, to enable information communication between different regions and capture global context, the cross-region rearrangement is proposed to circularly shift all tokens along spatial directions.

看来有些类似于SwinTransformer的处理,局部处理后,再进行全局的偏移来关联各个区域。

The proposed HireMLP architecture is built with simple channel-mixing MLPs and rearrangement operations, thus enjoys high flexibility and inference speed.

这里的处理似乎没有提到空间MLP,那么在局部区域内部是如何处理的?

Experiments show that our Hire-MLP achieves state-of-the-art performance on the ImageNet-1K benchmark. In particular, Hire-MLP achieves an 83.4% top-1 accuracy on ImageNet, which surpasses previous Transformer-based and MLP-based models with better trade-off for accuracy and throughput.

主要内容

image.png

可以看到,主要改动在于原始的空间MLP处,被替换为Hire Module。

Hierarchical Rearrangement

这里是基于区域进行的处理,所以模块中需要先按照H轴和W轴进行区块的划分。

image.png

这里的重排操作是从两个维度上展开的,一个是沿着高一个是沿着宽。

Inner-region Rearrangement

从图中可以看到,区域内的重排,高方向就是将H轴上的相邻层(局部条形区域)堆叠到通道C维度上,宽方向上处理也是类似,通过堆叠到通道上,直接就是用通道MLP就可以实现局部区域特征的处理。

这里的想法很有意思。

但是仔细思考的话,这里其实可以看做是对于卷积的一种分解。在pytorch中,使用nn.Unfold实现卷积操作的处理实际上和此类似。就是通过将局部窗口的数据堆叠到通道维度上,之后使用全连接层处理即可等价于更大的核的卷积操作。

而这里,可以看做是划窗无重叠的处理。说不定这篇后续的工作会尝试使用有重叠的形式。

不过这种方式,倒是更像卷积了。

>>> # Convolution is equivalent with Unfold + Matrix Multiplication + Fold (or view to output shape)
>>> inp = torch.randn(1, 3, 10, 12)
>>> w = torch.randn(2, 3, 4, 5)
>>> inp_unf = torch.nn.functional.unfold(inp, (4, 5))
>>> out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
>>> out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))
>>> # or equivalently (and avoiding a copy),
>>> # out = out_unf.view(1, 2, 7, 8)
>>> (torch.nn.functional.conv2d(inp, w) - out).abs().max()
tensor(1.9073e-06)

另外一点是,这里将局部方形窗口的处理拆分成了H和W两个不同方向的一维条形窗口。即将kxk拆分成1xk和kx1的形式。

看来卷积模型的各种设计几乎已经将模型结构的基础单元穷尽了吧(^__^) 。

Cross-region Rearrangement

而对于跨区域的重排,则是特征沿着H轴或者W轴进行整体的shift(torch.roll)来处理。该操作单独使用似乎没什么用,但是如果在重排后再前面设计的区域内的处理,那么就恰好实现了局部区域的跨窗口。

但是这里有个问题需要注意,可以看到,这里的局部区域的处理仅仅包含在窗口特征偏移之后,而没有考虑对偏移之前的特征进行局部处理。更合理的形式应该是窗口内部处理->窗口特征偏移->窗口内部处理->偏移的窗口特征位置恢复->窗口内部处理(可选),这样的交叉处理似乎才可以更好的覆盖到更广的空间范围,不像现在,窗口处理始终对应的都是固定的区域

实验结果

image.png

消融实验主要讨论了以下几点:

  • 划分的窗口数量:默认是沿着H轴和W轴划分的窗口宽度是一样的。更小的窗口会更加强调局部信息。并且实验中,经验性的在更浅层中使用更大的窗宽,来获得更大的感受野。

image.png

可以看到,当逐渐增大窗宽的时候,性能会有下降。作者推测随着区域大小的增加,瓶颈结构中可能会有一些信息丢失 。

  • 跨窗口偏移的步长s的讨论。

image.png

可以看到,浅层窗口偏移大一点效果会好。或许是因为增大浅层特征的感受野可以带来一定的益处。

  • 针对不同padding形式的讨论。这里是因为对于224的输入,stage4时特征大小为7x7,这不能实现窗口的均匀划分,所以在本文中无重叠的窗口的设定下,是需要加padding的。这里讨论了几种策略。

image.png

  • Hire Module中不同分支的重要性。

image.png

可以看到,区域内部的处理这个部分是非常重要的。其实也可以理解。如果没有了这样的局部操作。那么简单的偏移跨窗口偏移对于通道MLP而言也是没有什么意义的。因为本来就是点操作。

  • 不同的跨窗口通信的形式。

image.png

这里比较了偏移(可以保留一定的相邻关系,即相对位置信息)和shufflenet那样的组间shuffle。可以看到相对位置信息还是比较重要的。

链接