源 | 機器之心

---

導讀

考慮到 Transformer 對于機器學習最近一段時間的影響，這樣一個研究就顯得異常引人注目了。

Transformer 有著巨大的內存和算力需求，因為它構造了一個注意力矩陣，需求與輸入呈平方關系。谷歌大腦 Krzysztof Choromanski 等人最近提出的 Performer 模型因為隨機正正交特性為注意力矩陣構建了一個無偏的估計量，可以獲得線性增長的資源需求量。這一方法超越了注意力機制，甚至可以說為下一代深度學習架構打開了思路。

自面世以來，Transformer 模型已經在多個領域取得了 SOTA 結果，包括自然語言處理、圖像處理甚至是音樂處理。眾所周知，Transformer 架構的核心是注意力模塊，它計算輸入序列中所有位置對的相似度得分。然而，隨著輸入序列長度的增加，注意力機制本身的問題也越來越突出，因為它需要二次方的計算時間來產生所有的相似度得分，用來存儲這些得分的內存大小也是如此。

針對那些需要長距離注意力的應用，部分研究者已經提出了一些速度快、空間利用率高的方法，其中比較普遍的方法是稀疏注意力。

標準的稀疏化技術

然而，稀疏注意力方法也有一些局限。首先，它們需要高效的稀疏矩陣乘法運算，但這并不是所有加速器都能做到的；其次，它們通常不能為自己的表示能力提供嚴格的理論保證；再者，它們主要針對 Transformer 模型和生成預訓練進行優化；最后，它們通常會堆更多的注意力層來補償稀疏表示，這使其很難與其他預訓練好的模型一起使用，需要重新訓練，消耗大量能源。

此外，稀疏注意力機制通常不足以解決常規注意力方法應用時所面臨的所有問題，如指針網絡。還有一些運算是無法稀疏化的，比如常用的 softmax 運算。

為了解決這些問題，來自谷歌、劍橋大學、DeepMind、阿蘭 · 圖靈研究所的研究者提出了一種新的 Transformer 架構——Performer。它的注意力機制能夠線性擴展，因此能夠在處理長序列的同時縮短訓練時間。這點在 ImageNet64 等圖像數據集和 PG-19 文本數據集等序列的處理過程中都非常有用。

論文鏈接：
https://arxiv.org/pdf/2009.14794.pdf

Performer 使用一個高效的（線性）廣義注意力框架（generalized attention framework），允許基于不同相似性度量（核）的一類廣泛的注意力機制。該框架通過谷歌的新算法 FAVOR+（ Fast Attention Via Positive Orthogonal Random Features）來實現，后者能夠提供注意力機制的可擴展低方差、無偏估計，這可以通過隨機特征圖分解（常規 softmax-attention）來表達。該方法在保持線性空間和時間復雜度的同時準確率也很有保證，也可以應用到獨立的 softmax 運算。此外，該方法還可以和可逆層等其他技術進行互操作。

研究者表示，他們相信該研究為注意力、Transformer 架構和核方法提供了一種新的思維方式。

代碼地址：
https://github.com/google-research/google-research/tree/master/performer

論文公布之后，Youtube 知名深度學習頻道 Yannic Kilcher 對該文章進行了解讀。

廣義的注意力機制

在以往的注意力機制中，分別對應矩陣行與列的 query 和 key 輸入相乘，通過 softmax 計算形成一個注意力矩陣，以存儲相似度系數。值得注意的是，這種方法不能將 query-key 生成結果傳遞給非線性 softmax 計算之后再將其分解為原始的 query 和 key。然而，將注意力矩陣分解為原始 query 和 key 的隨機非線性函數的乘積是可以的，即所謂的隨機特征（random feature），這樣就可以更加高效地對相似度信息進行編碼。

標準注意力矩陣包括每一對 entry 的相似度系數，由 query 和 key 上的 softmax 計算組成，表示為 q 和 k。

常規的 softmax 注意力可以看作是由指數函數和高斯投影定義的非線性函數的一個特例。在這里我們也可以反向推理，首先實現一些更廣義的非線性函數，隱式定義 query-key 結果中其他類型的相似性度量或核函數。研究者基于早期的核方法（kernel method），將其定義為廣義注意力（generalized attention）。盡管對于大多核函數來說，閉式解并不存在，但這一機制仍然可以應用，因為它并不依賴于閉式解。

該研究首次證明了，任意注意力矩陣都可以通過隨機特征在下游 Transformer 應用中實現有效地近似。實現這一點的的新機制是使用正隨機特征，即原始 query 和 key 的正直非線性函數，這對于避免訓練過程中的不穩定性至關重要，并實現了對常規 softmax 注意力的更準確近似。

