蘋果讓Transformer拋棄注意力機制，一切只為效率，項目已開源丨華人一作

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

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

所有的這一切，都來源于這篇論文中的Transformer模型所采用的注意力（Attention）機制。

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

△“Attention Is All You Need作者：拳頭硬了”

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

為什么要“拋棄”注意力？

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

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

因此在大規(guī)模的語義環(huán)境里，Transformer的效率就顯得不高。

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

△對比現有的Transformer變體，AFT的時間復雜度是最低的

算法設計

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

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

這一新操作的內存復雜度、文本規(guī)模、特征維度都是線性的，這就使輸入大小和模型尺寸互相兼容。

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

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

觀察圖可以看到，標準Transformer表現出了比較強烈的局部模式，這就激發(fā)了AFT的一種變體：AFT-local。

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

在將空間權值共享（卷積）的理念也納入考慮后，研究者還得到了另一種變體：AFT-conv。

在將K的維度與頭（head）個數聯(lián)系起來后，就能使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”之風發(fā)出靈魂拷問：所以你們到底需要什么？

△“不需要卷積，不需要注意力……啥你都不需要”

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

一作為華人學者

這篇論文的一作翟雙飛曾在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/