[論文レビュー] Unlimiformer: Long-Range Transformers with Unlimited Length Input

Amanda Bertsch, Uri Alon+2 more|arXiv (Cornell University)|May 2, 2023

Handwritten Text Recognition Techniques被引用数 23

ひとこと要約

Unlimiformer は、クロスアテンションのための単一の k-NN インデックスを既存のエンコーダ-デコーダ・トランスフォーマーに追加し、試行時に追加学習なしで無制限の入力長を実現し、長文ドキュメントおよび書籍要約の性能を改善します。

ABSTRACT

Since the proposal of transformers, these models have been limited to bounded input lengths, because of their need to attend to every token in the input. In this work, we propose Unlimiformer: a general approach that wraps any existing pretrained encoder-decoder transformer, and offloads the cross-attention computation to a single k-nearest-neighbor (kNN) index, while the returned kNN distances are the attention dot-product scores. This kNN index can be kept on either the GPU or CPU memory and queried in sub-linear time; this way, we can index practically unlimited input sequences, while every attention head in every decoder layer retrieves its top-k keys, instead of attending to every key. We evaluate Unlimiformer on several long-document and book-summarization benchmarks, showing that it can process even 500k token-long inputs from the BookSum dataset, without any input truncation at test time. We demonstrate that Unlimiformer improves pretrained models such as BART and Longformer by extending them to unlimited inputs without additional learned weights and without modifying their code. We make our code and models publicly available at https://github.com/abertsch72/unlimiformer .

研究の動機と目的

• 標準的なコンテキストウィンドウを超える極端に長い入力を、再訓練せずに処理できるよう動機づけと実現を図る。
• 全入力トークンに対するフル・クロスアテンションを置換する、一般的でパラメータを持たないリトリーバル機構を提案する。
• 単一の k-NN インデックスが、全デコーダ層とヘッドのアテンションマスを近似するのに十分であることを示す。
• 複数のベースモデルと訓練方式にわたる長文ドキュメントおよび書籍要約ベンチマークで改善を示す。

提案手法

• 各デコーダ層のクロスアテンションの前に k-NN リトリーバル・ステップを挿入し、ヘッドごとに top-k キーを選択する。
• 長い入力をオーバーラップを持つチャンクに分割してエンコードし、各チャンクの中間半分の隠れ状態をインデックス化する。
• QWqWk^T による per-head 投影を介してエンコーダ隠れ状態の単一インデックスを参照できるよう、アテンション計算を改写し、層/ヘッドを跨ぐ共有インデックスを可能にする。
• デコード時にインデックスをクエリして、上位-k retrieved keys のみに対してアテンションを行い、ドット積距離をアテンションスコアとして利用する。
• 16-bit 隠れ状態を用いてメモリを抑制（例: 1,000,000 トークンあたり 2 GB）、必要に応じてインデックスを CPU/GPU にオフロードする。
• 低コストのテスト時バリエーション（+test Unlimiformer、+early stop w/ Unlimiformer）と、長距離訓練手法（Random-encoded、Retrieval、Alternating）を提供する。
• LLaMA-2 および HuggingFace Transformers に対応したコードベースとモデルリリースを提供する。

実験結果

リサーチクエスチョン

• RQ1 エンコーダ-デコーダ・トランスフォーマーのクロスアテンションを k-NN インデックスにオフロードして、試行時の入力長を無限大にサポートできるか。
• RQ2 全デコーダ層/ヘッドにまたがる単一の共有 k-NN インデックスは、効果的なリトリーバルを満たし、ほとんどのアテンションマスを保持できるか。
• RQ3 追加の学習パラメータなしで、無制限の入力長を扱えるよう、既存の事前学習済みモデルを拡張できるか。
• RQ4 Unlimiformer の使用による長距離要約と関連タスクで、精度と計算のトレードオフはどのようになるか。

主な発見

• Unlimiformer は追加の訓練なしで長文要約のベースモデルを改善（例: BART_base +test Unlimiformer が標準のファインチューニングより ROUGE/L や BERTScore が高い）。
• Unlimiformer を用いた早期停止は追加訓練コストなしに substantial な向上を提供（例: GovReport: ROUGE-1 が 48.7 から 51.0 へ）。
• Unlimiformer で訓練された場合、PRIMERA のような長距離ベースラインを上回るか同等になり、Unlimiformer はそれらをさらに改善可能（例: PRIMERA +test Unlimiformer は標準 PRIMERA より ROUGE/L と EntMent を改善）。
• BookSum の結果は EntMent の利得が見られ、Unlimiformer+PRIMERA で EntMent が baseline PRIMERA の 11.6 から 25.5 に上昇。
• Retrieval-focused training variants（Retrieval、Random-encoded、Alternating）はデータセット全体で競合的な利得をもたらし、最良の手法はモデルとデータによって異なる。