北大&華為諾亞提出Vision Transformer的後訓練量化方法

作者 | 小馬

編輯 | CV君

報道 | 我愛計算機視覺(微信id:aicvml)

【寫在前面】

近年來,Transformer模型在計算機視覺領域取得了很大的進展。視覺Transformer的快速發展主要得益於其從輸入影象中提取資訊特徵的能力。然而,主流Transformer模型採用深層架構設計,隨著深度的增加,特性多樣性將不斷減少,即

特徵崩潰(feature collapse)

在本文中,作者從理論上分析了特徵崩潰現象,並研究了這些Transformer模型中的shortcuts和特徵多樣性之間的關係。然後,作者提出了一個增強的shortcuts方案,該方案在原始shortcuts上並行插入具有可學習引數的附加路徑。

為了節省計算成本,作者進一步探索了一種有效的方法,使用

塊迴圈投影(block-circulant projection)

來實現增強的shortcuts。在基準資料集上進行的大量實驗證明了該方法的有效性,在不明顯增加引數和FLOPs的情況下,使SOTA的視覺Transformer的精度提高了1%。

1。 論文和程式碼地址

北大&華為諾亞提出Vision Transformer的後訓練量化方法

論文名稱:Augmented Shortcuts for Vision Transformers

論文地址:

https://arxiv.org/abs/2106.15941

程式碼地址:未開源

2。 Motivation

源於自然語言處理,Transformer模型最近在各種計算機視覺任務中取得了巨大進展,如影象分類、目標檢測和影象處理。其中,ViT模型將輸入影象劃分為視覺序列,並獲得88。36%的top-1精度,這與SOTA的卷積神經網路(CNN)模型(例如,EfficientNet)效能相似。與具有先驗知識(區域性性、平移等變性)的CNN相比,視覺Transformer引入的假設偏置更小,並且在不同視覺任務上實現更好效能的潛力更大。

視覺Transformer中除了自注意層外,通常還包括一個shortcut,可以直接將多個層用恆等對映連線起來。shortcut的引入源於CNN的殘差結構設計。目前有一系列關於解釋和理解shortcut作用的工作。

有的工作對深層網路的梯度進行了分析,結果表明,shortcut連線可以有效地緩解梯度消失和爆炸的問題;有的工作將ResNet視為具有不同長度的路徑的集合,可以透過shortcut解決梯度消失問題;還有的工作證明了shortcut連線可以避免網路引數陷入偽區域性最優,使網路引數收斂到全域性最優值。因此,除了CNN之外,shortcut連線也在其他深度網路中被應用,比如LSTM。

北大&華為諾亞提出Vision Transformer的後訓練量化方法

在視覺Transformer中,shortcut連線繞過了多頭自注意(MSA)和多層感知機(MLP)模組,這對視覺Transformer的成功起著關鍵作用。如上表所示,去掉shortcut的ViT效能非常低。

北大&華為諾亞提出Vision Transformer的後訓練量化方法

根據實驗,刪除shortcut會導致不同patch的特徵隨著網路的變深而變得難以區分(如上圖(a)所示),並且這些特徵對於下游任務預測的表示能力也有限,作者把這種現象稱為

特徵崩潰

。在Transformer中新增shortcut可以緩解這種現象(如上圖(b)所示),並使生成的特徵多樣化。然而,傳統的shortcut只是將輸入特徵複製到輸出,限制了其增強特徵多樣性的能力。

北大&華為諾亞提出Vision Transformer的後訓練量化方法

在這篇文章中,作者提出了一種增強的shortcut方案來提高視覺Transformer中的特徵多樣性(如上圖所示)。除了傳統的shortcut外,作者還提出將MSA模組與多個引數化shortcut並行,以提供更多的替代路徑來繞過注意機制。實現上,將增強的shortcut構造為具有可學習引數和非線性啟用函式的線性投影序列。

作者從理論上證明了,Transformer模型配備了增強的shortcut,可以避免特徵崩潰,併產生更多樣化的特徵。由於計算效率的原因,作者進一步用

塊迴圈矩陣(block-circulant matrices)

代替了原始的密集矩陣,這些矩陣在傅立葉頻域具有較低的計算複雜度,而在空間域具有較高的表示能力。由於迴圈投影的緊湊性,與MSA和MLP模組相比,由增強的shortcut引入的引數和計算成本可以忽略不計。

作者在ImageNet資料集上使用ViT模型及其SOTA變體對增強shortcut的有效性進行了評估。結果表明,配備增強shortcut後,這些模型的效能可提高約1%,且計算成本相當,如下圖所示。

北大&華為諾亞提出Vision Transformer的後訓練量化方法

3。 方法

3。1。 Preliminaries

