Transformer優化之稀疏注意力

張雨石發表於2020-10-01

最近要開始使用Transformer去做一些事情了,特地把與此相關的知識點記錄下來,構建相關的、完整的知識結構體系。

以下是要寫的文章,文章大部分都發布在公眾號【雨石記】上,歡迎關注公眾號獲取最新文章。

Overall

和Linformer,Longformer類似,Sparse Attention也是為了解決Transformer模型隨著長度的增加,Attention部分所佔用的記憶體和計算呈平方比增加的問題。

但Sparse Attention在時間上比Linformer,Longformer等要早,所以效果上並沒有它們好。而另一點與其他模型不同的是,論文[1]是在影像、音樂、文字上分別做的實驗,而非其他論文所使用的的GLUE task和QA task。

影像生成問題

Transformer的Decoder部分(應用在GPT上)是一個自迴歸模型,即基於前面的context去生成下一個位置上的值,然後生成的值也被納入context中去做下一次預測,以此類推。

依照這個模式,可以生成的不僅僅是文字,還可以是影像。可以把影像的畫素點按照從上到下從左到右的方式當成一個序列,然後在序列上去做自迴歸。

論文首先構造了一個128層的全注意力網路,並在Cifar10生成問題上進行了訓練。如下圖所示,底部的黑色部分表示還沒有生成的部分,白色凸顯的部分則是注意力權重高的地方。下圖是比較低的層次的注意力,可以看到,低層次的時候主要關注的還是區域性區域的部分。

在這裡插入圖片描述

而在第19層和20層,Attention學習到了橫向和縱向的規律。

在這裡插入圖片描述

進一步,還有可能學習到和資料本身相關的attention。比如下圖,第二列第二張學習到了鳥的邊緣。

在這裡插入圖片描述

但無論如何,注意力權重高的地方只佔一小部分,這就為稀疏注意力提供了資料上的支援。

Self-Attention的分解

為了實現注意力,將計算Attention的公式進行重寫。如下圖所示,重寫後的公式定義了一個集合S,裡面儲存的是需要關注的位置的索引。而同樣的,Wq, Wk, Wv是針對q,k,v的矩陣。

Transformer優化之稀疏注意力

當定義S={j: j<=i}的時候,和正常的Attention計算沒有區別,因為都是全注意力。

而分解後的注意力則是:

  • 有p個頭的注意力
  • 對於每個頭,都定義Ai(m) 是S={j: j<=i}的子集。

而對於A的定義,期望設計出一些模式,使得 |Ai|正比於sqrt(n, 1/p)。

同時,對於A的定義,雖然直接的attention是稀疏的,但是期望某個位置可以經過一些路徑attend到之前所有位置。同時路徑長度不超過p+1,這樣保證所有原本在全注意力上能夠傳遞的訊號在稀疏注意力的框架下仍然可以有效傳遞,也把時間和空間複雜度變為了: O(n*sqrt(n, 1/p))

稀疏注意力機制

論文幫助我們定義兩種符合上述條件的注意力機制。第一種是定義兩個這樣的注意力頭:

  • Ai(1) = {t, t+1, … , i} for t = max(0, i-l)
  • Ai(2) = {j: (i-j) mod l = 0}

其中l=sqrt(n)。

這個模式如下圖所示,被稱為strided attention。這種模式在有規律的資料上很有用,比如影像,比如音訊。

Transformer優化之稀疏注意力

第二種,類似的,定義兩個注意力頭:

  • Ai(1) = {j: floor(j/l) = floor(i/l)}
  • Ai(2) = {j: j mod l ∈ {t, t+1, …, l}},其中t=l-c且c是超引數。

這種被稱為fixed attention。一般情況下,l取值為{128, 256}, c取值為{8, 16, 32}。

Transformer優化之稀疏注意力

稀疏注意力的組合

一個直接的方法來使用稀疏注意力是在每個殘差塊使用同樣的稀疏機制,在不同的塊使用不同的。

另一種方式則是在每個殘差塊使用組合的稀疏注意力,組合的方法則是把不同的A合併起來。

稀疏注意力的計算

上述所提到的稀疏注意力可以被高效在GPU上進行計算,但是需要做一下操作:

  • 切片和聚合
  • 自實現的GPU kernel來進行快速計算。

深度殘差Transformer

深層次的Transformer訓練起來十分困難,因為使用殘差的方式會比較好。

Transformer優化之稀疏注意力

公式如下,其中第k層是由k-1層加上當前層結果得到的。

Transformer優化之稀疏注意力

殘差塊的內部則是這樣計算的。

Transformer優化之稀疏注意力

總結與思考

作為解決注意力平方問題的早期論文,論文從影像生成的問題上揭示了attention的原罪,那就是其實不需要那麼密集的注意力,Top-k的注意力已經足夠可以保證效果了。而無論是Reformer還是Longformer,都是在這一方向上的擴充套件。

勤思考,多提問是Engineer的良好品德。

  1. 試解釋GPU中如何實現並行化,並實現自定義的一個計算核心。
  2. 上述的兩種稀疏注意力,你認為哪種在文字上應用的好,為什麼?

後續答案將會發布在公眾號【雨石記】上,歡迎關注。

Transformer優化之稀疏注意力

參考文獻

  • [1]. Child, Rewon, et al. “Generating long sequences with sparse transformers.” arXiv preprint arXiv:1904.10509 (2019).

相關文章