首頁 > 科技

蘋果讓Transformer拋棄注意力機制,一切只為效率丨開源

2021-06-10 16:14:13

博雯 發自 凹非寺量子位 報道 | 公眾號 QbitAI

2017年,一句Attention Is All You Need引爆了整個NLP圈,甚至在近些年還有踢館計算機視覺領域的趨勢。

所有的這一切,都來源於這篇論文中的Transformer模型所採用的注意力(Attention)機制。

但現在,你能想象一個不再需要注意力的Transformer嗎?

△「Attention Is All You Need作者:拳頭硬了」

今天就來一起看看這項來自蘋果的研究:無注意力Transformer

為什麼要「拋棄」注意力?

自注意力機制的應用,使得Transformer突破了RNN不能平行計算的限制,並大大提高了捕捉長期依賴關係的能力。

但與此同時,這種機制也意味著需要額外執行具有時間和空間複雜性的注意力操作,這就為操作帶來了二次成本。

因此在大規模的語義環境裡,Transformer的效率就顯得不高。

而這種無注意力Transformer(AFT)則消除了點積自注意力(dot-product self attention)的需要,展現出了出色的效率。

△AFT的時間複雜度最低

演算法設計

這項研究的全名為Attention Free Transformer,它類似於標準的點積注意力演算法,同樣由查詢向量Q,被查向量K,內容向量V相互作用而成。

但不同的是,AFT中的K和V首先與一組學習得到的位置偏差(position bias)結合,然後再進行同位元素對應相乘(element-wise multiplication)。

這一新操作的記憶體複雜度、文字規模、特徵維度都是線性的,這就使輸入大小和模型尺寸互相相容。

同時,研究者還在AFT的基礎上提出了幾種模型變體。

比如,將一張在ImagenetNet驗證集上進行預訓練的Vit圖(由12層組成,每層有6個頭,空間大小為14×14)視覺化後,最終產生了一組尺寸為12×6×27×27的注意力圖。

觀察圖可以看到,標準Transformer表現出了比較強烈的局部模式,這就激發了AFT的一種變體:AFT-local

AFT-local在保持全局連通性的同時利用了局部性概念,將學習到的位置偏差固定在一個局部區域中。

在將空間權值共享(卷積)的理念也納入考慮後,研究者還得到了另一種變體:AFT-conv

在將K的維度與頭(head)個數聯絡起來後,就能使AFT-conv可依賴於深度可分離卷積、全局池化和元素操作來實現。

表現如何?

在影象自迴歸建模(Image Autoregressive Modeling)上,每個字元的位元數(bpc)越低越好。在這一點上,AFT做到了SOTA

而對於老本行語言建模(Language Modeling)問題,研究者則基於Enwik8進行字元級語言建模。

而結果是,AFT僅僅消耗了1/3的記憶體,就提供了提供了44%的速度提升,且在bpc上與完整的Transformer的距離只相差0.024,在參數、速度、記憶體和效能方面均取得了最佳平衡

研究者也將AFT測試了影象分類( Image Classification)的任務。

最終AFT的變體AFT-full在更好的記憶體佔用和相似的速度下,取得了與基線Transformer DeiT相當的效能。

而在參數數量相似或更少的情況下,AFT-conv的準確率更高

與Lambda網路相比,所有的AFT變體都達到了相當或更好的精度,且速度相當,記憶體佔用更小。

「所以我們到底需要什麼?」

對於這一丟棄了注意力機制的Transformer,有人表示深度學習研究果然就是個圈,指不定過幾年KNN也能再次偉大。

也有人對這種「Not Need」之風發出靈魂拷問:需要卷積,要注意力,啥你都不需要……所以你們到底需要什麼?

最終,機器學習的研究者們決定用魔法打敗魔法:

一作為華人學者

這篇論文的一作翟雙飛曾在IBM Research實習過9個月,2017年6月加入蘋果,並在其機器學習研究部門任職至今。

他本科就讀於中國科學技術大學的電子工程與資訊科學學院,並於2017年在賓漢姆頓大學獲得了電腦科學博士學位。

下載地址:https://github.com/rish-16/aft-pytorch論文地址:https://arxiv.org/abs/2105.14103參考連結:https://www.reddit.com/r/MachineLearning/comments/npmq5j/r_an_attention_free_transformer/


IT145.com E-mail:sddin#qq.com