這是近期的一篇視覺Transformer領域的工作，文章并沒有設計更加復雜的token mixer，而是通過簡單的池化算子驗證視覺Transformer的成功在于整體架構設計，即MetaFormer。

簡介

Transformer已經在計算機視覺中展現了巨大的潛力，一個常見的觀念是視覺Transformer之所以取得如此不錯的效果主要是由于基于self-attention的token mixer模塊。但是視覺MLP的近期工作證明將這個token mixer換為spatial MLP依然可以保持相當好的效果。作者并沒有在這方面做過多的探究，而是認為這些工作之所以成功的原因是因為他們模型結構采用MetaFormer這樣的通用架構（即token mixer+channel MLP（FFN）），至于具體采用哪種token mixer模塊，影響并不是那么大。

• 論文標題

  MetaFormer is Actually What You Need for Vision

• 論文地址

  https://arxiv.org/abs/2111.11418

• 論文源碼

  https://github.com/sail-sg/poolformer

介紹

如上圖第二幅圖所示，Transformer的編碼器一般包含兩個組件，即一個注意力模塊和channel MLP等后續組件，前者用于混合token之間的信息，稱為token mixer，后者包括channel MLP和殘差連接等。

忽略掉token mixer用注意力模塊實現的細節，可以將上述架構抽象為上圖第一幅圖的MetaFormer架構，后來的一些工作將注意力模塊換成了簡單的spatial MLP，如代表性的MLP-Mixer，也取得了不遑多讓的效果。這也促使了研究者不斷探索token mixer的形式，比如最近的工作采用傅里葉變換作為token mixer。

但是，這篇論文的作者并沒有在這個思路上越扎越深，而是回過頭了做了個總結，竟然各種token mixer效果都很好，那說明這些模型成功的原因應該是整個架構的設計啊，也就是MetaFormer可以保證模型總能獲得一個不錯的效果。那么自然就會產生一個新的問題了，這個token mixer可以簡潔到什么地步呢？

這篇論文中，作者將注意力模塊換為了一個沒有參數的空間池化算子，原論文對此的描述為簡單到讓人尷尬。

PoolFormer

MetaFormer

首先，作者對視覺Transformer和視覺MLP等架構做了一個歸納，得到了上圖最左側的通用架構MetaFormer。

首先，輸入$I$唄切分為patch并轉換為token，得到序列$X\in {\mathbb{R}}^{N\times C}$，其中$N$為序列長度，$C$是嵌入的維度。

$$X=\mathrm{InputEmb}(I)$$

然后，這個embedding tokens會被送入多個MetaFormer blocks中，每個block包含兩個子block。

第一個子blcok通過一個token mixer模塊來實現不同token之間的信息交互，$Norm$表示某種標準化方法，如LN、BN。$TokenMixer$的形式 就比較多樣了，可以是最近的視覺Transformer提出的注意力模塊，也可以是視覺MLP采用的spatialMLP結構。需要注意的是，盡管有些token mixer能夠混合通道之間的信息，但是token mixer的主要功能還是混合不同token之間的信息。

$$Y=\text{?TokenMixer?}(\mathrm{Norm}(X))+X$$

第二個子block主要由兩層MLP和非線性激活函數組成，形式如下。$W_1\in {\mathbb{R}}^{C\times rC}$和$W_2\in {\mathbb{R}}^{rC\times C}$是可學習參數。

$$Z=\sigma \left(\mathrm{Norm}(Y)W_1\right)W_2+Y$$

將上述過程的token mixer更換，就形成了主流的視覺Transformer和MLP模型。

PoolFormer

此前的很多工作要么改進了自注意力模塊，要么設計了更加精致的token mixer模塊，很少有人關注整體架構。這篇論文的作者認為，MetaFormer這種通用架構才是Transformer和MLP模型取得成功的主要原因。

為了驗證猜想，作者設計了一個沒有參數的空間池化算子來作為token mixer模塊，這個池化沒有任何可學習的參數，只是用于使得每個token平均聚合其附近的tokens的信息。若輸入為$T\in {\mathbb{R}}^{C\times H\times W}$，該池化算子數學上表示如下，$K$表示池化核尺寸，減去自身是因為后續有個殘差連接會再次加上（為了統一為MetaFormer形式）。

$$T_{:,i,j}^{\prime }=\frac{1}{K\times K}\sum _{p,q=1}^K{T}_{:,i+p?\frac{K+1}{2},i+q?\frac{K+1}{2}}?T_{:,i,j}$$

該算子的PyTorch風格代碼如下，其中池化核和padding的設置是為了輸入和輸出尺寸不變。

要知道，self-attention和spatial MLP 的計算復雜度與要混合的token數量成平方倍。 更糟糕的是，spatial MLP 在處理更長的序列時會帶來更多的參數。 因此，self-attention 和spatial MLP 通常只能處理數百個token。相比之下，池化的計算復雜度和序列的長度是線性關系，且不需要可學習參數。

因此，作者以池化為token mixer參考CNN結構和最近的層級Transformer結構，設計了如下圖所示的網絡，即PoolFormer。它由四個stage組成，每個stage下采樣加倍，具體配置見下標。

實驗

作者在圖像分類、目標檢測和實例分割以及語義分割等任務上驗證了PoolFormer的效果。

下表是在ImageNet上驗證圖像分類任務的效果。可以看到，PoolFormer-S24 和 PoolFormer-S36 這樣的小模型就可以分別達到$80.3\%$和$81.4\%$的 Top-1，它們僅僅需要3.6G 和5.2G 的 MACs，超過了幾種典型的 視覺Transformer 和 視覺MLP 模型。這說明，即使使用池化這種極其簡單的 token mixer，MetaFormer仍然具有很強的性能，說明 這整個架構才是我們在設計視覺模型時所真正需要的。

下表是在COCO上目標檢測和實例分割的結果，也很能說明問題。

最后，在ADE20K上進行語義分割的實驗，結果如下。

作者還進行了一些消融實驗，都在下表中列了出來，著重看最后一部分，作者將四個stage的池化換為了注意力模塊或者spatial MLP，發現這些結構混用也不會有什么問題，而且前兩個stage池化后兩個stage注意力這種設計效果尤其不錯，同比之下的ResMLP-B24需要7倍的參數量和8.5倍的MACs才能獲得同等精度。

總結

這篇文章中，作者獨樹一幟提出視覺Transformer及其變種的成功原因主要是架構的設計，并且將token mixer換為了簡單的池化獲得了相當好的效果。這也反映了視覺Transformer其實還有很多值得研究的地方，這篇論文的代碼也已經開源，代碼量不大，感興趣的可以通過源碼了解到更多的細節。本文也只是我本人從自身出發對這篇文章進行的解讀，想要更詳細理解的強烈推薦閱讀原論文。最后，如果我的文章對你有所幫助，歡迎一鍵三連，你的支持是我不懈創作的動力。

總結

以上是生活随笔為你收集整理的PoolFormer解读的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。