作者丨蘇劍林

單位丨追一科技

研究方向丨NLP，神經(jīng)網(wǎng)絡(luò)

個(gè)人主頁(yè)丨kexue.fm

如今 NLP 領(lǐng)域，Attention 大行其道，當(dāng)然也不止 NLP，在 CV 領(lǐng)域 Attention 也占有一席之地（Non Local、SAGAN 等）。在 18 年初一文讀懂「Attention is All You Need」| 附代碼實(shí)現(xiàn)一文中，我們就已經(jīng)討論過(guò) Attention 機(jī)制，Attention 的核心在于 Q,K,V 三個(gè)向量序列的交互和融合，其中 Q,K 的交互給出了兩兩向量之間的某種相關(guān)度（權(quán)重），而最后的輸出序列則是把 V 按照權(quán)重求和得到的。?

顯然，眾多 NLP & CV 的成果已經(jīng)充分肯定了 Attention 的有效性。本文我們將會(huì)介紹 Attention 的一些變體，這些變體的共同特點(diǎn)是——“為節(jié)約而生”——既節(jié)約時(shí)間，也節(jié)約顯存。

背景簡(jiǎn)述

Attention is All You Need?一文討論的我們稱之為“乘性 Attention”，目前用得比較廣泛的也就是這種 Attention：?

另外還有加性 Attention，但加性 Attention 并行不大容易實(shí)現(xiàn)（或者實(shí)現(xiàn)起來(lái)占用顯存多），所以一般只用來(lái)將變長(zhǎng)向量序列編碼為固定長(zhǎng)度的向量（取代簡(jiǎn)單的 Pooling），而很少用來(lái)做序列到序列的編碼。

而在乘性 Attention 中，用得最廣泛的當(dāng)數(shù) Self Attention 了，這種情況下 Q,K,V 都是同一個(gè) X 經(jīng)過(guò)線性變換之后的結(jié)果，這樣一來(lái)輸出結(jié)果就是跟 X 一樣長(zhǎng)的向量序列，并且能夠直接捕捉X中任意兩個(gè)向量的關(guān)聯(lián)，而且易于并行，這都是 Self Attention 的優(yōu)點(diǎn)。?

然而，從理論上來(lái)講，Self Attention 的計(jì)算時(shí)間和顯存占用量都是級(jí)別的（n 是序列長(zhǎng)度），這就意味著如果序列長(zhǎng)度變成原來(lái)的 2 倍，顯存占用量就是原來(lái)的 4 倍，計(jì)算時(shí)間也是原來(lái)的 4 倍。當(dāng)然，假設(shè)并行核心數(shù)足夠多的情況下，計(jì)算時(shí)間未必會(huì)增加到原來(lái)的 4 倍，但是顯存的 4 倍卻是實(shí)實(shí)在在的，無(wú)可避免，這也是微調(diào) Bert 的時(shí)候時(shí)不時(shí)就來(lái)個(gè) OOM 的原因了。

稀疏Attention

我們說(shuō) Self Attention 是的，那是因?yàn)樗獙?duì)序列中的任意兩個(gè)向量都要計(jì)算相關(guān)度，得到一個(gè)大小的相關(guān)度矩陣：

▲?標(biāo)準(zhǔn)Self Attention的注意力矩陣（左）和關(guān)聯(lián)圖示（右）

在上圖中，左邊顯示了注意力矩陣，右變顯示了關(guān)聯(lián)性，這表明每個(gè)元素都跟序列內(nèi)所有元素有關(guān)聯(lián)。?

所以，如果要節(jié)省顯存，加快計(jì)算速度，那么一個(gè)基本的思路就是減少關(guān)聯(lián)性的計(jì)算，也就是認(rèn)為每個(gè)元素只跟序列內(nèi)的一部分元素相關(guān)，這就是稀疏 Attention 的基本原理。

