作者 |?陳萍、杜偉

來源 |?機器之心

CNN 通過堆疊更多的卷積層來提高性能，而 transformer 在層次更深時會很快進入飽和。基于此，來自新加坡國立大學和字節(jié)跳動 AI Lab 的研究者引入了 Re-attention 機制，以很小的計算代價重新生成注意力圖以增強各層之間的多樣性。提出的 DeepViT 模型也實現(xiàn)了非常不錯的性能。

視覺 transformer (ViT) 現(xiàn)已成功地應用于圖像分類任務。近日，來自新加坡國立大學和字節(jié)跳動美國 AI Lab 的研究者表明，不同于卷積神經網絡通過堆疊更多的卷積層來提高性能，ViT 的性能在擴展至更深時會迅速飽和。

他們是如何得出這一結論的呢？

具體而言，研究者根據(jù)經驗觀察到，這種擴展困難是由注意力崩潰（attention collapse）引起的：隨著 Transformer 加深，注意力圖在某些層之后逐漸變得相似甚至幾乎相同。換句話說，在 deep ViT 模型的頂層中，特征圖趨于相同。這一事實表明，在更深層的 ViT 中，自注意力機制無法學習有效的表征學習概念，并且阻礙了模型獲得預期的性能提升。

基于以上觀察，研究者提出了一種簡單而有效的方法 Re-attention，它可以忽略計算和存儲成本重新生成注意力圖以增加其在不同層的多樣性。借助于該方法，我們可以通過對現(xiàn)有 ViT 模型的微小修改來訓練具有持續(xù)性能改進的更深的 ViT 模型。此外，當使用 32 個 transformer 塊訓練 DeepViT 模型時，在 ImageNet 數(shù)據(jù)集上實現(xiàn)了頗具競爭力的 Top-1 圖像分類準確率。相較于 ViT-32B， 變體模型 DeepViT-32B 的 Top-1 準確率提升了 1.6%。

論文鏈接：https://arxiv.org/pdf/2103.11886.pdf

ViT 模型

如下圖 2 所示，ViT 模型由三部分組成：用于 patch 嵌入的線性層、具有多頭自注意力和特征編碼前饋層的 transformer 塊、以及用于分類得分預測的線性層。研究者首先回顧了 transformer 塊的獨特性，特別是自注意力機制，然后研究了自注意力的崩潰問題。

具有 N 個 transformer 塊的原版 ViT 模型與該研究所提 DeepViT 模型的結構對比。

注意力崩潰

受深度 CNN 成功的啟發(fā)，研究者對 ViT 隨深度的增加帶來的性能改變進行了系統(tǒng)的研究。在不失一般性的情況下，他們首先將隱藏維數(shù)和注意力頭數(shù)分別固定為 384 和 12，然后堆疊不同數(shù)量的 transformer 塊（數(shù)量從 12 到 32），以建立對應不同深度的多個 ViT 模型。在 ImageNet 數(shù)據(jù)集上，原版 ViT 和 DeepViT 的圖像分類 Top-1 準確率變化曲線如下圖 1 所示：

結果表明，隨著模型深度的增加，分類準確率提升緩慢，飽和速度加快。更具體地說，在采用 24 個 transformer 塊時，性能停止提升。這種現(xiàn)象表明，現(xiàn)有的 ViT 很難在更深的架構中提高性能。

為了衡量注意力圖在不同層上的改變，研究者計算了來自不同層的注意力圖之間的跨層相似性：

其中和四個因素有關：p 和 q 是兩個不同的層、h 是注意力頭、t 是具體的輸入，結果如下圖 3 所示。圖（a）表明了隨著深度的增加，注意力圖和 k 個附近塊的注意力圖越來越相似；圖（c）表明了即使到了第 32 個 block，同層注意力頭之間的相似度還是比較低的，這說明主要的相似還是層之間的相似。

為了了解注意力崩潰如何影響 ViT 模型的性能，研究者進一步探索了它是如何影響更深層次的特征學習。對于特定的 32 塊 ViT 模型，通過研究它們的余弦相似性，將最終輸出特征與每個中間 transformer 塊的輸出進行比較。

