Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

[e4exp]

• https://arxiv.org/abs/2108.12409v1
• 2021

Vaswaniら（2017）による変換モデルの導入以来、基本的な問題は未解決のままである： トレーニング中に見られたよりも長いシーケンスへの推論時の外挿をどのように実現するか？ 我々はまず、既存の提案では効率的な外挿ができないことがわかったが、位置表現の方法を変えることで外挿が改善できることを示す。 我々は、外挿を可能にするシンプルで効率的な方法、Attention with Linear Biases (ALiBi)を紹介する。 ALiBiは、単語埋め込みに位置埋め込みを追加するのではなく、クエリキーの注目度スコアを、その距離に比例した項でバイアスする。 この方法により、長さ1024の入力配列に対して13億パラメータのモデルを学習し、長さ2048の入力配列に外挿することで、長さ2048の入力で学習した正弦波の位置埋め込みモデルと同じパープレキシティを達成し、11%の速度と11%のメモリ使用量を削減できることを示します。 ALiBiの再帰性への帰納的バイアスにより、WikiText-103ベンチマークにおいて、複数の強力な位置推定法を上回る結果を得ることができました。 最後に、なぜALiBiが優れた性能を発揮するのかを理解するために、ALiBiの分析を行います。

https://github.com/ofirpress/attention_with_linear_biases

[e4exp]

1 はじめに

現代の大規模な変換器ベースの言語モデルを構築する際に，主要な設計上の決定事項は，トレーニングシーケンスの長さ（これまでは推論シーケンスの長さにも相当していた）であり，本稿ではLと表記している。 Lを大きくすることで、より多くの文脈が得られ、推論時の予測精度が向上します。 しかし、より長い配列は学習コストが高い。2

トランスフォーマーが登場する前、RNN言語モデルはより短いL配列で学習され、推論時にはより長い文脈に一般化すると想定されていた（Mikolov et al.、2010; Mikolov & Zweig、2012; Zaremba et al.、2014）。 Vaswaniら（2017）は、トランスフォーマーを紹介しながら、"学習中に遭遇したものよりも長い配列長に外挿するかもしれない[...]と推測している。 " 我々は、外挿性とは、検証時の入力トークン数がモデルの学習時のトークン数を超えて増加しても、モデルが良好なパフォーマンスを継続する能力と定義している。 図1に示すように、正弦波の位置埋め込みを用いたトランスフォーマーLMは、外挿性が非常に低いことが分かりました。 この外挿の失敗は、位置埋め込み法に起因することを示している。 図1に見られるように、オリジナルの正弦波位置法に対する最近の代替法(Su et al., 2021; Raffel et al., 2020)は外挿性を改善している。 これらのうち優れたもの、T5で使用されているバイアスは、正弦波アプローチに比べて少なくとも2倍遅く、余分なメモリとパラメータを使用します（図2）。

私たちは、効率的に外挿を可能にするALiBi（Attention with Linear Biases）を導入しました。 同じ入力長で学習した正弦波モデルと比較して、我々の手法は追加のランタイムやパラメータを必要とせず、メモリの増加も無視できる程度（0-0.7%）で済む。 ALiBiは、既存の変換器のコードを数行変更するだけで実装することができます。 ALiBiを使用することで、変換器LMは短いL列で学習することができるため、はるかに低コストでありながら、実行時に長文にも確実に適用することができます。 例えば、ALiBiを用いてL = 1024トークンで学習した13億パラメータのLMは、L = 2048で学習した正弦波モデルと同じパープレキシティを達成していますが、我々のモデルは11%高速でメモリ使用量も11%少ないです。 モデルが学習されたトークン数の2倍程度で性能がピークに達しますが、ALiBiは長さ10,000のシーケンスでも強力な性能を維持することができます。 最近検討されている、NLPの学習例がLMの文脈として与えられる設定(Brown et al., 2020)では、我々のアプローチにより、より多くの例に触れることができます。 さらに、より長い出力を生成することができます。

また、なぜALiBiが効果的に働くのかを調査します。 §5では、ALiBiのパフォーマンスの向上は、初期のトークンの呪いをうまく回避したことで大きく説明できることがわかりました。 このことは、ALiBiをベースにした将来の研究が、より長い履歴をうまく利用することで、さらなる利益を得ることができることを示唆しています。

[e4exp]

2 既存の手法では効率的な外挿ができない

技術的には外挿できるはずの正弦波位置法ですが、実際には非常に限定的な外挿しかできないことを初めて明らかにしました。 回転位置法は、正弦波法よりも改善されていますが、それでも満足のいく結果は得られません。 他の条件はそのままに、T5バイアス位置法がこれらのいずれよりも優れた外挿性を示すことを初めて確認し、外挿性は位置の埋め込みに大きく依存するという結論に至りました。 残念ながら、この方法は計算コストがかかります（図2）。

