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

Hire-MLP: Vision MLP via Hierarchical Rearrangement

原始文档：https://www.yuque.com/lart/pa...

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

从摘要了解文章

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_.

主要内容

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

Hierarchical Rearrangement

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

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

Inner-region Rearrangement

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

这里的想法很有意思。

但是仔细思考的话，这里其实可以看做是对于卷积的一种分解。在pytorch中，使用[nn.Unfold](https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html?highlight=unfold#torch.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）来处理。该操作单独使用似乎没什么用，但是如果在重排后再前面设计的区域内的处理，那么就恰好实现了局部区域的跨窗口。

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

实验结果

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

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

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

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

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

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

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

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

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

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

链接

• Hire-MLP: Vision MLP via Hierarchical Rearrangement: https://arxiv.org/pdf/2108.13341.pdf