Performer:用隨機投影將Attention的複雜度線性化

©PaperWeekly 原創 · 作者|蘇劍林

單位|追一科技

研究方向|NLP、神經網絡

Attention 機制的 複雜度是一個老大難問題了,改變這一複雜度的思路主要有兩種:一是走稀疏化的思路,比如我們以往介紹過的 Sparse Attention 以及 Google 前幾個月搞出來的 Big Bird [1] ,等等;二是走線性化的思路,這部分工作我們之前總結在線性 Attention 的探索:Attention 必須有個 Softmax 嗎?中,讀者可以翻看一下。

本文則介紹一項新的改進工作 Performer,出自 Google 的文章 Rethinking Attention with Performers,它的目標相當霸氣:通過隨機投影,在不損失精度的情況下,將 Attention 的複雜度線性化。

論文標題:

Rethinking Attention with Performers

論文鏈接:

https://arxiv.org/abs/2009.14794

說直接點,就是理想情況下我們可以不用重新訓練模型,輸出結果也不會有明顯變化,但是複雜度降到了 !看起來真的是“天上掉餡餅”般的改進了,真的有這麼美好嗎?

Attention

我們知道,Attention 的一般定義爲:

對於標準的 Scaled-Dot Attention 來說, (有時候指數部分還會多個縮放因子,這裏我們就不顯式寫出來了),將整個序列的運算寫成矩陣形式就是:

我們主要關心 Self Attention 場景,所以一般有 。在上式中, 這一步相當於要對 個向量對做內積,得到 個實數,因此不管時間還是空間複雜度都是 的。

而對於線性 Attention 來說, ,其中 是值域非負的激活函數。這樣一來,Attention 的核心計算量(式(1)中的分子部分)就變成了:

上式左端的複雜度依然是 的,由於矩陣乘法滿足結合律,我們可以先算後面兩個矩陣的乘法,這樣複雜度就可以降爲 了,詳細介紹還是請讀者翻看之前的文章線性 Attention 的探索:Attention 必須有個 Softmax 嗎?,這裏就不過多展開了。

Performer

現在我們就可以進入到 Performer 的介紹了,開頭已經說了,Performer 的出發點還是標準的 Attention,所以在它那裏還是有 ,然後它希望將複雜度線性化,那就是需要找到新的 ,使得:

如果找到合理的從 到 的映射方案,便是該思路的最大難度了。

2.1 漂亮的隨機映射

Performer 的最大貢獻就在於,它找到了一個非常漂亮的映射方案:

我們先分析一下上式究竟說了什麼。第一個等號意味着左右兩端是恆等的,那也就意味着,只要我們從標準正態分佈 中採樣無窮無盡的 ,然後算出 的平均,結果就等於 ,寫成積分形式就是:

當然實際情況下我們只能採樣有限的 m 個,因此就得到了第二個約等號,它正好可以表示爲兩個 m 維向量的內積的形式,這正是我們需要的 !

所以,藉助這個近似,我們就可以將兩個 d 維向量的內積的指數,轉化爲了兩個 m 維向量的內積了,從而理論上來說,我們就可以將原來 head_size 爲 d 的標準 Attention,轉化爲 head_size 爲 m 的線性 Attention 了,這便是整篇論文的核心思路。

2.2 推導過程討論

可能有些讀者會比較關心式(5)的來源,這裏展開討論一下,當然如果不關心的讀者可以直接跳過這一節。儘管直接通過計算積分可以驗證式(5)是成立的,但對於任意定義的 ,是否可以找到類似的線性近似?

下面我們將會證明,類似的線性化方案可以通過一個一般化的流程找到,只不過得到的結果可能遠遠不如式(5)漂亮有效。

具體來說,對於任意的 ,我們可以改寫爲:

然後我們可以對 做傅立葉變換:

至於爲什麼要先乘以 ,那是因爲直接對 做傅立葉變換的結果可能不好看甚至不存在,乘一個函數之後可能就可以了,比如可以讓 ,只要 足夠大,就可以讓很多 都完成傅立葉變換了。

接着我們執行逆變換,並代回原式,就得到:

如果我們能算出 並完成歸一化,那麼它就可以成爲一個可以從中採樣的分佈,從中可以採樣出隨機向量 來,然後近似轉化爲 組成的向量序列的內積。

當然,這裏的運算可能涉及到虛數,而我們一般只是處理實數,但這問題不大,我們可以用歐拉公式 展開,整個運算過程只保留實部即可,形式不會有太大的變化。

