Transformer是一種強(qiáng)大的序列模型，但是它所需的時(shí)間和內(nèi)存會(huì)隨著序列長(zhǎng)度出現(xiàn)二階增長(zhǎng)。近日，OpenAI研究人員開發(fā)出了一種深度神經(jīng)網(wǎng)絡(luò)Sparse Transformer，該網(wǎng)絡(luò)在預(yù)測(cè)長(zhǎng)序列方面創(chuàng)造了新紀(jì)錄——無論預(yù)測(cè)的是文本、圖像還是聲音。該神經(jīng)網(wǎng)絡(luò)利用注意力機(jī)制中的一種改進(jìn)算法，可以從長(zhǎng)度可能是之前30倍的序列中提取模式。
現(xiàn)在，AI 研究中的一項(xiàng)挑戰(zhàn)是在圖像、視頻或聲音等復(fù)雜數(shù)據(jù)中進(jìn)行長(zhǎng)序列的精細(xì)相關(guān)性建模。Sparse Transformer 合并了 O（N^2）Transformer 自注意力機(jī)制的 O(N√N(yùn)) 重組以及其他一些改進(jìn)，從而直接用于這些豐富的數(shù)據(jù)類型。以前，這些數(shù)據(jù)上所使用的模型是專為某個(gè)領(lǐng)域制作的，或者很難將序列擴(kuò)展到包含幾千個(gè)元素。
相比之下，OpenAI 開發(fā)的模型通過使用數(shù)以百計(jì)的層可以對(duì)包含上萬個(gè)元素的序列進(jìn)行建模，在諸多領(lǐng)域都取得了當(dāng)前最佳的表現(xiàn)。OpenAI 研究人員利用該模型幫助創(chuàng)建能夠更好地理解世界的 AI 系統(tǒng)。
深度注意力
在 Transformer 中，每一個(gè)輸出元素與輸入元素相連接，同時(shí)根據(jù)具體情況對(duì)它們之間的權(quán)重進(jìn)行動(dòng)態(tài)計(jì)算，這一過程被稱為「注意力機(jī)制」。雖然人們相信這使得 Transformer 較那些具有固定連接模式的模型更為靈活，但實(shí)際操作中需要為每一層和注意力頭創(chuàng)建一個(gè) N×N 注意力矩陣，當(dāng)應(yīng)用于圖像或原始音頻等具有許多元素的數(shù)據(jù)類型時(shí)會(huì)消耗大量?jī)?nèi)存。
當(dāng)矩陣存儲(chǔ)在內(nèi)存或在逆推計(jì)算過程中進(jìn)行再計(jì)算時(shí)，深度 Transformer（64 層和 4 個(gè)頭）的注意力內(nèi)存使用情況。作為參考，用于深度學(xué)習(xí)的標(biāo)準(zhǔn) GPU 內(nèi)存通常是 12-32GB.
減少內(nèi)存消耗的一種方法是在反向傳播過程中從檢查點(diǎn)處重新計(jì)算注意力矩陣，這是深度學(xué)習(xí)中的一種成熟的方法，以更多的計(jì)算來減少內(nèi)存使用。
當(dāng) Transformer 中的注意力矩陣完成時(shí)，這意味著最大的內(nèi)存消耗將不受層數(shù)的支配，使研究人員訓(xùn)練網(wǎng)絡(luò)的深度大大超過從前。在實(shí)際操作中，研究人員發(fā)現(xiàn)在處理 CIFAR-10 等基準(zhǔn)測(cè)試任務(wù)時(shí)，深度達(dá) 128 層的 Transformer 表現(xiàn)出的性能優(yōu)于較淺的網(wǎng)絡(luò)。
為了訓(xùn)練深度更大的模型，研究人員對(duì) transformer 的操作順序進(jìn)行了幾次調(diào)整，修改了初始化方法。詳情參見論文。
稀疏注意力
然而，對(duì)于非常大的輸入來說，甚至計(jì)算單個(gè)注意力矩陣都是不現(xiàn)實(shí)的。因此，OpenAI 使用了稀疏注意力模式，在這種模式中，每個(gè)輸出位置僅從輸入位置子集中計(jì)算權(quán)重。當(dāng)子集相對(duì)于整個(gè)輸入集較小時(shí)（如元素?cái)?shù)量是√N(yùn) 而不是 N），即使對(duì)于非常長(zhǎng)的序列，注意力計(jì)算也會(huì)變得比較容易，算法復(fù)雜度為 O（N√N(yùn)）而不是 O（N^2）。
為了評(píng)估該方法的可行性，研究人員首先可視化并學(xué)習(xí)了圖像上深度 Transformer 的注意力模式，發(fā)現(xiàn)其中許多模式表現(xiàn)出了可解釋和結(jié)構(gòu)化的稀疏模式。以下每幅圖像都顯示了哪個(gè)輸入像素（白色高亮標(biāo)出）由一個(gè)給定的注意力頭處理，以預(yù)測(cè)圖像中的下一個(gè)值。當(dāng)輸入部分集中在小的子集上并顯示出高度規(guī)律性時(shí)，該層就易于稀疏化。以下是 CIFAR-10 圖像上 128 層模型的樣本：
左：Layer 19，右：Layer 20。為一個(gè) 128 層的 CIFAR-10 網(wǎng)絡(luò)的若干層學(xué)習(xí)注意力模式（白色高亮顯示）。這些層學(xué)會(huì)了在兩個(gè)維度上分割注意力。Layer 19 匯總每一行的信息，Layer 20 按列匯總這些信息，從而有效分解了全注意力運(yùn)算。
為獲取位置記憶而訓(xùn)練的層（左：Layer 6；右：Layer 36），它們通常關(guān)注類似的位置，不管輸入數(shù)據(jù)或時(shí)間步長(zhǎng)如何（Layer 6）。其他層學(xué)習(xí)高度依賴數(shù)據(jù)的訪問模式（Layer 36）。
雖然許多層顯示出稀疏的結(jié)構(gòu)，但有些層清晰地顯示出了動(dòng)態(tài)注意力，這種注意力延伸到整個(gè)圖像。為了保持網(wǎng)絡(luò)學(xué)習(xí)這種模式的能力，研究人員實(shí)現(xiàn)了注意力矩陣的二維分解，其中網(wǎng)絡(luò)可以通過兩步稀疏注意力關(guān)注到所有位置。
第一版 strided attention 大概等同于每個(gè)位置處理自己的行和列，它與以上網(wǎng)絡(luò)學(xué)得的注意力模式類似。（注意，列注意力可等同于處理轉(zhuǎn)置矩陣的行）。第二版 fixed attention 在最新的列元素之后處理固定列和元素，研究者認(rèn)為這個(gè)模式對(duì)于數(shù)據(jù)無法擬合二維結(jié)構(gòu)（如文本）的情況很有用。
實(shí)驗(yàn)結(jié)果
Sparse Transformer 在 CIFAR-10、Enwik8 和 Imagenet 64 數(shù)據(jù)集上刷新了當(dāng)前最優(yōu)密度估計(jì)分?jǐn)?shù)。
在 CIFAR-10、Enwik8 和 Imagenet 64 數(shù)據(jù)集上的密度估計(jì)性能（單位為 bits per byte／dim）。M 表示網(wǎng)絡(luò)中使用的參數(shù)（單位為百萬），W 表示網(wǎng)絡(luò)寬度，L 表示層數(shù)，H 表示頭數(shù)。
研究者還發(fā)現(xiàn)稀疏注意力比完整注意力的損失更低，且速度更快。這可能指向稀疏模式產(chǎn)生的有用歸納偏置，或者密集注意力的底層優(yōu)化問題。
生成圖像
使用了稀疏注意力的 Transformer 似乎有一種全局結(jié)構(gòu)的概念，這可以通過觀察圖像補(bǔ)全（image completion）進(jìn)行定性評(píng)估。下圖可視化了一個(gè)在 64×64 ImageNet 上訓(xùn)練的模型：
損壞原圖
修復(fù)圖像
真實(shí)圖像
研究人員還生成了完全無條件的樣本，其中未調(diào)整的 softmax 溫度為 1.0。這些模型使用最大似然目標(biāo)進(jìn)行訓(xùn)練，其覆蓋了所有的數(shù)據(jù)模式（其中包括可能不存在的數(shù)據(jù)），而不是增強(qiáng)較小部分?jǐn)?shù)據(jù)的保真度。從具有未調(diào)整溫度的模型中取樣，研究人員看到了該模型認(rèn)為世界上存在的圖像的完整分布。因此，一些樣本看起來奇奇怪怪的。
模型示例
生成原始音頻波形
通過簡(jiǎn)單改變位置嵌入，稀疏 Transformer 還能用來生成原始音頻，而非圖像。隨著深度學(xué)習(xí)擴(kuò)展到新的數(shù)據(jù)類型，用這類網(wǎng)絡(luò)來指定歸納偏置也很容易。
該模型是在原始的古典音樂片段上訓(xùn)練的，并使用了稀疏注意力來生成長(zhǎng)度為 65000 的序列。這相當(dāng)于大約 5 秒長(zhǎng)的原始音頻，研究人員在下面的每個(gè)片段中將幾個(gè)樣本連接在一起。
代碼公布
通常，實(shí)現(xiàn)稀疏注意力需要將查詢和關(guān)鍵矩陣分割成塊，因此為了簡(jiǎn)化實(shí)驗(yàn)，OpenAI 實(shí)現(xiàn)了一組塊稀疏核，這些核在 GPU 上高效地執(zhí)行這些操作。OpenAI 開源了這些核并提供了稀疏注意力函數(shù)的示例：
https://github.com/openai/sparse_attention
未來發(fā)展和限制
本文介紹的稀疏注意力模式只是對(duì)長(zhǎng)序列進(jìn)行高效建模的初步嘗試。研究人員認(rèn)為，探索稀疏注意力的不同模式和各種組合非常有用，而且學(xué)習(xí)稀疏模式對(duì)下一代神經(jīng)網(wǎng)絡(luò)架構(gòu)來說也是一個(gè)很重要的研究途徑。
即使有了上述改進(jìn)，自回歸序列生成對(duì)非常高分辨率圖像和音頻來說仍是不切實(shí)際的。但是，研究人員介紹的優(yōu)化注意力操作可能有用，將它與其它方法（如多尺度方法）結(jié)合，可以建模高維數(shù)據(jù)。
論文：Generating Long Sequences with Sparse Transformers
論文鏈接：https://d4mucfpksywv.cloudfront.net/Sparse_Transformer/sparse_transformers.pdf
摘要：Transformer 是一種強(qiáng)大的序列模型，但是它所需的時(shí)間和內(nèi)存會(huì)隨著序列長(zhǎng)度出現(xiàn)二階增長(zhǎng)。這篇論文介紹了注意力矩陣的稀疏因式分解，可以將其降低到 O(N√N(yùn))。該研究提出了 a）訓(xùn)練更深網(wǎng)絡(luò)的架構(gòu)和初始化變體；b）重新計(jì)算注意力矩陣以節(jié)省內(nèi)存；c）用于訓(xùn)練的快速注意力內(nèi)核。研究者將具備這些變化的網(wǎng)絡(luò)稱為 Sparse Transformer，并證明該網(wǎng)絡(luò)可以使用數(shù)百個(gè)層來建模成千上萬個(gè)時(shí)間步長(zhǎng)的序列。
該網(wǎng)絡(luò)在從原始字節(jié)中建模圖像、音頻和文本時(shí)使用的是同樣的架構(gòu)，在 Enwik8、CIFAR10 和 ImageNet-64 數(shù)據(jù)集上取得了當(dāng)前最優(yōu)的密度估計(jì)性能。研究者生成的無條件樣本展示了全局一致性和極大的多樣性，并證明原則上可以使用自注意力建模長(zhǎng)度超百萬的序列。