[19.04] Sparse Transformer
稀疏注意力機制
Generating Long Sequences with Sparse Transformers
Transformer 的自注意力機制的計算複雜度為 。
當任何一個演算法的計算複雜度出現這種情況時,都意味著其中有非常大量的優化空間。
OpenAI 在基礎研究上做了非常多的努力,其中一個就是本篇論文所提出的 Sparse Transformer。
定義問題
為了研究原本的 Transformer 模型在實際場景中,注意力機制是如何運作的,研究者首先做了一個影像生成模型。
這個模型是基於 Decoder-Only 的 Transformer 架構,將影像切分為 Patch,然後將這些 Patch 作為輸入序列,並透過 Transformer 生成影像。
模型有 128 層,在 CIFAR-10 資料集上訓練,實驗結果如上圖,這張圖在這篇論文中非常重要,這是本研究的起點。
在上圖中,你會看到一些白色的光斑,這是作者將模型的注意力圖可視化後得到的結果。黑色的部分是還沒有解碼的部分,所以你看到黑色的交界處,就是此刻的「正在解碼的 Token」,而白色的光斑這是這個「正在解碼的 Token」所關注的範圍。
這裡分為幾個部分進行分析:
最前面的幾層
在最前面的幾層,模型的注意力範圍環繞在「正在解碼的 Token」周圍,這像是卷積神經網絡的感受野。
不是我們擷取的影像解析度很低,而是因為 CIFAR-10 的影像解析度就是 32x32。
第 19~20 層
在第 19 層和第 20 層,網路學會了將注意力分散到行注意力和列注意力上,從而有效地分解全局注意力計算。
不特定幾層
在有些層,注意力範圍非常大,這意味著這些層有全局的視野。
第 64~128 層
這里佔據了模型一半的層數,在後半部分。
可以發現,這些層的注意力高度稀疏,幾乎沒有被激活。
OpenAI 有提供動態圖片展示:
這種實驗結果,帶給我們的啟示是:
- 在 Transformer 中,注意力機制並不是所有 Token 都需要關注所有 Token,不斷地堆疊全局注意力是非常浪費的行為。
解決問題
解決問題的思路就是:我們需要一個機制,讓模型能夠選擇性地關注某些 Token,而不是所有 Token。
在上圖 (a) 中,指的是原本的 Transformer,而 (b) 和 (c) 是作者嘗試使用的 Sparse Transformer 的架構。上面的小圖指的是在圖上的注意力分佈,大圖是把這些注意力攤平成 N x N 的矩陣。
你可能會注意到上三角的部分呈現空白,因為這是一個 Decoder-Only 的模型,所以每個 Token 只需要關注自己之前的 Token。
解題思路 1: 固定步長
也就是上圖 (b) 中的架構,這個架構參考了剛才所提到的實驗結果:
- 我們可能需要關注的是前面幾個 Token,即 Row 的部分。
- 我們可能需要關注的是前 個 Token,即 Column 的部分(例如前 6, 18, 24,... 個)。