比標準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工作五年。
論文地址：
https：//arxiv。org/abs/2209。07484
參考連結：
https：//www。reddit。com/r/MachineLearning/comments/xgqwvu/r_hydra_attention_efficient_attention_with_many/
— 完 —
量子位 QbitAI · 頭條號簽約