論文閱讀筆記（四）：AS-MLP AN AXIAL SHIFTED MLP ARCHITECTUREFOR VISION

1.摘要

本文提出了一種軸向移位的MLP體系結構(AS-MLP),更關注區域性特徵的互動，透過特徵圖的通道軸移動，AS-MLP能夠從不同的軸獲取資訊，這使得網路能夠捕捉區域性依賴（可以理解為cnn的空間不變性），這樣的操作使我們能夠利用一個純的MLP體系結構來實現與cnn類體系結構相同的區域性感受野。我們也可以設計AS-MLP的感受野的大小和塊的擴充。本文提出的AS-MLP架構在ImageNet-1K資料集上的表現優於所有基於MLP的架構，並且與基於transformer相比即使FLOPs稍低，也能獲得具有競爭力的效能。此外，AS-MLP也是第一個應用於下游任務(如物件檢測和語義分割)的基於mlp的體系結構。

2.引入

1. 動機

   1. 最近，Tolstikhin等人(2021)首先提出了基於MLP的架構，其中幾乎所有的網路引數都是從MLP(線性層)學習而來的。它取得了驚人的結果，可與cnn類模型相媲美。這些有希望的結果推動了我們對基於mlp架構的探索。
   2. MLP-Mixer 很少能充分利用區域性資訊，這在類cnn架構中是非常重要的因為並非所有畫素都需要長期依賴，區域性資訊更多地關注於提取底層特徵。Swin Transformer一個視窗(7×7)中計算自注意，而不是全域性接收域，這類似於直接使用一個大核心大小的卷積層(7×7)來覆蓋區域性接收域。在這些思想的驅動下，本文主要探討區域性性對基於mlp的體系結構的影響。
   3. 如果將視窗切分（比如說7*7）並在視窗內做token混合對映，那麼線性層中視窗間共享49 *49引數，這極大的限制了模型的容量，影響了引數的學習和最終結果。相反的，如果線性層沒有在視窗之間共享，用固定的影像大小訓練的模型權值不能適應不同輸入大小的下游任務，因為不固定的輸入大小會導致視窗數量的不匹配。因此，更理想的引入局域性的方法是直接對任意位置的特徵點與其周圍特徵點之間的關係建模，而不需要預先設定固定的視窗(以及視窗大小)

2. 主要工作

   1. 提出了一種基於mlp架構的軸向轉移策略，在該架構中，我們在水平和垂直方向上對特徵進行空間轉移。

   2. 之後，通道混合MLP融合（透過1*1卷積實現）這些特徵，使模型能夠獲得區域性依賴關係。

      

      1*1卷積圖示

3.結論

本文提出了一個軸向移動的MLP架構，命名為AS-MLP，用於視覺。與MLP-Mixer相比更注重區域性特徵的提取，透過簡單的特徵軸向移動，充分利用了不同空間位置之間的通道相互作用。透過所提出的AS-MLP，進一步提高了基於mlp的體系結構的效能，實驗結果令人印象深刻。本文模型在ImageNet-1K資料集上的Top-1精度為83.3%，引數為88M, GFLOPs為15.2。這種簡單而有效的方法優於所有基於mlp的體系結構，並且與基於變壓器的體系結構相比，即使FLOPs略低，也能獲得具有競爭力的效能。本文也是第一個將AS-MLP應用於下游任務(例如，物件檢測和語義分割)的研究。與基於變壓器的體系結構相比，結果也具有競爭力，甚至更好，這表明基於mlp的體系結構在處理下游任務方面的能力。

4.相關工作

MLP- mixer (Tolstikhin et al.， 2021)設計了一個非常簡潔的框架，利用矩陣變換和MLP在空間特徵之間傳遞資訊，並取得了很好的效能。併發工作FF (Melas-Kyriazi, 2021)也採用了類似的網路架構，並得出了類似的結論。隨後，本文提出了Res-MLP (Touvron et al.， 2021a)，該演算法僅在ImageNet-1K上訓練剩餘MLP，也獲得了令人印象深刻的效能。gMLP (Liu et al.， 2021a)和EA (Guo et al.， 2021)分別引入了空間門單元(SGU)和外部關注來改善純基於mlp的架構的效能。最近，Container (Gao等人，2021年)提出了一種統一卷積、變壓器和MLP-Mixer的通用網路。S2-MLP (Yu et al.， 2021)使用空間平移MLP進行特徵交換。ViP (Hou等人，2021年)。提出了一種用於空間資訊編碼的Permute-MLP層，以捕獲遠端依賴。與這些工作不同的是，我們專注於在空間維度上捕獲具有軸向移動特徵的區域性依賴關係，從而獲得更好的效能，並可應用於下游任務。此外，還提出了與我們最接近的並行工作CycleMLP (Chen et al.， 2021b)和S2-MLPv2 (Yu et al.， 2021)。S2-MLPv2改進了S2-MLP, CycleMLP設計了Cycle full - connected Layer (Cycle FC)，以獲得比Channel FC更大的接收域。