2.1 背景と実験的設定

変形言語モデルは、トークンのリストを受け取り、次のトークンの予測を表す確率分布を出力します。 通常、言語モデルへの入力は、(より長い)訓練や評価シーケンスの部分シーケンスであるため、ここでは、入力リストを現在の入力部分シーケンスと呼びます。 トレーニングとパープレキシティ評価（固定シーケンスのスコアリング）の両方において、各ポジションの予測がその左隣のトークンからのみ影響を受けることを保証する「因果マスク」を用いて、多くの予測を一度に計算することができます。 学習中の各入力サブシーケンスの長さをLとすると，平均してL+12トークンの（左）文脈にアクセスできるL個の予測が含まれる． モデルの外挿性を調べるために、評価時に長さLvalid > Lのシーケンスを考慮する場合に興味があります。 推論時にLが学習時と異なる場合、学習時に使用された部分配列の長さをLとし、検証時に使用された部分配列の長さをLvalidとしています。

重複しない推論

L個のトークンよりも長いシーケンスを学習・評価するには、通常、シーケンスをL個の長さの部分シーケンスに分割し、それらを独立に学習・評価します。 特に断りのない限り、perplexityスコアの報告にはnon-overlapping推論を使用します。

推論時の外挿

形式的には，変換層を定義する関数は，入力の長さに依存しない3． テキストのように本質的に連続しているデータにトランスフォーマを適用すると、位置情報がさまざまな方法で入力に注入される。 Vaswaniら（2017）は、単語埋め込みに追加するベクトルに位置を埋め込む方法として、特定の位置に対する学習埋め込みと、未学習の正弦波埋め込みの2つの選択肢を議論した。 彼らは、これら2つの間で同様のパフォーマンスを観察しましたが、推論中に長い入力シーケンスに外挿する可能性があると主張した正弦波アプローチを好んでいました。 このモデルでは、L以降の数十トークン以上には外挿できないことがわかりました。

実験設定

まず、Baevski & Auli (2018)のtransformer言語モデルを用いて、WikiText-103コーパス(Merity et al., 2016)上で、様々な位置法の外挿能力をテストします。 このモデルを使用するのは、最近の言語モデリングの開発において顕著な役割を果たしているからである（Khandelwal et al.、2020; Press et al.、2021）。 学習セットは、英語版Wikipediaの約1億300万トークン（半分のギガバイト）です。 このモデルは、8つのヘッドを持つ1024次元の16の変換層を持ち、4096のフィードフォワード内寸を持つ。 このモデルは、単語埋め込み行列とソフトマックス行列を結びつける（Press & Wolf, 2017; Inan et al. 実験では、位置法と学習サブシーケンス長を変化させる以外に、ランダムシードや学習エポック数（205）など、他のハイパーパラメータは変更していません。

2.2 SINUSOIDAL POSITION EMBEDDINGS

正弦波位置埋め込み(Vaswani et al. (2017); §3.5)は、変換器の第1層への入力でトークン埋め込みに追加される一定の非学習ベクトルである。 これらは、変換器の言語モデリング（Baevski & Auli, 2018; Lewis et al., 2021）や機械翻訳（Vaswani et al., 2017; Ott et al., 2018）モデルで頻繁に使用されています。 まず、正弦波の位置埋め込みを用いるBaevski & Auli (2018)の修正されていないモデルを、L = 512トークンで学習した後、L + kトークンの検証セットで、kが0から15,000までの範囲で、それを使った推論を実行します。 図1（左）と表2（付録）によると、このモデルはk = 20まではperplexityを向上させることができますが、k = 20からk = 50までは性能の向上が止まり、安定した状態が続き、その後、性能が低下し始めます。 L = 1024個のトークンで学習したモデルでも同様の結果が得られました（図1右、および付録の表3）。 このモデルでは、Lvalid = L + 50 tokensまでは改善することができますが、その後は性能が低下していきます。

2.3 ROTARY POSITION EMBEDDINGS

回転位置法はSuら(2021)によって導入され、最近ではオープンソースのGPT-3(Brownら、2020)実装のGPT-J(Wang & Komatsuzaki、2021)によって普及しています。 彼らは、変換ネットワークの下部に正弦波エンベディングを追加する代わりに、すべての注目層のキーとクエリに正弦波エンベディングを乗算します。 なお、正弦波埋め込みや位置学習埋め込みの手法とは異なり、ロータリー法では、最初の層だけでなく、すべての層でモデルに位置情報を注入します。 また、自己注目副層の値には位置情報を付加しません。 自己吸着副層の出力は、入力された値のベクトルを線形変換して加重加算したものなので、値に位置情報を挿入しないことで、各変換層の出力には、明示的な位置情報が入っていないことになります。 我々は、この位置情報の分離が外挿に有益ではないかと考え、我々の方法（§3）を設計する際に、この方法からヒントを得た。 パープレキシティの結果（図1、付録の表2、3）は、L = 512（L = 1024）のモデルでは、学習時に見たトークンよりもk = 200（k = 100）まで追加してパープレキシティを向上させることができるので、正弦波アプローチよりも優れていますが、これは学習と推論が遅くなるという代償を伴います（図2）。

