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

本文作者：楊曉凡

2019-04-26 09:59

導(dǎo)語：OpenAI 設(shè)計(jì)了一種新的Transformer模型：稀疏Transformer，它在長序列預(yù)測任務(wù)中取得了新的表現(xiàn)

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

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

深層注意力

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

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

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

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

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

稀疏注意力

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

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


第 19 層	第 20 層

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


第 6 層	第 36 層

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

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

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

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

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

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

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

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

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

用稀疏注意力生成圖像

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

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

待補(bǔ)全圖像

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

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

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

真實(shí)圖像

他們還生成了完全無限制的圖像，對應(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é)的溫度對模型進(jìn)行采樣，得出的圖像可以看作是「模型認(rèn)為這樣的圖像在世界上是存在的」。當(dāng)人類觀察的時候，自然會覺得有一些樣本很奇怪。

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

模型在無條件限制下生成的樣本

生成原始音頻波形

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

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