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