为节约而生：从标准Attention到稀疏Attention
作者丨蘇劍林
單位丨追一科技
研究方向丨NLP，神經(jīng)網(wǎng)絡(luò)
個人主頁丨kexue.fm
如今 NLP 領(lǐng)域，Attention 大行其道，當(dāng)然也不止 NLP，在 CV 領(lǐng)域 Attention 也占有一席之地（Non Local、SAGAN 等）。在 18 年初一文讀懂「Attention is All You Need」| 附代碼實現(xiàn)一文中，我們就已經(jīng)討論過 Attention 機制，Attention 的核心在于 Q,K,V 三個向量序列的交互和融合，其中 Q,K 的交互給出了兩兩向量之間的某種相關(guān)度（權(quán)重），而最后的輸出序列則是把 V 按照權(quán)重求和得到的。?
顯然，眾多 NLP & CV 的成果已經(jīng)充分肯定了 Attention 的有效性。本文我們將會介紹 Attention 的一些變體，這些變體的共同特點是——“為節(jié)約而生”——既節(jié)約時間，也節(jié)約顯存。
背景簡述
Attention is All You Need?一文討論的我們稱之為“乘性 Attention”，目前用得比較廣泛的也就是這種 Attention：?
另外還有加性 Attention，但加性 Attention 并行不大容易實現(xiàn)（或者實現(xiàn)起來占用顯存多），所以一般只用來將變長向量序列編碼為固定長度的向量（取代簡單的 Pooling），而很少用來做序列到序列的編碼。
而在乘性 Attention 中，用得最廣泛的當(dāng)數(shù) Self Attention 了，這種情況下 Q,K,V 都是同一個 X 經(jīng)過線性變換之后的結(jié)果，這樣一來輸出結(jié)果就是跟 X 一樣長的向量序列，并且能夠直接捕捉X中任意兩個向量的關(guān)聯(lián)，而且易于并行，這都是 Self Attention 的優(yōu)點。?
然而，從理論上來講，Self Attention 的計算時間和顯存占用量都是級別的（n 是序列長度），這就意味著如果序列長度變成原來的 2 倍，顯存占用量就是原來的 4 倍，計算時間也是原來的 4 倍。當(dāng)然，假設(shè)并行核心數(shù)足夠多的情況下，計算時間未必會增加到原來的 4 倍，但是顯存的 4 倍卻是實實在在的，無可避免，這也是微調(diào) Bert 的時候時不時就來個 OOM 的原因了。
稀疏Attention
我們說 Self Attention 是的，那是因為它要對序列中的任意兩個向量都要計算相關(guān)度，得到一個大小的相關(guān)度矩陣：
▲?標(biāo)準(zhǔn)Self Attention的注意力矩陣（左）和關(guān)聯(lián)圖示（右）
在上圖中，左邊顯示了注意力矩陣，右變顯示了關(guān)聯(lián)性，這表明每個元素都跟序列內(nèi)所有元素有關(guān)聯(lián)。?
所以，如果要節(jié)省顯存，加快計算速度，那么一個基本的思路就是減少關(guān)聯(lián)性的計算，也就是認(rèn)為每個元素只跟序列內(nèi)的一部分元素相關(guān)，這就是稀疏 Attention 的基本原理。
本文所要介紹的稀疏 Attention，源于 OpenAI 的論文 Generating Long Sequences with Sparse Transformers，但沒有按照原論文的方式來介紹，而是用一種筆者認(rèn)為更加自然的思路來介紹。?
Atrous Self Attention?
第一個要引入的概念是 Atrous Self Attention，中文可以稱之為“膨脹自注意力”、“空洞自注意力”、“帶孔自注意力”等。這個名稱跟后面的 Local Self Attention 一樣，都是筆者根據(jù)它的特性自行命名的，原論文 Generating Long Sequences with Sparse Transformers 并沒有出現(xiàn)過這兩個概念，但我認(rèn)為將它們單獨引出來是有意義的。?
很顯然，Atrous Self Attention 就是啟發(fā)于“膨脹卷積（Atrous Convolution）”，如下右圖所示，它對相關(guān)性進(jìn)行了約束，強行要求每個元素只跟它相對距離為 k,2k,3k,…的元素關(guān)聯(lián)，其中 k>1 是預(yù)先設(shè)定的超參數(shù)。從下左的注意力矩陣看，就是強行要求相對距離不是k的倍數(shù)的注意力為 0（白色代表 0）：
由于現(xiàn)在計算注意力是“跳著”來了，所以實際上每個元素只跟大約 n/k 個元素算相關(guān)性，這樣一來理想情況下運行效率和顯存占用都變成了，也就是說能直接降低到原來的 1/k。?
Local Self Attention
另一個要引入的過渡概念是 Local Self Attention，中文可稱之為“局部自注意力”。其實自注意力機制在 CV 領(lǐng)域統(tǒng)稱為“Non Local”，而顯然 Local Self Attention 則要放棄全局關(guān)聯(lián)，重新引入局部關(guān)聯(lián)。具體來說也很簡單，就是約束每個元素只與前后 k 個元素以及自身有關(guān)聯(lián)，如下圖所示：
從注意力矩陣來看，就是相對距離超過 k 的注意力都直接設(shè)為 0。?
其實 Local Self Attention 就跟普通卷積很像了，都是保留了一個 2k+1 大小的窗口，然后在窗口內(nèi)進(jìn)行一些運算，不同的是普通卷積是把窗口展平然后接一個全連接層得到輸出，而現(xiàn)在是窗口內(nèi)通過注意力來加權(quán)平均得到輸出。
對于 Local Self Attention 來說，每個元素只跟 2k+1 個元素算相關(guān)性，這樣一來理想情況下運行效率和顯存占用都變成了 ?((2k+1)n)～?(kn) 了，也就是說隨著 n 而線性增長，這是一個很理想的性質(zhì)——當(dāng)然也直接犧牲了長程關(guān)聯(lián)性。?
Sparse Self Attention?
到此，就可以很自然地引入 OpenAI 的 Sparse Self Attention 了。我們留意到， Atrous Self Attention 是帶有一些洞的，而 Local Self Attention 正好填補了這些洞，所以一個簡單的方式就是將 Local Self Attention 和 Atrous Self Attention 交替使用，兩者累積起來，理論上也可以學(xué)習(xí)到全局關(guān)聯(lián)性，也省了顯存。
簡單畫個草圖就可以知道，假如第一層用 Local Self Attention 的話，那么輸出的每個向量都融合了局部的幾個輸入向量，然后第二層用 Atrous Self Attention，雖然它是跳著來，但是因為第一層的輸出融合了局部的輸入向量，所以第二層的輸出理論上可以跟任意的輸入向量相關(guān)，也就是說實現(xiàn)了長程關(guān)聯(lián)。
但是 OpenAI 沒有這樣做，它直接將兩個 Atrous Self Attention 和 Local Self Attention 合并為一個，如下圖：
從注意力矩陣上看就很容易理解了，就是除了相對距離不超過 k 的、相對距離為 k,2k,3k,… 的注意力都設(shè)為 0，這樣一來 Attention 就具有“局部緊密相關(guān)和遠(yuǎn)程稀疏相關(guān)”的特性，這對很多任務(wù)來說可能是一個不錯的先驗，因為真正需要密集的長程關(guān)聯(lián)的任務(wù)事實上是很少的。
代碼實現(xiàn)
上面的 Atrous Self Attention、Local Self Attention、Sparse Self Attention 都算是稀疏 Attention，直觀上來看就是注意力矩陣變得很稀疏了。那怎么實現(xiàn)它們呢？如果直接在注意力矩陣中對為零的部分進(jìn)行 mask 的話，那在數(shù)學(xué)上（功能上）是沒有問題的，但這樣做并不能提速，也不能省顯存。?
官方實現(xiàn)
OpenAI 也開源了自己的實現(xiàn)，位于：
https://github.com/openai/sparse_attention?
這是基于 tensorflow 的，還用到了它們自己的一個稀疏矩陣庫 blocksparse。不過這玩意似乎封裝得很奇怪，我不知道怎么將它遷移到 Keras，而且它用了很多 Python 3 的特性，不能直接用于 Python 2。如果用Python 3和純Tensorflow的朋友可以試試。?
還有一個問題是 OpenAI 原論文主要是用稀疏 Attention 來生成超長序列，所以它在論文中和代碼中都把注意力矩陣的所有上三角部分都 mask 了（避免用到未來信息），但未必所有用到稀疏 Attention 的都是生成場景，而且對于基本概念的介紹來說，這是不必要的，這也是筆者不按原論文的思路來介紹的原因之一。?
Keras實現(xiàn)
對于 Keras，筆者根據(jù)自己構(gòu)思的寫法實現(xiàn)了上述三種稀疏 Attention，并且和原來寫過的 Attention 代碼統(tǒng)一規(guī)范化了一下，還是放在原來的位置：?
https://github.com/bojone/attention/blob/master/attention_keras.py?
經(jīng)過實驗，發(fā)現(xiàn)在筆者的寫法中，這三種稀疏 Attention 相比全量 Attention 確實能節(jié)省一些內(nèi)存，但遺憾的是，除了 Atrous Self Attention 外，剩下兩種 Attention 的實現(xiàn)都不能提速，反而降低了一點速度，這是因為實現(xiàn)過程中沒有充分利用稀疏性所致的，而 OpenAI 的 blocksparse 則是經(jīng)過高度優(yōu)化，而且是直接寫的 CUDA 代碼，這沒法比。但不管速度如何，三種稀疏 Attention 功能上應(yīng)該是沒毛病的。
代碼實現(xiàn)
也沒什么好總結(jié)的了，就介紹并實現(xiàn)了三種稀疏 Attention。除了省顯存外，稀疏的 Attention 應(yīng)該能夠更好地適應(yīng)一些任務(wù)，畢竟大多數(shù)任務(wù)的關(guān)聯(lián)主要都在局部的，而且是從局部到整體的形式。尤其是最后一個 Sparse Self Attention 所體現(xiàn)的“局部緊密相關(guān)和遠(yuǎn)程稀疏相關(guān)”，應(yīng)當(dāng)能滿足大多數(shù)任務(wù)的特點，如果有相應(yīng)任務(wù)的讀者，不妨試用一下。
點擊以下標(biāo)題查看作者其他文章：?
變分自編碼器VAE：原來是這么一回事 | 附源碼
再談變分自編碼器VAE：從貝葉斯觀點出發(fā)
變分自編碼器VAE：這樣做為什么能成？
簡單修改，讓GAN的判別器秒變編碼器
深度學(xué)習(xí)中的互信息：無監(jiān)督提取特征
全新視角：用變分推斷統(tǒng)一理解生成模型
細(xì)水長flow之NICE：流模型的基本概念與實現(xiàn)
細(xì)水長flow之f-VAEs：Glow與VAEs的聯(lián)姻
深度學(xué)習(xí)中的Lipschitz約束：泛化與生成模型
#投 稿 通 道#
?讓你的論文被更多人看到?
如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達(dá)讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認(rèn)識的人。
總有一些你不認(rèn)識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學(xué)者和學(xué)術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)習(xí)心得或技術(shù)干貨。我們的目的只有一個，讓知識真正流動起來。
??來稿標(biāo)準(zhǔn)：
? 稿件確系個人原創(chuàng)作品，來稿需注明作者個人信息（姓名+學(xué)校/工作單位+學(xué)歷/職位+研究方向）?
? 如果文章并非首發(fā)，請在投稿時提醒并附上所有已發(fā)布鏈接?
? PaperWeekly 默認(rèn)每篇文章都是首發(fā)，均會添加“原創(chuàng)”標(biāo)志
? 投稿郵箱：
? 投稿郵箱：hr@paperweekly.site?
? 所有文章配圖，請單獨在附件中發(fā)送?
? 請留下即時聯(lián)系方式（微信或手機），以便我們在編輯發(fā)布時和作者溝通
?
現(xiàn)在，在「知乎」也能找到我們了
進(jìn)入知乎首頁搜索「PaperWeekly」
點擊「關(guān)注」訂閱我們的專欄吧
關(guān)于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學(xué)術(shù)平臺。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
▽ 點擊 |?閱讀原文?| 查看作者博客
總結(jié)
以上是生活随笔為你收集整理的为节约而生：从标准Attention到稀疏Attention的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 线下沙龙 × 上海 | 小身材大能量！用
- 下一篇： 线下沙龙 | 小身材大能量！用英伟达智能