Longformer: 長いシーケンスへの対応を可能にしたトランスフォーマーの改良版

BERT 2023年09月04日

3つの要点
✔️ 長いシーケンスの効率的な処理という課題に対する解決策を提示した
✔️ Sliding Window Attenion、Dilated Sliding Window Attention、Global Attentionという3つのAttentionを使ってTransformernの計算量を削減した
✔️ 長い文章を入力とするタスクの精度が向上した

Longformer: The Long-Document Transformer
written by Iz Beltagy, Matthew E. Peters, Arman Cohan
(Submitted on 10 Apr 2020 (v1), last revised 2 Dec 2020 (this version, v2))
Comments: Version 2 introduces the Longformer-Encoder-Decoder (LED) model
Subjects: Computation and Language (cs.CL)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

概要

LongformerはTransformerのself-attentionの計算量が，$O(n^2)$であり，長い文章を入れると，メモリの使用量が入力の増加に伴って2次的に急増する問題に対して，attentionを工夫することで対応しようとしたものです．

Transformerの問題点とは

Transformerは計算量が入力シーケンスに伴って2次的に増加していきます．これにより，メモリの使用量も計算時間も非常に長くなります．

なぜ，Transformerの計算量は入力シーケンスに伴って2次的に増加するのかというと，Transformerの主要パーツである，Scaled Dot-Product Self-Attentionのせいです．そもそも，Scaled Dot-Product Self-Attentionとは，クエリとキー・バリューののペアを使ってAttentionを計算するものです．つまり，Scaled Dot-Product Self-Attentionの計算式は，クエリ，キー，バリュー($Q，K，V$)を用いて，

$$Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}}V)$$

で示されます．この時，クエリとバリュー（Q，V）の積は文書の長さ（$n$）の二乗になります．そのため，2046トークンが入力として与えられた時にAttentionの計算に用いる行列サイズは，2024*2024となり，約410万要素を持つ行列をAttentionの計算で処理しなければならなくなります．

これがバッチサイズ分の計算ともなると膨大になり，メモリ容量との相談となってきます．そこで，このTransformerの問題である入力シーケンスの二乗に伴って計算量が増加する問題に取り組んだのが本論文であるLong Transformerです．

Longformerとは

Scaled Dot-Product Self-Attentionでは，すべての単語からすべての単語へAttentionを向けていることによって，計算量もメモリの使用量も$n^2$となっていました．そのため，可能な限り，重要な単語から重要な単語へのみAttentionを向けるように工夫を提案したものです．

実際に下記で示すように，かなり計算時間もメモリの容量も抑えられています．

Longformerの学習の工夫

LongformerでのAttentionの工夫を説明します．具体的には三つの工夫したAttentionを提案しています．

（a）Full $ n^2$ attention

Trasformerで使われているScaled Dot-Product Self-Attentionのことです．そのため，すべての単語からすべての単語へとアテンションが向けられています．

（b）Sliding window attention

Sliding window attentionでは，自分自身の近くにしかAttentionを向けません．このときに attentionを向ける単語数によってメモリの使用量を計算することができます．具体的な使用メモリの変化量は，ウィンドウサイズ$w$として，自分自身から左右に$\frac{1}{2}w$個の単語へAttentionを向けていることになるため，メモリの使用量は，文書の長さに対して線形となり，文書の長さ$n$とした時に，計算量は$O(nw)$となります．

ウィンドウサイズ$w$の設定をレイヤーごとに変える工夫もしています．具体的にはウィンドウサイズ$w$を下流のレイヤーでは小さくし，上流のレイヤーでは大きくしています．このことにより，

- 下流レイヤーでは局所的な情報の収集
- 上流レイヤーでは全体的な情報の収集

のような効果が得られます．

（c）Dilated sliding window

Dilated sliding window は自分自身の近くだけではなく，単語をある一定間隔で飛ばすことで，計算量を抑えつつも，適度に遠くにもAttentionを向けるというものです．論文では，単語を飛ばす感覚をギャップサイズとして$d$と表しています．また，Attentionを向ける単語を$d$個ずつとばして，multi-head attentionごとに$d$の値を変更することで性能向上が見込めるようです．

ただ，Dilated sliding windowは下流レイヤーでは局所的な情報を収集するために行うことはしておらず，上位レイヤーの2つのHeadのみに限定しています．
これは，実験されており結果は以下のように，ウィンドウサイズを増加させた場合の方が精度が改善し，上位レイヤーの2つのHeadのみに限定した方が限定しないよりも精度が向上しているためです．

