Unlimiformer: Long-Range Transformers with Unlimited Length Input

要約

トランスフォーマーが提案されて以来、これらのモデルは入力内のすべてのトークンに注意する必要があるため、制限された入力長に制限されてきました。
この研究では、Unlimiformer を提案します。これは、既存の事前トレーニング済みエンコーダー デコーダー トランスフォーマーをラップし、クロス アテンションの計算を単一の k 近傍 (kNN) インデックスにオフロードし、返された kNN 距離がアテンション ドットである一般的なアプローチです。
-製品のスコア。
この kNN インデックスは GPU または CPU メモリに保持され、サブリニア時間でクエリできます。
このようにして、すべてのデコーダ層のすべてのアテンション ヘッドがすべてのキーに注目するのではなく、その上位 k 個のキーを取得しながら、事実上無制限の入力シーケンスにインデックスを付けることができます。
Unlimiformer をいくつかの長い文書と本の要約ベンチマークで評価し、BookSum データセットからの 500k トークン長の入力でも、テスト時に入力を切り捨てることなく処理できることを示しました。
Unlimiformer は、追加の学習重みやコードの変更を行わずに無制限の入力に拡張することで、BART や Longformer などの事前トレーニング済みモデルを改善することを示します。
私たちはコードとモデルを https://github.com/abertsch72/unlimiformer で公開しています。

要約(オリジナル)

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 .

arxiv情報

提供元, 利用サービス

arxiv.jp, Google