5.網路結構

對於一個RGB影像，I \(\in R^{3 \times H \times W}\),H和W代表影像的高度和寬度，同時，AS-MLP進行了patch切分操作，這一操作將原影像切分為大小\(4 \times 4\)的大小，這樣的操作，使得所有token（這裡token指的是每一個4* 4的影像塊）的組合將變為\(48 \times \frac{H}{4} \times \frac{W}{4}\)的大小（這樣的原因可以理解為原先輸入的是單個的畫素點，而現在是將4 *4的影像塊當作輸入），AS-MLP有四個階段，在第一階段，對每個token做線性嵌入和AS-MLP，輸出大小為\(C \times \frac{H}{4}\times \frac{W}{4}\)；在第二階段，首先對上一階段輸出的特徵進行patch merge，將相鄰的2 × 2個patch進行分組，得到\(4C \times \frac{H}{8} \frac{W}{8}\)，然後採用線性層將feature大小變形為\(2C×\frac{H}{8} ×\frac{W}{8}\)，然後是級聯的AS-MLP塊.階段3和階段4的結構與階段2相似，層次表示將在這些階段生成。下圖是一個AS-MLP塊內進行的操作。

• Linear Embedding

  PatchEmbed 是一個將輸入影像轉換為影像塊的模組。在模組初始化時，我們需要指定輸入影像的大小 img_size、影像塊的大小 patch_size、輸入影像的通道數 in_chans 和輸出特徵圖的通道數 embed_dim 等引數。此外，我們還可以指定一種歸一化方法（例如，nn.LayerNorm）作為可選引數 norm_layer。

  在 forward 方法中，PatchEmbed 首先透過 nn.Conv2d 將輸入影像 x 劃分成多個大小為 patch_size 的影像塊，然後將所有影像塊展平成一維向量，最終返回特徵向量 x。如果指定了歸一化方法，則在將影像塊展平之前對每個影像塊應用歸一化操作。經過這個操作之後，輸出的形狀是（B,embed_dim,H_patch,W_patch），其中H_patch和W_patch是輸入影像劃分成的patch數量，由輸入影像尺寸和patch_size計算得到。輸出的影像將傳入AS-MLP 塊

• AS-MLP Block

# norm: normalization layer
# proj: channel projection
# actn: activation layer
import torch
import torch.nn.functional as F
def shift(x, dim):
	x = F.pad(x, "constant", 0)
	x = torch.chunk(x, shift_size, 1)
	x = [ torch.roll(x_s, shift, dim) for x_s,
	shift in zip(x, range(-pad, pad+1))]
	x = torch.cat(x, 1)
	return x[:, :, pad:-pad, pad:-pad]

def as_mlp_block(x):
	shortcut = x
	x = norm(x)
	x = actn(norm(proj(x)))
	x_lr = actn(proj(shift(x, 3)))
	x_td = actn(proj(shift(x, 2)))
	x = x_lr + x_td
	x = proj(norm(x))
	return x + shortcut

為了理解這個過程，我們拿一個（3，5，5）的torch張量，通道1資料設為1，通道2資料設為2，通道3設為3做演示，這裡shift_size取5，因此pad=1，這裡的proj是1*1卷積。

經過填充之後，

切分之後

x = [ torch.roll(x_s, shift, dim) for x_s,shift in zip(x, range(-pad, pad+1))]，這裡演示對H維的操作

x = torch.cat(x, 0)

x=x[:, pad:-pad, pad:-pad]

這裡是對H維進行shift操作之後的torch張量，我們同樣可以得到W維進行shift操作的torch張量

對於得到的兩個維度經過shift操作之後的圖片，我們在進行1*1卷積進行通道資訊融合（proj操作）

