TRAIN SHORT, TEST LONG: ATTENTION WITH LINEAR BIASES ENABLES INPUT LENGTH EXTRAPOLATION

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

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）。

私たちは、効率的に外挿を可能にするAttention with Linear Biases（ALiBi）を導入しました。同じ入力長で学習した正弦波モデルと比較して、我々の手法は追加のランタイムやパラメータを必要とせず、メモリの増加も無視できる程度（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をベースにした将来の研究が、より長い履歴をうまく利用することで、さらなる利益を得られる可能性を示唆しています。

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

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

2.1 背景と実験的設定

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

重複しない推論

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

推論中の外挿

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

Experimental Setup

我々はまず、Baevski & Auli (2018)のトランスフォーマー言語モデルを用いて、WikiText-103コーパス(Merity et al., 2016)上で様々な位置法の外挿能力をテストする。このモデルを使用するのは、最近の言語モデリングの開発において顕著な役割を果たしているからである（Khandelwal et al.、2020; Press et al.、2021）。学習セットは、英語版Wikipediaの約1億300万トークン（半分のギガバイト）です。このモデルは、1024次元の16の変換層を持ち、8つのヘッドを持ち、フィードフォワードの内側の次元は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 bias" と呼ぶ。修正されていない変換器では、注目値を計算するために、すべてのクエリとすべての関連するキーのドット積を計算し、これらの生の注目値をソフトマックスします。 T5バイアス法では、先ほどと同様に注目値を計算しますが、学習した共有バイアスを各クエリ-キーのスコアに加えます。つまり、クエリとキーの距離がゼロの場合（つまり、クエリは自分自身に注目している）、クエリとキーの距離が1単語の場合は別のバイアスを得ることになります。なお、ロータリー方式と同様に、T5バイアスは各層でモデルに位置情報を注入しており、自己注目値ベクトルに明示的な位置情報を統合していません。

Raffelら(2020)は、T5バイアスが外挿を可能にするのではないかと提案しているが、この特性を検証する実験は報告していない。ここでは、T5バイアス法が実際に言語モデルの外挿を可能にすることを示します。これは、Baevski & Auliモデルを再度修正し、今度はT5バイアスを挿入することで実現しています6。

図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バイアス法で短いシーケンスを学習するメリットは実際にはありません。

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）によってハイパーパラメータ（波長の幾何学的進行の開始と終了）が一度設定され、その後、異なるデータセット上の異なるサイズのモデルで再利用された正弦波アプローチに似ています。 ALiBiは、遠く離れたクエリキーペア間の注目度スコアにペナルティを課すことで、再帰性への帰納的バイアスをかけており、キーとクエリの距離が大きくなるにつれてペナルティが増加します。異なるヘッドは、傾きの大きさに応じて、異なる割合でペナルティを増加させます。私たちは当初、スロープを学習可能にする実験を行いましたが、強力な外挿結果には至りませんでした8。約10個のスロープセットを手作業で簡単に調査した結果、最終的に選んだスロープセットを発見しました。この調査から得られた主な洞察は、最も効果的な傾きのセットは、傾きが(0, 1)の範囲にあり、0に近づくにつれて傾きの密度が高くなるものであるということです。この調査の中で、私たちの手法は傾きの選択に対してロバストであることがわかりました。指数分布からのランダムなサンプリングでも、場合によってはうまくいくことがあります（ただし、この方法は分散が大きい）。私たちのALiBi方式は相対位置法であり、T5バイアスやロータリー方式のように、すべての層で位置情報をキーとクエリに追加しますが、値には追加しないことに注意してください。これらの特性が外挿に有利に働くのではないかと仮説を立てています。

実装

ALiBiの実装は簡単で、数行のコードですべての変更を実現しています。 ALiBiの実装は、マスクマトリックスに線形バイアスを追加することで行います。実際には、変換言語モデルを学習する際、クエリqiは1からiまでのキーのみを対象としますが、これはソフトマックス演算を適用する前に、クエリ-キーのドット積にマスク行列を追加することで実装されています。つまり、我々の手法は、ネットワークに何の演算も追加せず、各注目副層の注目スコアに追加されるマスクを変更するだけなので、実行時のペナルティはありません。同じ入力長で学習された正弦波モデルと比較すると、修正されていない変換器ではマスクがL×Lのサイズであるため、ALiBiを使用した場合にはメモリの増加（一部の実験では最大100MB）がありますが、ALiBiを使用した場合には、各ヘッドに追加される線形バイアスが異なる勾配を使用しなければならないため、マスクはn×L×L（ここでnはヘッドの数）とわずかに大きくなります。しかし、ALiBiを使えば、より小さい配列での学習が可能になる一方で、より長い配列で正弦波エンベッディングを使って得られた結果を達成（時には上回る）することができ、数ギガバイトのメモリを節約することができます。

4 RESULTS

