Transformer是一種強大的序列模型，但是它所需的時間和內存會隨著序列長度出現二階增長。近日，OpenAI研究人員開發出了一種深度神經網絡Sparse Transformer，該網絡在預測長序列方面創造了新紀錄——無論預測的是文本、圖像還是聲音。該神經網絡利用注意力機制中的一種改進算法，可以從長度可能是之前30倍的序列中提取模式。

現在，AI 研究中的一項挑戰是在圖像、視頻或聲音等復雜數據中進行長序列的精細相關性建模。Sparse Transformer 合并了 O（N^2）Transformer 自注意力機制的 O(N√N) 重組以及其他一些改進，從而直接用于這些豐富的數據類型。以前，這些數據上所使用的模型是專為某個領域制作的，或者很難將序列擴展到包含幾千個元素。

相比之下，OpenAI 開發的模型通過使用數以百計的層可以對包含上萬個元素的序列進行建模，在諸多領域都取得了當前最佳的表現。OpenAI 研究人員利用該模型幫助創建能夠更好地理解世界的 AI 系統。

深度注意力

在 Transformer 中，每一個輸出元素與輸入元素相連接，同時根據具體情況對它們之間的權重進行動態計算，這一過程被稱為「注意力機制」。雖然人們相信這使得 Transformer 較那些具有固定連接模式的模型更為靈活，但實際操作中需要為每一層和注意力頭創建一個 N×N 注意力矩陣，當應用于圖像或原始音頻等具有許多元素的數據類型時會消耗大量內存。

當矩陣存儲在內存或在逆推計算過程中進行再計算時，深度 Transformer（64 層和 4 個頭）的注意力內存使用情況。作為參考，用于深度學習的標準 GPU 內存通常是 12-32GB.

減少內存消耗的一種方法是在反向傳播過程中從檢查點處重新計算注意力矩陣，這是深度學習中的一種成熟的方法，以更多的計算來減少內存使用。

當 Transformer 中的注意力矩陣完成時，這意味著最大的內存消耗將不受層數的支配，使研究人員訓練網絡的深度大大超過從前。在實際操作中，研究人員發現在處理 CIFAR-10 等基準測試任務時，深度達 128 層的 Transformer 表現出的性能優于較淺的網絡。

為了訓練深度更大的模型，研究人員對 transformer 的操作順序進行了幾次調整，修改了初始化方法。詳情參見論文。

稀疏注意力

然而，對于非常大的輸入來說，甚至計算單個注意力矩陣都是不現實的。因此，OpenAI 使用了稀疏注意力模式，在這種模式中，每個輸出位置僅從輸入位置子集中計算權重。當子集相對于整個輸入集較小時（如元素數量是√N 而不是 N），即使對于非常長的序列，注意力計算也會變得比較容易，算法復雜度為 O（N√N）而不是 O（N^2）。

為了評估該方法的可行性，研究人員首先可視化并學習了圖像上深度 Transformer 的注意力模式，發現其中許多模式表現出了可解釋和結構化的稀疏模式。以下每幅圖像都顯示了哪個輸入像素（白色高亮標出）由一個給定的注意力頭處理，以預測圖像中的下一個值。當輸入部分集中在小的子集上并顯示出高度規律性時，該層就易于稀疏化。以下是 CIFAR-10 圖像上 128 層模型的樣本：

左：Layer 19，右：Layer 20。為一個 128 層的 CIFAR-10 網絡的若干層學習注意力模式（白色高亮顯示）。這些層學會了在兩個維度上分割注意力。Layer 19 匯總每一行的信息，Layer 20 按列匯總這些信息，從而有效分解了全注意力運算。

為獲取位置記憶而訓練的層（左：Layer 6；右：Layer 36），它們通常關注類似的位置，不管輸入數據或時間步長如何（Layer 6）。其他層學習高度依賴數據的訪問模式（Layer 36）。

雖然許多層顯示出稀疏的結構，但有些層清晰地顯示出了動態注意力，這種注意力延伸到整個圖像。為了保持網絡學習這種模式的能力，研究人員實現了注意力矩陣的二維分解，其中網絡可以通過兩步稀疏注意力關注到所有位置。