視覺Transformer(ViT)採用Transformer中提出的架構,MSA和MLP模組交替堆疊。對於影象,ViT將其拆分為N個patch,並將每個patch投影到d維向量中。因此,視覺Transformer的輸入可以看作是矩陣

。帶有H個頭的MSA模組定義為:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中,

表示第

層的特徵,

分別是第h個頭中相應的注意圖和value投影矩陣。

表示H個頭特徵的拼接,

是輸出投影矩陣。注意力矩陣

由自注意機制計算得到,即,

北大&華為諾亞提出Vision Transformer的後訓練量化方法

分別是query和value的投影矩陣,注意矩陣

反映了不同patch之間的關係,

值越大表明patch

和patch

之間的關係越強。MLP模組獨立地從每個patch中提取特徵,這是由兩個線性投影層(引數為

)和一個非線性啟用函式構成的,表示為

。Transformer透過交替堆疊的MLP和MSA模組來構建。

北大&華為諾亞提出Vision Transformer的後訓練量化方法

除了MSA和MLP之外,shortcut連線對於實現高效能也起著至關重要的作用,刪除shortcut會嚴重損害ViT模型的精度。為了探究背後的原因,作者分析ViT類模型中的中間特徵,並計算了不同patch特徵之間的餘弦相似性,並在上圖中顯示了低層(第1層)、中層(第6層)和深層(第12層)的相似矩陣。隨著網路深度的增加,不同patch的特徵很快變得難以區分。作者稱這種現象為

特徵崩潰

,它極大地限制了特徵表示能力,進而阻礙了更高的效能。

特徵崩潰主要是由注意機制引起的,注意機制將不同patch的特徵逐層聚集在一起。將

表示為輸出特徵,注意機制可以表示為:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中,來自第i個patch的特徵

是來自所有其他patch的特徵的加權平均值,權重是注意圖

中的值。雖然注意機制捕獲了全域性關係,但patch特徵的多樣性也減少了。理論上,特徵

的多樣性可以透過特徵和Rank-1矩陣之間的差異來衡量,即:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中

表示矩陣範數,

是向量,

是一個全1的向量。矩陣

的秩為1,直觀上,如果

可以表示N個patch的

,則生成的特徵是冗餘的。較大

意味著特徵

較強的多樣性,因此,作者用

作為多樣性指標。由於MSA模組會導致特徵崩潰,因此作者將重點放在由注意力模組堆疊的模型上,並具有以下定理:

定理1:給定由MSA模組堆疊的模型,第

層特徵的多樣性

可以由輸入資料

限定:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中

是head的數量,

是特徵維度,

是與MSA模組中權重矩陣範數相關的常數。

通常小於1,因此特徵多樣性

會隨著網路深度的增加而迅速減小。此外,作者也從實驗上證明了定理1,如下圖所示:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

可以看出,新增一個與MSA模組並行的shortcut連線可以保持特徵多樣性,具有shortcut的MSA模組可表述為:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中殘差對映(即

) 與MSA模組平行。shortcut連線繞過MSA模組,並提供另一條替代路徑,在該路徑中,特徵可以直接複製到下一層,而不受其他patch的干擾。shortcut的成功表明,使用額外的路徑繞過注意層是增強特徵多樣性和改進transformer模型效能的有效方法。

然而,在一般的vision transformer模型中,每個MSA模組只有一個shortcut連線,它僅將輸入特徵複製到輸出。這種簡單的對映可能沒有足夠的表示能力來最大限度地提高特徵多樣性。因此,作者重新定義了視覺Transformer中現有的shortcut連線,並探索有效且強大的增強shortcut,以產生具有更高多樣性的視覺特徵。

3。2。 Augmented Shortcuts

作者提出了增強的Shortcuts,透過新增更多並行的引數化的投影來減輕特徵崩潰。配備T個增強Shortcuts的MSA模組可以表述為:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中,

是第l層的第i個增強的shortcut,

表示其引數。除了原始的shortcut外,增強的shortcuts還提供了更多的替代路徑來繞過注意力機制。與直連對映將輸入patch複製到相應的輸出不同,引數化投影

可以將輸入特徵轉換為另一個特徵空間。實際上,投影

將對輸入特徵進行不同的變換,只要它們的權重矩陣

不同,並行更多的增強shortcuts可以豐富特徵空間。

的簡單公式可以是線性投影和啟用函式的序列,即:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中

為權重矩陣,

為非線性啟用函式。在沒有shortcuts的Transformer的模型中,特徵多樣性

的上限隨著網路深度的增加而急劇下降 (見定理1)。

定理2:給定由AugMSA模組堆疊的模型,第

層特徵的多樣性

可以由輸入資料

限制,即:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中

