Transformer | 詳細(xì)解讀Transformer怎樣從零訓(xùn)練并超越ResNet？

本文證明了在沒有大規(guī)模預(yù)訓(xùn)練或強(qiáng)數(shù)據(jù)增廣的情況下，在ImageNet上從頭開始訓(xùn)練時(shí)，所得ViT的性能優(yōu)于類似大小和吞吐量的ResNet！而且還擁有更敏銳的注意力圖。
作者單位：谷歌,UCLA

1簡介

Vision Transformers(ViTs)和MLPs標(biāo)志著在用通用神經(jīng)架構(gòu)替換手動(dòng)特征或歸納偏置方面的進(jìn)一步努力。現(xiàn)有工作通過大量數(shù)據(jù)為模型賦能，例如大規(guī)模預(yù)訓(xùn)練和/或重復(fù)的強(qiáng)數(shù)據(jù)增廣，并且還報(bào)告了與優(yōu)化相關(guān)的問題（例如，對初始化和學(xué)習(xí)率的敏感性）。

因此，本文從損失幾何的角度研究了ViTs和MLP-Mixer，旨在提高模型在訓(xùn)練和推理時(shí)的泛化效率?？梢暬虷essian揭示了收斂模型極其敏感的局部最小值。

同時(shí)通過使用最近提出的銳度感知優(yōu)化器提高平滑度，進(jìn)而大大提高了ViT和MLP-Mixer在跨越監(jiān)督、對抗、對比和遷移學(xué)習(xí)（例如，+5.3% 和 +11.0%）的各種任務(wù)上的準(zhǔn)確性和魯棒性使用簡單的Inception進(jìn)行預(yù)處理，ViT-B/16和Mixer-B/16在ImageNet上的準(zhǔn)確率分別為Top-1）。

作者研究表明，改進(jìn)的平滑度歸因于前幾層中較稀疏的活動(dòng)神經(jīng)元。在沒有大規(guī)模預(yù)訓(xùn)練或強(qiáng)數(shù)據(jù)增強(qiáng)的情況下，在ImageNet上從頭開始訓(xùn)練時(shí)，所得 ViT的性能優(yōu)于類似大小和吞吐量的ResNet。同時(shí)還擁有更敏銳的注意力圖。

2Background和Related Work

最近的研究發(fā)現(xiàn)，ViT中的self-attention對性能并不是至關(guān)重要的，因此出現(xiàn)了一些專門基于mlp的架構(gòu)。這里作者以MLP-Mixer為例。MLP-Mixer與ViT共享相同的輸入層;也就是說，它將一個(gè)圖像分割成一系列不重疊的Patches/Toekns。然后，它在torkn mlp和channel mlp之間交替使用，其中前者允許來自不同空間位置的特征融合。

3ViTs和MLP-Mixers收斂到銳局部極小值

目前的ViTs、mlp-mixer和相關(guān)的無卷積架構(gòu)的訓(xùn)練方法很大程度上依賴于大量的預(yù)訓(xùn)練或強(qiáng)數(shù)據(jù)增強(qiáng)。它對數(shù)據(jù)和計(jì)算有很高的要求，并導(dǎo)致許多超參數(shù)需要調(diào)整。

現(xiàn)有的研究表明，當(dāng)在ImageNet上從頭開始訓(xùn)練時(shí)，如果不結(jié)合那些先進(jìn)的數(shù)據(jù)增強(qiáng)，盡管使用了各種正則化技術(shù)(例如，權(quán)重衰減，Dropout等)ViTs的精度依然低于類似大小和吞吐量的卷積網(wǎng)絡(luò)。同時(shí)在魯棒性測試方面，vit和resnet之間也存在較大的差距。

此外，Chen等人發(fā)現(xiàn)，在訓(xùn)練vit時(shí)，梯度會出現(xiàn)峰值，導(dǎo)致精確度突然下降，Touvron等人也發(fā)現(xiàn)初始化和超參數(shù)對訓(xùn)練很敏感。這些問題其實(shí)都可以歸咎于優(yōu)化問題。

在本文中，作者研究了ViTs和mlp-mixer的損失情況，從優(yōu)化的角度理解它們，旨在減少它們對大規(guī)模預(yù)訓(xùn)練或強(qiáng)數(shù)據(jù)增強(qiáng)的依賴。

3.1 ViTs和MLP-Mixers收斂到極sharp局部極小值

