Longformer

Longformerは長いドキュメント向けのTransformerで、すべての自己注意を、各トークンの周囲にあるスライディングウィンドウ注意と、タスクに応じて選んだ少数のトークンに対するグローバル注意に置き換えます。そのため、計算コストは入力長の二乗ではなく線形に増加します。Allen Institute for AIのIz Beltagy、Matthew Peters、Arman Cohanが、2020年の論文Longformer: The Long-Document Transformerで発表しました。BERTやRoBERTaなどのエンコーダーは最大512トークンまでしか読み取れないため、それより長いテキストは分割する必要があり、分割間のつながりが失われていました。公開されたLongformerモデルは4,096トークンを一度に読み取るため、コンテキストウィンドウが8倍に拡大しました。

Longformerの注意機構の仕組み#

各層は1つの注意機構内で、局所的なパターンとグローバルなパターンを組み合わせます。

• スライディングウィンドウ注意: 各トークンは、公開モデルでは両側それぞれ256個の固定範囲内にある隣接トークンに注意を向けます。
• 拡張スライディングウィンドウ: 注意を向ける位置の間隔を空けることで、同じ計算コストのまま、ヘッドがより遠くまで見渡せるようにします。著者らは、文字レベルの言語モデルの上位層にある一部のヘッドで拡張を使用しました。
• グローバル注意: 選択されたトークンはすべての位置に注意を向け、すべての位置から選択されたトークンに注意を向けます。この処理には、独立したクエリー、キー、バリューの射影を使います。対象はタスクに応じて選択されます。分類では[CLS]トークン、質問応答では質問のトークンが選ばれます。

層を重ねると、積み重ねた畳み込みによって畳み込みニューラルネットワーク（CNN）の受容野が広がるのと同様に、ウィンドウの範囲が広がります。ウィンドウ幅がwのl層では、l × w個のトークンをカバーします。各層で計算する注意スコアはn × nではなく、およそn × wです。Longformerのマスク内の組み合わせ数を数えると、その違いが分かります。

import torch

window = 512  # Longformer-base: 256 tokens on each side of every token

for n in (1024, 2048, 4096):
    pos = torch.arange(n)
    # Sliding window: each token attends to neighbors within window / 2 positions
    mask = (pos[:, None] - pos[None, :]).abs() <= window // 2
    # Global attention on token 0 ([CLS]): it sees every token and every token sees it
    mask[0, :] = True
    mask[:, 0] = True
    print(f"{n} tokens: full attention {n * n / 1e6:.1f}M scores, Longformer {mask.sum().item() / 1e6:.1f}M")

Longformerでは0.5M、1.0M、2.0Mのスコアが出力される一方、フル注意では1.0M、4.2M、16.8Mとなります。入力を2倍にするたびにLongformerの注意計算量は2倍になりますが、フル注意では4倍になります。実際の実装ではスコアの帯状部分だけを計算し、これによってメモリを節約します。

訓練と結果#

LongformerはRoBERTaのチェックポイントから始めます。著者らは学習済み位置埋め込み512個を8回コピーして4,096位置に拡張し、長文書を使ってマスク言語モデルの事前学習を続けました。ファインチューニング後は、長文書タスクで一貫してRoBERTaを上回り、WikiHopとTriviaQAの質問応答ベンチマークで当時の最高性能を達成しました。また、文字レベルの言語モデルは、text8とenwik8で新たな記録を樹立しました。

Longformer-Encoder-Decoder（LED）は、テキスト要約やその他のテキスト生成タスク向けにデコーダーを追加したモデルで、arXivの科学論文全文の要約で実証されました。Hugging Face Transformersは両方のモデルを実装しており、グローバルトークンをglobal_attention_maskとして受け取ります。大規模な臨床テキストコーパスで事前学習され、固有表現認識、質問応答、ドキュメント分類など10個のタスクで評価されたClinical-Longformerなど、分野別モデルも登場しました。

Longformerと関連モデル#

• Transformer-XL: 長いテキストをセグメントごとに読み込み、各セグメントが前のセグメントのキャッシュ済み隠れ状態に注意を向けられるようにします。再帰によってコンテキストを拡張し、左から右への生成を行います。
• BigBird: 同じ年に登場したGoogleのモデルで、ウィンドウ注意とグローバル注意にランダムな注意リンクを加え、4,096トークンも読み取ります。
• スパース注意: Longformerが属する、固定された注意パターンのより広いカテゴリです。
• Flash Attention: オンチップメモリのタイル上で正確なフル注意を計算します。そのため、メモリ使用量は入力長に対して線形に増加しますが、スコアの数は二次関数的に増加します。

ビジョンにおける局所注意とグローバル注意#

高解像度画像でも同じ問題が生じます。ビジョントランスフォーマー（ViT）は画像を数千個のパッチトークンに変換するためです。Swin Transformerはシフトされた局所ウィンドウでこの問題に対処します。Ultralyticsパッケージで実行できるMetaのSegment Anything Model（SAM）は、ブロックごとにウィンドウ注意とグローバル注意を組み合わせ、画像セグメンテーションを行います。ViT-B画像エンコーダーは、64 × 64のパッチグリッドを重なりのない14 × 14のウィンドウに分割し、12個のブロックのうち8個ではウィンドウ内に注意を向け、ブロック2、5、8、11では4,096個すべてのパッチに対してグローバルに注意を向けます。