是第m層第i個增強shortcut中的權重矩陣,

是啟用函式

的Lipschitz常數。透過比較定理1和定理2,可以看出增強的shortcut多樣性引入了額外的一項

。因為

通常大於1,所以多樣性會隨著深度增加而增大,這就可以抑制注意力機制引起的多樣性衰減。

考慮到在MSA和MLP模組中都存在shortcut連線,因此增強shortcut也可以嵌入到MLP中,即:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中

是第l層中MLP模組的輸入特徵,

表示增強shortcut的引數。

3。3。 Efficient Implementation via Circulant Projection

如上所述,在vision transformer中使用MSA和MLP模組並行多個增強shortcut可以提高特徵多樣性以獲得更高的效能。然而,直接實現

涉及大量的矩陣乘法,計算成本很高。例如,給定特徵

和權重矩陣

,矩陣乘法

消耗

FLOPs,並且d通常在視覺Transformer中很大。因此,作者提出使用

塊迴圈矩陣

來實現增強的shortcut,與視覺Transformer中的其他模組相比,其計算成本可以忽略不計。

迴圈矩陣是結構化矩陣,在傅立葉域上具有極少的引數和極低的計算複雜度。迴圈矩陣

只有

個引數,並且

與向量之間的乘積透過快速傅立葉變換 (FFT) 僅具有

的計算複雜度。因此,作者利用迴圈矩陣來實現增強的shortcut。首先,將原始矩陣

拆分為

個子矩陣

,即:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中,

是子矩陣的大小,

。每個子矩陣

是透過迴圈一個

維向量中的元素而生成的迴圈矩陣

,即:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

為了有效實現投影

,作者首先將輸入

拆分為

個切片

,即

,然後將每個切片

乘以迴圈矩陣

。原定義域中的迴圈矩陣和向量之間的乘積等價於傅立葉定義域中的元素乘,即,

,其中

是離散傅立葉變換,

表示

中的每一行與

中的相應元素相乘。離散傅立葉變換及其逆變換都可以用快速傅立葉變換FFT和它的逆IFFT來有效地計算,並且只有

計算複雜度。輸出的計算方式為:

北大&華為諾亞提出Vision Transformer的後訓練量化方法

其中

的切片。最後,透過將切片

拼接得到

, 即,

4。實驗

4。1。 Experimental Results

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了不同baseline網路,加入增強shortcut之後的對比結果,可以看出,本文的方法對於不同的baseline網路都是有效的。

4。2。 The number of augmented shortcuts

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了不同增強shortcuts數量對結果的影響,可以看出,當路徑數量為2時,效能趨於飽和。

4。3。 Location for implementing the augmented shortcuts

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了不同位置上實現增強shortcuts對結果的影響,可以看出,在MLP和MSA上都加增強shortcut效能更好。

4。4。 Efficiency of the block-circulant projection

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了不同引數b下塊迴圈卷積實現和不使用塊迴圈卷積實現的結果,可以看出塊迴圈卷積可以降低計算量,並且效能降低有限。b=4時能夠在效能和計算量之間達到比較好的平衡。

4。5。 Formulation of the augmented shortcuts

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了不同增強shortcut實現的效能對比。

4。6。 Feature Visualization

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了ViT中三種不同設定下,底層、中層和高層特徵圖的視覺化結果。可以看出,Aug-ViT的特徵更具有可辨別性。

4。7。 Object Detection with Pure Transformer

北大&華為諾亞提出Vision Transformer的後訓練量化方法

上表展示了本文方法和baseline在目標檢測任務上的效能對比。

4。8。 Transfer Learning

北大&華為諾亞提出Vision Transformer的後訓練量化方法

為了驗證本文方法的泛化能力,作者也在其實資料集上做了實驗,結果如上表所示。

5。 總結

在本文中,作者提出了增強的shortcut,用於解決視覺Transformer中的特徵崩潰問題。增強的shortcut與原來的殘差連線平行,並且每個連線都有其可學習的引數,以對輸入特徵進行各種轉換。作者採用有效的迴圈投影用於實現增強的shortcut,其視訊記憶體和計算成本與視覺Transformer中的其他元件相比可以忽略不計。與殘差連線類似,增強的shortcut也不依賴於特定的結構設計,可以將其快速嵌入到視覺Transformer的各種變體中, 用於增強它們在不同任務上的效能。

在本文中,作者的Motivation是,注意力計算是基於全域性的,所以越往深層,特徵多樣性就會越小,殘差連線可以緩解這一現象,但是不能生成更多多樣性的特徵。因此作者加入了多個並行的增強shortcut來獲得更多多樣性的特徵來解決問題。為了計算的高效,作者引用了塊迴圈矩陣來實現這一方法。