OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

機器之心發表於2019-04-25

選自 openai,機器之心編譯,機器之心編輯部。

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

現在,AI 研究中的一項挑戰是在影像、視訊或聲音等複雜資料中進行長序列的精細相關性建模。Sparse Transformer 合併了 O(N^2)Transformer 自注意力機制的 O(N√N) 重組以及其他一些改進,從而直接用於這些豐富的資料型別。以前,這些資料上所使用的模型是專為某個領域製作的,或者很難將序列擴充套件到包含幾千個元素。

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

深度注意力

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

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

當矩陣儲存在記憶體或在逆推計算過程中進行再計算時,深度 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 層模型的樣本:

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

左:
Layer 19,右:
Layer 20。
為一個 128 層的 CIFAR-10 網路的若干層學習注意力模式(白色高亮顯示)。
這些層學會了在兩個維度上分割注意力。
Layer 19 彙總每一行的資訊,Layer 20 按列彙總這些資訊,從而有效分解了全注意力運算。

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

為獲取位置記憶而訓練的層(左:
Layer 6;
右:
Layer 36),它們通常關注類似的位置,不管輸入資料或時間步長如何(Layer 6)。
其他層學習高度依賴資料的訪問模式(Layer 36)。

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

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

第一版 strided attention 大概等同於每個位置處理自己的行和列,它與以上網路學得的注意力模式類似。(注意,列注意力可等同於處理轉置矩陣的行)。第二版 fixed attention 在最新的列元素之後處理固定列和元素,研究者認為這個模式對於資料無法擬合二維結構(如文字)的情況很有用。

實驗結果

Sparse Transformer 在 CIFAR-10、Enwik8 和 Imagenet 64 資料集上重新整理了當前最優密度估計分數。

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

在 CIFAR-10、Enwik8 和 Imagenet 64 資料集上的密度估計效能(單位為 bits per byte/dim)。
M 表示網路中使用的引數(單位為百萬),W 表示網路寬度,L 表示層數,H 表示頭數。

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

生成影像

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

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

損壞原圖

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

修復影像

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

真實影像

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

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

模型示例

生成原始音訊波形

通過簡單改變位置嵌入,稀疏 Transformer 還能用來生成原始音訊,而非影像。隨著深度學習擴充套件到新的資料型別,用這類網路來指定歸納偏置也很容易。

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

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

程式碼公佈

通常,實現稀疏注意力需要將查詢和關鍵矩陣分割成塊,因此為了簡化實驗,OpenAI 實現了一組塊稀疏核,這些核在 GPU 上高效地執行這些操作。OpenAI 開源了這些核並提供了稀疏注意力函式的示例:github.com/openai/spar…

未來發展和限制

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

即使有了上述改進,自迴歸序列生成對非常高解析度影像和音訊來說仍是不切實際的。但是,研究人員介紹的優化注意力操作可能有用,將它與其它方法(如多尺度方法)結合,可以建模高維資料。

論文:Generating Long Sequences with Sparse Transformers

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍

論文連結:d4mucfpksywv.cloudfront.net/Sparse_Tran…

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

該網路在從原始位元組中建模影像、音訊和文字時使用的是同樣的架構,在 Enwik8、CIFAR10 和 ImageNet-64 資料集上取得了當前最優的密度估計效能。研究者生成的無條件樣本展示了全域性一致性和極大的多樣性,並證明原則上可以使用自注意力建模長度超百萬的序列。

OpenAI新研究補齊Transformer短板,將可預測序列長度提高30倍


相關文章