Vision Transformer 必读系列之图像分类综述(二): Attention-based
文 @ 000007
號外號外:awesome-vit 上新啦,歡迎大家 Star Star Star ~
https://github.com/open-mmlab/awesome-vit?github.com/open-mmlab/awesome-vit
0 前言
在Vision Transformer 必讀系列之圖像分類綜述(一):概述一文中對 Vision Transformer 在圖像分類中的發展進行了概述性總結,本文則對其中涉及的 Attention-based 部分進行詳細說明。下一篇文章則會對概述中涉及的其他部分進行說明。
ViT 進展匯總思維導圖如下圖所示:
注意:文中涉及到的思維導圖,可以通過?https://github.com/open-mmlab/awesome-vit?下載。
1 Transformer
論文題目: Attention is All You Need
論文地址:?https://arxiv.org/abs/1706.03762
Transformer 結構是 Google 在 2017 年為解決機器翻譯任務(例如英文翻譯為中文)而提出,從題目中可以看出主要是靠 Attention 注意力機制,其最大特點是拋棄了傳統的 CNN 和 RNN,整個網絡結構完全是由 Attention 機制組成。為此需要先解釋何為注意力機制,然后再分析模型結構。
1.1 Attention 注意力機制
人生來就有注意力機制,看任何畫面,我們會自動聚焦到特定位置特定物體上。此處的 Attention 機制也是同一個含義,對于需要的任何模態,不管是圖像、文本、點云還是其他,我們都希望網絡通過訓練能夠自動聚焦到有意義的位置,例如圖像分類和檢測任務,網絡通過訓練能夠自動聚焦到待分類物體和待檢測物體上。
注意力機制不是啥新鮮概念,視覺算法中早已廣泛應用,典型的如?SENet。
利用 Squeeze-and-Excitation 模塊計算注意力權重概率分布,然后作用于特征圖上實現對每個通道重加權功能。
可以舉一個更簡單的例子,假設有一個訓練好的分類網絡,輸入一張圖片,訓練好的分類網絡權重 W 和圖片 X 進行注意力計算,從 X 中提取能夠有助于分類的特征,該特征最終可以作為類別分類依據。W 和 X 都是矩陣,要想利用 W 矩陣來達到重加權 X 目的,等價于計算 W 和 X 的相似度(點乘),然后將該相似度變換為權重概率分布,再次作用于 X 上就可以。
以一個簡單貓狗二分類例子說明。網絡最終輸出是 2x1 的向量,第一個數大則表示貓類別,否則為狗類別,假設網絡已經訓練好了,其 W 為 shape 為 2x1 的向量,值為?[[0.1, 0.5]],X 表示輸入圖片 shape 也是 2x1,其值為?[[0.1, 0.8]],可以看出其類別是狗,采用如下的計算步驟即可正確分類:
- W 和 X 的轉置相乘,即計算 W 中每個值和 X 中每個值的相似度,得到 2x2 矩陣,值為 [[0.01,0.08], [0.05,0.4]]。
- 第二個維度進行 Softmax,將其轉化為概率權重圖即為 [[0.4825, 0.5175], [0.4134, 0.5866]]。
- 將上述概率權重乘以 X,得到 shape 為 2x1 輸出,值為 [[0.4622, 0.5106]]。
- 此時由于第二個值大,所以正確分類為狗。
X 是含有狗的圖片矩陣,能夠正確分類的前提是訓練好的 W 矩陣中第二個數大于第一個數。可以簡單理解上述過程是計算 W 和 X 的相似度,如果兩個向量相似(都是第二個比第一個數大),那么就分類為狗,否則就分類為貓。
上述計算過程可以用如下公式表示:
對應到上面例子,Q 就是訓練好的 W 矩陣,K 是圖片輸入,V 和 K 相等,其通用解釋為利用 Q 查詢矩陣和 K 矩陣進行相似度計算,然后轉換為概率分布,此時概率值大的位置表示兩者相似度大的部分,然后將概率分布乘上 V 值矩陣,從而用注意力權重分布加權了 V 矩陣,也就改變了 V 矩陣本身的分布。如果注意力機制訓練的很好,那么提取的 V 應該就是我們想要的信息。分母 d_k 的平方根是為了避免梯度消失,當向量值非常大的時候,Softmax 函數會將幾乎全部的概率分布都分配給了最大值對應的位置,也就是說所謂的銳化,通過除以分母可以有效避免梯度消失問題,穩定訓練過程。
上述公式就是論文中提出的最重要的?Scaled Dot-Product Attention?計算公式,先利用點乘計算 QK 矩陣的相似度,除以分母 d_k 平方根進行 Scaled 操作,然后 Softmax 操作將其轉換為概率乘以 V 實現 Attention 功能。
需要注意: 為了讓上面公式不會報錯,其 Shape 關系必須為?Q - (N, M),K - (P, M),V - (P, G), 一般來說 K 和 V Shape 相同,但是 Q Shape 不一定和 K 相同。通過靈活地改變這些維度就可以控制注意力層的計算復雜度,后續大部分改進算法都有利用這一點。
1.2 Transformer 結構分析
Transformer 是為了解決機器翻譯任務而提出。機器翻譯是一個歷史悠久的問題,可以理解為序列轉序列問題,也就是我們常說的 seq2seq 結構,解決這類問題一般是采用 encoder-decoder 結構,Transformer 也沿用了這種結構。翻譯任務一個常規的解決方案如下所示:
對應到 Transformer 中的一個更具體的結構為:
主要包括編碼器和解碼器組件,編碼器包括自注意力模塊(QKV 來自同一個輸入)和前向網絡,解碼器和編碼器類似,只不過內部多了編碼器和解碼器交互的交叉注意力模塊。
通常來說,標準的 Transformer 包括 6 個編碼器和 6 個解碼器串行。
Transformer 完整結構如下所示:
編碼器基本組件包括:?源句子詞嵌入模塊 Input Embedding、位置編碼模塊 Positional Encoding、多頭自注意力模塊 Muti-Head Attention、前向網絡模塊 Feed Forward 以及必要的 Norm、Dropout 和殘差模塊。
解碼器基本組件類似包括:?目標句子詞嵌入模塊 Output Embedding、位置編碼模塊 Positional Encoding、帶 mask 的自注意力模塊 Masked Muti-Head Attention、交叉互注意力模塊 Muti-Head Attention、前向網絡模塊 Feed Forward 、分類頭模塊 Linear+Softmax 以及必要的 Norm、Dropout 和殘差模塊。
由于本文重點是分析視覺方面的 Transformer,故沒有必要對機器翻譯過程進行深入解析,讀者只需要理解每個模塊的作用即可,而且視覺分類 Transformer 任務和 NLP 機器翻譯任務不一樣,實際上也不需要解碼器模塊,相比 NLP 任務會簡單很多。
1.2.1 編碼器基本組件
(1) 源句子詞嵌入模塊 Input Embedding
機器翻譯是句子輸入,句子輸出,每個句子由單詞構成,將句子編碼成程序可以理解的向量過程就叫做詞嵌入過程,也就是常說的 Word2Vec,對應到圖像中稱為 Token 化過程即如何將圖像轉換為更具語義的 Token,Token 概念會在 ViT 中詳細描述。
(2) 多頭自注意力模塊 Muti-Head Attention
在 1.1 小節已經詳細說明了注意力計算過程。左邊是最簡單的 Scaled Dot-Product Attention,單純看上圖你可以發現沒有任何可學習參數,那么其存在的意義是啥? 實際上可學習參數在 QKV 映射矩陣中,在自注意力模塊中會對輸入的向量分別乘上可學習映射矩陣 W_Q、W_K 和 W_V,得到真正的 Q、K 和 V 輸入,然后再進行 Scaled Dot-Product Attention 計算。
為了增加注意力特征提取的豐富性,不會陷入某種局部特性中,一般會在注意力層基礎上(單頭注意力層)引入多個投影頭,將 QKV 特征維度平均切分為多個部分(一般分成 8 部分),每個部分單獨進行自注意力計算,計算結果進行拼接 。在特征維度平均切分,然后單獨投影、計算,最后拼接可以迫使提取的注意力特征更加豐富。也就是上面的多頭注意力模塊 Multi-Head Attention。
(3) 前向網絡模塊 Feed Forward
前向網絡模塊主要是目的是對特征進行變換,其單獨作用于每個序列(只對最后一個特征維度進行變換)。由于沒有結構圖,故直接貼相關代碼,包括兩個 Position-wise FC 層、激活層、Dropout層和 LayerNorm 層。
(4) Norm、Dropout 和殘差模塊
在每個注意力層后面和前向網絡后都會接入 Dropout 、殘差模塊和 Layer Norm 模塊。這些必要的措施對整個算法的性能提升非常關鍵。至于為啥用 Layer Norm 而不是 Batch Norm ,原因是機器翻譯任務輸入的句子不一定是同樣長的,組成 Batch 訓練時候會存在大量 Padding 操作,如果在 Batch 這個維度進行 Norm 會出現大量無效統計,導致 Norm 值不穩定,而 Layer Norm 是對每個序列單獨計算,不考慮 Batch 影響,這樣比較符合不定長序列任務學習任務。當然如果換成圖像分類任務,則可以考慮使用 BN 層,后續有算法是直接采用 BN 的。
(5) 位置編碼 Positional Encoding
考慮一個分類任務,輸入一段句子判斷是疑問句還是非疑問句?現在有兩條語句分別是:
- 不準在地鐵上吃東西
- 在地鐵上吃東西準不
自注意力層的計算不會考慮字符間的順序,因為每個字符都是單獨和全局向量算相似度,也就是說上面兩個句子輸入進行注意力計算,輸出的向量值是相同的,只不過相對位置有變化。如果我們對輸出向量求和后值大于 0 還是小于 0 作為分類依據,那么上面兩個句子輸出相加值是完全相同的,那就始終無法區分到底是疑問句還是非疑問句,這就是我們常說的 Transformer 具有位置不變性。要解決這個問題,只需要讓模型知道輸入語句是有先后順序的,位置編碼可以解決這個問題。
加入位置信息的方式非常多,最簡單的可以是直接將輸入序列中的每個詞按照絕對坐標 0,1,2 編碼成相同長度的向量,然后和詞向量相加即可。作者實際上提出了兩種方式:
- 網絡自動學習,直接全 0 初始化向量,然后和詞向量相加,通過網絡學習來學習位置信息。
- 自己定義規則,規則自己定,只要能夠區分輸入詞順序即可,常用的是 sincos 編碼。
實際訓練選擇哪一種位置編碼方式發現效果一致,但是不管哪一種位置編碼方式都應該充分考慮在測試時候序列不定長問題,可能會出現測試時候非常長的訓練沒有見過的長度序列,后面會詳細說明。
1.2.2 解碼器基本組件
其大部分組件都和編碼器相同,唯一不同的是自注意力模塊帶有 mask,還額外引入了一個交叉注意力模塊以及分類頭模塊。
(1) 帶 mask 的自注意力模塊
注意這個模塊的輸入是目標序列轉化為詞向量后進行自注意力計算。機器翻譯是一個 seq2seq 任務,其真正預測是:最開始輸入開始解碼 token 代表解碼開始,解碼出第一個詞后,將前面已經解碼出的詞再次輸入到解碼器中,按照順序一個詞一個詞解碼,最后輸出解碼結束 token,表示翻譯結束。
也就是當解碼時,在解碼當前詞的時候實際上不知道下一個詞是啥,但在訓練時,是將整個目標序列一起輸入,然而注意力計算是全局的,每個目標單詞都會和整個目標句子計算自注意力,這種訓練和測試階段的不一致性無法直接用于預測。為此我們需要在訓練過程中計算當前詞自注意力時候手動屏蔽掉后面的詞,讓模型不知道后面詞。具體實現就是輸入一個 mask 來覆蓋掉后面的詞。
由于這種特性是只存在于 NLP 領域,圖片中不存在,故不再進行更深入分析。
(2) 交叉注意力
交叉注意力模塊和自注意力模塊相同,只不過其 QKV 來源不同,Q 來自解碼器,KV 來自編碼器,交叉注意力模塊會利用 Q 來提取編碼器提取的特征 KV,然后進行分類。
(3) 分類頭
分類頭就是普通的線性映射,轉換輸出維度為分類個數,然后采用 CE Loss 進行訓練即可。
1.3 總結
Transformer 結構內部存在多個組件,但是最核心的還是注意力模塊,在原始論文中作者也引入了大量的可視化分析來證明注意力模塊的作用,有興趣的建議閱讀原文。可能作者自己也沒有想到這篇論文會在視覺領域引起另一個全新的風尚,開辟出一條新的看起來前途一片光明的道路。
圖片來自?A Survey of Visual Transformers。
2 Vision Transformer
論文題目: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
論文地址:?https://arxiv.org/abs/2010.11929
ViT 是第一篇成功將 Transformer 引入到視覺領域且成功的嘗試,開辟了視覺 Transformer 先河。其結構圖如下所示:
其做法非常簡單,簡要概況為:
- 將圖片分成無重疊的固定大小 Patch (例如 16x16),然后將每個 Patch 拉成一維向量, n 個 Patch 相當于 NLP 中的輸入序列長度(假設輸入圖片是 224x224,每個 patch 大小是 16x16,則 n 是 196),而一維向量長度等價于詞向量編碼長度(假設圖片通道是 3, 則每個序列的向量長度是 768) 。
- 考慮到一維向量維度較大,需要將拉伸后的 Patch 序列經過線性投影 (nn.Linear) 壓縮維度,同時也起到特征變換功能,這兩個步驟可以稱為圖片 Token 化過程 (Patch Embedding) 。
- 為了方便后續分類,作者還額外引入一個可學習的 Class Token,該 Token 插入到圖片 token 化后所得序列的開始位置 。
- 將上述序列加上可學習的位置編碼輸入到 N 個串行的 Transformer 編碼器中進行全局注意力計算和特征提取,其中內部的多頭自注意模塊用于進行 Patch 間或者序列間特征提取,而后面的 Feed Forward(Linear+ GELU+Dropout+ Linear+ Dropout) 模塊對每個 Patch 或者序列進行特征變換。
- 將最后一個 Transformer 編碼器輸出序列的第 0 位置( Class Token 位置對應輸出)提取出來,然后后面接 MLP 分類后,然后正常分類即可。
可以看出圖片分類無需 Transformer 解碼器,且編碼器幾乎沒有做任何改動,針對圖像分類任務,只需單獨引入一個 Image to Token 操作和 Class Token 的概念即可。
如何理解 Token? 個人覺得任何包括圖片更加高級的語義向量都可以叫做 Token,這個概念在 NLP 中應用非常廣泛,表征離散化后的高級單詞語義,在圖像中則可以認為是將圖像轉化為離散的含更高級語義的向量。
ViT 證明純 Transformer 也可以取得非常好的效果,相比 CNN 在數據量越大的情況下優勢更加明顯,但是 ViT 也存在如下問題:
- 不采用超大的 JFT-300M 數據集進行預訓練,則效果無法和 CNN 媲美,原因應該是 Transformer 天然的全局注意力計算,沒有 CNN 這種 Inductive Bias 能力,需要大數據才能發揮其最大潛力。
- ViT 無法直接適用于不同尺寸圖片輸入,因為 Patch 大小是固定的,當圖片大小改變,此時序列長度就會改變,位置編碼就無法直接適用了,ViT 解決辦法是通過插值,這種做法一般會造成性能損失,需要通過 Finetune 模型來解決,有點麻煩。
- 因為其直筒輸出結構,無法直接應用于下游密集任務。
后面的文章對上述缺點采用了各種各樣的改進,并提出了越來越先進的處理手段,推動了視覺 Transformer 的巨大進步。
3 全局概述
由于內容非常多,為了更容易理解,我在拆解模塊的基礎上對每個模塊進行分析,而不是對某篇文章進行概括。綜述部分的分析流程按照結構圖順序描述,我將近期圖像分類視覺 Transformer 發展按照 ViT 中是否包括自注意層模塊來劃分,包括:
在上述三個方向中,Attention-based 是目前改進最多,最熱門的,也是本綜述的核心。本文按照 3 個分類順序依次分析,最后進行通用架構分析。通過 General architecture analysis 部分可以深化Attention-based、MLP-based 和 ConvMixer-based 三者的聯系和區別。本文僅僅涉及 Attention-based 部分。
3.1 Attention-based
Attention-based 表示這類算法必然包括注意力模塊,我們將按照廣度優先順序進行一次分析。
繼 ViT 后,我們將其發展分成兩條線路:訓練策略和模型改進。
其中訓練策略表示目前主流對 ViT 模型的訓練改進方式,而模型改進則是對各個部件進行改進。
- 訓練策略包括兩篇論文:DeiT 和 Token Labeling。兩者提出的出發點一致,都是為了克服 ViT 需要 JFT-300M 大數據集進行預訓練的缺點。DeiT 是通過引入蒸餾學習解決,而 Token Labeling 通過引入顯著圖然后施加密集監督解決。后續發展中大部分算法都是參考了 DeiT 的訓練策略和超參設置,具有非常大的參考價值。
- 模型改進方面,我將其分成了 6 個組件以及其他方面的改進,6 個組件包括:
下面按照訓練策略和模型改進順序分析。
3.1.1 訓練策略
訓練策略解決 ViT 需要大數據先預訓練問題以及超參有待優化問題。
3.1.1.1 DeiT
如果說 ViT 開創了 Transformer 在視覺任務上面的先河,那么?DeiT?的出現則解決了 ViT 中最重要的問題:如果不采用超大的 JFT-300M 數據集進行預訓練,則效果無法和 CNN 媲美。在單個節點 8 張 V100 且無需額外數據的情況下,用不到 3 天的時間訓練所提的 ViT(86M 參數),在 ImageNet 上單尺度測試達到了 83.1% 的 top-1 準確率。
DeiT 核心是引入蒸餾手段加上更強的 Aug 和更優異的超參設置。其蒸餾的核心做法如下所示:
ViT 的 Class Token 是加到圖片輸入序列的前面,那么蒸餾 Token 可以插到輸入序列的后面,當然插入到哪個位置其實無所謂,你也可以插入到 Class Token 后面,經過 Transformer 編碼器輸出的序列相比 ViT 也會多一個,然后額外的一個輸出 Token 經過線性層輸出相同類別通道,最后進行蒸餾學習。
對于蒸餾學習來說,做法通常有兩個:
- Soft 蒸餾,即學生模型和教師模型預測的 Softmax 概率分布值計算 KL Loss。
- Hard 蒸餾,即教師模型預測的 Softmax 概率分布值中,值最大對應的類別作為標簽,然后和學生模型預測的 Softmax 概率分布值計算 CE Loss。
蒸餾學習中,通常教師模型會選擇一個比學生模型性能更強的且已經提前訓練好的模型,教師模型不需要訓練,通過蒸餾 loss 將教師模型知識以一種歸納偏置的方式轉移給學生模型,從而達到提升學生模型性能的目的。因為引入了額外的蒸餾 Token,而且該 Token 訓練任務也是分類,所以實際上 DeiT 在推理時,是將 Class Token 和 Distillation Token 的預測向量求平均,再轉換為概率分布。
為了證明 Distillation Token 的有效性,而不是只由于多了一個 Token 或者說多了一個可學習參數導致的,作者還做了對比試驗,不加 Distillation Token,而是再加一個 Class Token,相當于有兩個分類頭,兩個 Token 獨立且隨機初始化,實驗發現他們最終收斂后兩個分類 Token 的相似度達到 0.999,并且性能更弱,這樣證明了加入 Distillation Token 的意義。
通過大量實驗,作者總結了如下結論:
- 蒸餾做法確實有效,且 Hard 蒸餾方式效果會更好,泛化性能也不錯。
- 使用 RegNet 作為教師網絡可以取得更好的性能表現,也就是說相比 Transformer,采用卷積類型的教師網絡效果會更好。
除了上述蒸餾策略,還需要特別注意 DeiT 引入了非常多的 Aug 并且提供了一套更加優異的超參,這套參數也是后續大部分分類模型直接使用的訓練參數,非常值得學習,如下所示:
總而言之 DeiT 算法非常優異,實驗也非常多(建議去閱讀下),最大貢獻是通過蒸餾策略省掉了 ViT 需要 JFT-300M 數據集進行預訓練這個步驟,并且提供了一套非常魯棒且實用的超參配置,深深地影響了后續的大部分圖像分類視覺 Transformer 模型。
3.1.1.2 Token Labeling
DeiT 不是唯一一個解決 ViT 需要大數據量問題的算法,典型的還有?Token Labeling,其在 ViT 的 Class Token 監督學習基礎上,還對編碼器輸出的每個序列進行額外監督,相當于將圖片分類任務轉化為了多個輸出 Token 識別問題,并為每個輸入 Patch 的預測 Token 分配由算法自動生成的基于特定位置的監督信號,簡要圖如下所示:
從上圖明顯可以看出,相比 ViT 額外多了輸出 Token 的監督過程,這些監督可以當做中間監督,監督信息是通過 EfficientNet 或者 NFNet ( F6 86.3% Top-1 accuracy) 這類高性能網絡對訓練圖片提前生成的顯著圖,每個顯著圖維度是和類別一樣長的 C 維,輔助 Loss 和分類一樣也是 CE Loss。當然最終實驗結果表明性能比 DeiT 更優異,而且由于這種密集監督任務,對于下游密集預測任務泛化性也更好。
在此基礎上 DeiT 已經證明通過對 ViT 引入更多的強 Aug 可以提升性能,例如引入 CutMix,但是本文的做法無法直接簡單增加 CutMix,為此作者還專門設計了一個 MixToken,大概做法是在 Pathc Embedding 后,對 Token 進行了相應的 CutMix 操作。性能表如下所示:
LV-ViT 即為本文所提模型。相比 DeiT,作者認為本文做法更加優異,體現在:
- 不需要額外的教師模型,是一個更加廉價的做法。
- 相比于單向量監督,以密集的形式監督可以幫助訓練模型輕松發現目標物體,提高識別準確率,實驗也證明了對下游密集預測任務(例如語義分割)更友好。
下表是對訓練技術的詳細分析:
簡而言之,Token Labeling 的核心做法是通過引入額外的顯著圖來監督每個 patch 輸出的預測 token,雖然不需要教師模型,但是依然需要利用更優異的模型對所有訓練圖片生成顯著圖。
3.1.2 模型改進
在 DeiT 提出后,后續基于它提出了大量的改進模型,涉及到 ViT 的方方面面。前面說過 ViT 模型主要涉及到的模塊包括:
3.1.2.1 Token 模塊
Token 模塊包括兩個部分:
下面是完整結構圖:
3.1.2.1.1 Image to Token 模塊
首先需要明確:Patch Embedding 通常包括圖片窗口切分和線性嵌入兩個模塊,本小結主要是說圖片窗口切分方式,而具體實現不重要,常用的 2 種實現包括 nn.Conv 和 nn.Unfold,只要設置其 kernel 和 stride 值相同,則為非重疊 Patch Embedding,如果 stride 小于 kernel 則為重疊 Patch Embedding。
(1) 非重疊 Patch Embedding
ViT 和目前主流模型例如 PVT 和 Swin Transformer 等都是采用了非重疊 Patch Embedding,其簡要代碼為:
通過設置 16x16 的 kernel 和 stride 可以將圖片在空間維度進行切割,每個 patch 長度為 16x16x3,然后將每個 Patch 重排拉伸為一維向量后,經過線性層維度變換,輸出 shape 為 (B, Num_Seq, Dim)。
在?TNT?中作者提出了一種更精細的非重疊 Patch Embedding 模塊,如下圖所示:
他的基本觀點是自然圖像的復雜度相較于自然文本更高,細節和顏色信息更豐富,而 ViT 的非重疊 Patch Embedding 做法過于粗糙,因為后續自注意力計算都是在不同 Patch 間,這會導致 Patch 內部的局部自注意力信息沒有充分提取,而這些信息在圖像中也是包含了不同的尺度和位置的物體特征,是非常關鍵的。故我們不僅要考慮 Patch 間自注意力,還要考慮 Patch 內自注意力,為此作者在 外層 Transformer 中又嵌入了一個內層 Transformer,相應的非重疊 Patch Embedding 也分成兩步:整圖的非重疊 Patch Embedding 過程和 Patch 內部更細粒度的非重疊 Patch Embedding 過程。
通過上圖大概可以看出其具體做法,內部相當于有兩個 Transformer,第一個 Transformer (Outer Transformer )和 ViT 完全一樣,處理句子 Sentences 信息即圖片 Patch 級別信息,第二個 Transformer (Inner Transformer,也需要額外加上 Inner Transformer 所需要的位置編碼) 處理更細粒度的 Words 信息即圖片 Patch 內再切分為 Patch,為了能夠將兩個分支信息融合,內部會將 Inner Transformer 信息和 Outer Transformer 相加。將上述 Transformer block 嵌入到 PVT 模型中驗證了其對下游任務的適用性,通過進一步的可視化分析側面驗證了分類任務上 TNT 相比 DeiT 的優異性。
(2) 重疊 Patch Embedding
在常規的 CNN 網絡中一般都是采用重疊計算方式,為此是否采用重疊 Patch Embedding 會得到更好的性能?
直接將非重疊 Patch Embedding 通過修改 Unfold 或者 Conv 參數來實現重疊 Patch Embedding 功能的典型算法包括?T2T-ViT?和 PVTv2,這兩個算法的出發點都是非重疊 Patch Embedding 可以加強圖片 Patch 之間的連續性,不至于出現信息斷層,性能應該會比重疊 Patch Embedding 高。PVTv2 內部采用 Conv 實現,而 T2T ViT 是通過 Unfold 方式實現(論文中稱為 soft split)。
前面說過 CNN 網絡中一般都是采用重疊計算方式,那么是否可以用 ResNet Stem 替換非重疊 Patch Embedding過程,性能是否會更好?
在?Early Convolutions Help Transformers See Better?論文中作者進行了深度分析,雖然作者只是簡單的將圖片 Token 化的 Patch Embedding 替換為 ResNet Conv Stem,但是作者是從優化穩定性角度入手,通過大量的實驗驗證上述做法的有效性。作者指出 Patch Embedding 之所以不穩定,是因為該模塊是用一個大型卷積核以及步長等于卷積核的卷積層來實現的,往往這個卷積核大小為 16*16,這樣的卷積核參數量很大,而且隨機性很高,從某種程度上造成了 Transformer 的不穩定,如果用多個小的卷積來代替則可以有效緩解。結構如下所示:
考慮了和 ViT 公平對比,新引入的 Conv Stem 計算量約等于一個 transformer block,故后續僅僅需要 L-1 個 transformer block。作者通過大量分析實驗得到一些經驗性看法:
(a) ViT 這類算法對 lr 和 wd 超參的選擇非常敏感,而替換 Stem 后會魯棒很多。
(b) ViT 這類算法收斂比較慢,而本算法會快很多,例如都在 100 epoch 處本文性能遠優于 ViT。
ViT_p 即為 Patch Embedding 模式,ViT_c 即為 Conv Stem 模式,可以看出在不同 flops 下模型收斂速度都是 ViT_c 快于 ViT_p,雖然到 400 epoch 時候性能都非常接近。
(c) ViT 這類算法只能采用 AdamW 訓練,而本文更加通用,采用 SGD 后性能沒有顯著下降。
眾所周知,ViT 類模型都只能用 AdamW 訓練,其占據顯存是 SGD 的 3 倍,所以一般在 CNN 網絡中都是用過 SGD 模型,性能通常不錯,而通過替換 Patch Embedding 后也可以用 SGD 訓練了
(d) 僅僅采用 ImageNet 訓練 ViT 性能難以超越 CNN,而本文可以進一步提升 ViT 性能。
與上述論文持相同觀點的也包括?ResT?、Token Learner、CSWin Transformer 等算法,他們都采用了完全相同的做法。 更進一步在?PS-ViT?中為了能夠方便后續的漸進采樣模塊穩定提取更好的特征點,作者在 Image to Token 模塊中不僅僅引入了 ResNet 的 Conv Stem 模塊,還在后面再使用了 ResNet 第一個 stage 的前兩個殘差 block,在 Token to Token 模塊中會詳細說明 PS-ViT。
在?CeiT?中作者出發點是 CNN 中的諸多特性已經被證明是很成功的,純粹的 Transformer 需要大量的數據、額外的監督才能達到和 CNN 相同的精度,出現這種問題的原因可能是 NLP 中的 Transformer 直接搬到圖像任務中可能不是最合適的,應該考慮部分引入 CNN 來增強 Transformer。具體來說,在圖片轉 Token 方案中提出 Image-to-Tokens (I2T) 模塊,不再是從圖片中直接進行 Patch Emeding ,而是對 CNN 和 Pool 層所提取的底層特征進行 Patch Embedding,借助圖像特征會比直接使用圖片像素更好。
上圖的上半部分是 ViT 的 Patch Embedding 過程,下圖是 CeiT 所提出的做法,核心就是引入卷積操作提取底層特征,然后在底層特征上進行 Patch Embedding 操作。
既然采用 Conv Stem 可以解決很多問題,那么理論上經過精心設計的 Conv 結構也必然是有效的,例如?ViTAE?中就采用了空洞卷積做法,本質上是希望能夠利用卷積提供多尺度上下文信息,這有助于后續模塊信息提取,如下圖所示:
對圖片或者特征圖應用多個不同空洞率的卷積提取信息后,進行拼接和 GeLU 激活函數后,直接拉伸為一維向量,從而轉換為序列,并且由于空洞卷積可以實現下采樣功能,故也可以有效地減少后續注意力模塊計算量。
3.1.2.1.2 Token to Token 模塊
大部分模型的 Token to Token 方案和 Image to Token 做法相同,但是也有些算法進行了相應改造。經過整理,我們將其分成兩種做法:
固定窗口是指 Token 化過程是固定或者預定義的規則,典型的重疊和非重疊 Patch Embedding 就是固定窗口,因為其窗口劃分都是提前訂好的規則,不會隨著輸入圖片的不同而不同,而動態窗口是指窗口劃分和輸入圖片語義相關,不同圖片不一樣,是一個動態過程。
(1) 固定窗口 Token 化
這個做法通常和 Image to Token 模塊完全一樣,也可以分成非重疊 Patch Embedding 和重疊 Patch Embedding,大部分算法都屬于這一類,例如 PVT、Swin Transformer 等。
(2) 動態窗口 Token 化
動態窗口 Token 化過程典型代表是 PS-ViT 和 TokenLearner。
前面說過,Vision Transformer with Progressive Sampling (PS-ViT)?中為了方便后續的漸進采樣模塊能夠穩定提取更好的特征點,在 Image to Token 模塊中不僅僅引入了 ResNet 的 Conv Stem 模塊,還在后面再使用了 ResNet 第一個 stage 的前兩個殘差 block。在特征圖 F 后,作者在 Token to Token 環節引入了一個漸進式采樣模塊,其出發點是?ViT 采用固定窗口劃分機制,然后對每個窗口進行 Token 化,這種做法首先不夠靈活,而且因為圖片本身就是密集像素,冗余度非常高,采用固定劃分方法對于分類來說可能就某幾個窗口內的 Token 實際上才是有意義的,假設物體居中那么物體四周的 Token 可能是沒有作用的,只會增加無效計算而已。基于此作者設計一個自適應采樣的 Token 機制,不再是固定的窗口采樣,而是先初始化固定采樣點,如下圖紅色點所示,然后通過 refine 機制不斷調整這些采樣點位置,最終得到的采樣點所對應的 Token 就是最有代表力的。其完整分類網絡結構圖如下所示:
得到特征圖 F 后,經過漸進采樣模塊,不斷 refine 采樣點,最終輸出和采樣點個數個序列,將該序列作為 ViT 模型的輸入即可。簡單來看漸進采樣模塊起到了 Token to Token 作用。其中的漸進采樣模塊結構圖如下所示:
詳細計算過程如下:
可以發現,和 ViT 的主要差異就在于其采樣點不是固定的均勻間隔,而是基于語義圖自適應,從而能夠在減少計算量的前提下進一步提升性能。PS-ViT 在 top-1 精度方面比普通 ViT 高 3.8%,參數減少約 4 倍,FLOP 減少約 10 倍,性能比較優異。
基于類似出發點,TokenLearner?提出可以基于空間注意力自適應地學習出更具有代表性的 token,從而可以將 ViT 的 1024 個 token 縮減到 8-16 個 token,計算量減少了一倍,性能依然可以保持一致,這也側面說明了 ViT 所采樣的固定窗口 token 有大量冗余,徒增計算量而已。其核心示意圖如下所示:
假設想僅僅采樣出 8 個 token,首先采用 Conv Stem 提取圖片特征,然后分別輸入到 8 個空間注意力分支中,空間注意力分支首先會應用一系列卷積生成空間 attention 圖,然后逐點和輸入特征相乘進行特征加權,最后通過空間全局 pool 層生成 1x1xC 的 Token,這樣就將 HXWXC 的特征圖轉換為了 8 個通道為 C 的 Token。
為了進一步提高信息,作者還額外提出一個 TokenFuser 模塊,加強 Token 和 Token 之間的聯系以及恢復空間結構,整個分類網絡的結構如下所示:(a) 為不包括 TokenFuser 的改進 ViT 結構,(b) 為包括 TokenFuser 的改進 ViT 結構。
從上述結構可以發現 TokenLearner 模塊起到了自適應提取更具語義信息的 Token,并且還能夠極大地減少計算量,而 TokenFuser 可以起到加強 Token 和 Token 之間的聯系以及恢復空間結構的功能,TokenLearner+Transformer+ TokenFuser 構成 Bottleneck 結構。其中 TokenFuser 示意圖如下所示:
其接收兩個輸入,一個是 TokenLearner 前的保持了空間信息的 1024 個 token 特征,一個是 TokenLearner 后經過自適應采樣的 8 個 token 特征,然后以注意力模式兩者進行乘加操作,融合特征以及恢復空間結構。
作者的分類實驗依然采用了 JFT-300M 數據集進行預訓練,然后在 ImageNet 1k上面微調,也就是說和最原始的 ViT 進行比較。
TokenFuser 也進行了相應的對比實驗。
at 6 表示 TokenLearner 插入到第 6 個 Transformer Encoder 后。
3.1.2.2 位置編碼模塊
位置編碼模塊是為 Transformer 模塊提供 Patch 和 Patch 之間的相對關系,非常關鍵。在通用任務的 Transformer 模型中認為一個好的位置編碼應該要滿足以下特性:
- 不同位置的位置編碼向量應該是唯一的
- 不能因為不同位置位置編碼的值大小導致網絡學習有傾向性
- 必須是確定性的
- 最好能夠泛化到任意長度的序列輸入
ViT 位置編碼模塊滿足前 3 條特性,但是最后一條不滿足,當輸入圖片改變時候需要微調,比較麻煩。基于此也出現了不少的算法改進,結構圖如下所示:
按照是否顯式的設置位置編碼向量,可以分成:
隱式位置編碼對于圖片長度改變場景更加有效,因為其是自適應圖片語義而生成。
3.1.2.2.1 顯式位置編碼
顯式位置編碼,可以分成絕對位置編碼和相對位置編碼。
(1) 絕對位置編碼
絕對位置編碼表示在 Patch 的每個位置都加上一個不同的編碼向量,其又可以分成固定位置編碼即無需學習直接基于特定規則生成,常用的是 Attention is all you need 中采用的 sincos 編碼,這種編碼方式可以支持任意長度序列輸入。還有一種是可學習絕對位置編碼,即初始化設置為全 0 可學習參數,然后加到序列上一起通過網絡訓練學習,典型的例如 ViT、PVT 等等。
(2) 相對位置編碼
相對位置編碼考慮為相鄰 Patch 位置編碼,其實現一般是設置為可學習,例如 Swin Transformer 中采用的可學習相對位置編碼,其做法是在 QK 矩陣計完相似度后,引入一個額外的可學習 bias 矩陣,其公式為:
Swin Transformer 這種做法依然無法解決圖片尺寸改變時候對相對位置編碼插值帶來的性能下降的問題。在 Swin Transformer v2 中作者做了相關實驗,在直接使用了在 256 * 256 分辨率大小,8 * 8 windows 大小下訓練好的 Swin-Transformer 模型權重,載入到不同尺度的大模型下,在不同數據集上進行了測試,性能如下所示 (Parameterized position bias 這行):
每個表格中的兩列表示沒有 fintune 和有 fintune,可以看出如果直接對相對位置編碼插值而不進行 fintune,性能下降比較嚴重。故在 Swin Transformer v2 中引入了對數空間連續相對位置編碼 log-spaced continuous position bias,其主要目的是可以更有效地從低分辨權重遷移到高分辨率下游任務。
相比于直接應用可學習的相對位置編碼,v2 中先引入了 Continuous relative position bias (CPB),
B 矩陣來自一個小型的網絡,用來預測相對位置,該模塊的輸入依然是 Patch 間的相對位置,這個小型網絡可以是一個 2 層 MLP,然后接中間采用激活函數連接。
其性能如上表的 Linear-Spaced CPB 所示,可以發現相比原先的相對位置編碼性能有所提升,但是當模型尺度繼續增加,圖片尺寸繼續擴大后性能依然會下降比較多,原因是預測目標空間是一個線性的空間。當 Windows 尺寸增大的時候,比如當載入的是 8*8 大小 windows 下訓練好的模型權重,要在 16*16 大小的 windows 下進行 fine-tune,此時預測相對位置范圍就會從 [-7,7] 增大到 [-15,15],整個預測范圍的擴大了不少,這可能會出現網絡不適應性。為此作者將預測的相對位置坐標從 linear space 改進到 log space 下,這樣擴大范圍就縮小了不少, 可以提供更加平滑的預測范圍,這會增加穩定性,提升泛化能力,性能表如上的 Log-Spaced CPB 所示。
在 Swin Transformer 中相對位置編碼矩陣 shape 和 QK矩陣計算后的矩陣一樣大,其計算復雜度是 O(HW),當圖片很大或者再引入 T 時間軸,那么計算量非常大, 故在?Improved MViT?,作者進行了分解設計,分成 H 軸相對位置編碼,W 軸相對位置編碼,然后相加,從而將復雜度降低為 O(H+W)。
關于絕對位置編碼和相對位置編碼到底哪個是最好的,目前還沒有定論,在不同的論文實驗中有不同的結論,暫時來看難分勝負。但是從上面可以分析來看,在 ViT 中不管是絕對位置編碼和相對位置編碼,當圖片大小改變時候都需要對編碼向量進行插值,性能都有不同程度的下降 ( Swin Transformer v2 在一定程度上解決了)。
3.1.2.2.2 隱式位置編碼
當圖片尺寸改變時候,隱式位置編碼可以很好地避免顯式位置編碼需要對對編碼向量進行插值的弊端。其核心做法都是基于圖片語義自適應生成位置編碼。
在論文 How much position information do convolutional neural networks encode? 中已經證明 CNN 不僅可以編碼位置信息,而且越深的層所包含的位置信息越多,而位置信息是通過 zero-padding 透露的。既然 Conv 自帶位置信息,那么可以利用這個特性來隱式的編碼位置向量。大部分算法都直接借鑒了這一結論來增強位置編碼,典型代表有 CPVT、PVTv2 和 CSwin Transformer 等。
CPVT?指出基于之前 CNN 分類經驗,分類網絡通常需要平移不變性,但是絕對位置編碼會在一定程度打破這個特性,因為每個位置都會加上一個獨一無二的位置編碼。看起來似乎相對位置編碼可以避免這個問題,其天然就可以適應不同長度輸入,但是由于相對位置編碼在圖像分類任務中無法提供任何絕對位置信息(實際上相對位置編碼也需要插值),而絕對位置信息被證明非常重要。以 DeiT-Tiny 模型為例,作者通過簡單的對比實驗讓用戶直觀的感受不同位置編碼的效果:
2D PRE 是指 2D 相對位置編碼,Top-1@224 表示測試時候采用 224 圖片大小,這個尺度和訓練保持一致,Top-1@384 表示測試時候采用 384 圖片大小,由于圖片大小不一致,故需要對位置編碼進行插值。從上表可以得出:
- 位置編碼還是很重要,不使用位置編碼性能很差
- 2D 相對位置編碼性能比其他兩個差,可學習位置編碼和 sin-cos 策略效果非常接近,相對來說可學習絕對位置編碼效果更好一些(和其他論文結論不一致)
- 在需要對位置編碼進行插值時候,性能都有下降
基于上述描述,作者認為在視覺任務中一個好的位置編碼應滿足如下條件:
- 模型應該具有 permutation-variant 和 translation-equivariance 特性,即對位置敏感但同時具有平移不變性
- 能夠自然地處理變長的圖片序列
- 能夠一定程度上編碼絕對位置信息
基于這三個原則,CPVT 引入了一個帶有 zero-padding 的卷積 ( kernel size k ≥ 3) 來隱式地編碼位置信息,并提出了 Positional Encoding Generator (PEG) 模塊,如下所示:
將輸入序列 reshape 成圖像空間維度,然后通過一個 kernel size 為 k ≥ 3, (k?1)/2 zero paddings 的 2D 卷積操作,最后再 reshape 成 token 序列。這個 PEG 模塊因為引入了卷積,在計算位置編碼時候會考慮鄰近的 token,當圖片尺度改變時候,這個特性可以避免性能下降問題。算法的整體結構圖如下所示:
基于 CPVT 的做法,PVTv2?將 zero-padding 卷積思想引入到 FFN 模塊中。
通過在常規 FFN 模塊中引入 zero-padding 的逐深度卷積來引入隱式的位置編碼信息(稱為 Convolutional Feed-Forward)。
同樣的,在?CSWin Transformer?中作者也引入了 3x3 DW 卷積來增強位置信息,結構圖如下所示:
APE 是 ViT 中的絕對位置編碼,CPE 是 CPVT 中的條件位置編碼,其做法是和輸入序列 X 相加,而 RPE 是 Swin Transformer 中所采用的相對位置編碼,其是加到 QK 矩陣計算后輸出中,而本文所提的 Locally-Enhanced Positional Encoding (LePE),是在自注意力計算完成后額外加上 DW 卷積值,計算量比 RPE 小。
LePE 做法對于下游密集預測任務中圖片尺寸變化情況比較友好,性能下降比較少。
除了上述分析的諸多加法隱式位置編碼改進,?ResT?提出了另一個非常相似的,但是是乘法的改進策略,結構圖如下所示:
對 Patch Embedding 后的序列應用先恢復空間結構,然后應用一個 3×3 depth-wise padding 1的卷積來提供位置注意力信息,然后通過 sigmoid 操作變成注意力權重和原始輸入相乘。代碼如下所示:
作者還指出這個Pixel Attention (PA) 模塊可以替換為任意空間注意力模塊,性能優異,明顯比位置編碼更加靈活好用。
3.1.2.3 自注意力模塊
Transformer 的最核心模塊是自注意力模塊,也就是我們常說的多頭注意力模塊,如下圖所示:
注意力機制的最大優勢是沒有任何先驗偏置,只要輸入足夠的數據就可以利用全局注意力學到泛化性能不錯的特征。當數據量足夠大的時候,注意力機制是 Transformer 模型的最大優勢,但是一旦數據量不夠就會變成逆勢,后續很多算法改進方向都是希望能夠引入部分先驗偏置輔助模塊,在減少對數據量的依賴情況下加快收斂,并進一步提升性能。同時注意力機制還有一個比較大的缺點:因為其全局注意力計算,當輸入高分辨率圖時候計算量非常巨大,這也是目前一大改進方向。
簡單總結,可以將目前自注意力模塊分成 2 個大方向:
如果整個 Transformer 模型不含局部注意力模塊,那么其主要改進方向就是如何減少空間全局注意力的計算量,而引入額外的局部注意力自然可以很好地解決空間全局注意力計算量過大的問題,但是如果僅僅包括局部注意力,則會導致性能下降嚴重,因為局部注意力沒有考慮窗口間的信息交互,因此引入額外的局部注意力的意思是在引入局部注意力基礎上,還需要存在窗口間交互模塊,這個模塊可以是全局注意力模塊,也可以是任何可以實現這個功能的模塊。其結構圖如下所示:
3.1.2.3.1 僅包括全局注意力
標準的多頭注意力就是典型的空間全局注意力模塊,當輸入圖片比較大的時候,會導致序列個數非常多,此時注意力計算就會消耗大量計算量和顯存。以常規的 COCO 目標檢測下游任務為例,輸入圖片大小一般是 800x1333,此時 Transformer 中的自注意力模塊計算量和內存占用會難以承受。其改進方向可以歸納為兩類:減少全局注意力計算量以及采用廣義線性注意力計算方式。
(1) 減少全局注意力計算量
全局注意力計算量主要是在 QK 矩陣和 Softmax 后和 V 相乘部分,想減少這部分計算量,那自然可以采用如下策略:
(a) 降低 KV 維度
降低 KV 維度做法的典型代碼是 PVT,其設計了空間 Reduction 注意力層 (SRA) ,如下所示:
其做法比較簡單,核心就是通過 Spatial Reduction 模塊縮減 KV 的輸入序列長度,KV 是空間圖片轉化為 Token 后的序列,可以考慮先還原出空間結構,然后通過卷積縮減維度,再次轉化為序列結構,最后再算注意力。假設 QKV shape 是完全相同,其詳細計算過程如下:
- 在暫時不考慮 batch 的情況下,KV 的 shape 是 (H'W', C),既然叫做空間維度縮減,那么肯定是作用在空間維度上,故首先利用 reshape 函數恢復空間維度變成 (H', W', C)。
- 然后在這個 shape 下應用 kernel_size 和 stride 為指定縮放率例如 8 的二維卷積,實現空間維度縮減,變成 (H/R, W/R, C), R 是縮放倍數。
- 然后再次反向 reshape 變成 (HW/(R平方), C),此時第一維(序列長度)就縮減了 R 平方倍數。
- 然后采用標準的多頭注意力層進行注意力加權計算,輸出維度依然是 (H'W', C)。
而在?Twins?中提出了所謂的 GSA,其實就是 PVT 中的空間縮減模塊。
同時基于最新進展,在 PVTV2 算法中甚至可以將 PVTv1 的 Spatial Attention 直接換成無任何學習參數的 Average Pooling 模塊,也就是所謂的 Linear SRA,如下所示:
同樣參考 PVT 設計,在?P2T?也提出一種改進版本的金字塔 Pool 結構,如下所示:
(b) 即為改進的 Spatial Attention 結構,對 KV 值應用不同大小的 kernel 進行池化操作,最后 cat 拼接到一起,輸入到 MHSA 中進行計算,通過控制 pool 的 kernel 就可以改變 KV 的輸出序列長度,從而減少計算量,同時金字塔池化結構可以進一步提升性能(不過由于其 FFN 中引入了 DW 卷積,也有一定性能提升)。
從降低 KV 空間維度角度出發,ResT 算法中也提出了一個內存高效的注意力模塊 E-MSA,相比 PVT 做法更近一步,不僅僅縮減 KV 空間維度,還同時加強各個 head 之間的信息交互,如下所示:
其出發點有兩個:
- 當序列比較長或者維度比較高的時候,全局注意力計算量過大。
- 當多頭注意力計算時,各個頭是按照 D 維度切分,然后獨立計算最后拼接輸出,各個頭之間沒有交互,當 X 維度較少時,性能可能不太行。
基于上述兩點,作者引入 DWConv 縮放 KV 的空間維度來減少全局注意力計算量,然后在 QK 點乘后引入 1x1 Conv 模塊進行多頭信息交互。其詳細做法如下:
其核心代碼如下所示:
(b) 降低 QKV 維度
Multiscale Vision Transformers (MViT) 也考慮引入 Pool 算子來減少全局注意力計算量。MViT 主要為了視頻識別而設計,因為視頻任務本身就會消耗太多的顯存和內存,如果不進行針對性設計,則難以實際應用。正如其名字所言,其主要是想參考 CNN 中的多尺度特性(淺層空間分辨率大通道數少,深層空間分辨率小通道數多)設計出適合 Transformer 的多尺度 ViT。自注意力模塊如下所示:
相比 Transformer 自注意力模塊,其主要改變是多了 Pool 模塊,該模塊的主要作用是通過控制其 Stride 參數來縮放輸入的序列個數,而序列個數對應的是圖片空間尺度 THW。以圖像分類任務為例,
- 任意維度的序列 X 輸入,首先和 3 個獨立的線性層得到 QKV,維度都是 (HW, D)。
- 將QKV 恢復空間維度,變成 (H, W, D),然后經過獨立的 3 個 Pool 模塊,通過控制 stride 參數可以改變輸出序列長度,變成 (H', W', D),設置 3 個 Pool 模塊不同的 Stride 值可以實現不同大小的輸出。
- 將輸入都拉伸為序列格式,然后采用自注意力計算方式輸出 (H'W‘, D)。
- 為了保證輸出序列長度改變而無法直接應用殘差連接,需要在 X 側同時引入一個 Pool 模塊將序列長度和維度變成一致。
由于 MViT 出色的性能,作者將該思想推廣到更多的下游任務中(例如目標檢測),提出了改進版本的 Imporved MViT,其重新設計的結構圖如下所示:
Imporved MViT 在不同的下游任務提升顯著。
(2) 廣義線性注意力計算方式
基于 NLP 中 Transformer 進展,我們可以考慮將其引入到 ViT 中,典型的包括 Performer,其可以通過分解獲得一個線性時間注意力機制,并重新排列矩陣乘法,以對常規注意力機制的結果進行近似,而不需要顯示構建大小呈平方增長的注意力矩陣。在 T2T-ViT 算法中則直接使用了高效的 Performer。
在 NLP 領域類似的近似計算方式也有很多,由于本文主要關注 ViT 方面的改進,故這部分不再展開分析
3.1.2.3.2 引入額外局部注意力
引入額外局部注意力的典型代表是 Swin Transformer,但是卷積模塊工作方式也可以等價為局部注意力計算方式,所以從目前發展來看,主要可以分成 3 個大類:
- 局部窗口計算模式,例如 Swin Transformer 這種局部窗口內計算。
- 引入卷積局部歸納偏置增強,這種做法通常是引入或多或少的卷積來明確提供局部注意力功能。
- 稀疏注意力。
結構圖如下所示:
需要特別注意的是:
(1) 局部窗口計算模式
局部注意力的典型算法是?Swin Transformer,其將自注意力計算過程限制在每個提前劃分的窗口內部,稱為窗口注意力 Window based Self-Attention (W-MSA),相比全局計算自注意力,明顯可以減少計算量,但是這種做法沒法讓不同窗口進行交互,此時就退化成了 CNN,所以作者又提出移位窗口注意力模塊 Shifted window based Self-Attention (SW-MSA),示意圖如下所示,具體是將窗口進行右下移位,此時窗口數和窗口的空間切分方式就不一樣了,然后將 W-MSA 和 SW-MSA 在不同 stage 之間交替使用,即可實現窗口內局部注意力計算和跨窗口的局部注意力計算,同時其要求 stage 個數必須是偶數。
大概流程為:
上述只是原理概述,實際上為了保證上述操作非常高效,作者對代碼進行了非常多的優化,相對來說是比較復雜的。值得注意的是 Swin Transformer 相比其他算法(例如 PVT )非常高效,因為整個算法中始終不存在全局注意力計算模塊( SW-MSA 起到類似功能),當圖片分辨率非常高的時候,也依然是線性復雜度的,這是其突出優點。憑借其局部窗口注意力機制,刷新了很多下游任務的 SOTA,影響非常深遠。
在 Swin Transformer v2 中探討了模型尺度和輸入圖片增大時候,整個架構的適應性和性能。在大模型實驗中作者觀察到某些 block 或者 head 中的 attention map 會被某些特征主導,產生這個的原因是原始 self-attention 中對于兩兩特征之間的相似度衡量是用的內積,可能會出現某些特征 pair 內積過大。為了改善這個問題,作者將內積相似度替換為了余弦相似度,因為余弦函數本身的取值范圍本身就相當于是被歸一化后的結果,可以改善因為些特征 pair 內積過大,主導了 attention 的情況,結構圖如下所示:
Swin Transformer 算法在解決圖片尺度增加帶來的巨大計算量問題上有不錯的解決方案,但是 SW-MSA 這個結構被后續諸多文章吐槽,主要包括:
基于這三個問題,后續學者提出了大量的針對性改進,可以歸納為兩個方向:
(a) 拋棄 SW-MSA,依然需要全局注意力計算模塊
Imporved MViT?在改進的 Pool Attention 基礎上,參考 Swin Transformer 在不同 stage 間混合局部注意力 W-MSA 和 SW-MSA 設計,提出 HSwin 結構,在 4 個 stage 中的最后三個 stage 的最后一個 block 用全局注意力 Pool Attention 模塊(具體間 3.1.2.3.1 小節),其余 stage 的 block 使用 W-MSA ,實驗表明這種設計比 Swin Transformer 性能更強,也更簡單。
同樣 Twins 也借鑒了 W-MSA 做法,只不過由于其位置編碼是隱式生成,故不再需要相對位置編碼,而且 SW-MSA 這種 Shift 算子不好部署,所以作者的做法是在每個 Encoder 中分別嵌入 Locally-grouped self-attention (LSA) 模塊即不帶相對位置編碼的 W-MSA 以及 GSA 模塊,GSA 就是 PVT 中使用的帶空間縮減的全局自注意力模塊,通過 LSA 計算局部窗口內的注意力,然后通過全局自注意力模塊 GSA 計算窗口間的注意力,結構圖如下所示:
(b) 拋棄 SW-MSA,跨窗口信息交互由特定模塊提供
參考 CNN 網絡設計思想,可以設計跨窗口信息交互模塊,典型的論文包括 MSG-T 、Glance-and-Gaze Transformer 和 Shuffle Transformer。
MSG-Transformer?基于 W-MSA,通過引入一個 MSG Token 來加強局部窗口之間的信息交互即在每個窗口內額外引入一個 MSG Token,該 Token 主要作用就是進行窗口間信息傳遞,所設計的模塊優點包括對 CPU 設備友好,推理速度比 SWin 快,性能也更好一些。結構圖如下所示:
在第 3 步的 W-MSA 計算中,可以認為同一個窗口內會進行信息流通,從而 2x2 個 MSG Token 都已經融合了對應窗口內的信息,然后經過第 4 步驟 MSG Token 交換后就實現了局部窗口間信息的交互。MSG Token 信息交互模塊完成了 Swin Transformer 中 SW-MSA ,相比 SW-MSA 算子,不管是計算復雜度還是實現難度都非常小。Shuffle 計算過程和 ShuffleNet 做法完全一樣,如下所示:
將 Swin Transformer 中的 block 全部換成 MSG-Transformer block ,通過實驗驗證了本結構的優異性。
Shuffle Transformer?也是從效率角度對 Swin Transformer 的 SW-MSA 進行改進,正如其名字,其是通過 Shuffle 操作來加強窗口間信息交流,而不再需要 SW-MSA,由于其做法和 ShuffleNet 一致就不再詳細說明,核心思想如下所示 (c):
將 Swin Transformer 中的 SW-MSA 全部換成 Shuffle W-MSA,在此基礎上還引入了額外的 NWC 模塊,其是一個 DW Conv,其 kernel size 和 window size 一樣,用于增強鄰近窗口信息交互,Shuffle Transformer 的 block 結構如下所示:
在 ImageNet 數據集上,同等條件上 Shuffle Transformer相比 Swin 有明顯提升,在 COCO 數據集上,基于 Mask R-CNN,Shuffle Transformer 和 Swin 性能不相上下。
因為 Swin Transformer 不存在 NWC 模塊,作者也進行了相應的對比實驗:
這也進一步驗證了引入適當的 CNN 局部算子可以在幾乎不增加計算量的前提下顯著提升性能。
MSG-Transformer 和 Shuffle Transformer 都是通過直接替換 SW-MSA 模塊來實現的,Glance-and-Gaze Transformer (GG-Transformer)?則認為沒有必要分成兩個獨立的模塊,只需要通過同一個模塊的兩個分支融合就可以同時實現 W-MSA 和 SW-MSA 功能。結構圖如下所示:
其提出一種新的局部窗口注意力計算機制,相比常規的近鄰劃分窗口,其采用了自適應空洞率窗口劃分方式,上圖是假設空洞率是 2 即每隔 1 個位置,這樣就可以將圖片劃分為 4 個窗口,由于其采樣劃分方式會橫跨多個像素位置,相比 Swin Transofrmer 劃分方式具有更大的感受野,不斷 stage 堆疊就可以實現全局感受野。在 Glance 分支中采用 MSA 局部窗口計算方法計算局部注意力,同時為了增強窗口之間的交互,其將 V 值還原為原先劃分模式,然后應用 depth-wise conv 來提取局部信息,再通過自適應空洞劃分操作的逆操作還原,再加上 Attention 后的特征。
Glance 分支用于在劃分窗口內單獨計算窗口內的局部注意力,由于其自適應空洞率窗口劃分方式,使其能夠具備全局注意力提取能力,而 Gaze分支用于在劃分的窗口間進行信息融合,具備窗口間局部特征提取能力。將 Swin Transformer 中的 block 全部換成 GG-Transformer block ,通過實驗驗證了其性能優于 Swin Transformer 。
在改進 Swin Transformer 的窗口注意力計算方式這方面,CSWin Transformer?相比其余改進更加獨特,其提出了十字架形狀的局部窗口劃分方式,如下圖所示:
假設一共將圖片劃分成了 9 個窗口,本文所提注意力的計算只會涉及到上下左右中 5 個窗口,同時為了進一步減少計算量,又分成 horizontal stripes self-attention 和 vertical stripes self-attention,每個自注意力模塊都只涉及到其中 3 個窗口,這種計算方式同時兼顧了局部窗口計算和跨窗口計算,一步到位。所謂 horizontal stripes self-attention 是指沿著 H 維度將 Tokens 分成水平條狀 windows,假設一共包括 k 個頭,則前 k/2 個頭用于計算 horizontal stripes self-attention,后面 k/2 個頭用于計算 vertical stripes self-attention。
兩組self-attention是并行的,計算完成后將 Tokens 的特征 concat 在一起,這樣就構成了CSW self-attention,最終效果就是在十字形窗口內做 Attention,可以看到 CSW self-attention 的感受野要比常規的 Window Attention 的感受野更大。可以通過控制每個條紋的寬度來控制自注意力模塊的復雜度,默認 4 個 stage 的條紋寬度分別設為 1, 2, 7, 7(圖片空間維度比較大的時候采用較小的條紋寬度,減少計算量)。
(2) 引入卷積的局部歸納偏置能力
上述都是屬于 Swin Transformer 改進,在引入卷積局部歸納偏置增強方面,典型算法為?ViTAE?和?ELSA,ViTAE 結構圖如下所示:
其包括兩個核心模塊:reduction cell (RC) 和 normal cell (NC)。RC 用于對輸入圖像進行下采樣并將其嵌入到具有豐富多尺度上下文的 token 中,而 NC 旨在對 token 序列中的局部性和全局依賴性進行聯合建模,可以看到,這兩種類型的結構共享一個簡單的基本結構。
對于 RC 模塊,分成兩個分支,第一條分支首先將特征圖輸入到不同空洞率并行的卷積中,提取多尺度特征的同時也減少分辨率,輸出特征圖拼接+ GeLU 激活,然后輸入到注意力模塊中,第二條分支是純粹的 Conv 局部特征提取,用于加強局部歸納偏置,兩個分支內容相加,然后輸入到 FFN 模塊中。
對于 NC 模塊,類似分成兩個分支,第一條是注意力分支,第二條是 Conv 局部特征提取,用于加強局部歸納偏置,兩個分支內容相加,然后輸入到 FFN 模塊中。
基于上述模塊,構建了兩個典型網絡,如下所示:
至于為何要如此設置以及各個模塊的前后位置,作者進行了大量的實驗研究:
ELSA: Enhanced Local Self-Attention for Vision Transformer?基于一個現狀:Swin Transformer 中所提的局部自注意力(LSA)的性能與卷積不相上下,甚至不如動態過濾器。如果是這樣,那么 LSA 的重要性就值得懷疑了。
最近也有很多學者發現了個奇怪的現象,例如?Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight?中深入分析了注意力和 Conv 的關系,特別是 DW Conv,但是大部分都沒有深入探討 LSA 性能如此平庸的具體原因,本文從這個方面入手,并提出了針對性改進。
奇怪現象
作者以 Swin Tiny 版本為例,將其中的局部窗口注意力模塊 LSA 替換為 DW Conv、decoupled dynamic filter (DDF),從上圖可以看出 DWConv 和 DDF 性能都比 LSA 強的,特別是 DW Conv,在參數量和 FLOPs 更小的情況下性能會比 Swin Transformer 高。
原因分析
作者試圖從 2 個角度來統一分析 LSA、DWConv 和 DDF,分別是通道數設置和空間處理方式 spatial processing。
通道數設置 channel setting 可以用于統一建模 LSA、DWConv 在通道數上的差異。DWConv 是逐通道(深度)上計算 Conv,也就是說在不同的通道上應用了不同的濾波器參數,而 DDF 和 LSA 是將通道數分組,然后在同一個組內采用相同的濾波器參數,LSA 的通道數分組實際上就是常說的多頭機制即DwConv 可以視作一種特殊的多頭策略。如果將 DWConv 的通道數設置為 LSA 的頭個數,那么通道數設置原則上就是相同效果了。對比結果如下所示:
從上述曲線可以看出:
- 在相同通道配置下(例如1x、2x,這里是特指4 個 stage 中 head 個數,典型的 1x 是 3, 6, 12, 24),DwConv 版本仍與 LSA 版本具有相似性能。
- 從上述 1x, 2x 結果可以看出,通常來說 DWConv 的通道數肯定比 LSA 的 Head 個數多,這可能是 DWConv 性能比 LSA 高的原因,但是當設置為相同的 Head 個數時,DWConv 性能甚至比 LSA 類型更好一些。,這說明通道配置并非導致前述現象的主要原因
- 當 Head 個數配置大于 1x 時,LSA 性能反而下降,但是 DWConv 性能確實能夠提升,這說明兩個問題:(1)?LSA 中直接提升頭數并不能改善通道容量與性能,但是 DWConv 可以; (2) LSA 中需要一種新的策略以進一步提升通道容量和性能。
空間處理方式 spatial processing 是指如何得到濾波器或者注意力圖并對空域進行信息聚合。DwConv采用靜態的濾波器即一旦訓練完成,不管輸入啥圖片都是采用固定 kernel,而其他兩者則采用動態濾波器。為了方便統一建模分析,作者采用統一的表示式來說明,如下所示:
Conv 和 DwConv 計算公式為:
當僅僅使用 $$r_{j-i}^b$$相對位置偏置,Norm 設置為恒等變換,$$\Theta$$表示滑動窗口計算方式則可以表示 Conv 和 DwConv,j-i 表示相對偏移。
Dynamic filters 計算公式為:
W 是 1x1 卷積,當僅僅使用$$q_ir^k_{j-i}$$($$r_{j-i}^k$$表示 k 的相對位置嵌入向量), Norm 設置為恒等變換則可以表示 Dynamic filters,注意 Dynamic filters 有非常多種做法,以上寫的是其中一種。
LSA 計算公式為:
我們可以很容易地將公式 1 退化為 LSA 模式,$$\Omega$$表示非重疊局部窗口計算方式。自從可以將 DwConv 、Dynamic filters 和 LSA 統一起來,并且將其分成三個核心不同部分:參數形式 parameterization, 規范化方式 normalization 和濾波器應用方式 filter application。
參數形式 parameterization
從上表可以看出:
- 動態濾波器的參數策略要比標準 LSA 策略具有更優的性能(Net2 vs Net1)
- 動態濾波器變種策略 (Net6) 具有與 SwinT 相當的性能
- LSA 參數策略與動態濾波器參數策略的組合 (Net7) 可以進一步提升模型性能
規范化方式 Normalization
- 當采用 Net7 的參數形式組合 Identity 時,模型訓練崩潰
- 相比 FilterNorm,Softmax 規范化具有更優的性能
- 規范化方式并非 LSA 平庸的原因
濾波器應用方式 filter application
當將濾波器采用滑動窗口形式時,Net6 與 Net7 均得到了顯著性能提升?。這意味著:近鄰處理方式是空域處理的關鍵。
基于上述實驗,使 LSA 變平庸的因素可以分為兩個因素:
- 相對位置嵌入是影響性能的一個關鍵因素
- 另外一個關鍵因素是濾波器使用方式,即滑動窗口 vs 非重疊窗口
DwConv 能夠與 LSA 性能相媲美的原因在于:它采用了滑動窗口處理機制。當其采用非重疊窗口機制時,性能明顯弱于 LSA(見 Table1 中的 Net4)。
動態濾波器性能優于 LSA 的原因在于相對位置嵌入與近鄰濾波器使用方式。兩者的集成 (Net7) 取得了最佳的性能。
對比非重疊局部窗口與滑動窗口,局部重疊的峰值性能要弱于滑動窗口。局部窗口的一個缺點在于:窗口間缺乏信息交互,限制了其性能;而滑動窗口的缺陷在于低吞吐量。那么,如何避免點乘同時保持高性能就成了新的挑戰?。
基于上述分析,作者提出一種增強的局部自注意力模塊 ELSA,性能可以超越了 SwinT 中的 LSA 與動態濾波器。
將該模塊替換掉 Swin Transformer 中的前三個 stage 中的 LSA 模塊即可。ELSA 的關鍵模塊為Hadamard 注意力與 Ghost 頭模塊,表達式如下所示:
相比 Dot product,Hadamard product ( Pytorch 中的 A*B )可以有效地提取高階信息,而 ghost 模塊是受啟發于 GhostNet,可以在有限的容量下提取更豐富的通道特征。作者在分類、目標檢測和語義分割任務上都進行了驗證。
ELSA 的解決辦法看起來過于復雜,整個分析過程也感覺有點復雜,可能會存在更簡單的改進策略。
(3) 稀疏注意力(局部緊密相關和遠程稀疏)
和 Performer 一樣,NLP 領域內也有很多對自注意力模塊的改進中是引入局部注意力的,典型的例如 Sparse Transformers,其出發點是實際場景中或許用不到真正的全局注意力,只要提供稍微等價的全局效果就行,不一定真的理論上達到全局。基于此可以實現局部緊密相關和遠程稀疏的 Sparse Transformers,后續改進版本也非常多,例如 LongFormer。同樣的,本文對 NLP 領域發展不展開描述。
3.1.2.4 FFN 模塊
FFN 模塊比較簡單,主要是進行特征變換,ViT 中代碼如下所示:
FFN 的改進方向和前面類似,也是希望引入 CNN 局部特征提取能力,加速收斂、提升性能,主要可以歸納為如下所示:
在引入 Conv 增強局部信息特征信息提取方面,也有不少論文進行了嘗試。
LocalViT?中在其他模塊都不改動情況下,僅僅在 FFN 中引入 1x1 卷積和帶 padding 的 3x3 DW 卷積來增強局部特征提取能力,實驗結果能夠帶來不少的性能提升,如下所示:
(a) 是簡單的卷積 forward 模塊,(b) 是反轉殘差塊,(c) 是本文所提結構,僅僅需要通過替換 FFN 中間的 MLP 層為卷積層。同樣采用類似做法的有 PVTv2 和 HRFormer 等等。
CeiT?中也提出了非常類似的結構,其出發點也是希望引入 CNN 來加強局部特征能力,結構如下所示:
主要包括線性投影、恢復空間結構、3x3 DW 卷積、flatten 和線性投影。
在 3.1.2.2 位置編碼模塊中說過 PVTv2 在 FFN 中引入了零填充的逐深度通道卷積來自動感知位置信息,從而可以省掉位置編碼。從 FFN 角度來看,除了有省掉位置編碼的好處外, 還能夠加強局部特征提取能力,相比 LocalViT 和 CeiT ,直接拋棄位置編碼的做法更加徹底,更加合理。
3.1.2.5 Norm 位置改動
Norm 通常是 Layer Norm,按照該模型放在自注意力和 FFN 模塊的前面還是后面,可以分成 pre norm 和 post norm 方式,如下所示:
絕大部分模型例如 ViT、Swin Transformer 等都是 pre norm 方式,如下所示:
但是 Swin Transformer v2 卻改成了 post norm。Swin Transformer V2?主要探討如何 scale up 視覺模型,并最終在 4 個數據集上重新刷到了新的 SOTA。將視覺 Transformer 的模型參數調大,作者認為會面臨以下幾個挑戰:
- 增大視覺模型時候可能會帶來很大的訓練不穩定性
- 在很多需要高分辨率的下游任務上,還沒有很好地探索出對低分辨率下訓練好的模型遷移到更大scale 模型上的方法
- GPU 內存代價太大
作者發現,將 Swin Transformer 模型從 small size 增大到 large size 后,網絡深層的激活值會變得很大,與淺層特征的激活值有很大的 gap,如下圖所示,可以看出來隨著模型 scale 的增大,這個現象會變得很嚴重。
作者發現使用 post-norm 操作后,上面所觀察到的問題可以得到很明顯的改善,并且為了更進一步穩定 largest Swin V2 的訓練,在每 6 個 transformer block 后還額外加了一層 layer normalization。也就是說在大模型 Transformer 中使用 post Norm 能夠更加穩定訓練過程。
3.1.2.6 分類預測頭模塊
在 ViT 中通過追加額外一個 Class Token,將該 Token 對應的編碼器輸出輸入到 MLP 分類頭(實際上是一個線性投影層)進行分類,如下所示:
為何不能和我們常規的圖像分類一樣,直接聚合所有特征,而需要單獨引入一個 Class Token ?這個問題自然有人進一步探索,經過簡單總結,可以歸納為如下結構圖:
從目前來看主要包括 2 種做法: Class Token 和常規分類一樣直接聚合特征。
Class Token 方面的改進典型代碼是?CeiT?,作者的出發點是 ViT 僅僅使用最后一個 Transformer 輸出做分類,但是實驗觀察不同的 Transformer 層輸出特性不一樣,如果能夠聚合多個輸出層的特征進行分類,也許會更好,具體做法如下所示:
首先收集不同 Transformer 層的 Class Token 對應的輸出,然后通過一個 Layer-wise Class-token Attention (LCA) 模塊來通過注意力自適應聚合所需要的 token。LCA 可以看出一個簡單的 Transformer 模塊,其包括多頭注意力模塊 MSA 和前向網絡 FFN,在聚合信息后,取輸出序列的最后一個輸出進行分類訓練。為了節省計算量, MSA 只計算最后一個 Class Token 和其他 Class Token 之間的單向注意力,而沒有必要算全局。
隨著視覺 Transformer 的快速發展,后人發現其實直接采用 Avg Pool 聚合特征性能也很不錯,甚至會更好,典型的代表是 Swin Transformer 和 CPVT-GAP,這種做法無需引入額外的 Class Token,和我們 CNN 中常用的分類套路一致,如下所示:
3.1.2.7 其他
這里說的其他包括兩個內容部分,如下所示:
3.1.2.7.1 多尺度特征圖輸出
在 CNN 時代的各種下游任務(例如目標檢測、語義分割)中已經被廣泛證明多分辨率多尺度特征非常重要,不同尺度特征可以提供不同的感受野,適合提取不同物體尺度的特征,然而 ViT 僅僅是為圖像分類而設計,無法很好地應用于下游任務,這嚴重制約了視覺 Transformer 的廣泛應用,故迫切需要一種能夠類似 ResNet 在不同 stage 輸出不同尺度的金字塔特征做法。
ViT 要輸出多尺度特征圖,最常見做法是 Patch Merging。Patch Mergeing 含義是對不同窗口的 Patch 進行合并,在目前主流的 PVT、Twins、Swin Transformer 和 ResT 中都有廣泛的應用,以?PVT?為例詳細說明,結構圖如下所示:
假設圖片大小是 (H, W, 3),暫時不考慮 batch。
CSWin Transformer 則更加徹底,直接將上述 stride 下采樣過程替換為 kernel 為 3x3,stride為 2 的卷積,做法和 ResNet 中下采樣一致。我們將上述改變 stride 導致序列長度變少從而縮減空間特征圖的過程統稱為 Patch Merging。
除了上述這種相對樸素的做法,還有一些其他做法。例如?MViT?,其不存在專門的 Patch Merging 模塊,而是在注意力模塊中同時嵌入下采樣功能,如下所示:
只要在每個 stage 中改進 Pool 模塊的 stride 就可以控制實現 ResNet 一樣的多尺度輸出,從而實現多分辨率金字塔特征輸出。
ViTAE 則直接引入空洞卷積,然后設置不同的空洞率來實現下采樣從而輸出不同尺度的特征圖。
3.1.2.7.2 訓練深層 Transformer
探討如何訓練更深的 Transformer 典型算法是 CaiT 和 DeepViT。前面的諸多 ViT 改進都是在編碼層為 6 的基礎上進行設計的,是否可以類似 CNN 分類網絡設計更深的 Transformer,性能也能出現一致性提升,不至于過早飽和?
(1) CaiT
在?CaiT?算法中,作者從 Transformer 架構和優化關系會相互影響相互作用角度出發進行探討。在 ResNet 中,作者也說過較深的網絡能不能順利的優化才是其成功的關鍵,而不是在于特征提取能力多強,對應到 ViT 中依然是首先要考慮如何更容易優化,這會涉及到 Normalize 選擇、權重初始化、殘差塊初始化等模塊,不能小看這些模塊,一個良好的初始化設計都可能避免深模型的過早飽和問題。
除了上述影響,作者發現 ViT 的 class token 似乎有點不合理,ViT 在一開始就將 class token 和 patch embedding 輸出序列拼接,然后輸入給自注意力模塊,這個自注意力模塊需要同時完成兩個任務:
- 引導 attention 過程,幫助得到 attention map。
- token 最后輸入到 classifier 中,完成分類任務。
也就是說將 class token 和 patch embedding 過早融合,但是兩者優化方向不一樣,可能存在矛盾。同時說到更順利的優化,Normalize 的重要性應該是第一個能想到的。基于這兩個出發點,作者進行了如下改進:
- 提出?LayerScale,更合理的 Norm 策略使深層 Transformer 易于收斂,并提高精度。
- 提出?class-attention layers,class token 和 patch embedding 在最后融合,并且通過 CA 模塊來更加高效地將 patch embedding 信息融合到 class embedding 中,從而提升性能。
結構圖如下所示:
(a) LayerScale
- (a) 是 ViT 做法即先進行 Layer Normalization,再進行 Self-attention 或者 FFN,然后結果與 block 輸入相加。
- (b) 是 ReZero、Skipinit 和 Fixup 算法的做法,引入了一個可學習的參數 alpha 作用在 residual block 的輸出,并移除了 Layer Normalization 操作,這個可學習參數初始化可以從 0 開始,也可以從 1 開始,實驗表示都無法很好地解決深層 Transformer 優化問題。
- (c) 是一個組合做法,實驗發現這種組合做法有一定效果。
- (d) 是本文所提 LayerScale,效果最好。
LayerScale 的做法是保留 Layer Normalization,并對 Self-attention 或者 FFN 的輸出乘上一個對角矩陣,由于其對角矩陣,這個實際上就等效于通道注意力(對不同通道乘上不同的系數),這些系數的設置比較有講究,在 18 層之前,它們初始化為 0.1,若網絡更深,則在 24 層之前初始化為 10^-5,若網絡更深,則在之后更深的網絡中初始化為 10^-6,這樣設計的原因是希望越深的 block 在一開始的時候更接近恒等映射,在訓練的過程中逐漸地學習到模塊自身所需要的特征,這個理論和 ResNet 中初始化設計非常類似。
(b) Class-Attention Layers
從上述結構圖中可以看出, class token 是在后面融合,前面的 Transformer 層專注于 patch embedding 信息提取。CA 模塊做法為:
Q 僅僅來自 class token,而 KV 來自 z=[class token, patch embedding],上述做法是常規的交叉注意力計算方式。作者實驗發現 CA layer 使用 2 層就夠了,前面依舊是 N 個正常的 Transformer Block。如果直接看代碼可能會更清晰一些:
上述就是 CaiT 算法主要改進,但是其實要想成功訓練出來還需要諸多 trick,作者在論文中也貼了大量對比實驗,有興趣的建議直接閱讀原文,本文僅僅描述了最重要的改進而已。更貼心的是,作者還列出了 DeiT-S 到 CaiT-36 的性能提升 trick 表,這個表格同樣重要。
(2) DeepViT
CaiT 算法是從整體架構和優化聯系角度入手,而?DeepViT?不一樣,他通過分析得出深層 Transformer 性能飽和的原因是注意力崩塌,即深層的 Transformer 學到的 attention 非常相似。這意味著隨著 ViT 的層次加深,self-attention 模塊在生成不同注意力以捕獲多樣性特征方面變得低效。如下圖所示,上面一行是 ViT 訓練結果,隨著模型深入,自注意力層學到的自注意力權重圖趨向均勻且都非常類似,下面一類是作者改進后的可視化結果,根據多樣性。
DeepViT 中作者也做了很多前期實驗,來驗證注意力崩塌現象。那么現在核心問題就是要使得各個注意力層所提取的注意力權重圖根據多樣性,獨特性。
為了解決這個問題,作者提出了 2 種解決辦法:
(a) 增加 self-attention 模塊的嵌入維度
作者認為增加嵌入維度,可以增加表征容量從而避免注意力崩塌。實驗結果如上所示,隨著 Embedding Dimension 的增長,相似的 Block 的數量在下降, Acc 在上升,說明這個策略實際上確實是有效的。但增加 Embedding Dimension 也會顯著增加計算成本,帶來的性能改進往往會減少,且需要更大的數據量來訓練,增加了過擬合的風險,這個也不是啥實質性的解決辦法。為此作者提出了第二個改進 Re-attention。
(b) Re-attention
只需要替換 Self-Attention 為 Re-Attention 即可,其結構圖如下右圖所示:
作者通過實驗觀察到同一個層的不同 head 之間的相似度比較小,這表明來自同一自注意力層的不同head 關注到了輸入 token 的不同方面,基于此作者提出可以把不同的 head 的輸出信息融合,然后利用它們再生出一個新的 attention map。具體為引入一個可學習 Linear Transformation,在注意力圖生成后乘上 Linear Transformation,從而生成一個新的 attention map,然后進行 Norm 后乘上 V。
因為我們目的是加強同一個自注意力層間不同 head (通常是8,也就是 Linear Transformation 矩陣大小是 8x8) 之間的交流,所以這個可學習 Linear Transformation 是作用在 head 維度,同時 Norm 是 BatchNorm 而不是傳統的 LayerNorm。
CaiT 和 DeepViT 都是關注深層 Transformer 出現的過早飽和問題,不過關注的地方不一樣,解決辦法也完全不同,但是效果類似,這或許說明還有很大的改進空間。
4 總結
ViT 的核心在于 Attention,但是整個架構也包括多個組件,每個組件都比較關鍵,有諸多學者對多個組件進行了改進。我們可以簡單將 ViT 結構分成 6 個部分:
隨著研究的不斷深入,大家發現 Attention 可能不是最重要的,進而提出了 MLP-based 和 ConvMixer-based 類算法,這些算法都是為了說明自注意力模塊可以采用 MLP 或者 Conv 層代替,這說明 Transformer 的成功可能來自整個架構設計。MLP-based 和 ConvMixer-based 部分將會在下一篇文章中進行說明。
對 Vision Transformer 系列內容感興趣的朋友,不要忘記 star 啦~
https://github.com/open-mmlab/awesome-vit?github.com/open-mmlab/awesome-vit
總結
以上是生活随笔為你收集整理的Vision Transformer 必读系列之图像分类综述(二): Attention-based的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 猎头职场:真正城府深的人都不会做这些
- 下一篇: 主动降噪技术