本文所要介紹的稀疏 Attention，源于 OpenAI 的論文 Generating Long Sequences with Sparse Transformers，但沒(méi)有按照原論文的方式來(lái)介紹，而是用一種筆者認(rèn)為更加自然的思路來(lái)介紹。?

Atrous Self Attention?

第一個(gè)要引入的概念是 Atrous Self Attention，中文可以稱之為“膨脹自注意力”、“空洞自注意力”、“帶孔自注意力”等。這個(gè)名稱跟后面的 Local Self Attention 一樣，都是筆者根據(jù)它的特性自行命名的，原論文 Generating Long Sequences with Sparse Transformers 并沒(méi)有出現(xiàn)過(guò)這兩個(gè)概念，但我認(rèn)為將它們單獨(dú)引出來(lái)是有意義的。?

很顯然，Atrous Self Attention 就是啟發(fā)于“膨脹卷積（Atrous Convolution）”，如下右圖所示，它對(duì)相關(guān)性進(jìn)行了約束，強(qiáng)行要求每個(gè)元素只跟它相對(duì)距離為 k,2k,3k,…的元素關(guān)聯(lián)，其中 k>1 是預(yù)先設(shè)定的超參數(shù)。從下左的注意力矩陣看，就是強(qiáng)行要求相對(duì)距離不是k的倍數(shù)的注意力為 0（白色代表 0）：

▲?Atrous Self Attention的注意力矩陣（左）和關(guān)聯(lián)圖示（右）

由于現(xiàn)在計(jì)算注意力是“跳著”來(lái)了，所以實(shí)際上每個(gè)元素只跟大約 n/k 個(gè)元素算相關(guān)性，這樣一來(lái)理想情況下運(yùn)行效率和顯存占用都變成了，也就是說(shuō)能直接降低到原來(lái)的 1/k。?

Local Self Attention

另一個(gè)要引入的過(guò)渡概念是 Local Self Attention，中文可稱之為“局部自注意力”。其實(shí)自注意力機(jī)制在 CV 領(lǐng)域統(tǒng)稱為“Non Local”，而顯然 Local Self Attention 則要放棄全局關(guān)聯(lián)，重新引入局部關(guān)聯(lián)。具體來(lái)說(shuō)也很簡(jiǎn)單，就是約束每個(gè)元素只與前后 k 個(gè)元素以及自身有關(guān)聯(lián)，如下圖所示：

▲?Local Self Attention的注意力矩陣（左）和關(guān)聯(lián)圖示（右）

從注意力矩陣來(lái)看，就是相對(duì)距離超過(guò) k 的注意力都直接設(shè)為 0。?

其實(shí) Local Self Attention 就跟普通卷積很像了，都是保留了一個(gè) 2k+1 大小的窗口，然后在窗口內(nèi)進(jìn)行一些運(yùn)算，不同的是普通卷積是把窗口展平然后接一個(gè)全連接層得到輸出，而現(xiàn)在是窗口內(nèi)通過(guò)注意力來(lái)加權(quán)平均得到輸出。

對(duì)于 Local Self Attention 來(lái)說(shuō)，每個(gè)元素只跟 2k+1 個(gè)元素算相關(guān)性，這樣一來(lái)理想情況下運(yùn)行效率和顯存占用都變成了 ?((2k+1)n)～?(kn) 了，也就是說(shuō)隨著 n 而線性增長(zhǎng)，這是一個(gè)很理想的性質(zhì)——當(dāng)然也直接犧牲了長(zhǎng)程關(guān)聯(lián)性。?

Sparse Self Attention?

到此，就可以很自然地引入 OpenAI 的 Sparse Self Attention 了。我們留意到， Atrous Self Attention 是帶有一些洞的，而 Local Self Attention 正好填補(bǔ)了這些洞，所以一個(gè)簡(jiǎn)單的方式就是將 Local Self Attention 和 Atrous Self Attention 交替使用，兩者累積起來(lái)，理論上也可以學(xué)習(xí)到全局關(guān)聯(lián)性，也省了顯存。