眾所周知，當(dāng)模型收斂到曲率小的平坦區(qū)域時(shí)模型會具有更好的泛化性能。在[36]之后，當(dāng)resnet、vit和MLP-Mixers在ImageNet上使用基本的初始風(fēng)格預(yù)處理從頭開始訓(xùn)練時(shí)，作者繪制損失圖：

如圖1(a)到1(c)所示，ViTs和mlp-mixer比ResNets收斂到更清晰的區(qū)域。

在表1中，通過計(jì)算主要的Hessian特征值 進(jìn)一步驗(yàn)證了結(jié)果。ViT和MLP-Mixer的 值比ResNet大一個(gè)數(shù)量級，并且MLP-Mixer的曲率在3種中是最大的(具體分析見4.4節(jié))。

3.2 Small training errors

這種向sharp區(qū)域的收斂與圖2(左)所示的訓(xùn)練動(dòng)態(tài)一致。盡管Mixer-B/16參數(shù)少于ViT-B/16(59M vs 87M)，同時(shí)它有一個(gè)小的訓(xùn)練誤差，但測試性能還是比較差的，這意味著使用cross-token MLP學(xué)習(xí)的相互作用比ViTs’ self-attention機(jī)制更容易過度擬合。這種差異可能解釋了mlp-mixer更容易陷入尖銳的局部最小值。

3.3 ViTs和MLP-Mixers的可訓(xùn)練性較差

此外，作者還發(fā)現(xiàn)ViTs和MLP-Mixers的可訓(xùn)練性較差，可訓(xùn)練性定義為通過梯度下降優(yōu)化的網(wǎng)絡(luò)的有效性。Xiao等人的研究表明，神經(jīng)網(wǎng)絡(luò)的可訓(xùn)練性可以用相關(guān)的神經(jīng)切線核(NTK)的條件數(shù)來表征:

其中 是雅可比矩陣。

用 表示NTK 的特征值，最小的特征值 以條件數(shù)κ 的速率指數(shù)收斂。如果κ是發(fā)散的，那么網(wǎng)絡(luò)將變得不可訓(xùn)練。如表1所示，ResNets的κ是相當(dāng)穩(wěn)定的，這與之前的研究結(jié)果一致，即ResNets無論深度如何都具有優(yōu)越的可訓(xùn)練性。然而，當(dāng)涉及到ViT和時(shí)，條件數(shù)是不同的MLP-Mixer，證實(shí)了對ViTs的訓(xùn)練需要額外的輔助。

4CNN-Free視覺架構(gòu)優(yōu)化器原理

常用的一階優(yōu)化器(如SGD,Adam)只尋求最小化訓(xùn)練損失。它們通常會忽略與泛化相關(guān)的高階信息，如曲率。然而，深度神經(jīng)網(wǎng)絡(luò)的損失具有高度非凸性，在評估時(shí)容易達(dá)到接近0的訓(xùn)練誤差，但泛化誤差較高，更談不上在測試集具有不同分布時(shí)的魯棒性。

由于對視覺數(shù)據(jù)缺乏歸納偏差ViTs和MLPs放大了一階優(yōu)化器的這種缺陷，導(dǎo)致過度急劇的損失scene和較差的泛化性能，如前一節(jié)所示。假設(shè)平滑收斂時(shí)的損失scene可以顯著提高那些無卷積架構(gòu)的泛化能力，那么最近提出的銳度感知最小化(SAM)可以很好的避免銳度最小值。

4.1 SAM:Overview

從直覺上看，SAM尋找的是可以使整個(gè)鄰近訓(xùn)練損失最低的參數(shù)w，訓(xùn)練損失 通過構(gòu)造極小極大目標(biāo):

其中 是neighbourhood ball的大小。在不失一般性的情況下，這里使用 范數(shù)作為其強(qiáng)經(jīng)驗(yàn)結(jié)果，這里為了簡單起見省略了正則化項(xiàng)。

由于內(nèi)部最大化下式的確切解很難獲得：

因此，這里采用了一個(gè)有效的一階近似:

在 范數(shù)下， 是當(dāng)前權(quán)值 的縮放梯度。計(jì)算 后，SAM基于銳度感知梯度更新w。

4.2 SAM優(yōu)化器實(shí)質(zhì)上改進(jìn)了ViTs和MLP-Mixers

作者在沒有大規(guī)模的預(yù)訓(xùn)練或強(qiáng)大的數(shù)據(jù)增強(qiáng)的情況下訓(xùn)練了vit和MLP-Mixers。直接將SAM應(yīng)用于vit的原始ImageNet訓(xùn)練pipeline，而不改變?nèi)魏纬瑓?shù)。pipeline使用了基本的Inception-style的預(yù)處理。最初的mlp-mixer的訓(xùn)練設(shè)置包括強(qiáng)數(shù)據(jù)增強(qiáng)的組合;也用同樣的Inception-style的預(yù)處理來替換它，以便進(jìn)行公平的比較。

