OpenAI提出Sparse Transformer，文本、圖像、聲音都能預測，序列長度提高30倍

新聞 04-24

新智元報道

編輯：大明

【新智元導讀】OpenAI提出新的神經網路模型「稀疏Transformer」，能夠預測文本、圖像和聲音等序列的後續內容，該模型是對注意力機制的一個改進，預測長度達到之前最佳水平的30倍。

目前人工智慧研究的一大挑戰是對複雜數據（如圖像，視頻或聲音）中的大範圍微妙的相互依賴性進行建模。稀疏Transformer降低了傳統注意力機制模型的計算複雜度，將其直接應用於不同的數據類型中。以前，在這些數據上使用的模型是針對某個專門領域設計的，難以擴展到超過幾千個元素的序列規模上應用。

此次OpenAI提出的模型可以使用數百個層對數萬個元素的序列進行建模，在多個域中實現最先進的性能。稀疏Transformer能夠幫助我們構建具有更強的理解世界能力的AI系統。

深度注意力機制

在稀疏Transformer中，每個輸出元素都與每個輸入元素相連，它們之間的權重是根據環境動態計算的，這個過程稱為注意力。雖然這樣會讓模型比固定連接模式的模型更加靈活，但在實踐中需要為每個層和注意力頭N×N注意力矩陣，面對元素數量眾多的數據類型時會消耗大量的內存，比如圖像或原始音頻數據。

當矩陣存儲在內存中或在後向傳遞期間重新計算時，深度Transformer的內存消耗情況（64層、4個注意力頭）。作為參考，用於深度學習的標準GPU通常配備12-32GB的內存

減少內存消耗一種方法是在反向傳播期間從檢查點重新計算注意力矩陣，這是深度學習中的一種成熟技術，以增加計算量為代價來減少內存使用。在計算Transformer的注意力矩陣時，意味著最大的內存成本與層數無關，這使我們能夠以比以前更大的深度訓練神經網路。

實際上，我們發現深度達128層的Transformer在常用數據集基準任務（如CIFAR-10）上的表現優於較淺層的網路。

為了更深入地訓練這些模型，我們對Transformer中的操作順序進行了幾次調整，並修改了初始方案。

稀疏注意力機制：顯著降低計算複雜度

然而，即使是計算單個注意力矩陣，對於非常大的輸入也是不切實際。因此我們使用稀疏注意力模式，即每個輸出位置僅計算來自輸入位置子集的權重。當子集相對於整個輸入集較小時，即使對於非常長的序列，所得到的注意力計算也是容易處理的，演算法複雜度為O（N *sqrt ），而不是O（N^2）。

為了評估該方法的可行性，我們首先將深度Transformer在圖像上的學習注意模式進行可視化，發現許多模型表現出可解釋和結構化的稀疏模式。下面的每個圖像顯示給定的注意頭處理哪些輸入像素（以白色突出顯示）以便預測圖像中的下一個值。

當輸入部分聚焦在小的子集上並顯示出高度的規則性時，該層就是易於稀疏化的。下圖為CIFAR-10圖像上的128層模型示例。

左圖為19層，右圖為20層

學習後的128層CIFAR-10網路的多個層的注意力模式（白色高亮部分）。這些層學會將注意力分散在兩個維度上。其中第19層總結了每一行的信息，第20層則按列聚合這些信息，從而能夠對全面注意力操作進行有效分解。

左圖為第6層，右圖為第36層

一些層學會了訪問位置存儲器，無論輸入數據或時間步長如何，通常都會訪問類似的位置（第6層）。還有的層學習了高度依賴數據的訪問模式（第36層）。

雖然許多圖層顯示出了稀疏結構，某些層還清晰地顯示出在整個圖像上延伸的動態注意力。為了讓網路保持學習這些模式的能力，我們進行了注意力矩陣的二維分解，網路可以通過兩個稀疏注意力步驟來關注所有位置。

（左）普通transformer，（中）範圍注意力，（右）固定注意力

第一個版本，大範圍注意力，大致相當於參與其行和列的每個位置，並且類似於上面的網路學習的注意力模式。（注意，列注意力可以等效地表示成轉置矩陣的行注意力）。第二個版本是固定注意力，注意固定列和最新列元素之後的元素，我們發現這種模式在數據不適合二維結構（如文本）時很有用。

實驗結果：創造多個數據集上的新紀錄

稀疏Transformer在CIFAR-10，Enwik8和Imagenet 64上創造了密度估計的最新記錄。如下表所示：

在一系列數據集上的密度建模表現，M為網路中使用的參數數量（百萬），W為網路寬度，L為層數，H為注意力頭數量。

我們還發現，除了速度明顯更快之外，稀疏注意力模型的損失也要低於完全注意力模型。這可能表明我們的稀疏模式存在有用的歸納偏差，或是密集關注的潛在優化問題。

使用稀疏注意力的Transformer似乎有一個全局結構的概念，可以通過查看圖像完成來定性評估。我們對64×64 ImageNet上訓練的模型進行了可視化，如下圖所示：

Prompt

Completions

Ground truth

我們還利用未調整的softmax temperature 1.0下生成了完全無條件的樣圖。這些模型使用最大似然目標進行訓練，眾所周知，這類訓練的目標是覆蓋所有數據模式（包括可能不存在的數據），而不是增加小部分數據的保真度。從這些具有未調整溫度的模型中生成樣圖，可以讓我們看到模型認為存在於真實世界中圖像的完整分布。結果，一些樣本看起來很奇怪。