簡(jiǎn)單畫個(gè)草圖就可以知道，假如第一層用 Local Self Attention 的話，那么輸出的每個(gè)向量都融合了局部的幾個(gè)輸入向量，然后第二層用 Atrous Self Attention，雖然它是跳著來(lái)，但是因?yàn)榈谝粚拥妮敵鋈诤狭司植康妮斎胂蛄?#xff0c;所以第二層的輸出理論上可以跟任意的輸入向量相關(guān)，也就是說(shuō)實(shí)現(xiàn)了長(zhǎng)程關(guān)聯(lián)。

但是 OpenAI 沒(méi)有這樣做，它直接將兩個(gè) Atrous Self Attention 和 Local Self Attention 合并為一個(gè)，如下圖：

▲?Sparse Self Attention的注意力矩陣（左）和關(guān)聯(lián)圖示（右）

從注意力矩陣上看就很容易理解了，就是除了相對(duì)距離不超過(guò) k 的、相對(duì)距離為 k,2k,3k,… 的注意力都設(shè)為 0，這樣一來(lái) Attention 就具有“局部緊密相關(guān)和遠(yuǎn)程稀疏相關(guān)”的特性，這對(duì)很多任務(wù)來(lái)說(shuō)可能是一個(gè)不錯(cuò)的先驗(yàn)，因?yàn)檎嬲枰芗拈L(zhǎng)程關(guān)聯(lián)的任務(wù)事實(shí)上是很少的。

代碼實(shí)現(xiàn)

上面的 Atrous Self Attention、Local Self Attention、Sparse Self Attention 都算是稀疏 Attention，直觀上來(lái)看就是注意力矩陣變得很稀疏了。那怎么實(shí)現(xiàn)它們呢？如果直接在注意力矩陣中對(duì)為零的部分進(jìn)行 mask 的話，那在數(shù)學(xué)上（功能上）是沒(méi)有問(wèn)題的，但這樣做并不能提速，也不能省顯存。?

官方實(shí)現(xiàn)

OpenAI 也開(kāi)源了自己的實(shí)現(xiàn)，位于：

https://github.com/openai/sparse_attention?

這是基于 tensorflow 的，還用到了它們自己的一個(gè)稀疏矩陣庫(kù) blocksparse。不過(guò)這玩意似乎封裝得很奇怪，我不知道怎么將它遷移到 Keras，而且它用了很多 Python 3 的特性，不能直接用于 Python 2。如果用Python 3和純Tensorflow的朋友可以試試。?

還有一個(gè)問(wèn)題是 OpenAI 原論文主要是用稀疏 Attention 來(lái)生成超長(zhǎng)序列，所以它在論文中和代碼中都把注意力矩陣的所有上三角部分都 mask 了（避免用到未來(lái)信息），但未必所有用到稀疏 Attention 的都是生成場(chǎng)景，而且對(duì)于基本概念的介紹來(lái)說(shuō)，這是不必要的，這也是筆者不按原論文的思路來(lái)介紹的原因之一。?

Keras實(shí)現(xiàn)

對(duì)于 Keras，筆者根據(jù)自己構(gòu)思的寫法實(shí)現(xiàn)了上述三種稀疏 Attention，并且和原來(lái)寫過(guò)的 Attention 代碼統(tǒng)一規(guī)范化了一下，還是放在原來(lái)的位置：?

https://github.com/bojone/attention/blob/master/attention_keras.py?

經(jīng)過(guò)實(shí)驗(yàn)，發(fā)現(xiàn)在筆者的寫法中，這三種稀疏 Attention 相比全量 Attention 確實(shí)能節(jié)省一些內(nèi)存，但遺憾的是，除了 Atrous Self Attention 外，剩下兩種 Attention 的實(shí)現(xiàn)都不能提速，反而降低了一點(diǎn)速度，這是因?yàn)閷?shí)現(xiàn)過(guò)程中沒(méi)有充分利用稀疏性所致的，而 OpenAI 的 blocksparse 則是經(jīng)過(guò)高度優(yōu)化，而且是直接寫的 CUDA 代碼，這沒(méi)法比。但不管速度如何，三種稀疏 Attention 功能上應(yīng)該是沒(méi)毛病的。