新算法 FAVOR+

上文描述的分解允許我們以線性而非二次內存復雜度的方式存儲隱式注意力矩陣。我們還可以通過分解獲得一個線性時間注意力機制。雖然在分解注意力矩陣之后，原始注意力機制與具有值輸入的存儲注意力矩陣相乘以獲得最終結果，我們可以重新排列矩陣乘法以近似常規注意力機制的結果，并且不需要顯式地構建二次方大小的注意力矩陣。最終生成了新算法 FAVOR+。

左：標準注意力模塊計算，其中通過執行帶有矩陣 A 和值張量 V 的矩陣乘法來計算最終的預期結果；右：通過解耦低秩分解 A 中使用的矩陣 Q′和 K′以及按照虛線框中指示的順序執行矩陣乘法，研究者獲得了一個線性注意力矩陣，同時不用顯式地構建 A 或其近似。

上述分析與雙向注意力（即非因果注意力）相關，其中沒有 past 和 future 的概念。對于輸入序列中沒有注意前后 token 的單向（即因果）注意力而言，研究者稍微修改方法以使用前綴和計算（prefix-sum computation），它們只存儲矩陣計算的運行總數，而不存儲顯式的下三角常規注意力矩陣。

左：標準單向注意力需要 mask 注意力矩陣以獲得其下三角部分；右：LHS 上的無偏近似可以通過前綴和獲得，其中用于 key 和值向量的隨機特征圖的外積（outer-product）前綴和實現動態構建，并通過 query 隨機特征向量進行左乘計算，以在最終矩陣中獲得新行（new row）。

性能

研究者首先對 Performer 的空間和時間復雜度進行基準測試，結果表明，注意力的加速比和內存減少在實證的角度上近乎最優，也就是說，這非常接近在模型中根本不使用注意力機制的情況。

在以時間（T）和長度（L）為度量的雙對數坐標軸中，常規 Transformer 模型的雙向 timing。

研究者進一步證明，使用無偏 softmax 近似，該 Performer 模型在稍微進行微調之后可以向后兼容預訓練 Transformer 模型，從而在提升推理速度的同時降低能耗，并且不需要從頭訓練預先存在的模型。

在 One Billion Word Benchmark \(LM1B\) 數據集上，研究者將原始預訓練 Transformer 的權重遷移至 Performer 模型，使得初始非零準確度為 0.07（橙色虛線）。但在微調之后，Performer 的準確度在很少的梯度步數之后迅速恢復。

應用示例：蛋白質建模

蛋白質具有復雜的 3D 結構，是生命必不可少的擁有特定功能的大分子。和單詞一樣，蛋白質可以被看做線性序列，每個字符代表一種氨基酸。將 Transformers 應用于大型未標記的蛋白質序列語料庫，生成的模型可用于精確預測折疊功能大分子。正如該研究理論結果所預測的那樣，Performer-ReLU 在蛋白質序列數據建模方面表現良好，而 Performer-Softmax 與 Transformer 性能相媲美。

Performer 在蛋白質序列建模時的性能。

下面可視化一個蛋白質 Performer 模型，該模型使用基于 ReLU 的近似注意力機制進行訓練。研究者發現，Performer 的密集注意力近似有可能捕捉到跨多個蛋白質序列的全局相互作用。作為概念的證明，研究者在串聯蛋白長序列上訓練模型，這使得常規的 Transformer 模型內存過載。但由于具有良好的空間利用效率，Performer 不會出現這一問題。

左：從注意力權重估計氨基酸相似性矩陣。該模型可以識別高度相似的氨基酸對，例如 \(D,E\) 和 \(F,Y\)。

Performer 和 Transformer 在長度為 8192 的蛋白質序列上的性能。

隨著 Transformer 的頻繁跨界，越來越多的研究者開始關注其內存占用和計算效率的問題，比如機器之心前段時間介紹的《拋棄注意力，比 EfficientNet 快 3.5 倍，類 Transformer 新模型跨界視覺任務實現新 SOTA 》。在那篇文章中，研究者提出了一種名為「lambda」的層，這些層提供了一種捕獲輸入和一組結構化上下文元素之間長程交互的通用框架。類似的改進還在不斷涌現，我們也將持續關注。

參考鏈接：
http://ai.googleblog.com/2020/10/rethinking-attention-with-performers.html

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

有頂會審稿人、大廠研究員、知乎大V和妹紙

等你來撩哦~

總結

以上是生活随笔為你收集整理的谷歌大改Transformer注意力，速度大涨，显存大降！的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。