比標(biāo)準(zhǔn)Attention快197倍！Meta推出多頭注意力機制“九頭蛇”
思路來自線性注意力的一個“矛盾點”
豐色 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
盡管Transformer已經(jīng)開始在諸多視覺任務(wù)上“大展身手”，但還有一個問題。
那就是在處理大圖像上計算比較費勁。
比如面對一個1080p的圖時，它會有超過60%的計算量都耗費在了創(chuàng)建和應(yīng)用注意矩陣上。
究其原因，主要是因為自注意力頭的數(shù)量是token的平方，而token的數(shù)量又與圖形大小呈二次方的關(guān)系。
那能怎么辦呢？
好消息是——
現(xiàn)在Meta搗鼓出了一種多頭注意力操作方法，可以做到足足比標(biāo)準(zhǔn)注意力快197倍！
而且在提高計算速度的同時，它也不會犧牲準(zhǔn)確率，有時甚至還能將準(zhǔn)確率提高1-2個點。
具體怎么回事兒？
思路來源一個“矛盾點”
這個方法名叫Hydra Attention，主要針對Vision Transformer。
（“Hydra”有“九頭蛇”之義，來自希臘神話。）
Hydra Attention的思路源于線性注意力中的一種有點矛盾的點：
使用標(biāo)準(zhǔn)的多頭自注意力，再向模型中添加更多頭可以保持計算量不變。
而在線性注意力中改變操作順序后，增加更多的頭實際上還會降低層的計算成本。
于是，作者通過將模型中的注意力頭數(shù)量設(shè)置成特征（feature）數(shù)，創(chuàng)建出一個token和feature的計算都是線性的注意力模塊，從而把上述特性發(fā)揮到極致。
具體來說：
當(dāng)標(biāo)準(zhǔn)自注意力頭是token數(shù)的平方（O(T2D)）時，通過使用可分解核（decomposable kernel），我們重新安排操作順序，讓注意力頭的數(shù)量變?yōu)樘卣鱀的平方。
然后再使用Hydra Trick，最大化注意力頭H的數(shù)量，讓H=D，最終就可以化為一個在空間和時間上的O（TD）簡化操作了。
其中，Hydra Trick的依據(jù)見下圖：
作者在ImageNet-1k上訓(xùn)練了具有不同頭數(shù)的DeiT-B模型，包括使用標(biāo)準(zhǔn)自注意力（藍(lán)線，基于softmax）和使用多頭線性注意（紅線，基于余弦相似性）的。
前者在H>96、后者在H<3時出現(xiàn)內(nèi)存不足的情況。
當(dāng)他們往模型中添加更多的頭時，Softmax注意力模型的準(zhǔn)確性似乎會崩潰，而多頭線性注意力仍可以保持一致，因此就有了上述操作。
（需要注意的是，H必須除以D=768。）
速度快197倍，準(zhǔn)確率還能更上層樓
來看看Hydra Attention交出的成績單。
可以看到，Hydra注意力比標(biāo)準(zhǔn)注意力快197倍（T=197）。
隨著圖像大小的增加，它顯著提高了模型的FLOP數(shù)，在創(chuàng)建和應(yīng)用注意力矩陣所占的計算量上也始終只有0.02%。
除此之外，作者發(fā)現(xiàn)，使用不同的內(nèi)核，Hydra Attention還能讓模型的準(zhǔn)確率提高大約兩個百分點。
或者用Hydra Attention替換特定的注意力層，也能將模型的精度在ImageNet上提高1%或者與基線維持不變。
當(dāng)然，最多可替換8層。
另外，作者表示，這個方法應(yīng)該可以擴展到NLP領(lǐng)域，不過他們還沒試。
作者介紹
這篇成果已入選ECCV 2022 Workshop。
作者一共5位，分別來自Meta AI和佐治亞理工學(xué)院。
其中華人3名，分別是：
Cheng-Yang Fu，本碩畢業(yè)于清華大學(xué)，博士畢業(yè)于美國北卡羅來納大學(xué)教堂山分校，現(xiàn)在是Meta計算機視覺方面的研究科學(xué)家。
Xiaoliang Dai，本科畢業(yè)于北大，博士畢業(yè)于普林斯頓大學(xué)，同就職于Meta。
Peizhao Zhang，本碩畢業(yè)于中山大學(xué)，博士于美國德克薩斯A&M大學(xué)，已在Meta工作五年。
論文地址：
https://arxiv.org/abs/2209.07484
參考鏈接：
https://www.reddit.com/r/MachineLearning/comments/xgqwvu/r_hydra_attention_efficient_attention_with_many/
- 北大開源最強aiXcoder-7B代碼大模型！聚焦真實開發(fā)場景，專為企業(yè)私有部署設(shè)計2024-04-09
- 剛剛，圖靈獎揭曉！史上首位數(shù)學(xué)和計算機最高獎“雙料王”出現(xiàn)了2024-04-10
- 8.3K Stars!《多模態(tài)大語言模型綜述》重大升級2024-04-10
- 谷歌最強大模型免費開放了！長音頻理解功能獨一份，100萬上下文敞開用2024-04-10
相關(guān)閱讀
谷歌大模型研究陷重大爭議：訓(xùn)練數(shù)據(jù)之外完全無法泛化？網(wǎng)友：AGI奇點推遲了
網(wǎng)友找出論文中更多關(guān)鍵卻被忽略的細(xì)節(jié)，比如只做了GPT-2規(guī)模的試驗等