該公式計算的是某一個畫素點在經過AS-MLP 塊之後的畫素值，\(W^{as-h},W^{as-v} \in R^C\)通道投影的可學習權值在水平和垂直方向上。對於一個三通道的影像，\(Y_{i,j}\)也是三通道的，這樣需要對X使用三個三通道的1*1卷積，對於每一個通道的\(Y_{i，j}\)，我們可以透過上面所舉的例子總結規律，\(Y_{i，j}\)的計算實際上是透過\(X_{i',j}\)和\(X_{i,j'}\)計算的，這裡\(X_{i',j}\)是經過H維度shift操作之後與\(Y_{i，j}\)對應的畫素點，\(X_{i,j'}\)是經過W維度shift操作之後與\(Y_{i，j}\)對應的畫素點，以上面舉例的經過H維度shift操作的特徵圖為例，對於通道一，當前位置為(3,3)的畫素點，在經過shift操作之前位置是（5，3），對於通道二，當前位置為（3，3）的畫素點，在經過shift操作之前的位置為（3，4），對於通道三，當前位置為（3，3）的畫素點，在經過shift操作之前位置為（3，3），由此，我們可以知道，\(X_{i,j}\)位置的畫素點，在經過shift操作之後的位置座標，與通道數、shift_size、膨脹率d有關，我們可以使用不同的通道時、shift_size,膨脹率舉例，可以總結出上式的規律。經過AS-MLP後，特徵圖形狀不變，為\(C \times \frac{h}{4}\times \frac{W}{4}\)

由上式也可推得某一畫素點的感受野。

• LayerNorm：Layer Normalization是一種常見的神經網路正則化方法，它可以用來加速訓練和提高模型的效能。相比於Batch Normalization，Layer Normalization不需要對一個批次內的資料進行統計計算，而是對每個樣本進行獨立的歸一化處理。假設有一個大小為\((n,m)\)的張量\(X\)，其中\(n\)表示樣本數，\(m\)表示特徵數。對於每個樣本\(x_i\)，可以計算出該樣本所有特徵值的均值和方差：

  \[\mu_i=\frac{1}{m} \sum_{j=1}^{m}x_{ij} \]
  \[\sigma_i^2=\frac{1}{m}\sum_{j=1}^{m}{(x_{ij}-\mu_i)^2} \]

  然後可以對每個樣本進行歸一化處理：

  \[\check{x}_{ij}=\frac{x_{ij}-\mu_i}{\sqrt{\sigma_i^2+\epsilon}} \]

  最後，可以將歸一化後的特徵值\(\hat{x}_i\)和一個可學習的權重向量\(w\)進行線性變換，得到輸出特徵值：

  \[y_i=\omega \cdot \check{x}_i+b \]

  其中，\(w\)和\(b\)是可學習的引數，用來調整輸出特徵值的幅度和偏置。

  與BN的區別：

  相比於Batch Normalization，Layer Normalization的主要區別在於計算統計量的方式。Batch Normalization是對一個批次內的所有樣本進行統計計算，然後對每個樣本進行歸一化處理；而Layer Normalization是對每個樣本單獨進行統計計算和歸一化處理，不需要考慮批次的大小和分佈。因此，Layer Normalization更加適合處理樣本數較少或者樣本分佈不均勻的情況，而Batch Normalization則更加適合處理大批次的資料。此外，Layer Normalization可以更好地處理RNN等序列模型，因為它不需要對序列進行批次劃分。但是，相比於Batch Normalization，Layer Normalization的訓練速度和收斂速度較慢，且需要更多的記憶體空間。

  另外，由於Batch Normalization是對整個批次進行歸一化處理，因此它可以有效地減少內部協變數移位（Internal Covariate Shift）的影響，從而加速模型訓練和提高模型的穩定性。而Layer Normalization在處理協變數移位時的效果相對較差。但是，相比於Batch Normalization，Layer Normalization的輸出更加穩定，因為它不會受到批次內樣本數量的影響。

• Patch Merging實現了將一個尺寸為 \(B\times C\times H\times W\) 的輸入特徵 x，轉化為一個尺寸為 \(B\times 4C\times \frac{H}{2}\times \frac{W}{2}\) 的輸出特徵，其中輸入的 \(H\) 和 \(W\) 需要是偶數，轉化方式如下實際上實現了下采樣的過程，通道數加倍，長寬減半：

  1. 將 \(x\) 按照 \(H\) 和 \(W\) 方向上每隔兩個元素取一個進行切片，即將 \(H\times W\) 的矩陣轉化為 \(\frac{H}{2}\times\frac{W}{2}\times 4\) 的四維張量，四個維度對應四個切片；
  2. 將四個切片沿著第二個維度進行拼接，即將 \(\frac{H}{2}\times\frac{W}{2}\times C\) 的四個張量拼接成 \(\frac{H}{2}\times\frac{W}{2}\times 4C\) 的張量；
  3. 對拼接後的張量在第二個維度上應用 LayerNorm；
  4. 對應用 LayerNorm 後的張量應用 \(1\times 1\) 的卷積變換，將通道數從 \(4C\) 減少為 \(2C\)。

