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

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

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

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

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

深度注意力架構

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

當矩陣存儲在內存或在逆推計算過程當中進行再計算時,深度 Transformer(64 層和 4 個頭)的注意力內存使用狀況。
做爲參考,用於深度學習的標準 GPU 內存一般是 12-32GB.

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

當 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-十、Enwik8 和 Imagenet 64 數據集上刷新了當前最優密度估計分數。

在 CIFAR-十、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 開源了這些核並提供了稀疏注意力函數的示例:github.com/openai/spar…

將來發展和限制

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

即便有了上述改進,自迴歸序列生成對很是高分辨率圖像和音頻來講還是不切實際的。可是,研究人員介紹的優化注意力操做可能有用,將它與其它方法(如多尺度方法)結合,能夠建模高維數據。

論文:Generating Long Sequences with Sparse Transformers

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

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

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

參考連接: openai.com/blog/sparse…
相關文章
相關標籤/搜索