Hugging Face Reads: Long-range Transformers

TL;DR

Hugging Face 分析了四種主要的架構方法——Longformer、Compressive Transformer、Linformer 和 Performer——以解決標準 Transformer 的二次記憶體與時間複雜度瓶頸。這些方法能夠處理遠超傳統 512 或 1024 token 限制的序列，這對於文件級 NLP、語音和蛋白質建模至關重要。

Overcoming the Quadratic Bottleneck

標準 Transformer 的自注意力機制會隨著序列長度 $n$ 以二次方 ($O(n^2)$) 比例擴展，這為長文件帶來了顯著的記憶體與運算瓶頸。為了應對這一點，研究人員開發了「高效 Transformer」（Efficient Transformers），旨在將此複雜度降低至線性 ($O(n)$)。這些方法通常分為四類：自定義注意力模式、遞歸、低秩近似和核函數近似。

Architectural Approaches to Long-Range Modeling

Longformer: Custom Attention Patterns

Longformer 使用結合了窗口（局部）與全域注意力的組合來取代標準自注意力，使其能夠隨序列長度線性擴展。

  • Mechanism: 它使用擴張窗口自注意力（dilated windowed self-attention）進行自回歸語言建模，並使用局部窗口與全域雙向注意力的混合模式進行編碼器預訓練。全域注意力被應用於特定任務的 token（例如 [CLS] token 或問答系統中的問題 token），以允許資訊在整個序列中流動。
  • Key Advantage: 自注意力層是一個即插即用的替換方案，這意味著預訓練權重可以被調整以適應長範圍輸入，而不需要從頭開始進行昂貴的預訓練。
  • Trade-off: 滑動窗口注意力依賴於索引操作，這在某些硬體（如 TPU）上可能會很慢。

Compressive Transformer: Recurrence

基於 Transformer-XL，Compressive Transformer 引入了壓縮記憶體來儲存原本會被捨棄的過去激活值。

  • Mechanism: 它使用壓縮函數（例如 max/mean pooling 或 1D convolution）將過去的激活值壓縮 $c$ 倍。這使得模型能夠同時關注近期 token 的常規記憶體與更久遠 token 的壓縮記憶體。
  • Key Advantage: 它顯著提升了在長範圍語言建模基準測試（如 enwik8 和 WikiText-103）上的困惑度（perplexity），特別是對於在長距離出現的罕見詞彙。
  • Trade-off: 訓練過程較為脆弱，需要特殊的優化策略，其中有效批次大小（effective batch size）會逐漸增加。

Linformer: Low-Rank Approximations

Linformer 基於自注意力矩陣是低秩的這一觀察，透過將序列長度投影到較小的維度來降低複雜度。

  • Mechanism: 利用 Johnson-Lindenstrauss lemma，Linformer 學習注意力上下文矩陣的低秩分解。這確保了不會計算或儲存任何 $n imes n$ 矩陣。
  • Key Advantage: 推論速度（time-clock）不會受到序列長度增加的影響，且與標準 Transformer 相比，收斂速度保持穩定。
  • Trade-off: 該分解是針對訓練時確定的固定上下文長度設計的，若無調整，則無法泛化到更長的序列。

Performer: Kernel Approximations

Performer 使用 FAVOR+ (Fast Attention Via Orthogonal Random positive features) 演算法來近似 softmax 注意力核函數，而不依賴於稀疏性或低秩先驗。

  • Mechanism: 它使用隨機特徵映射來近似 softmax 函數，允許在進行 query 乘法之前先執行矩陣乘法 $K imes V$。這有效地繞過了 $n imes n$ 注意力矩陣的計算。
  • Key Advantage: 因為它不對注意力矩陣的結構做任何假設，所以它在不同模態（包括語音和蛋白質序列）中都具有高度的適用性。
  • Trade-off: 微小的近似誤差可能會在多個 Transformer 層中傳播，進而可能影響預訓練網路微調時的穩定性。

Comparative Analysis and Discussion

Inductive Biases and Trade-offs

架構的選擇取決於特定任務的需求與可用數據：

  • Longformer: Longformer 使用固定稀疏模式，而 Linformer 學習低秩分解。Longformer 通常比 Linformer 效率低，但在序列長度方面更具靈活性（儘管 Linformer 限制於其訓練時的上下文長度）。
  • Performer: 與前兩者不同，它透過近似核函數本身來達成，使其成為一種多功能的即插即用替換方案，且不假設注意力矩陣是稀疏或低秩的。

Positional Embeddings

位置編碼是長範圍效率的關鍵因素：

  • Relative Positional Embeddings: 用於 Transformer-Xl 與 Compressive Transformers；它們可以輕易地擴展到未見過的序列長度，但運算成本較高。
  • Absolute Positional Embeddings: 用於 Longformer 與 Linformer；它們在運算上更有效率，但在處理比訓練時見過的序列更長時，靈活性較低。
  • Position-Infused Attention: 由 Shortformer 引入的一種替代方案，它將位置資訊加入到 queries 與 keys，而非 token embeddings。

Training Strategies

來自 Shortformer 與 Longformer 等模型的證據表明，在短序列上進行訓練並逐漸增加長度，可以實現更快的訓練速度與更強的下游任務表現，並能防止模型依賴數據中的虛假相關性。

Sources