BigBirdのブロック疎なアテンションの理解

BigBirdは、標準的なTransformerモデルの$O(n^2)$の時間およびメモリ計算量を解決し、最大4096トークンのシーケンスを扱うことを可能にします。フルアテンションをブロック疎なアテンションメカニズムに置き換えることで、BigBirdは長い文書の要約や長いコンテキストを伴う質問応答などの長シーケンスタスクにおいて、最先端の結果を達成します。

ブロック疎なアテンションメカニズム

BigBirdのブロック疎なアテンションは、優位性よりも効率性を目的として設計されたBERTのフルアテンションの近似です。これは、以下の3つの異なるタイプのアテンションを組み合わせることで、クエリトークンが注目すべきトークンの数を削減します：

  • Sliding Attention: トークンは自身の直近の隣接トークンに注目します。これは、シーケンス内の単語が通常、隣接するトークンに強く依存しているため、局所的な依存関係を捉えます。
  • Global Tokens: 少数のトークンセットがシーケンス内の他のすべてのトークンに注目し、また他のすべてのトークンからも注目されます。これにより、スライディングアテンション単独よりも迅速に情報がシーケンス全体に伝播します。
  • Random Tokens: いくつかのトークンがランダムに選択され、他のトークンに注目します。これにより、シーケンスグラフ内の離れたノード間の情報の移動距離をさらに短縮します。

フルアテンションは単一のレイヤーで任意の2つのトークン間で情報を転送できますが、ブロック疎なアテンションでは、情報がシーケンス全体を移動するために複数のレイヤーが必要になる場合があります。GlobalおよびRandomの接続を組み合わせることで、わずか数レイヤーで情報を迅速に伝達できることが保証されます。

技術的実装

GPUおよびTPU上でこれを効率的に実装するために、BigBirdはブロックベースのアプローチを使用します。シーケンスはサイズ$b$のブロックに分割されます。

ブロックによるアテンション計算

  • First and Last Blocks: 最初のブロック($q_1$)と最後のブロック($q_n$)は、シーケンス内の他のすべてのトークンに対して通常のアテンション操作を実行します。
  • Middle Blocks: ブロック$q_{3:n-2}$のトークンについては、モデルはglobal、sliding、およびrandomのキーを集め、選択されたそれらのキーに対してのみアテンションを計算します。
  • Boundary Blocks: ブロック$q_2$および$q_{n-1}$は、疎な構造を維持するために、特定のキーのサブセット（最初のブロック、最後のブロック、および近接するスライディングブロックを含む）を集めます。

計算量比較

BigBirdは、シーケンス長が増加するにつれて計算負荷を大幅に軽減します。BERTの二次関数的なスケーリングと比較して、BigBirdは線形にスケーリングします：

Attention Type Sequence Length Time & Memory Complexity
Original Full (BERT) 512 $T$
Original Full (BERT) 1024 $4 ✕ T$
Original Full (BERT) 4096 $64 ✕ T$
Block Sparse (BigBird) 1024 $2 ✕ T$
Block Sparse (BigBird) 4096 $8 ✕ T$

トレーニング戦略: ITC vs. ETC

BigBirdは、2つの異なる戦略を使用してトレーニングできます：

  • Internal Transformer Construction (ITC): 少数のglobalトークンを使用します。計算量が少なく、多くのタスクにおいて十分です。
  • Extended Transformer Construction (ETC): 追加のglobalトークンを追加します。これは、質問応答のような、コンテキストトークンによって質問全体がグローバルに注目される必要があるタスクにおいて特に有用です。

Hugging Face Transformersとの統合

BigBirdは、transformersライブラリを通じてBigBirdModelとして利用可能です。主な実装の詳細は以下の通りです：

  • Sequence Length Requirements: シーケンス長はblock_sizeの倍数である必要があります。Hugging Faceは、シーケンスをブロックサイズの最小の倍数に自動的にパディングします。
  • Automatic Fallback: シーケンス長が、必要なglobal、random、およびslidingトークンをサポートするために短すぎる場合、またはモデルがデコーダー(BigBirdForCausalLM)として使用される場合、ライブラリはattention_typeを自動的にoriginal_fullに切り替えます。
  • Recommended Usage: 著者は、1024トークン未満のシーケンスに対してはattention_type="original_full"に設定することを推奨しています。

利用可能なチェックポイントは、bigbird-roberta-base、bigbird-roberta-large、およびbigbird-base-trivia-itcです。

Sources