• 隨後將進入下一個AS-MLP塊，第三四階段過程和第二階段相同

6.實驗

1. imagenet-1k資料集上的影像分類：為了評估AS-MLP的有效性，在ImageNet-1K基準上進行了影像分類的實驗，該實驗收集於(Deng et al.， 2009)。它包含1.28萬張訓練影像和來自1000個類的20K張驗證影像。我們報告了單作物Top-1精度的實驗結果。我們使用的初始學習率為0.001，帶有餘弦衰減和20個時期的線性預熱。AdamW (Loshchilov & Hutter, 2019)最佳化器被用於訓練300個時代的整個模型，批處理大小為1024。遵循Swin Transformer (Liu et al.， 2021b)的訓練策略，我們還使用平滑比為0.1的標籤平滑(Szegedy et al.， 2016)和DropPath (Huang et al.， 2016)策略。我們將所有的網路架構劃分為基於cnn、基於transformer和基於mlp的架構。輸入解析度為224 × 224。當保持類似的引數和FLOPs時，我們提出的AS-MLP的效能優於其他基於mlp的體系結構。例如，與Mixer-B/16 (Tolstikhin et al.， 2021)(76.4%)和ViP-Medium/7 (Hou et al.， 2021)(82.7%)相比，AS-MLP-S以更少的引數獲得更高的top-1精度(83.1%)。此外，與基於變壓器的架構相比，它獲得了具有競爭力的效能，例如AS-MLP- b(83.3%)與swi -b (Liu等人，2021b)(83.3%)，這表明了我們的AS-MLP架構的有效性。

2. as-mlp塊的選擇和影響：

   1. AS-MLP塊的不同配置：

      • 零填充比其他填充方式更合適
      • 增大膨脹率回事AS-MLP效能有所下降，通常使用影像膨脹進行語義分割而不是影像分類
      • 當增大shift_size時，精度先增大後減小，可能的原因是接收域被放大時，更關注全域性依賴關係，這會使得對區域性特徵的提取被忽略，因此在實驗中使用的是shift size = 5，零填充，膨脹率= 1

      

   2. 連線型別：在不同的shift_size上，並聯方式的效能始終優於串聯方式，說明了並聯方式的有效性。當位移大小為1時，序列連線較好，但由於只使用通道混合的MLP，序列連線不具有代表性。

      

   3. AS-MLP塊的影響。我們還評估了表4中AS-MLP塊的影響。我們設計了5條基線:i) Global-MLP;ii) AxialMLP;iii) Window-MLP;Iv)位移大小(5,1);v)位移大小(1,5)。前三條基線從如何利用MLP進行不同位置的特徵融合的角度設計，後兩條基線從單一方向的軸向位移的角度設計。具體設定請參見附錄A.5。表4中的結果顯示，我們的移動大小為(5,5)的AS-MLP塊的表現優於其他基線。

      

3. coco上的物件檢測

4. 在ade20k上進行語義分割

7.討論

在本文中，我們提出了一個軸向移動的MLP架構，命名為AS-MLP，用於視覺。與MLP-Mixer相比，我們更注重區域性特徵的提取，透過簡單的特徵軸向移動，充分利用了不同空間位置之間的通道相互作用。透過所提出的AS-MLP，我們進一步提高了基於mlp的體系結構的效能，實驗結果令人印象深刻。我們的模型在ImageNet-1K資料集上的Top-1精度為83.3%，引數為88M, GFLOPs為15.2。這種簡單而有效的方法優於所有基於mlp的體系結構，並且與基於變壓器的體系結構相比，即使FLOPs略低，也能獲得具有競爭力的效能。我們也是第一個將AS-MLP應用於下游任務(例如，物件檢測和語義分割)的研究。與基於變壓器的體系結構相比，結果也具有競爭力，甚至更好，這表明基於mlp的體系結構在處理下游任務方面的能力。在未來的工作中，我們將研究AS-MLP在自然語言處理中的有效性，並進一步探討AS-MLP在下游任務中的效能。