理論上來說,整套流程可以走下來,不會有什麼困難,但是相比式(5),存在的問題是:1)現在需要採樣兩組隨機變量 ,會擴大估計的方差;2)最終保留實部後,得到的將會是 的組合形式,其結果無法保證非負性,需要額外的 clip 來處理。

而式(5)的特別之處在於, 可以改寫爲:

所以只需要轉化爲單個變量 的問題,而 的傅立葉變換正好是 ,所以做逆變換我們有:

到這裏,如果直接取實部展開,得到的也是 的組合,這就是原論文說的 形式的投影方案。不過,有一個更加巧妙的性質可以改變這一切!注意到上式是一個恆等式,所以我們可以左右兩邊都置換 ,結果是:

這便是式(5)。置換 使得上述左邊保持不變,並且右邊完全脫離虛數還保持了非負性,真的是集衆多巧合於一身,“只此一家,別無分號”的感覺~

2.3 正交化降低方差

除了提出式(5)來對標準 Attention 進行線性化外,原論文還做了進一步的增強。在式(5), 是獨立重複地從 中採樣出來的,而原論文則指出,如果將各個 正交化,能有效地降低估算的方差,提高單次估算的平均精度。

注意,這裏的正交化指的是保持 的模長不變,僅僅是對其方向進行施密特正交化。這個操作首先提出在同樣是 Google 的論文 The Unreasonable Effectiveness of Structured Random Orthogonal Embeddings [2] 中,而 Performer 在其論文附錄裏邊,足足花了 6 頁紙來論證這個事情。這裏照搬 6 頁證明顯然是不適合的,那對於我們來說,該怎麼去理解這個策略比較好呢?

其實,這個策略有效的最根本原因,是採樣分佈 的各向同性,即其概率密度函數 只依賴於 的模長 ,所以它在方向上是均勻的。

而如果我們要降低估算的方差,那麼就應該要降低採樣的隨機性,使得采樣的結果更爲均勻一些。而各個向量正交化,是方向上均勻的一種實現方式,換句話說,將各個  正交化促進了採樣結果的均勻化,從而降低估算的方差。此外,正交化操作一般只對 有效,如果 ,原論文的處理方式是每 d 個向量爲一組分別進行正交化。

我們可以聯想到,正交化操作只是讓採樣的方向更加均勻化,如果要做得更加徹底些,可以讓採樣的模長也均勻化。具體來說,將標準正態分佈變換爲 d 維球座標 [3] 得到其概率微元爲:

其中 代表在 d 維球面上的積分微元。上式就顯示出,標準正態分佈在方向是均勻的,模長的概率密度函數正比於 ,我們定義它的累積概率函數:

如果要採樣 m 個樣本,那麼讓 ,從中解出 m 個 作爲模長就行了。

性能與效果

理論的介紹就到這裏了,其實已經展開得都多了,一般來說只要對線性 Attention 有所瞭解的話,那麼只要看到式(5),就能很快理解 Performer 的要點了,剩下的都是錦上添花的內容,不影響全文的主線。接下來我們就主要看對 Performer 的評測。

3.1 原論文的評測

我們先來看看原論文的評測,其實它的評測還是很豐富的,但是有點亂,對於主要關心 NLP 的讀者來說可能還有點莫名其妙。

一開始是速度上的評測,這個不意外,反正就是序列長了之後 Performer 比標準的 Transformer 有明顯優勢:

▲ Performer與標準Transformer的速度對比(實線:Performer,虛線:標準Transformer)

接着,是近似程度的比較,說明了採樣正交化的有效性,以及所提的式(5)相比舊的基於 函數的形式的精確性:

▲ 左圖:比較採樣向量正交化的有效性;有圖:比較Performer所用的近似與舊的基於三角函數的近似的精確性

那能不能達到我們一開始的預期目標——不用重新訓練已訓練好的模型呢?很遺憾,不行,原論文做了兩個實驗,顯示 Performer 直接加載 Transformer 的權重不能恢復已有的結果,但經過 finetune 後可以迅速恢復回覆。至於爲什麼一開始沒有很好地回覆,論文沒有做展開分析。

▲ Performer加載已訓練好的Transformer權重實驗

最後,論文還做了蛋白質序列和圖像的實驗,證明 Performer 對於長序列是很有效的,特別地,至少比 Reformer 有效,全文的實驗差不多就是這樣,內容很多,但有點找不到北的感覺。

3.2 其他論文的評測

也許是自己的同事都看不下去了,後來 Google 又出了兩篇論文 Efficient Transformers: A Survey [4] 和 Long Range Arena: A Benchmark for Efficient Transformers [5] ,系統地評測和比較了已有的一些改進 Transformer 效率的方法,其中就包含了 Performer。

