比標準Attention快197倍！Meta推出多頭注意力機制“九頭蛇”

思路來自線性注意力的一個“矛盾點”

豐色 發自 凹非寺

量子位 | 公眾號 QbitAI

儘管Transformer已經開始在諸多視覺任務上“大展身手”，但還有一個問題。

那就是在處理 大影象 上計算比較費勁。

比如面對一個1080p的圖時，它會有 超過60%的計算量 都耗費在了建立和應用注意矩陣上。

究其原因，主要是因為自注意力頭的數量是token的平方，而token的數量又與圖形大小呈二次方的關係。

那能怎麼辦呢？

好訊息是——

現在Meta搗鼓出了一種 多頭注意力 操作方法，可以做到足足 比標準注意力快197倍 ！

而且在提高計算速度的同時，它也不會犧牲準確率，有時甚至還能將準確率提高1-2個點。

具體怎麼回事兒？

思路來源一個“矛盾點”

這個方法名叫 Hydra Attention ，主要針對Vision Transformer。

（“Hydra”有“九頭蛇”之義，來自希臘神話。）

Hydra Attention的思路源於線性注意力中的一種 有點矛盾 的點：

使用標準的多頭自注意力，再向模型中新增更多頭可以保持計算量不變。

而線上性注意力中改變操作順序後，增加更多的頭實際上還會降低層的計算成本。

於是，作者通過將模型中的 注意力頭數量設定成特徵 （feature） 數 ，創建出一個token和feature的計算都是線性的注意力模組，從而把上述特性發揮到極致。

具體來說：

當標準自注意力頭是token數的平方（O(T2D)）時，通過使用可分解核（decomposable kernel），我們重新安排操作順序，讓注意力頭的數量變為特徵D的平方。

然後再使用Hydra Trick，最大化注意力頭H的數量，讓H=D，最終就可以化為一個在空間和時間上的O（TD）簡化操作了。

其中，Hydra Trick的依據見下圖：

作者在ImageNet-1k上訓練了具有不同頭數的DeiT-B模型，包括使用標準自注意力（藍線，基於softmax）和使用多頭線性注意（紅線，基於餘弦相似性）的。

前者在H>96、後者在H<3時出現記憶體不足的情況。

當他們往模型中新增更多的頭時，Softmax注意力模型的準確性似乎會崩潰，而多頭線性注意力仍可以保持一致，因此就有了上述操作。

（需要注意的是，H必須除以D=768。）

速度快197倍，準確率還能更上層樓

來看看Hydra Attention交出的成績單。

可以看到，Hydra注意力比標準注意力快197倍（T=197）。

隨著影象大小的增加，它顯著提高了模型的FLOP數，在建立和應用注意力矩陣所佔的計算量上也 始終只有0.02% 。

除此之外，作者發現，使用不同的核心，Hydra Attention還能讓模型的準確率提高大約兩個百分點。

或者用Hydra Attention替換特定的注意力層，也能將模型的精度在ImageNet上提高1%或者與基線維持不變。

當然，最多可替換8層。

另外，作者表示，這個方法應該可以擴充套件到NLP領域，不過他們還沒試。

作者介紹

這篇成果已入選 ECCV 2022 Workshop 。

作者一共5位，分別來自Meta AI和佐治亞理工學院。

其中華人3名，分別是：

Cheng-Yang Fu，本碩畢業於清華大學，博士畢業於美國北卡羅來納大學教堂山分校，現在是Meta計算機視覺方面的研究科學家。

Xiaoliang Dai，本科畢業於北大，博士畢業於普林斯頓大學，同就職於Meta。

Peizhao Zhang，本碩畢業於中山大學，博士於美國德克薩斯A&M大學，已在Meta工作五年。

論文地址：

http://arxiv.org/abs/2209.07484

參考連結：

http://www.reddit.com/r/MachineLearning/comments/xgqwvu/r_hydra_attention_efficient_attention_with_many/

版權所有，未經授權不得以任何形式轉載及使用，違者必究。