今回の論文
今回取り上げるのは、Ofir Press、Noah A. Smith、Mike Lewisによる「Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation」です。2021年8月27日にarXivで初版が公開され、ICLR 2022の会議論文になりました。研究分野は、Transformer言語モデルの位置表現と長さ外挿です。arXivの書誌情報と論文から確認できます。本記事は2022年4月22日のv2を参照します。
選んだ理由は、「長い入力を扱うには、長い入力で学習するしかないのか」という設計上の疑問に、シンプルな計算規則で取り組んでいるからです。既存モデルをそのまま改造する万能な方法ではありませんが、学習時の計算予算と提供する入力長を分けて考えるヒントになります。
どんな技術か
ALiBiはAttention with Linear Biasesの略で、Attentionが各トークンを参照する強さに、距離に比例したペナルティを加える技術です。遠い位置ほどスコアを下げ、近い位置を参照しやすくします。位置番号のベクトルを入力へ足す代わりに、参照関係そのものへ位置の情報を入れます。
たとえば、文の続きを予測するとき、直前の語と100語前の語を同じ条件では扱いません。ただし、遠い語を参照禁止にするわけではありません。内容に基づくスコアと距離の補正を組み合わせて、最終的な重みを決めます。
この規則は学習中に現れなかった距離でも計算できます。その性質を使い、短い系列で学習したモデルを、追加の長文ファインチューニングなしで長い系列へ適用するのが論文の狙いです。論文・第3節
課題
Transformerでは、単語の内容に加えて並び順を伝える必要があります。「会社が顧客を訪ねる」と「顧客が会社を訪ねる」では、単語が似ていても意味が変わるからです。一方、その位置表現が学習時の入力長に強く依存すると、入力を長くした途端に予測が不安定になる可能性があります。
学習する位置埋め込みでは、学習範囲の外に対応するベクトルがありません。数式で任意の位置を表せる方式でも、その数式をモデルが未知の位置でうまく使えるとは限りません。「位置の値を作れる」と「その位置で推論品質を保てる」は別の問題です。
最初から長い系列で学習すればよいように見えますが、密なAttentionでは系列長に応じて比較する組み合わせが二乗で増えます。長文入力を必要とするアプリを作る際、学習時から最大長を使い続ける負担は小さくありません。ALiBiは、位置表現を変えることで、この学習側の制約を緩められないかを検証します。
用語解説
- 位置表現
- トークンの順序や距離をモデルへ伝える方法です。ALiBiでは入力に位置ベクトルを加えず、Attentionスコアの距離補正として各層へ組み込みます。
- Query・Key・Value
- Attentionで使う三つの表現です。QueryとKeyの内積から参照スコアを計算し、その重みに従ってValueを集約します。ALiBiが直接変更するのは、Valueの内容ではなく集約前のスコアです。
- 因果マスク
- 次のトークンを予測する際に、未来のトークンを参照できないようにする制約です。距離による減点とは役割が違い、ALiBiでも未来の参照は禁止します。
- 長さ外挿
- 学習した系列長を超える入力でも性能を維持することです。単に実行時の長さ上限を変更するだけではなく、未知の長さで予測品質を確かめる必要があります。
- Perplexity（PPL）
- 正解トークンへ割り当てた確率から計算する、言語モデルの予測の難しさを示す指標です。低いほどよい値ですが、文書QAで必要な根拠を見つける正答率を直接測っているわけではありません。
技術の仕組み
距離をSoftmax前のスコアから引く
因果的な自己Attentionで、現在のQuery位置をi、参照するKey位置をjとします。jはi以下です。通常のスケーリングを明示すると、ヘッドhのスコアは次のように書けます。
score(h, i, j) = dot(q[h, i], k[h, j]) / sqrt(d) - m[h] * (i - j)
weight(h, i, :) = softmax(score(h, i, :))dはヘッドの次元、mは正の傾きです。未来の位置は別途マスクします。距離補正は内積をスケーリングした後に加えるため、補正までsqrt(d)で割らないことに注意します。
傾きが0.5なら、距離1の位置は0.5、距離4の位置は2だけ減点します。これは仕組みを説明する仮例です。Softmax後の重みを一定量引く処理ではなく、確率へ変換する前のスコアを補正します。
数学的には、補正前の非正規化の重みに「exp(-m×距離)」を掛ける形です。スコア上では線形の減点ですが、重みへの作用は指数的な減衰になります。この違いを押さえると、「少し遠いだけでどれくらい不利になるか」を考えやすくなります。
ヘッドごとに距離への感度を変える
全ヘッドを同じ強さで近距離へ偏らせるのではなく、ヘッドごとに異なる固定の傾きを使います。8ヘッドの例では、1/2、1/4、…、1/256です。大きい傾きは近距離を強く優先し、小さい傾きは遠い位置を比較的残しやすくします。
傾きは学習パラメータではありません。モデルは固定の距離規則がある状態で、Query・Key・Valueや他のネットワーク重みを学習します。「どのヘッドも近くしか見ない」とも、「特定ヘッドが必ず長距離の推論を担当する」とも断定できません。論文・第3節の傾きの設定
学習と推論で同じ規則を使う
学習では短い系列を使い、通常の次トークン予測の損失でモデルを更新します。ALiBi専用の教師ラベルや追加の生成器は必要ありません。推論では入力長に合わせて距離補正を計算し、同じ学習済み重みを使います。
著者の公開実装の説明では、位置埋め込みを外し、ヘッド別のバイアスを因果マスクへ加え、層のループより前でマスクを用意する変更を示しています。固定の補正を各層で使う構成です。
自社実装へ落とすなら、KVキャッシュを使う生成時にも、QueryとKeyの系列全体での位置をそろえる必要があります。1トークンのQuery配列内の番号と、会話全体での位置を混同すると距離の意味が変わります。独自補正を高速な処理へ組み込む設計については、FlexAttentionのスコア補正とカーネル生成が補足になります。
実験と結果
短い学習から長い評価へ外挿できるか
論文はWikiText-103、Toronto BookCorpus、CC100とRoBERTaのコーパスを組み合わせたデータで言語モデリングを評価しています。以下はv2の付録表6・12の値です。主な指標はPPLであり、QAの正答率や会話の指示追従率ではありません。
WikiText-103の約247Mパラメータのモデルでは、次の結果でした。評価長はいずれも3,072です。
| 位置表現 | 学習長 | テストPPL ↓ |
|---|---|---|
| Sinusoidal | 3,072 | 19.38 |
| ALiBi | 512 | 19.08 |
| ALiBi | 3,072 | 18.30 |
512で学習したALiBiは、6倍の長さへ外挿しても、3,072で学習したSinusoidalを上回りました。ただし、ALiBi同士なら長く学習した条件のほうがよく、「長文で学習する価値がなくなる」という結果ではありません。論文・付録表6
計算予算をそろえたときの意味
1.3Bパラメータ、CC100+RoBERTaの実験では、同じ50,000更新を行ったとき、学習長1,024のALiBiを2,048で評価したPPLは8.92でした。学習・評価とも2,048のSinusoidalは8.83で、ALiBi側の学習は約11%高速、メモリ使用量も約11%少ないと報告しています。近いPPLを、短い学習系列で実現した比較です。
一方、同じ学習時間を割り当てた付録表11では、長さ1,024のALiBiが50,000更新を完了し、長さ2,048のSinusoidalは約44,200更新でした。2,048でのPPLはそれぞれ8.92と9.01です。同じ更新数の比較と、同じ時間の比較は区別する必要があります。論文・付録表11〜12
長文の改善をどう解釈するか
付録Bでは、評価データを独立した短い区間に切ると、区間の先頭付近で過去文脈が不足する影響を分析しています。長い区間で評価すると、その不利を受ける位置の割合が減ります。ALiBiの外挿によるPPL改善の多くは、この効果で説明されるとしています。
そのため、PPLが改善したことだけから、何千トークンも前の根拠を積極的に読み取れるとは言えません。長文処理を評価するときは、系列の区切り方までそろえる必要があります。論文・付録B
何に使える？
以下は論文を踏まえた応用上の考察です。業務アプリで直接実証された成果ではありません。
自社言語モデルの学習予算を設計する
業務文書やログ向けに言語モデルを学習する場合、短い系列で学習してから、必要な推論長で評価する候補を持てます。最大入力長に合わせて学習を重くする前に、位置表現と学習長の組み合わせを比較できることが利点です。
ただし、長い文書の離れた条件を結び付ける用途では、その関係を学ぶデータも必要です。短文だけで訓練すれば長距離の業務推論まで身に付く、という意味ではありません。
長い履歴を受け取る生成処理
モデル内部を管理できる環境なら、学習長を超える会話履歴やコードを受け取る生成処理の候補になります。位置表現に起因する不安定さを抑えられれば、入力の切り捨てを減らす余地があります。
一方、距離のペナルティは古い情報を不利にします。会話冒頭の制約やコード先頭の定義を利用できるかは別に測ります。外部LLM APIへプロンプトを書き足すだけでは、ALiBiへ変更できません。
長文対応モデルを選ぶ際の評価軸
API中心の小規模アプリでも、「入力上限」「予測品質」「根拠利用」を分ける考え方は役立ちます。同じ根拠を先頭・中央・末尾へ動かす評価を加えれば、入力容量だけでは見えない弱さを調べられます。Lost in the Middleの根拠位置を変える評価は、その具体的な方法です。
開発や事業へのヒント
何をそろえて比較したかを記録する
モデル比較では、更新数、総学習トークン数、実時間、ピークメモリを別々に残すと判断しやすくなります。同じ時間なら多く学べる方式と、同じ更新数なら品質が高い方式は、予算の制約によって評価が変わるからです。
小規模な検証でも、学習長と評価長を別の設定として保存できます。短い入力だけで改善を確認して終わらず、提供予定の長さで、PPLと実タスクの品質を測る手順が重要です。
長文機能の価値を業務の成功率で決める
長い入力を受け付けられるだけでは、高価格の長文プランを支える根拠として不十分です。たとえば契約書QAなら、離れた条項の例外を拾えるか、文書要約なら冒頭の条件を末尾まで保てるかを検証する必要があります。
ALiBiの結果から持ち帰れるのは、モデルの構造と学習予算を工夫する余地です。自社の業務品質が上がるかは、実際の長さと情報配置を持つ評価セットで確認します。
補正の計算と実行基盤を一緒に見る
距離規則が簡単でも、ヘッド数×系列長×系列長のバイアス行列をそのまま保存するとメモリを使います。導入時は、補正をカーネル内で計算できるか、学習・生成の両方で同じ規則を使えるかを確認します。これは実装上の検討事項であり、論文の速度差が別のGPUやライブラリでも再現するという保証ではありません。
限界
ALiBiは、Attentionの比較対象を減らす疎化手法ではありません。密なAttentionの計算量は長さに対して二乗で増え、生成時のKVキャッシュも長さとともに増えます。短く学習できる利点と、長く推論するときのコストは分けて考えます。
また、既存のRoPEモデルや学習済み位置埋め込みのモデルへ、推論時だけALiBiを足して品質を保てることは本論文では示されていません。中心となる実験は、最初からその位置表現を使って学習したモデルです。
データ面では英語の言語モデリングが中心です。日本語の社内文書QA、複数資料を使う推論、エージェントの長期記憶での効果は別途検証が必要です。近距離への偏りが、遠くの重要情報を使いたい用途に常に適しているとも限りません。
外挿できることは、入力を無制限に伸ばして品質が改善し続けることを意味しません。付録表12では、512で学習したALiBiのPPLは評価長1,024で9.30、2,048で9.54でした。実装の最大長、メモリ、数値精度、評価データに応じて提供範囲を決める必要があります。論文・付録表12
よくある質問
Q. ALiBiを使えば既存のLLMを追加学習なしで長文化できますか？
A. 論文で追加学習なしに外挿しているのは、ALiBiを使って学習したモデルです。別の位置表現で学習したモデルを推論時だけ置き換える効果は、この結果からは判断できません。
Q. 遠いトークンを切り捨てる技術ですか？
A. 参照は禁止せず、Softmax前のスコアを距離に応じて下げます。内容のスコアも残るため、ウィンドウ外を一律に見えなくする局所Attentionとは異なります。計算量を自動で減らす処理でもありません。
Q. 長文RAGの正答率も上がりますか？
A. PPLの改善だけでは判断できません。検索で拾った根拠を長い入力に含め、その位置を動かして正答率と引用を確認してください。とくに遠い根拠への偏りを測る必要があります。
Q. 11%の高速化は推論の数字ですか？
A. 本記事の11%は、1.3Bモデルで学習長1,024のALiBiと学習長2,048のSinusoidalを比べた学習側の結果です。同じ長さの推論を常に11%速くするという意味ではありません。
Q. 小さな開発チームは何から試せますか？
A. 自前の学習基盤があれば、短い学習長と長い評価長を組み合わせ、費用と実タスク品質を比較します。外部API中心なら、まず入力長と根拠位置を変える回帰評価を作ると、この論文の問題意識を活かせます。
今日の学び
この論文は、学習時より長い系列で言語モデルの予測品質が崩れる課題を扱いました。位置ベクトルを入力へ加える代わりに、固定の距離ペナルティをAttentionスコアへ入れるALiBiで解こうとしています。
得られるヒントは、学習時の長さと推論時の長さを分けて設計し、計算予算を比較することです。そのうえで、長文を入力できる能力と、必要な情報を使って業務を成功させる能力を、それぞれ測る必要があります。