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

博雯 發自 凹非寺
量子位 報道 | 公眾號 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）的需要，展現出了出色的效率。

△對比現有的Transformer變體，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/