ViP解读:视觉MLP结构新作
最近的一篇新的視覺MLP方面的工作,天津大學程明明組參與其中,通過充分利用2D特征圖的位置信息設計的一個MLP結構ViP(Vision Permutator),獲得了相當不錯的性能。
簡介
論文提出一種簡單高效的類MLP結構,Vision Permutator,用于視覺識別任務。通過意識到2D特征圖所攜帶的位置信息的重要性,因此,不同于之前的MLP模型總是沿著展平的空間維度編碼空間信息,ViP通過線性投影的方式沿著高和寬兩個維度編碼特征表達。這使得ViP可以沿著某個空間方向捕獲長程依賴并且同時在另一個方向保留精確的位置信息。通過這種方式獲得的位置感知的輸出特征以一個互相補充的方式聚合以形成感興趣目標的有效表示。實驗表明,ViP具有比肩CNN和視覺Transformer的能力。不需要2D卷積也不需要注意力機制,不使用額外的大尺度數據集,ViP僅僅通過25M的參數就取得了ImageNet上81.5的top-1精度,同等模型大小的約束下,這已經超越了大多數CNN和Transformer,更大版本的ViP精度更高,達到SOTA水準。
-
論文標題
Vision Permutator: A Permutable MLP-Like Architecture for Visual Recognition
-
論文地址
http://arxiv.org/abs/2106.12368
-
論文源碼
https://github.com/Andrew-Qibin/VisionPermutator
介紹
最近的MLP-Mixer和ResMLP證明了純MLP結構也可以在圖像分類等任務上表現得非常好。相比于卷積神經網絡和視覺Transformer分別使用空域卷積和自注意力層來編碼空間信息,類MLP網絡只使用全連接層(或者1x1卷積),因此無論是訓練還是測試都更加高效。但是,MLPs在圖像分類上的優異表現往往得益于大尺度數據集的訓練(如ImageNet-22K和JFT-300M),沒有這些數據的支撐,其性能距離CNNs和Vision Transformers還有不小的差距。
這篇論文中,作者希望探索一個MLP結構僅僅使用ImageNet-1k數據集訓練即可達到不錯的效果,因此,提出了Vision Permutator結構(ViP)。具體而言是,Vision Permutator 通過提出一種新的層來創新現有的 MLP 架構,該結構可以基于基本矩陣乘法更有效地編碼空間信息。 不同于現有的MLP模型,如MLP-Mixer和ResMLP,他們通過展平空間維度然后再線性投影來編碼空間信息的,也就是說作用于尺寸為tokens x channels的張量上,這導致了2D特征圖的位置信息丟失,ViP則保持了原始輸入token的空間維度并沿著寬和高兩個維度編碼空間信息,這樣可以保留充分的位置信息。
具體而言,和很多視覺Transformer結構類似,ViP首先通過對圖像切塊進行token化,將切得得小patch通過線性投影映射為token embedding,如上圖所示。此時獲得的token embedding的維度是height x width x channels,它們被送入一系列的Permutator block中,這個Permutator block包含一個用來編碼空間信息的Permute-MLP和一個用來混合通道信息的Channel-MLP,這兩個結構下文會詳細介紹。隨后經過這些block獲得的特征會送入GAP和Linear層中用于圖像分類任務。和現有的很多MLP結構混合兩個空間方向信息的方式相比,ViP沿著不同的方向獨立處理輸入特征,保證token具有特定方向的信息,這已經被視覺中的很多方法證明是很重要的(如CoordAttention)。
ViP
ViP的整體結構如下圖所示,該網絡以224x224的圖像作為輸入并且將其均分為多個圖像塊(image patches),如14x14或7x7,這些圖像塊隨后被映射為linear embedding(也叫token),這個映射通過一個共享的線性層完成,這和MLP-Mixer是一致的。
接著,這些token送入一序列的Permutator來編碼空間和通道信息,產生的新token沿著空間維度做平均再用全連接層做分類。下文將詳細描述核心的Permutator block以及網絡的一些設置。
Permutator
Permutator的結構如下圖所示,如果不考慮用于標準化的LN和跳躍鏈接,這個結構的核心其實是兩個組件:Permute-MLP和Channel-MLP,從名字上也知道,它們分別負責編碼空間信息和通道信息。Channel-MLP和Transformer原論文中的前饋層結構類似,由兩個全連接層中間夾著一個GELU激活函數。不過,關于空間信息的編碼,最近的MLP-Mixer采用對輸入token沿著空間維度進行線性投影的方式來處理,但是ViP則選擇沿著高和寬兩個方向分別處理。 數學上,對一個CCC維的輸入tokens X∈RH×W×C\mathbf{X} \in \mathbb{R}^{H \times W \times C}X∈RH×W×C,Permutator的運算過程可以總結如下式,其中LNLNLN表示LayerNorm,輸出的Z\mathbf{Z}Z則用于下一層block的輸入。
Y=Permute-MLP?(LN(X))+XZ=Channel-MLP?(LN?(Y))+Y\begin{aligned} &\mathbf{Y}=\text { Permute-MLP }(\mathrm{LN}(\mathbf{X}))+\mathbf{X} \\ &\mathbf{Z}=\text { Channel-MLP }(\operatorname{LN}(\mathbf{Y}))+\mathbf{Y} \end{aligned} ?Y=?Permute-MLP?(LN(X))+XZ=?Channel-MLP?(LN(Y))+Y?
由于Channel-MLP就是原始Transformer中的FFN結構,這里就不贅述了,下面來看這篇論文最核心的Permute-MLP。它的詳細結構如下圖所示,可以發現和之前視覺Transformer結構以及MLP-Mixer這類工作不一樣的地方,它的輸入其實不是二維的(即tokens×channels=HW×Ctokens \times channels = HW \times Ctokens×channels=HW×C),而是三維的。
如上圖所示,Permute-MLP包含三個分支,每個分支負責沿高度、寬度或通道維度對信息進行編碼。其中通道信息的編碼比較簡單,只需要使用一個權重為WC∈RC×C\mathbf{W}_{C} \in \mathbb{R}^{C \times C}WC?∈RC×C的全連接層對輸入進行X\mathbf{X}X進行線性投影即可,得到XC\mathbf{X}_{C}XC?,這就是上圖最下面的分支。
下面介紹如何編碼空間信息,其本質上是一個逐head的permute操作。假定隱層維度為384,輸入圖像為224x224的分辨率,為了沿著高度這個維度編碼空間信息,首先需要對height和channel這兩個維度進行調換(PyTorch中即為Permute操作)。那么具體如何做呢?對輸入X∈RH×W×C\mathbf{X} \in \mathbb{R}^{H \times W \times C}X∈RH×W×C,首先沿著通道對其分成S組(每組為一個segment),即分為[XH1,XH2,?,XHS]\left[\mathbf{X}_{H_{1}}, \mathbf{X}_{H_{2}}, \cdots, \mathbf{X}_{H_{S}}\right][XH1??,XH2??,?,XHS??],每組包含NNN個通道,故有C=N?SC=N * SC=N?S。若patch size為14×1414 \times 1414×14,NNN為16,則XHi∈RH×W×N,(i∈{1,?,S})\mathbf{X}_{H_{i}} \in \mathbb{R}^{H \times W \times N},(i \in\{1, \cdots, S\})XHi??∈RH×W×N,(i∈{1,?,S})。接著,對每個segment XHi\mathbf{X}_{H_{i}}XHi??進行height和channel的permute操作,產生的結果記為[XH1?,XH2?,?,XHS?]\left[\mathbf{X}_{H_{1}}^{\top}, \mathbf{X}_{H_{2}}^{\top}, \cdots, \mathbf{X}_{H_{S}}^{\top}\right][XH1???,XH2???,?,XHS???],它們最終沿著通道維度concat到一起作為輸出,接著這個輸出送入一個參數為WH∈RC×C\mathbf{W}_{H} \in \mathbb{R}^{C \times C}WH?∈RC×C的全連接層中進行height信息的混合,然后再進行一次permute即可恢復維度,記這個輸出為XH\mathbf{X}_{H}XH?,這就是上圖的第一個分支。第二個分支是和第一個分支一樣的操作不過針對于width這個維度,其輸出為XW\mathbf{X}_{W}XW?。最后,這三個分支的輸出加到一起然后通過一個全連接層得到最終的Permute-MLP的輸出,可以用下式表示,其中FC(?)FC(\cdot)FC(?)表示一個參數為WP∈RC×C\mathbf{W}_{P} \in \mathbb{R}^{C \times C}WP?∈RC×C的全連接層。
X^=FC(XH+XW+XC)\hat{\mathbf{X}}=\mathrm{FC}\left(\mathbf{X}_{H}+\mathbf{X}_{W}+\mathbf{X}_{C}\right) X^=FC(XH?+XW?+XC?)
作者也在論文中給出了Permute-MLP的PyTorch風格的偽代碼,如下所示。
接著,作者發現上面這種相加融合三個分支的信息的方式采用的這種逐元素相加的操作效果不是很好,因此進一步提出了重校準不同分支重要性的加權融合方式,并順勢提出了Weighted Permute-MLP,和split attention操作類似,只不過split attention是對組卷積的一組tensor進行,這里是對XH,XW\mathbf{X}_{H}, \mathbf{X}_{W}XH?,XW?和XC\mathbf{X}_{C}XC?進行的。這個過程并不復雜,我這里直接貼上作者的源碼了,需要注意的是,下文所說的ViP默認均采用這種Weight Permute-MLP。
class WeightedPermuteMLP(nn.Module):def __init__(self, dim, segment_dim=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):super().__init__()self.segment_dim = segment_dimself.mlp_c = nn.Linear(dim, dim, bias=qkv_bias)self.mlp_h = nn.Linear(dim, dim, bias=qkv_bias)self.mlp_w = nn.Linear(dim, dim, bias=qkv_bias)self.reweight = Mlp(dim, dim // 4, dim *3)self.proj = nn.Linear(dim, dim)self.proj_drop = nn.Dropout(proj_drop)def forward(self, x):B, H, W, C = x.shapeS = C // self.segment_dimh = x.reshape(B, H, W, self.segment_dim, S).permute(0, 3, 2, 1, 4).reshape(B, self.segment_dim, W, H*S)h = self.mlp_h(h).reshape(B, self.segment_dim, W, H, S).permute(0, 3, 2, 1, 4).reshape(B, H, W, C)w = x.reshape(B, H, W, self.segment_dim, S).permute(0, 1, 3, 2, 4).reshape(B, H, self.segment_dim, W*S)w = self.mlp_w(w).reshape(B, H, self.segment_dim, W, S).permute(0, 1, 3, 2, 4).reshape(B, H, W, C)c = self.mlp_c(x)a = (h + w + c).permute(0, 3, 1, 2).flatten(2).mean(2)a = self.reweight(a).reshape(B, C, 3).permute(2, 0, 1).softmax(dim=0).unsqueeze(2).unsqueeze(2)x = h * a[0] + w * a[1] + c * a[2]x = self.proj(x)x = self.proj_drop(x)return xViPs
下圖是ViP的各種變種形式,采用不同的配置。patch size越大模型越小,ViP-Small/14與ViP-Small/16僅有一個塊嵌入模塊后接16個Permutators,而ViP-Small/7、ViP-Medium/7以及ViP-Large/7則具有兩個階段,每個階段有一個塊嵌入模塊,對于這些模型,作者添加了一些 Permutator,用于編碼細粒度表示,作者發現這對模型性能有益。
實驗
下表是和近期MLP結構的對比,不難看到,ViP-Small/7憑借25M參數取得了81.5%的top1精度,這優于大部分MLP結構,ViP-Medium/7則憑借55M參數取得了82.7%top1精度,超越其他所有MLP結構,ViP-Large/7則憑借88M參數將精度進一步提升到了83.2%,為當前MLP結構中最高精度。
下表則是ViP和SOTA的CNN和ViT的對比結果,雖然比有的方法效果好,但是距離SOTA的CNN和視覺Transformer還有不小距離,也就是說,視覺MLP模型和視覺Transformer結構一樣還有很大的改進空間。
消融實驗部分感興趣的可以查看原論文,這里不多贅述了。
總結
位置信息的利用是這篇文章的核心工作,從這點出發設計了ViP結構,該結構通過不同方向編碼空間信息達到了充分利用二維視覺特征的目的,大大改進了目前MLP模型的性能,為視覺MLP模型的發展提供了一個思路。本文也只是我本人從自身出發對這篇文章進行的解讀,想要更詳細理解的強烈推薦閱讀原論文。最后,如果讀到了這里并且我的文章對你有所幫助,歡迎一鍵三連,你的支持是我不懈創作的動力。
總結
以上是生活随笔為你收集整理的ViP解读:视觉MLP结构新作的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: GCT解读
- 下一篇: 0179-Largest Number(