まず、WikiText-103での結果を示し、我々のALiBi手法が効率的であり、短い入力サブシーケンスでモデルを学習することができ、学習したトークン数の6倍以上に外挿しても強力なベースラインを上回ることを示します。次に、WikiText-103で成功したのと同じハイパーパラメータ（スロープのセット）を用いて、WikiText-103とは全く異なるドメインである書籍のデータセットで、何の変更もなく強力な結果が得られることを示しました。最後に、より大きな（461GB）データセットで、より多くの計算機を使って学習された1.3Bパラメータモデルが、我々のALiBi法により、大幅に少ないメモリ使用量で同様のパープレキシティスコアを達成し、正弦波法の優れた代替手段となることを示します。 Vaswaniら（2017）で提示された位置法の複数の代替案が提案されていますが、その設定は小規模な実験よりもはるかに困難であるため、大規模（1B以上のパラメータ）な言語モデルで採用されたものはほとんどありません。 GPT-3とJurassic-1(Lieber et al., 2021)ではVaswani et al.の学習済み位置埋め込み法を採用し、GPT-Jではロータリー法を採用しています。 1.3Bパラメータモデルでの結果は、ハイパーパラメータを再調整することなく、我々の手法がより大きなモデル、データセットサイズ、トレーニング期間に一般化できることを示しています。

4.1 WIKITEXT-103での結果

最初に、Baevski & Auli (2018)の言語モデルの正弦波位置埋め込みをALiBiに置き換えて、WikiText-103コーパス(Merity et al., 2016)上で我々の手法を開発する。図4（および付録の対応する表5）は、入力サブシーケンストークン（L）の数を変化させて訓練したモデルについて、検証データセット上でより長いサブシーケンスの長さに外挿した結果を示しています。最初の観察結果は、外挿を行わない場合、どのLに対しても、我々のモデルは正弦波法を用いたモデルに勝っており、時にはかなりの量の差があります。例えば、Baevski & Auliモデルは、L = 3072の入力トークンで学習した場合、18.67±0.24 (std. dev.) perplexityを達成しましたが、我々のL = 3072モデルは17.60 perplexityを達成しました（両モデルともLvalid = 3072で評価した場合）。 2つ目の観察結果は，すべてのモデルが外挿性を持ち，学習時に観察したよりも多くのトークンを扱った場合に，より良いperplexityスコアを得ることができるということです．例えば、512個のトークンで学習したモデル（開発セットで512個の長さの部分列を評価したときに19.73のperplexityを達成）は、開発セットで3072個の長さの部分列に外挿したときに18.40のperplexityスコアを達成しました。驚くべきことに，このスコアは，L = 3072の正弦波モデルが開発セットで獲得したスコアを，統計的に有意な差で上回っています． L = 512からL = 2048で学習したすべてのモデルが、Lvalid = 3072に外挿する際に、L = 3072で学習した正弦波ベースラインを上回ることができましたが、これらのモデルはすべて、より短い部分配列で学習しているため、学習時間がはるかに短いにもかかわらずです。 L = 512モデルの学習速度は1.84倍ですが、Lvalid = 3072に外挿する際にはL = 3072に勝っています。さらに、L = 3072の正弦波モデルを学習するには、大きなアテンションマトリクスに対応するために16GB以上のメモリを搭載したGPUが必要ですが、L = 512は、アテンションマトリクスがはるかに小さいため、より少ないメモリのGPUでも学習可能であるにもかかわらず、その性能を上回っています。

さらに、表5（付録）では、Lが1024と3072の場合、Lvalid = Lであっても、我々の手法がロータリーバイアスモデルやT5バイアスモデルよりも優れた性能を発揮することも示されています。図1（および対応する付録の表2と3）には、他の位置の手法に対する我々の手法の幅広い調査結果が含まれています。これによると、T5バイアス（ベースラインのベスト）はLvalidが約2Lになるまでperplexityを向上させることができますが、WikiText-103データセットでは、我々の手法は少なくとも約3Lになるまでperplexityを継続的に向上させることができ、L=512モデルはLvalidが12kトークン以上の場合でもperplexityを向上させることができます。 L = 512モデルでは、Lvalidが12kトークン以上になっても、パープレキシティを向上させることができました。

図5は、図4の断面図であり、異なる訓練長のモデルと正弦波ベースラインを示している。これを見ると，512≦L＜3072のモデルは，L＝3072の正弦波モデルよりも学習速度が速いですが，検証セットでのパープレキシティスコアはすべて勝っています。 L = 3072のモデルは、正弦波モデルと同じくらいの速度で学習しますが、そのスコアには1 perplexity point以上の差があります（L = 3072の正弦波モデルの標準偏差は0.24）。付録の表6は、検証セットでの結果が、WikiText103のテストセットにも反映されていることを示しています。

現在、WikiText-103の結果を発表している他のほとんどのモデルは、スライドウィンドウ評価（§5で定義）を用いて、perplexityを計算しています。付録表7の我々のモデル（および、正弦波、ロータリー、T5バイアスのベースライン）にその方法を適用したところ、我々のL = 3072モデルは、Transformer-XL (Dai et al., 2019)、Sandwich (Press et al., 2020)、Shortformer (Press et al., 2021)モデルの性能を上回ることがわかりました。我々の結果は、段階的なトレーニングで得られたものと似ていますが（Pressら、2021年）、Routing Transformer Royら（2020年）とkNN-LM Khandelwalら（2020年）で得られた結果よりも遅れています。これらのモデルで使用されている手法は我々のものとは直交しており、我々のものと組み合わせることで、さらに大きな性能向上につながるのではないかと考えています。

e4exp / paper_manager_abstract