Hire-MLP: Vision MLP via Hierarchical Rearrangement

原始文件：https://www.yuque.com/lart/papers/lbhadn

這篇文章非常好讀，沒有複雜的單詞和句式，從頭看到尾非常順暢。非常喜歡這樣的寫作風格。

從摘要了解文章

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

區域性和全域性空間資訊的整合：區域性可能依賴卷積操作；在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