作者 |?陳萍、杜偉

來源 |?機(jī)器之心

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

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

他們是如何得出這一結(jié)論的呢？

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

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

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

ViT 模型

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

具有 N 個(gè) transformer 塊的原版 ViT 模型與該研究所提 DeepViT 模型的結(jié)構(gòu)對比。

注意力崩潰

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

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

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

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

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

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

DeepViT 中的 Re-attention

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

高維空間中的自注意力

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

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

Re-attention

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

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

實(shí)驗(yàn)

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

注意力崩潰分析

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

Re-attention 機(jī)制

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

表 4：ImageNet Top-1 準(zhǔn)確率。

與 SOTA 模型對比

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

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

????

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

進(jìn)入知乎首頁搜索「PaperWeekly」

點(diǎn)擊「關(guān)注」訂閱我們的專欄吧

關(guān)于PaperWeekly

PaperWeekly 是一個(gè)推薦、解讀、討論、報(bào)道人工智能前沿論文成果的學(xué)術(shù)平臺。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號后臺點(diǎn)擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。