注意，在應(yīng)用SAM之前，我們對學(xué)習(xí)速率、權(quán)重衰減、Dropout和隨機(jī)深度進(jìn)行網(wǎng)格搜索。

1 局部極小值周圍的平滑區(qū)域

由于SAM, ViTs和mlp-mixer都匯聚在更平滑的區(qū)域，如圖1(d)和1(e)所示。

曲率測量，即Hessian矩陣的最大特征值 ，也減小到一個(gè)小值(見表1)。

2 Higher accuracy

隨之而來的是對泛化性能的極大改進(jìn)。在ImageNet驗(yàn)證集上，SAM將ViT-B/16的top-1精度從74.6%提高到79.9%，將Mixer-B/16的top-1精度從66.4%提高到77.4%。

相比之下，類似規(guī)模的ResNet-152的性能提高了0.8%。根據(jù)經(jīng)驗(yàn)，改進(jìn)的程度與架構(gòu)中內(nèi)置的歸納偏差水平呈負(fù)相關(guān)。與基于注意力的ViTs相比，具有inherent translation equivalence和locality的ResNets從landscape smoothing中獲益較少。MLP-Mixers從平滑的loss geometry中獲得最多。

此外，SAM對更大容量(例如:+4.1%的Mixer-S/16 vs. +11.0%的Mixer-B/16)和更長的patch序列(例如:+2.1%的vits/32 vs. +5.3%的vits /8)的模型帶來了更大的改進(jìn)。

3 Better robustness

作者還使用ImageNet-R和ImageNetC評估了模型的魯棒性，并發(fā)現(xiàn)了smoothed loss landscapes的更大影響。在ImageNet-C上，它通過噪音、惡劣天氣、模糊等來破壞圖像，實(shí)驗(yàn)了5種嚴(yán)重程度上19種破壞的平均精度。如表1和表2所示， ViT-B/16和Mixer-B/16的精度分別增加了9.9%和15.0%。

4.3 無需預(yù)訓(xùn)練或強(qiáng)大的數(shù)據(jù)增強(qiáng)ViTs優(yōu)于ResNets

模型體系結(jié)構(gòu)的性能通常與訓(xùn)練策略合并，其中數(shù)據(jù)增強(qiáng)起著關(guān)鍵作用。然而，數(shù)據(jù)增廣的設(shè)計(jì)需要大量的領(lǐng)域?qū)I(yè)知識，而且可能無法在圖像和視頻之間進(jìn)行轉(zhuǎn)換。由于有了銳度感知優(yōu)化器SAM，可以刪除高級的數(shù)據(jù)增強(qiáng)，并專注于體系結(jié)構(gòu)本身(使用基本的Inception-style的預(yù)處理)。

當(dāng)使用SAM在ImageNet上從0開始訓(xùn)練時(shí)，ViT的準(zhǔn)確性(在ImageNet、ImageNet-Real和ImageNet V2上)和健壯性(在ImageNet-R和ImageNet-R上)方面都優(yōu)于類似和更大的ResNet(在推理時(shí)也具有相當(dāng)?shù)耐掏铝?。

ViT-B/16在ImageNet、ImageNet-r和ImageNet-C上分別達(dá)到79.9%、26.4%和56.6%的top精度，而對應(yīng)的ResNet-152則分別達(dá)到79.3%、25.7%和52.2%(見表2)。對于小型架構(gòu)，vit和resnet之間的差距甚至更大。在ImageNet上，ViT-S/16的表現(xiàn)比同樣大小的ResNet-50好1.4%，在ImageNet-C上好6.5%。SAM還顯著改善了MLP-Mixers的結(jié)果。

4.4 SAM后的內(nèi)在變化

作者對模型進(jìn)行了更深入的研究，以理解它們?nèi)绾螐谋举|(zhì)上改變以減少Hessian的特征值 以及除了增強(qiáng)泛化之外的變化意味著什么。

結(jié)論1：每個(gè)網(wǎng)絡(luò)組件具有Smoother loss landscapes

在表3中，將整個(gè)體系結(jié)構(gòu)的Hessian分解成與每一組參數(shù)相關(guān)的小的斜對角Hessian塊，試圖分析在沒有SAM訓(xùn)練的模型中，是什么特定的成分導(dǎo)致 爆炸。