第一版 strided attention 大概等同于每個位置處理自己的行和列，它與以上網絡學得的注意力模式類似。（注意，列注意力可等同于處理轉置矩陣的行）。第二版 fixed attention 在最新的列元素之后處理固定列和元素，研究者認為這個模式對于數據無法擬合二維結構（如文本）的情況很有用。

實驗結果

Sparse Transformer 在 CIFAR-10、Enwik8 和 Imagenet 64 數據集上刷新了當前最優密度估計分數。

在 CIFAR-10、Enwik8 和 Imagenet 64 數據集上的密度估計性能（單位為 bits per byte／dim）。M 表示網絡中使用的參數（單位為百萬），W 表示網絡寬度，L 表示層數，H 表示頭數。

研究者還發現稀疏注意力比完整注意力的損失更低，且速度更快。這可能指向稀疏模式產生的有用歸納偏置，或者密集注意力的底層優化問題。

生成圖像

使用了稀疏注意力的 Transformer 似乎有一種全局結構的概念，這可以通過觀察圖像補全（image completion）進行定性評估。下圖可視化了一個在 64×64 ImageNet 上訓練的模型：

損壞原圖

修復圖像

真實圖像

研究人員還生成了完全無條件的樣本，其中未調整的 softmax 溫度為 1.0。這些模型使用最大似然目標進行訓練，其覆蓋了所有的數據模式（其中包括可能不存在的數據），而不是增強較小部分數據的保真度。從具有未調整溫度的模型中取樣，研究人員看到了該模型認為世界上存在的圖像的完整分布。因此，一些樣本看起來奇奇怪怪的。

模型示例

生成原始音頻波形

通過簡單改變位置嵌入，稀疏 Transformer 還能用來生成原始音頻，而非圖像。隨著深度學習擴展到新的數據類型，用這類網絡來指定歸納偏置也很容易。

該模型是在原始的古典音樂片段上訓練的，并使用了稀疏注意力來生成長度為 65000 的序列。這相當于大約 5 秒長的原始音頻，研究人員在下面的每個片段中將幾個樣本連接在一起。

代碼公布

通常，實現稀疏注意力需要將查詢和關鍵矩陣分割成塊，因此為了簡化實驗，OpenAI 實現了一組塊稀疏核，這些核在 GPU 上高效地執行這些操作。OpenAI 開源了這些核并提供了稀疏注意力函數的示例：

https://github.com/openai/sparse_attention

未來發展和限制

本文介紹的稀疏注意力模式只是對長序列進行高效建模的初步嘗試。研究人員認為，探索稀疏注意力的不同模式和各種組合非常有用，而且學習稀疏模式對下一代神經網絡架構來說也是一個很重要的研究途徑。

即使有了上述改進，自回歸序列生成對非常高分辨率圖像和音頻來說仍是不切實際的。但是，研究人員介紹的優化注意力操作可能有用，將它與其它方法（如多尺度方法）結合，可以建模高維數據。

論文：Generating Long Sequences with Sparse Transformers

論文鏈接：https://d4mucfpksywv.cloudfront.net/Sparse_Transformer/sparse_transformers.pdf

摘要：Transformer 是一種強大的序列模型，但是它所需的時間和內存會隨著序列長度出現二階增長。這篇論文介紹了注意力矩陣的稀疏因式分解，可以將其降低到 O(N√N)。該研究提出了 a）訓練更深網絡的架構和初始化變體；b）重新計算注意力矩陣以節省內存；c）用于訓練的快速注意力內核。研究者將具備這些變化的網絡稱為 Sparse Transformer，并證明該網絡可以使用數百個層來建模成千上萬個時間步長的序列。

該網絡在從原始字節中建模圖像、音頻和文本時使用的是同樣的架構，在 Enwik8、CIFAR10 和 ImageNet-64 數據集上取得了當前最優的密度估計性能。研究者生成的無條件樣本展示了全局一致性和極大的多樣性，并證明原則上可以使用自注意力建模長度超百萬的序列。