代碼實(shí)現(xiàn)

也沒(méi)什么好總結(jié)的了，就介紹并實(shí)現(xiàn)了三種稀疏 Attention。除了省顯存外，稀疏的 Attention 應(yīng)該能夠更好地適應(yīng)一些任務(wù)，畢竟大多數(shù)任務(wù)的關(guān)聯(lián)主要都在局部的，而且是從局部到整體的形式。尤其是最后一個(gè) Sparse Self Attention 所體現(xiàn)的“局部緊密相關(guān)和遠(yuǎn)程稀疏相關(guān)”，應(yīng)當(dāng)能滿足大多數(shù)任務(wù)的特點(diǎn)，如果有相應(yīng)任務(wù)的讀者，不妨試用一下。

點(diǎn)擊以下標(biāo)題查看作者其他文章：?

• 變分自編碼器VAE：原來(lái)是這么一回事 | 附源碼

• 再談變分自編碼器VAE：從貝葉斯觀點(diǎn)出發(fā)

• 變分自編碼器VAE：這樣做為什么能成？

• 簡(jiǎn)單修改，讓GAN的判別器秒變編碼器
  

• 深度學(xué)習(xí)中的互信息：無(wú)監(jiān)督提取特征

• 全新視角：用變分推斷統(tǒng)一理解生成模型

• 細(xì)水長(zhǎng)flow之NICE：流模型的基本概念與實(shí)現(xiàn)

• 細(xì)水長(zhǎng)flow之f-VAEs：Glow與VAEs的聯(lián)姻

• 深度學(xué)習(xí)中的Lipschitz約束：泛化與生成模型
  

#投 稿 通 道#

?讓你的論文被更多人看到?

如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達(dá)讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認(rèn)識(shí)的人。

總有一些你不認(rèn)識(shí)的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學(xué)者和學(xué)術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。

PaperWeekly 鼓勵(lì)高校實(shí)驗(yàn)室或個(gè)人，在我們的平臺(tái)上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)習(xí)心得或技術(shù)干貨。我們的目的只有一個(gè)，讓知識(shí)真正流動(dòng)起來(lái)。

??來(lái)稿標(biāo)準(zhǔn)：

? 稿件確系個(gè)人原創(chuàng)作品，來(lái)稿需注明作者個(gè)人信息（姓名+學(xué)校/工作單位+學(xué)歷/職位+研究方向）?

? 如果文章并非首發(fā)，請(qǐng)?jiān)谕陡鍟r(shí)提醒并附上所有已發(fā)布鏈接?

? PaperWeekly 默認(rèn)每篇文章都是首發(fā)，均會(huì)添加“原創(chuàng)”標(biāo)志

? 投稿郵箱：

? 投稿郵箱：hr@paperweekly.site?

? 所有文章配圖，請(qǐng)單獨(dú)在附件中發(fā)送?

? 請(qǐng)留下即時(shí)聯(lián)系方式（微信或手機(jī)），以便我們?cè)诰庉嫲l(fā)布時(shí)和作者溝通

?

現(xiàn)在，在「知乎」也能找到我們了

進(jìn)入知乎首頁(yè)搜索「PaperWeekly」

點(diǎn)擊「關(guān)注」訂閱我們的專欄吧

關(guān)于PaperWeekly

PaperWeekly 是一個(gè)推薦、解讀、討論、報(bào)道人工智能前沿論文成果的學(xué)術(shù)平臺(tái)。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號(hào)后臺(tái)點(diǎn)擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

▽ 點(diǎn)擊 |?閱讀原文?| 查看作者博客

總結(jié)

以上是生活随笔為你收集整理的为节约而生：从标准Attention到稀疏Attention的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。