相比之下,這兩篇論文給出的結果就直觀多了,簡單幾個圖表,就把各個模型的位置給表達清楚了。

▲ 各種高效Transformer分門別類

▲ 各類改進版Transformer比較。其中decode那一欄指的是能否mask掉未來信息,用於做語言模型

▲ 各個Transformer模型的“效果-速度-顯存”圖,縱軸是效果,橫軸是速度,圓圈的大小代表所需要的顯存。理論上來說,越靠近右上方的模型越好,圓圈越小的模型越好

更詳細的評測信息,大家自行去看這兩篇論文就好。

問題與思考

看起來 Performer 是挺不錯的,那是不是說我們就可以用它來替代 Transformer 了?並不是,縱然 Performer 有諸多可取之處,但仍然存在一些無法解決的問題。

首先,爲了更好地近似標準 Transformer,Performer 的 m 必須取得比較大,至少是 m > d,一般是 d 的幾倍,這也就是說,Performer 的 head_size 要比標準 Transformer 要明顯大。雖然理論上來說,不管 m 多大,只要它固定了,那麼 Performer 關於序列長度的複雜度是線性的,但是 m 變大了,在序列長度比較短時計算量是明顯增加的。

換句話說,短序列用 Performer 性能應該是下降的,根據筆者的估計,只有序列長度在 5000 以上,Performer 纔會有比較明顯的優勢。

其次,目前看來 Performer(包括其他的線性 Attention)跟相對位置編碼是不兼容的,因爲相對位置編碼是直接加在 Attention 矩陣裏邊的,Performer 連 Attention 矩陣都沒有,自然是加不了的。此外,像 UniLM 這種特殊的 Seq2Seq 做法也做不到了,不過普通的單向語言模型還是可以做到的。

總之, 的 Attention 矩陣實際上也帶來了很大的靈活性,而線性 Attention 放棄了這個 Attention 矩陣,也就放棄了這種靈活性了。

最後,也是筆者覺得最大的問題,就是 Performer 的思想是將標準的 Attention 線性化,所以爲什麼不乾脆直接訓練一個線性 Attention 模型,而是要想標準 Attention 靠齊呢?

從前面的最後一張圖來看,Performer 並沒有比 Linear Transformer 有大的優勢(而且筆者覺得最後一張圖的比較可能有問題,Performer 效果可能比 Linear Transformer 要好,但是速度怎麼可能還超過了 Linear Transformer?Performer 也是轉化爲 Linear Transformer 的,多了轉化這一步,速度怎能更快?),因此 Performer 的價值就要打上個問號了,畢竟線性 Attention 的實現容易得多,而且是通用於長序列/短序列,Performer 實現起來則麻煩得多,只適用於長序列。

全文的總結

本文主要介紹了 Google 的新模型 Performer,這是一個通過隨機投影將標準 Attention 的複雜度線性化的工作,裏邊有不少值得我們學習的地方,最後彙總了一下各個改進版 Transformer 的評測結果,以及分享了筆者對 Performer 的思考。

參考文獻

[1] https://arxiv.org/abs/2007.14062

[2] https://arxiv.org/abs/1703.00864

[3] https://en.wikipedia.org/wiki/N-sphere

[4] https://arxiv.org/abs/2009.06732

[5] https://arxiv.org/abs/2011.04006

更多閱讀

#投 稿 通 道#

 讓你的論文被更多人看到 

如何才能讓更多的優質內容以更短路徑到達讀者羣體,縮短讀者尋找優質內容的成本呢?答案就是:你不認識的人。

總有一些你不認識的人,知道你想知道的東西。PaperWeekly 或許可以成爲一座橋樑,促使不同背景、不同方向的學者和學術靈感相互碰撞,迸發出更多的可能性。 

PaperWeekly 鼓勵高校實驗室或個人,在我們的平臺上分享各類優質內容,可以是最新論文解讀,也可以是學習心得技術乾貨。我們的目的只有一個,讓知識真正流動起來。

???? 來稿標準:

• 稿件確係個人原創作品,來稿需註明作者個人信息(姓名+學校/工作單位+學歷/職位+研究方向) 

• 如果文章並非首發,請在投稿時提醒並附上所有已發佈鏈接 

• PaperWeekly 默認每篇文章都是首發,均會添加“原創”標誌

???? 投稿郵箱:

• 投稿郵箱:[email protected] 

• 所有文章配圖,請單獨在附件中發送 

• 請留下即時聯繫方式(微信或手機),以便我們在編輯發佈時和作者溝通

????

現在,在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

關於PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域,歡迎在公衆號後臺點擊「交流羣」,小助手將把你帶入 PaperWeekly 的交流羣裏。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章