（d）Global+sliding window

これは，Global AttentionとSliding window attentionの組み合わせで，Sliding window attentionではすべての単語にAttentionは向けず，自分自身の周りだけでしたが，Global Attentionでは特定の単語の時はすべての単語にAttentionを向けて，任意の単語の時はその特定の単語にAttentionを向けるようにするというものです．

具体的にいうと，BERTで考える場合，[CLS]というスペシャルトークンは文書のクラス分類に用いられます．そのため，[CLS]の位置の時はすべての単語にAttentionを向け，[CLS]以外の単語は[CLS]に対してAttentionを向けるようにするというものです．このことで，[CLS]はすべての単語へAttentionを向ける・向けられることができ文書分類を行う際に文全体の特徴を得ることができます．このGlobal Attentionによる計算量はあくまで特定の単語位置のみの話であるため$O(n)$です．

このGlobal AttentionはSliding window attentionと一緒に使うことが前提のように主張されています．

また，Global AttentionとSliding window attentionでは，Transformerのように，一つの$(Q, K, V)$でAttentionを計算するのではなく，それぞれの線型写像を用いて計算するようです．このことによりモデルがそれぞれのAttentionに柔軟に対応することができて，下流のタスクにおいてより良い性能を示せるようになります．

タスクごとの精度改善

評価のためのモデルの訓練

評価のためにモデルを訓練する際に5段階に分けて訓練しています．分ける理由ですが，モデルを訓練する際の理想はGPUメモリで最大のウィンドウサイズとシーケンス長で訓練することです．しかし，長いコンテキスト利用して学習する際はそうもいきません．また，学習初期の段階では多くの勾配の更新を行う必要があり，短いシーケンスから入力をして学習レートを低く設定し，ウィンドウサイズを小さくすることで学習が効率的に行えます．そのため本研究では，開始シーケンス長2,048、終了シーケンス長23,040として学習を5段階のフェーズに分け行っています．

最高性能を達成したモデルの学習の各段階における各パラメータの設定値の詳細は以下のようになっています．

学習の結果

学習の結果は旧来のTransformerやTransformerの改善を試みたReformerなどよりも性能が向上しています．

BERTでの検証

画期的なことで，Reformerなどでは行われていない，BERTでのタスクごとの精度比較を行っています．これによると，ほぼ全てのタスクで高精度を実現しています．

検証にはRoBERTaを使ってファインチューニングを行った結果を用いています．ただ，RoBERTaにはPosition Embeddingが512個しかないので，4096単語に対応させるためにPosition Embeddingを8個コピーして対応しています．単純にもかかわらず非常に有効であることが主張されていて，これは，コピーによってパーティションの境界が無くなるためのようです．

BERTでの検証の結果

BERT関連のモデルでの各タスクの評価は以下のようになっています．

それぞれのタスクは，

WikiHop:
複数文書を読んで情報をつなぎ合わせて回答するタスク

TriviaQA：
1つのインスタンスが質問，答え，根拠からなっているデータセット

HotpotQA：
複数文書を読んで情報をつなぎ合わせて回答するタスク

OntoNotes(Coreference Resolution)：
同一のモノを指す２つ以上の語句を見つけ対応づけるタスク

IMDB：
映画レビューデータセットで比較的短い文章の分類タスク

Hyperpartisan：
ニュースデータセットで比較的長い文章の分類タスク

それぞれのタスクでどのデータセットにおいても，Longformerはうまく処理できており，Longformerの有用性が示されています．

まとめ

Longformerは，長い文書を処理するためのスケーラブルなTransformerベースのモデルであり，長い入力をチャンク化や短縮することなく，またチャンク間の情報を結合するための複雑なアーキテクチャを用いることなく，幅広い文書レベルのNLPタスクを簡単に実行することができる仕組みです．Longformerは，局所的な情報と大域的な情報を組み合わせると同時に，シーケンスの長さに応じて線形にスケーリングするために3つのアテンション（Sliding Window Attention、Dilated Sliding Window Attention、Global Attention）を採用しています．また，Longformerはtext8とenwik8のタスクで最高性能を達成しました．さらに，Longformerは事前に学習させることで，長文文書タスクにおいてRoBERTaを常に上回り，WikiHopとTriviaQAにおいて最高性能を達成しました．