下圖 4 中的結果表明：特征圖與注意力圖相似性非常高，并且學習的特征在第 20 個塊之后停止變化。注意力相似度的增加與特征相似度之間存在著密切的相關性。這一觀察表明注意力崩潰是 ViT 不可擴展的原因。

DeepViT 中的 Re-attention

如上所述，將 ViT 擴展到更深層的一個主要障礙是注意力崩潰問題。因此，研究者提出了兩種解決方法，一種是增加計算自注意力的隱藏維度，另一種是新的 re-attention 機制。

高維空間中的自注意力

克服注意力崩潰的一個解決方案是增加每個 token 的嵌入維度，這將提升每個 token 嵌入的表征能力，以編碼更多信息。因此，所得到的注意力圖更具多樣性，降低每個塊注意力圖之間的相似性。在不喪失一般性的前提下，該研究通過一組基于 ViT 模型的實驗驗證了這種方法，其中有 12 個塊用于快速實驗。根據(jù)先前基于 transformer 的工作，研究者選擇了四個嵌入維度，范圍從 256 到 768。詳細配置結果如下表所示：

下圖 5 展示了嵌入維數(shù)對生成的跨層自注意力圖相似性的影響。可見，隨著嵌入維數(shù)的增加，相似注意力圖的數(shù)目減少。然而，模型尺寸也迅速增加。

Re-attention

上文證明了在高維、尤其在深層網絡中，不同 transformer 塊之間注意力圖存在相似性。然而，研究者發(fā)現(xiàn)同一 transformer 塊不同 head 的注意力圖相似性相當小，如上圖 3(c) 所示。顯然，來自同一自注意力層的不同 head 關注輸入 token 的不同方面。基于這一觀察，研究者建議建立 cross-head 通信來重新生成注意力圖，訓練 deep ViT 性能會更好。

具體而言，該研究以 head 的注意力圖為基礎，通過動態(tài)地聚合它們來生成一組新的注意力圖。采用一個變換矩陣和 multi-head attention maps 相乘來得到新的 map，這個變換矩陣是可學習的。公式如下：

實驗

在實驗部分，研究者首先通過實驗進一步證明注意力崩潰問題，然后通過大量的控制變量實驗來證明 Re-attention 方法的優(yōu)勢。該研究通過將 Re-attention 融入到 transformer 中，設計了兩個改進版 ViT，并命名為 DeepViT。最后將 DeepViT 與具有 SOTA 性能的模型進行了比較。

注意力崩潰分析

當模型越深時，越深區(qū)塊的注意力圖就越相似。這意味著在 deep ViT 模型上添加更多的塊可能不會提高模型性能。為了進一步驗證這一說法，研究者設計了一個實驗，以重用在 ViT 早期塊上計算的注意力圖并替換它之后的注意力圖。結果如下表 3 所示：

Re-attention 機制

Re-attention 與 Self-attention 對比。研究者首先評估了 Re-attention 有效性，直接將 ViT 中的 self-attention 替換為 Re-attention 模塊。下表 4 展示了在 ImageNet 數(shù)據(jù)集上，不用數(shù)量 transformer 塊時的 Top-1 準確率對比：

表 4：ImageNet Top-1 準確率。

與 SOTA 模型對比

研究者在 Re-attention 基礎上設計了兩個 ViT 變體，即 DeepViT-S 與 DeepViT-L，兩個變體分別具有 16 和 32 個 transformer 塊。對于這兩個模型，Re-attention 均替代 self-attention。

為了得到和其他 ViT 模型相似的參數(shù)量，研究者相應地調整了嵌入維度，結果如下表 6 所示：DeepViT 模型在參數(shù)量更少的情況下實現(xiàn)了比最近 CNN 和 ViT 模型更好的 Top-1 準確率。

????

現(xiàn)在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的如何做深ViT模型？NUS、字节：引入Re-attention机制，实现强大性能的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。