除了生成文本，還可以補(bǔ)全圖像、生成音頻序列的稀疏 Transformers

雷鋒網(wǎng) AI 科技評(píng)論按：OpenAI 設(shè)計(jì)了一種新的 Transformer 模型：稀疏 Transformer（Sparse Transformer），它在序列預(yù)測(cè)任務(wù)中取得了新的表現(xiàn)記錄，不論是預(yù)測(cè)文本、像素還是聲音序列。它在注意力機(jī)制的基礎(chǔ)上做了算法改進(jìn)，可以在比以前長(zhǎng) 30 倍的序列中提取模式。

AI 科研中現(xiàn)存挑戰(zhàn)之一就是對(duì)復(fù)雜數(shù)據(jù)中的長(zhǎng)距離、細(xì)微的互相關(guān)聯(lián)性做建模，比如圖像、視頻、音頻序列內(nèi)部存在的關(guān)聯(lián)性。原始的 Transformer 中的自我注意力機(jī)制有著 O(N²) 的計(jì)算復(fù)雜度，OpenAI 新提出的稀疏 Transformer 經(jīng)過(guò)重新設(shè)計(jì)，把計(jì)算復(fù)雜度降低到了 O(N√N(yùn))，以及加入了其它一些改進(jìn)，讓它可以處理更長(zhǎng)的序列甚至更復(fù)雜的數(shù)據(jù)類型。在此之前，處理這些數(shù)據(jù)的模型要么是針對(duì)單獨(dú)一種數(shù)據(jù)類型專門設(shè)計(jì)的，要么很難處理幾千個(gè)元素以及更長(zhǎng)的序列。而 OpenAI 的稀疏 Transformer 可以用數(shù)百層的模型處理上萬(wàn)個(gè)元素長(zhǎng)的序列額，并且在多個(gè)不同的任務(wù)中取得最佳表現(xiàn)。雷鋒網(wǎng) AI 科技評(píng)論根據(jù) OpenAI 技術(shù)博客介紹如下。

深層注意力

在 Transformer 模型中，每一個(gè)輸出元素都與每一個(gè)輸入元素相連接，每個(gè)連接的權(quán)重是根據(jù)不同的狀況動(dòng)態(tài)計(jì)算的，這個(gè)過(guò)程就叫做注意力。相比連接方式固定的做法，研究人員們相信這樣能夠讓 Transformer 更靈活，但代價(jià)就是在實(shí)際使用中它需要為每一層、每一個(gè)注意力頭建立一個(gè) N x N 大小的注意力矩陣，當(dāng)輸入圖像、音頻波形等含有大量元素的序列時(shí)會(huì)帶來(lái)大量的內(nèi)存開(kāi)銷。

一個(gè)具有 64 層、4 個(gè)注意力頭的深度 Transformer 模型的注意力內(nèi)存消耗。作為參考，用于深度學(xué)習(xí)的 GPU 的顯存大小一般在 12GB 到 32GB 的水平。

應(yīng)對(duì)這個(gè)問(wèn)題的一種做法是在反向傳播中根據(jù)建立的檢查點(diǎn)重新計(jì)算注意力矩陣，這是深度學(xué)習(xí)中的一種常見(jiàn)做法，可以用消耗更多計(jì)算為代價(jià)減小內(nèi)存的使用。對(duì) Transformer 模型中的注意力矩陣使用這種做法之后，它的效果是讓最大的內(nèi)存消耗不再和網(wǎng)絡(luò)的層數(shù)相關(guān)，就讓訓(xùn)練比以往深得多的模型變得可能。在實(shí)際嘗試中，OpenAI 的研究人員們訓(xùn)練了最深達(dá) 128 層的 Transformer，也在 CIFAR-10 之類的測(cè)試任務(wù)中擊敗了較淺的模型。

為了訓(xùn)練更深的模型，OpenAI 的研究人員們還對(duì) Transformer 模型中的運(yùn)算順序做了多種調(diào)整、修改了最初的注意力機(jī)制。詳細(xì)介紹可以參見(jiàn)論文原文。

