OpenAI Sparse Transformer

OpenAI 開發了 Sparse Transformer，這是一種旨在預測文本、圖像和聲音序列中下一個元素的深度神經網絡。透過對注意力機制（attention mechanism）進行演算法改進，該模型可以從比以往長 30 倍的序列中提取模式，從而能夠使用數百層來對數萬個元素進行建模。

使用深度注意力減少記憶體限制

標準的 Transformer 架構在每一層和每個注意力頭（attention head）都需要一個 $N \times N$ 的注意力矩陣。當處理高維數據（如原始音訊或大型圖像）時，這會造成顯著的記憶體瓶頸。例如，一張 64x64x3 像素的 ImageNet 圖像在深度 Transformer（64 層，4 個頭）中存儲矩陣需要 154 GB 的記憶體，遠遠超過標準 GPU 的 12-32 GB 容量。

為了應對這一點，OpenAI 實施了兩項主要優化：

  1. Checkpointing: 透過在反向傳播期間從檢查點（checkpoints）重新計算注意力矩陣，最大的記憶體成本變得與層數無關。這允許訓練具有實質上更大深度的網絡；OpenAI 發現，在 CIFAR-10 等基準測試中，多達 128 層的模型表現優於較淺的網絡。
  2. Operational Adjustments: 團隊修改了初始化方案和 Transformer 內部的操作順序，以支持增加的深度。

稀疏注意力與演算法複雜度

為了處理即使是單個注意力矩陣都變得不切實際的超大型輸入，Sparse Transformer 利用了稀疏注意力模式（sparse attention patterns）。在這種方法中，每個輸出位置僅從輸入位置的一個子集（例如 $\sqrt{N}$ 個元素而非 $N$ 個）計算權重。這將演算法複雜度從 $O(N^2)$ 降低到 $O(N \sqrt{N})$。

分解注意力模式

為了確保模型保留學習動態、全局模式的能力，OpenAI 對注意力矩陣進行了二維分解。這使得網絡能夠透過兩個步驟的稀疏注意力來關注所有位置：

  • Strided Attention: 每個位置關注其自身的行和列，有效地分解了完整的注意力操作。這對於具有二維結構的數據（如圖像）特別有用。
  • Fixed Attention: 網絡關注一個固定的列以及緊隨最新列元素之後的元素。這種模式針對不符合二維結構的數據（如文本）進行了優化。

實驗結果與性能

Sparse Transformer 在以下基準數據集上實現了密度估計的新 SOTA 分數：

  • CIFAR-10
  • Enwik8
  • Imagenet 64

OpenAI 觀察到，稀疏注意力不僅比全注意力（full attention）運行得更快，而且還實現了更低的損失（loss）。研究人員認為這可能是由於密集注意力（dense attention）潛在的優化問題，或者是稀疏模式提供的有益歸納偏置（inductive bias）。

跨模態的生成能力

圖像生成

該模型透過在 64x64 ImageNet 數據上的圖像補全任務展示了對全局結構的掌握。由於模型是使用最大似然目標進行訓練的，使用未調整的 softmax 溫度為 1.0 的無條件樣本，反映了模型認為存在的圖像完整分佈，這偶爾會導致看起來很奇怪的樣本。

原始音訊生成

透過修改位置嵌入（position embeddings），Sparse Transformer 可以生成原始音訊波形。OpenAI 在原始古典音樂片段上訓練了該模型，使其能夠生成 65,000 個元素的序列，這大約相當於 5 秒的原始音訊。

實施與限制

為了促進稀疏注意力的使用，OpenAI 開源了一組 block-sparse GPU kernels，可以高效地對 query 和 key 矩陣進行必要的切片操作。

儘管有這些進步，研究人員指出，對於非常高解析度的圖像或視頻，自回歸序列生成仍然是不切實際的的。他們建議，這些優化的注意力操作未來可以作為多尺度建模方法的基礎組件。

Sources