2.4 T5 BIAS

多くのモデルが訓練された位置エンベッディングまたは正弦波の位置エンベッディングを使用しているのに対し、Raffelら（2020）のT5モデルでは、相対位置法（Shaw et al., 2018; Huang et al., 2019）を使用しており、（以前の方法のように）単語エンベッディングに何らかの位置情報を追加するのではなく、注目値の計算方法を変更している。 我々はこの方法を "T5バイアス "と呼ぶ。 修正されていない変換器では、注目値を計算するために、すべてのクエリとすべての関連キーのドット積を計算し、これらの生の注目値をソフトマックスします。 T5バイアス法では、先ほどと同様に注目値を計算しますが、学習した共有バイアスを各クエリ-キーのスコアに加えます。 つまり、クエリとキーの距離がゼロ（すなわち、クエリとキーが同じトークンを表す）のすべてのクエリ・キーのスコアには特定のバイアスがかかり、クエリとキーが1単語離れているすべてのスコアには別のバイアスがかかる、というように、あるポイントまでは、複数の異なる距離が同じバイアスを共有することになります（これは外挿に有益かもしれません）。 なお、ロータリー方式と同様に、T5バイアスは各層でモデルに位置情報を注入しており、自己注目値ベクトルに明示的な位置情報を統合していません。

Raffelら(2020)は、T5バイアスが外挿を可能にするのではないかと提案しているが、この特性を検証する実験は報告していない。 ここでは、T5バイアス法が実際に言語モデルの外挿を可能にすることを示します。 図1が示すように、T5バイアスは、L = 512 (L = 1024)個の入力トークンで学習したモデルに対して、k = 600 (k = 800)個のトークンを追加して、学習したものよりも長いシーケンスのパープレキシティを改善することができます。 残念ながら、この素晴らしい性能には代償が伴います。 バッチ学習は正弦波モデルに比べて少なくとも2倍の時間がかかるため、このモデルの外挿性は実際には役に立ちません。 例えば、1024個のトークンで推論を行いたい場合、L = 1024で正弦波モデルを学習するか、L = 512のトークンでT5バイアスモデルを学習し、推論のために1024個に外挿することができます。 しかし、L = 1024の正弦波モデルは28.5k words per second (WPS)で実行されるのに対し、L = 512のT5バイアスモデルは14.4k WPSで実行されるため（付録表1）、T5バイアス法で短いシーケンスを学習するメリットは実際にはありません。

[e4exp]

3 ATTENTION WITH LINEAR BIASES (ALIBI)

Vaswaniら(2017)の変換モデルでは、ネットワークの下部で単語埋め込みに位置埋め込みが追加される。 長さLの入力サブシーケンスに対して、注目サブレイヤーは、最初のi個のキーK∈R i×d（dはヘッド次元）が与えられたときに、各ヘッドにおけるi番目のクエリqi∈R 1×d 、（1≦i≦L）の注目スコアを計算する

ALiBiを使用する際には、ネットワークのどの位置にも位置埋め込みを追加しません。 唯一の修正点は、クエリ・キーのドット積の後に、静的で非学習のバイアスを加算する

ここで、スカラーmは、トレーニング前に固定された頭部固有の傾きです。 図3にその様子を示します。 8つのヘッドを持つモデルでは，1/2から始まり，各要素に1 /2を掛けて次の要素を計算する幾何学的な配列の傾きを使用しました． 1/ 2^ 1 , 1 /2 ^2 , ..., 1 2 8 . 16個のヘッドを必要とするモデルでは，連続する各ペアを幾何学的に平均して8つのスロープを補うことで，√1 2から始まり，√1 2 : 1 2 0.5 , 1 2 1 , 1 2 1.5 , ..., 1 2 8の比率を持つ幾何学的な配列が得られました．一般に，n個の頭に対して，我々の傾きの集合は，2 -2 (-log2 n+3)で始まり，その同じ値を比とする幾何学的な数列である．§4では、この傾きのセットが様々なテキストドメインとモデルサイズで動作することを観察しており、新しいデータセットで新しいモデルを学習するたびに、これらの傾きの値を調整する必要はないと考えています。このことから、我々の手法は、Vaswaniら（2017）によってハイパーパラメータ（波長の幾何学的進行の開始と終了）が一度設定され、その後、異なるデータセット上の異なるサイズのモデルで再利用された正弦波アプローチに似ている。

[e4exp]

duplicated issue. https://github.com/e4exp/paper_manager_abstract/issues/622