稀疏注意力

即便已經(jīng)有了節(jié)省內(nèi)存的方法，面對(duì)很長(zhǎng)的序列輸入時(shí)把每個(gè)注意力矩陣都計(jì)算出來(lái)仍然是不實(shí)際的。OpenAI 轉(zhuǎn)而使用稀疏的注意力模式，也就是說(shuō)，每個(gè)輸出位置在計(jì)算時(shí)只考慮一部分輸入位置的權(quán)重。當(dāng)取的這一部分相比于全部輸入位置很小的時(shí)候（比如一共 N 個(gè)輸入元素，取 √N(yùn) 個(gè)），即便對(duì)于很長(zhǎng)的序列也是可以計(jì)算注意力的了，計(jì)算復(fù)雜度也就響應(yīng)降低為 O(N√N(yùn))，不再是完全計(jì)算時(shí)的 O(N²)。

這種做法的可行性需要加以評(píng)估。OpenAI 的研究人員們首先對(duì)原本的深度 Transformer 模型學(xué)到的注意力模式進(jìn)行了可視化，他們發(fā)現(xiàn)許多樣本中都體現(xiàn)出了可解釋的、結(jié)構(gòu)化的稀疏模式。在下面的四張圖中，用高亮顯示出了預(yù)測(cè)圖像中的下一個(gè)像素時(shí)有哪些已有的像素是被注意力頭所使用的。當(dāng)用到的輸入像素?cái)?shù)量不多而且表現(xiàn)出高度的規(guī)律性的時(shí)候，就可以認(rèn)為網(wǎng)絡(luò)中的這一層是可以做稀疏化改造的。圖中展示的是 128 層的網(wǎng)絡(luò)模型在 CIFAR-10 數(shù)據(jù)集中圖像上運(yùn)行時(shí)的分析結(jié)果。

第 19 層	第 20 層

網(wǎng)絡(luò)中的某些層學(xué)到的是單個(gè)維度上注意力模式。19 層關(guān)注的是每個(gè)橫排，20 層關(guān)注的是每個(gè)豎列，可以看做是完整的注意力操作的分解。

第 6 層	第 36 層

某些層學(xué)習(xí)到的是使用某些空間位置，在不同的輸入數(shù)據(jù)或者時(shí)間步驟中這些位置都是相似的。還有一些其他的層學(xué)到的訪問(wèn)模式會(huì)隨著數(shù)據(jù)的不同而變化。

雖然許多層都確實(shí)表現(xiàn)出了稀疏結(jié)構(gòu)，不過(guò)也有一些層明顯表現(xiàn)出了可以覆蓋整個(gè)圖像范圍的動(dòng)態(tài)注意力模式。為了仍然在網(wǎng)絡(luò)中保留學(xué)到這種模式的能力，OpenAI 的研究人員們實(shí)現(xiàn)了一種注意力矩陣的二維分解，網(wǎng)絡(luò)通過(guò)兩步稀疏注意力操作就仍然可以訪問(wèn)圖像中的所有位置。

其中的第一種，條狀注意力，基本就和前面例子中關(guān)注單行、單列一樣；第二種，固定注意力，關(guān)注的是某個(gè)固定的列以及其中最后一個(gè)列元素之后的元素，研究人員們發(fā)現(xiàn)這種模式對(duì)于無(wú)法用前一種模式覆蓋的數(shù)據(jù)結(jié)構(gòu)（比如文本）非常有用。更多的細(xì)節(jié)可以參見(jiàn)論文原文。

實(shí)驗(yàn)結(jié)果

稀疏 Transformers 模型在 CIFAR-10、Enwik8、Imagenet 64 數(shù)據(jù)集上的密集估計(jì)任務(wù)中取得了目前的最好成績(jī)。

對(duì)于稀疏 Transformer 模型的參數(shù)，W 表示網(wǎng)絡(luò)寬度、L 表示網(wǎng)絡(luò)層數(shù)、H 表示注意力頭的數(shù)量。