作者觀察到較淺的層具有較大的Hessian特征值 ，并且第1個(gè)linear embedding layer產(chǎn)生sharpest的幾何形狀。

此外，ViTs中的多頭自注意(MSA)和MLP-Mixers中的token mlp(Token mlp)跨空間位置混合信息，其 相對較低。SAM一致地降低了所有網(wǎng)絡(luò)塊的 。

可以通過遞歸mlp的Hessian矩陣得到上述發(fā)現(xiàn)。設(shè) 和 分別為第k層激活前的值和激活后的值。它們滿足，其中 為權(quán)值矩陣， 為激活函數(shù)(mlp-mixer中的GELU)。為了簡單起見，在這里省略偏置項(xiàng)。Hessian矩陣 相對于 的對角塊可遞歸計(jì)算為:

其中?為Kronecker product， 為第 層的預(yù)激活Hessian，L為目標(biāo)函數(shù)。因此，當(dāng)遞歸公式反向傳播到淺層時(shí)，Hessian范數(shù)累積，這也解釋了為什么表3中第一個(gè)塊的 比最后一個(gè)塊大得多。

結(jié)論2：Greater weight norms

應(yīng)用SAM后，作者發(fā)現(xiàn)激活后的值 的范數(shù)和權(quán)重 的范數(shù)變得更大(見表3)，說明常用的權(quán)重衰減可能不能有效地正則化ViTs和MLP-Mixers。

結(jié)論3：MLP-Mixers中較稀疏的active neurons

根據(jù)遞歸公式(3)到(4)，作者確定了另一個(gè)影響Hessian的MLP-Mixers的內(nèi)在度量:激活神經(jīng)元的數(shù)量。

事實(shí)上， 是由大于零的被激活神經(jīng)元決定的，因?yàn)楫?dāng)輸入為負(fù)時(shí)，GELU的一階導(dǎo)數(shù)變得非常小。因此，活躍的GELU神經(jīng)元的數(shù)量直接與Hessian規(guī)范相連。

圖2(右)顯示了每個(gè)塊中被激活的神經(jīng)元的比例，使用ImageNet訓(xùn)練集的10%進(jìn)行計(jì)算。可以看到，SAM極大地減少了前幾層被激活神經(jīng)元的比例，使它們處于更稀疏的狀態(tài)。這一結(jié)果也說明了圖像patch的潛在冗余性。

結(jié)論4：ViTs的active neurons高度稀疏

雖然公式(3)和(4)只涉及mlp，但仍然可以觀察到vit的第1層激活神經(jīng)元的減少(但不如MLP-Mixers顯著)。更有趣的是，作者發(fā)現(xiàn)ViT中被激活神經(jīng)元的比例比ResNets或MLP-Mixers中要小得多——在大多數(shù)ViT層中，只有不到5%的神經(jīng)元的值大于零。換句話說，ViT為網(wǎng)絡(luò)修剪提供了巨大的潛力。

這種稀疏性也可以解釋為什么一個(gè)Transformer可以處理多模態(tài)信號(視覺、文本和音頻)?

結(jié)論5：ViTs中有更多的感知注意力Maps

在圖3中可視化了classification token的attention map。有趣的是，經(jīng)過SAM優(yōu)化的ViT模型能夠編碼合理的分割信息，比傳統(tǒng)SGD優(yōu)化訓(xùn)練的模型具有更好的可解釋性。

結(jié)論6：Higher training errors

如圖2(左)所示，使用SAM的ViT-B/16比使用vanilla SGD的訓(xùn)練誤差更高。當(dāng)在訓(xùn)練中使用強(qiáng)數(shù)據(jù)增強(qiáng)時(shí)，這種正則化效應(yīng)也會發(fā)生，它迫使網(wǎng)絡(luò)顯式地學(xué)習(xí)RandAugment中的旋轉(zhuǎn)平移等方差和mixup中的線性插值等先驗(yàn)。然而，增益對不同的訓(xùn)練設(shè)置很敏感(第5.2節(jié))，并導(dǎo)致高噪聲損失曲線(圖2(中間))。

5實(shí)驗(yàn)

具有smoother loss geometry的ViTs和MLP-Mixers可以更好地遷移到下游任務(wù)。

6參考

[1].When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations.

本文論文原文獲取方式，掃描下方二維碼

掃描下方二維碼關(guān)注【集智書童】公眾號，獲取更多實(shí)踐項(xiàng)目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行！