字節(jié)跳動AI Lab最新的工作捡遍,研究了如何把ViT做的更深的問題。之前的ViT采用的都是固定的12層,所不同的是head的數(shù)目和embedding的通道數(shù)攻旦,由此構(gòu)成small、medium以及l(fā)arge版本生逸。但是關(guān)于ViT層數(shù)的問題牢屋,則很少有探討。
本文首先將ViT由12層加深到32層槽袄,發(fā)現(xiàn)隨著網(wǎng)絡(luò)層數(shù)的加深烙无,性能并沒有獲得相應(yīng)幅度的上升,甚至還有下降:
類似的網(wǎng)絡(luò)退化問題在cnn中也有出現(xiàn)遍尺,cnn中通過加入residual connection來解決網(wǎng)絡(luò)退化截酷。本文中則觀察發(fā)現(xiàn)ViT的退化是由于attention collapse,即隨著層數(shù)的加深狮鸭,attention map的變化越來越小合搅,相應(yīng)的feature map也開始停止學(xué)習(xí),因而隨著層數(shù)加深性能增長逐漸停滯歧蕉。
作者的解決方案其實很簡單灾部,觀察發(fā)現(xiàn)雖然同一個token在不同層的attention map差別小,但是同一層不同head之間差距還是很大惯退,因此作者額外加入一個HxH大小的矩陣赌髓,來對attention map進(jìn)行re-attention,利用其他head的attention map來對當(dāng)前attention map進(jìn)行增強催跪。
反應(yīng)在公式上的變化就是在與V相乘之前先與theta矩陣相乘锁蠕,看起來很簡潔了。消融實驗結(jié)果顯示懊蒸,在32層ViT上荣倾,本文所提的解決方案(DeepViT)可以比vanilla ViT有1.6的提升,在深層ViT上的提升還是比較顯著的骑丸。
作者解釋re-attention起作用的原因在于在不同head之間進(jìn)行了interaction舌仍,而不是單純的解決over smoothing妒貌,因而可以編碼更豐富的信息,當(dāng)然通過可視化結(jié)果也可以看到铸豁,加入re-attention之后灌曙,block的相似性也顯著降低了。