OpenAI 的研究人員們還發(fā)現(xiàn)稀疏注意力不僅比完整的注意力快多了，損失也要更低（具體對(duì)比見(jiàn)論文）。這表明他們的稀疏注意力設(shè)計(jì)可能附帶了有好處的先驗(yàn)偏倚，或者是密集注意力自身隱含著一定的優(yōu)化問(wèn)題。

用稀疏注意力生成圖像

使用稀疏注意力的 Transformer 模型在測(cè)試中似乎體現(xiàn)出了全局結(jié)構(gòu)的概念，這在圖像補(bǔ)全任務(wù)中可以進(jìn)行量化測(cè)量。下面展示了對(duì)于 64x64 尺寸 ImageNet 圖像的補(bǔ)全演示。

待補(bǔ)全圖像

補(bǔ)全結(jié)果

真實(shí)圖像

他們還生成了完全無(wú)限制的圖像，對(duì)應(yīng)的未調(diào)節(jié) softmax 溫度為 1。模型是使用最大似然作為目標(biāo)訓(xùn)練的，這種方法可以覆蓋數(shù)據(jù)的所有模式（包括訓(xùn)練數(shù)據(jù)中未出現(xiàn)甚至不存在的模式），而不是某一小部分?jǐn)?shù)據(jù)中繼續(xù)增加保真度而已。取未調(diào)節(jié)的溫度對(duì)模型進(jìn)行采樣，得出的圖像可以看作是「模型認(rèn)為這樣的圖像在世界上是存在的」。當(dāng)人類觀察的時(shí)候，自然會(huì)覺(jué)得有一些樣本很奇怪。

模型在無(wú)條件限制下生成的樣本

生成原始音頻波形

只要簡(jiǎn)單地改變位置嵌入，稀疏 Transformer 可以用來(lái)生成音頻波形。這樣可以讓深度學(xué)習(xí)模型輕松地在不同的數(shù)據(jù)類型上工作，OpenAI 的研究人員們認(rèn)為這減小了之前的研究中處理不同類型數(shù)據(jù)時(shí)分別引入的不同先驗(yàn)偏倚，對(duì)科研會(huì)有所幫助。

他們?cè)诠诺湟魳?lè)片段（鋼琴曲）上訓(xùn)練了模型，然后通過(guò)稀疏注意力生成長(zhǎng)度為 65000 的音頻波形序列，對(duì)應(yīng)的播放時(shí)長(zhǎng)大約為 5 秒鐘。他們把多組生成結(jié)果拼成了三個(gè)音頻文件，試聽(tīng)可以訪問(wèn) https://soundcloud.com/openai_audio/sample-set-1。雷鋒網(wǎng) AI 科技評(píng)論試聽(tīng)后的評(píng)價(jià)是，有些片段較為激烈雜亂，也有一些較為簡(jiǎn)單安靜，能感覺(jué)到細(xì)微的節(jié)奏感（輕重節(jié)拍變換）；總的來(lái)說(shuō)，能輕松辨別出是鋼琴演奏，但是旋律比較難以欣賞。

現(xiàn)階段的不足以及未來(lái)方向

論文中設(shè)計(jì)的稀疏注意力模式僅僅是對(duì)高效地建模長(zhǎng)序列的新方法的非常初步的研究。作者們認(rèn)為探究不同的稀疏模式以及它們的組合是有用的研究，而且對(duì)于下一代神經(jīng)網(wǎng)絡(luò)架構(gòu)的發(fā)展也可以起到尤其有幫助的探路作用。

即便作者們提出的方法已經(jīng)帶來(lái)了明顯的改進(jìn)，非常高分辨率的圖像或者視頻的自回歸序列生成仍然是不可行的。論文中提出的注意力機(jī)制的優(yōu)化方法可以成為多尺度方法之類的高維數(shù)據(jù)建模方法的有效補(bǔ)充。

論文原文見(jiàn)：https://arxiv.org/abs/1904.10509

GitHub 地址：https://github.com/openai/sparse_attention

via openai.com，雷鋒網(wǎng) AI 科技評(píng)論編譯