論文提出CeiT混合網絡，結合了CNN在提取低維特徵方面的局部性優勢以及Transformer在建立長距離依賴關係方面的優勢。CeiT在ImageNet和各種下游任務中達到了SOTA，收斂速度更快，而且不需要大量的預訓練數據和額外的CNN蒸餾監督，值得借鑑

來源：曉飛的算法工程筆記公衆號

論文: Incorporating Convolution Designs into Visual Transformers

論文地址：https://arxiv.org/abs/2103.11816

Introduction

在視覺領域中，純Transformer架構往往需要大量的訓練數據或額外的監督來達到與CNN相當的性能。爲了克服這些限制，論文對直接使用Transformer架構的潛在缺點進行了分析，發現Transformer主要缺乏了CNN的平移不變性以及局部性。於是，論文將CNN在提取低維特徵方面的局部性優勢以及Transformer在建立長距離依賴關係方面的優勢進行結合，提出了Convolution-enhanced image Transformer(CeiT)混合網絡。

論文對原生Transformer做了三處修改：

設計了Image-to-Tokens（I2T）模塊，從生成的低維特徵中提取token序列，而不是將原始輸入圖像直接分割成token序列。
提出Locally-enchanced Feed-Forward(LeFF)層替換每個encoder中的feed-forward層，LeFF能夠促進相鄰token之間的相關性。
在Transformer的頂部附加Layer-wise Class token Attention(LCA)，能夠綜合多層特徵作爲最終輸出。

在ImageNet和七個下游任務的實驗結果表明，CeiT的性能和泛化能力比之前的Transformer和CNN更優，而且不需要大量的訓練數據和額外的CNN蒸餾。此外，CeiT模型的收斂性更好，訓練迭代次數減少了3倍，極大地降低了訓練成本。

Methodology

Image-to-Tokens with Low-level Features

爲了優化初始token序列的生成，論文提出了簡單而有效的Imageto-Tokens（I2T）模塊，從生成的低維特徵中提取token序列，而不是將原始輸入圖像直接分割。如圖2所示，I2T模塊是由卷積層和最大池化層組成的輕量級stem結構，卷積層後面會進行BN操作。整個模塊可表示爲：

其中\(x^{'}\in \mathbb{R}^{\frac{H}{S}\times \frac{W}{S}\times D}\)，\(S\)爲卷積的stride參數，\(D\)爲卷積輸出的通道數。

在得到輸出特徵圖後，根據空間維度從中切割圖像塊序列。爲了保持生成的標記數量與ViT一致，論文將圖像塊的分辨率縮減爲\(（\frac{P}{S} ,\frac{P}{S}）\)，在實踐中設定\(S = 4\)。最後，通過embedding操作將圖像塊序列轉換爲token序列。

I2T模塊能夠充分發揮CNN在提取低層次特徵方面的優勢，並且能夠通過縮小圖像塊的大小來降低embedding的訓練難度。與用ResNet-50來提取後兩個階段的高層特徵的混合類型Transformer對比，I2T模塊要輕量得多。

Locally-Enhanced Feed-Forward Network

爲了將CNN提取局部信息的優勢與Transformer建立長距離依賴關係的能力相結合，論文提出了Locally-enhanced FeedForward Network（LeFF）層。在每個encoder模塊中，保持MHSA模塊不變來保留捕捉token間全局相似性的能力，將原來的前饋網絡層用LeFF取代，LeFF的結構如圖3。

LeFF模塊的執行如公式5-11所示，每條公式對應以下一條處理：

定義MSA模塊生成的輸出爲\(x^h_t \in\mathbb{R}^{(N+1)\times C}\)，將其區分爲圖像token序列\(x^h_p\in \mathbb{R}^{N\times C}\)和一個class token \(x^h_c\in \mathbb{R}^C\)。
對圖像token序列進行線性投影，擴展到更高維度的\(x^{l1}_p\in \mathbb{R}^{N\times (e×C)}\)，其中\(e\)是擴展率。
根據相對於原始圖像的位置，將圖像token序列進行空間維度的還原，得到還原特徵圖\(x^s_p\in \mathbb{R}^{\sqrt{N}\times \sqrt{N}\times(e\times C)}\)。
對還原的特徵圖進行內核大小爲\(k\)的深度卷積處理，增強每個token與相鄰的\(k^2 - 1\)個token的特徵相關性，得到增強特徵圖\(x^d_p\in \mathbb{R}^{\sqrt{N}\times \sqrt{N}\times(e\times C)}\)。
將還原特徵圖中拉平爲\(x^f_p\in \mathbb{R}^{N\times (e\times C)}\)的序列。
將序列中的token映射回初始維度，得到最終的token序列\(x^{l2}_p\in \mathbb{R}^{N\times C}\)，
將最終的token序列與class tken連接，得到最終輸出\(x^{h+1}_t\in \mathbb{R}^{(N+1)\times C}\)。

需要注意，在每次線性投影和深度卷積之後，都會增加進行BatchNorm和GELU處理。

Layer-wise Class-Token Attention

在CNN中，特徵圖的感受域隨着網絡的加深而增加。在ViT中也有類似的現象，自注意計算範圍隨深度增加而增加。因此，特徵的表達在不同層會有所不同。爲了整合不同層的信息，論文設計了Layer-wise Class-token Attention（LCA）模塊。標準的ViT只使用第\(L\)層（最後）的class token \(x^{(L)}_c\)作爲最終特徵，而LCA則綜合不同層的class token作爲最終特徵。

如圖4所示，LCA將一串class token \(X_c = [x^{(1)}_c,\cdots,x^{(l)}_c,\cdots,x^{(L)}_c]\)作爲輸入，其中\(l\)表示層深度。LCA遵循Transformer block的標準實現，包含一個MSA和一個FFN層。LCA的MSA層只計算第\(L\)個class token \(x^{(L)}_c\)和其他class token之間的單向相似性，這樣可以將計算複雜度從\(O(n^2)\)降低到\(O(n)\)。聚合後的\(x^{(L)}_c\)的對應值被送入FFN層，從而得到最終特徵\(x^{(L)^{'}}_c\)。