天天看點

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

機器之心報道

機器之心編輯部

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

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

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

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

标準的稀疏化技術。

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

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

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

論文連結: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 對該文章進行了解讀。

00:00/00:00倍速

廣義的注意力機制

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

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

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

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

新算法 FAVOR+:通過矩陣相關性實作快速注意力

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

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

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

左:标準單向注意力需要 mask 注意力矩陣以獲得其下三角部分;右:LHS 上的無偏近似可以通過字首和獲得,其中用于 key 和值向量的随機特征圖的外積(outer-product)字首和實作動态建構,并通過 query 随機特征向量進行左乘計算,以在最終矩陣中獲得新行(new row)。

性能

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

在以時間(T)和長度(L)為度量的雙對數坐标軸中,正常 Transformer 模型的雙向 timing。

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

在 One Billion Word Benchmark (LM1B) 資料集上,研究者将原始預訓練 Transformer 的權重遷移至 Performer 模型,使得初始非零準确度為 0.07(橙色虛線)。但在微調之後,Performer 的準确度在很少的梯度步數之後迅速恢複。

應用示例:蛋白質模組化

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

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

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

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

空間注意力機制sam_自己挖坑自己填,谷歌大改Transformer注意力

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

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

參考連結:http://ai.googleblog.com/2020/10/rethinking-attention-with-performers.html