[論文レビュー] Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

Ofir Press, Noah A. Smith+1 more|arXiv (Cornell University)|Aug 27, 2021

Topic Modeling参考文献 30被引用数 30

ひとこと要約

本論文はALiBiを導入する。注意機構における単純な線形バイアスで、学習時に見た長さを超える長さの入力へ外挿することを可能にし、追加のパラメータや実行時コストを要しない。シ sinusoidal-positionモデルと同等かそれ以上のパープレキシティを達成しつつ、トレーニングはより速く、メモリ使用量も少なくなる。

ABSTRACT

Since the introduction of the transformer model by Vaswani et al. (2017), a fundamental question has yet to be answered: how does a model achieve extrapolation at inference time for sequences that are longer than it saw during training? We first show that extrapolation can be enabled by simply changing the position representation method, though we find that current methods do not allow for efficient extrapolation. We therefore introduce a simpler and more efficient position method, Attention with Linear Biases (ALiBi). ALiBi does not add positional embeddings to word embeddings; instead, it biases query-key attention scores with a penalty that is proportional to their distance. We show that this method trains a 1.3 billion parameter model on input sequences of length 1024 that extrapolates to input sequences of length 2048, achieving the same perplexity as a sinusoidal position embedding model trained on inputs of length 2048 but training 11% faster and using 11% less memory. ALiBi's inductive bias towards recency also leads it to outperform multiple strong position methods on the WikiText-103 benchmark.

研究の動機と目的

• 標準的な位置エンコーディングを用いた場合、トランスフォーマーが学習時より長いシーケンスへ外挿する能力が低い理由を調査する。
• 外挿の効率と性能の観点から、既存の位置埋め込み法（sinusoidal、rotary、T5 bias）を評価する。
• 追加のパラメータやランタイムコストなしに長いシーケンスへ外挿を可能にする、軽量で効果的な位置エンコーディング法を提案する。
• データセット・モデルサイズ・トレーニング予算を跨いだALiBiの有効性を示す。

提案手法

• ポジショナル埋め込みを、クエリ-キーのスコアに対してトークン距離に比例するヘッド固有の勾配バイアスを加える、線形にバイアスされた注意機構（ALiBi）に置換する。
• 他のトランスフォーマーの構成要素はすべて変更せずに保つ。外挿は更新済み埋め込みよりも距離ベースのバイアスから生じる。
• ヘッドごとに勾配を幾何級数として設定する（例：8ヘッドの場合は1/2, 1/4, ..., 1/256）、他のヘッド数には内挿する。勾配は学習前に固定される。
• 言語モデルを短い入力部分列(L)で訓練し、より長い有効な部分列(L_valid > L)への外挿を評価する。
• パープレキシティ、トレーニング速度、メモリ使用量の観点でALiBiをsinusoidal、rotary、T5 biasベースラインと比較する。

実験結果

リサーチクエスチョン

• RQ1 注意機構における単純な距離ベースのバイアスは、推論時により長いシーケンスへ信頼できる外挿を可能にするか。
• RQ2 ALiBiは、データセットとモデル規模を横断して、パープレキシティ、トレーニング速度、メモリの観点でsinusoidal、rotary、およびT5 bias法とどう比較するか。
• RQ3 ALiBiは、Wikipedia風テキストや書籍などのドメイン、および大規模コーパスに対して頑健か。
• RQ4 大規模LMにALiBiを採用する場合の実用的なトレーニング時間とメモリの影響は何か。
• RQ5 外挿を行わない場合にも性能を維持するか、および1.3Bパラメータモデル以上へどのようにスケールするか。

主な発見

• ALiBiは外挿を可能にする：L=1024トークンで訓練された1.3BパラメータのLMは2048トークンへ外挿し、パープレキシティは2048トークンで訓練されたsinusoidalモデルと同等となり、かつ11%速く、メモリも11%少なくなる。
• ALiBiは入力長が増加しても性能を維持し、トレーニング長の2倍を超えても利点が続き、非常に長いシーケンス（例：10,000トークン）まで競争力を維持する。
• ALiBiはWikiText-103で複数の位置法を上回り、他のドメイン（例：Toronto Book Corpus）へもハイパーパラメータを再調整せずに転移する。
• T5 biasと比較して、ALiBiは同等またはそれ以上のパープレキシティをほとんどランタイムのオーバーヘッドや追加パラメータなしに提供し、一般により速く訓練され、メモリも少なくなる。
• CC100+RoBERTaコーパスで1.3Bモデルの場合、ALiBiはsinusoidalベースラインと同等のパープレキシティを達成する一方でメモリを6-11%削減し、訓練を7%高速化する。2L–2L+トークンへの外挿はとりわけ効果的。
• 総じて、ALiBiは追加の longer sequences の訓練なしで、入力長の外挿へ向けた単純で頑健かつ効率的な道を提供する。