Hugging Face Reads: Long-range Transformers

TL;DR

Hugging Faceは、標準的なTransformerのメモリと時間の計算量が二次関数的（quadratic）になるボトルネックを解決するために、Longformer、Compressive Transformer、Linformer、Performerという4つの主要なアーキテクチャ的アプローチを分析しました。これらの手法は、従来の512または1024トークンの制限をはるかに超えるシーケンスの処理を可能にし、これはドキュメントレベルのNLP、音声、およびタンパク質モデリングにおいて極めて重要です。

二次関数のボトルネックの克服

標準的なTransformerの自己注意（self-attention）は、シーケンス長 $n$ に対して二次関数的（$O(n^2)$）にスケールするため、長いドキュメントに対して重大なメモリと計算のボトルネックを生じさせます。これに対処するため、研究者たちは、この複雑さを線形（$O(n)$）に削減することを目指す「Efficient Transformers」を開発してきました。これらのアプローチは、一般的に、カスタム注意パターン、再帰、低ランク近似、およびカーネル近似の4つのカテゴリに分類されます。

長距離モデリングへのアーキテクチャ的アプローチ

Longformer: カスタム注意パターン

Longformerは、標準的な自己注意をウィンドウ化（local）およびグローバルな注意の組み合わせに置き換えることで、シーケンス長に対して線形にスケールすることを可能にします。

Mechanism: 拡張されたウィンドウ化自己注意（dilated windowed self-attention）を自己回帰的な言語モデリングに使用し、ローカルなウィンドウ化注意とグローバルな双方向注意を組み合わせてエンコーダーの事前学習に使用します。グローバルな注意は、タスク固有のトークン（例：[CLS]トークンやQAにおける質問トークン）に適用され、情報がシーケンス全体に流れるようにします。
Key Advantage: 自己注意層はドロップイン・リプレイスメント（そのまま置き換え可能）であり、これは、高コストな一からの事前学習を必要とせずに、事前学習済みチェックポイントを長距離入力に適応させることができることを意味します。
Trade-off: スライディングウィンドウ注意は、TPUなどの特定のハードウェア上でインデックス操作が遅くなる可能性があるインデックス操作に依存しています。

Compressive Transformer: 再帰

Transformer-XLをベースに、Compressive Transformerは、通常であれば破棄されてしまう過去のアクティベーションを保存するための圧縮メモリを導入します。

Mechanism: 圧縮関数（max/mean poolingや1D convolutionなど）を使用して、過去のアクティベーションを係数 $c$ で圧縮します。これにより、モデルは、最近のトークンの通常のメモリと、はるかに古いトークンの圧縮メモリの両方に注意を向けることができます。
Key Advantage: enwik8 や WikiText-103 のような長距離言語モデリングのベンチマークにおいて、特に長い距離で発生する稀な単語に対して、パープレキシティを大幅に改善します。
Trade-off: 学習は不安定（brittle）であり、有効なバッチサイズを段階的に増やしていく特殊な最適化スケジュールが必要です。

Linformer: 低ランク近似

Linformerは、自己注意行列が低ランクであることを観察に基づき、シーケンス長をより小さな次元に投影することで複雑さを削減します。

Mechanism: Johnson-Lindenstraussの補題を利用して、Linformerは注意コンテキスト行列の低ランク分解を学習します。これにより、$n imes n$ 行列が計算または保存されないことが保証されます。
Key Advantage: 推論速度（タイムクロック）はシーケンス長の増加によって影響を受けず、収束速度も標準的なTransformerと比較して安定しています。
Trade-off: 分解は、学習時に決定される固定のコンテキスト長向けに設計されており、適応なしにはより長いシーケンスには汎用化できません。

Performer: カーネル近似

Performerは、疎性（sparsity）や低ランクの事前知識に頼ることなく、softmax attentionカーネルを近似するためにFAVOR+ (Fast Attention Via Orthogonal Random positive features) アルゴリズムを使用します。

Mechanism: ランダムな特徴マップを使用してsoftmax関数を近似し、$K imes V$ の行列乗算をクエリ乗算の前に実行できるようにします。これにより、$n imes n$ の注意行列の計算を事実上回避できます。
Key Advantage: 注意行列の構造に関する仮定を置かないため、音声やタンパク質シーケンスを含む、異なるモダリティに非常に適用可能です。
Trade-off: 小さな近似誤差が複数のTransformer層を通じて伝播し、事前学習済みネットワークのファレクトライニングの安定性に影響を与える可能性があります。

比較分析と議論

帰納バイアスとトレードオフ

アーキテクチャの選択は、タスクの具体的な要件と利用可能なデータに依存します：

Longformer vs. Linformer: Longformerは固定の疎なパターンを使用しますが、Linformerは低ランク分解を学習します。Longformerは一般的にLinformerよりも効率が低いですが、シーケンス長に関してはより柔軟です（ただし、Linformerは学習コンテキスト長に制限があります）。
Performer: カーネル自体を近似することで、両者とは異なります。注意行列が疎であるか、あるいは低ランクであるかという仮定を置かないため、多用途なドロップイン・リプレイスメントとなります。

位置エンコーディング

位置エンコーディングは、長距離の効率性における重要な要因です：

Relative Positional Embeddings: Transformer-XLやCompressive Transformersで使用されます。これらは未知のシーケンス長に容易に拡張可能ですが、計算コストが高いです。
Absolute Positional Embeddings: LongformerやLinformerで使用されます。これらは計算効率が高いですが、学習時よりも長いシーケンスに対しては柔軟性が低くなります。
Position-Infused Attention: Shortformerによって導入された代替案で、トークン埋め込みではなく、クエリとキーに位置情報を追加します。

学習戦略

ShortformerやLongformerのようなモデルから得られる証拠は、短いシーケンスで学習を開始し、徐々に長さを増やすことが、より高速な学習とより強力なダウンストリーム性能をもたらし、モデルがデータ内の偽の相関関係に依存することを防ぐことを示唆しています。