ALiBi(Attention with Linear Biases)の理論と実装 — 学習不要の位置エンコーディングで外挿性能を向上させる

大規模言語モデル(LLM)を使っていて、こんな疑問を持ったことはないでしょうか。「このモデルは2048トークンで学習されたのに、なぜ4096トークンの文章を処理できないのだろう?」あるいは逆に、「なぜ学習時より長い文章を入力すると、急に出力の品質が劣化するのだろう?」

この問題の根源は位置エンコーディングにあります。Transformerは入力トークンの順序情報をSelf-Attentionの外から注入する必要がありますが、多くの手法は学習時に見た系列長の範囲でしか正しく機能しません。学習時に1024トークンまでしか見ていないモデルに2048トークンを入力すると、モデルは「1025番目の位置」をどう解釈すればよいかわからず、性能が崩壊します。

2021年、Press、Smith、Lewisは驚くほどシンプルな解決策を提案しました。それがALiBi(Attention with Linear Biases)です。位置エンコーディングを入力に加えるのではなく、Attentionスコアに「距離が遠いほど大きなペナルティ」を加算するだけ。学習パラメータはゼロ、実装はわずか数行、それでいて訓練時の系列長を大幅に超える入力に対しても安定した性能を発揮します。

ALiBiを理解すると、以下のことが可能になります。

  • 長文処理モデルの設計: 訓練コストを抑えつつ、推論時に長い系列を処理できるモデルの構築
  • 位置エンコーディング手法の選択: sin/cos、RoPE、ALiBiの特性を理解し、タスクに応じた最適な手法を選択
  • 最新LLMアーキテクチャの理解: BLOOM、MPTなどALiBiを採用したモデルの設計思想の把握
  • 外挿性の定量的評価: 異なる系列長でのPerplexity比較による手法の評価方法の習得

位置エンコーディング3手法の比較。絶対位置は入力埋め込みに位置ベクトルを加算、RoPEはQ,Kを回転、ALiBiはAttentionスコアに距離ペナルティを加算する

位置情報を「どこに」注入するかが3手法の本質的な違いです。絶対位置(sin/cos)は入力埋め込みに位置ベクトルを足し込み、RoPEはQuery/Keyを回転させます。ALiBiだけは入力・Q・K・Vに一切触れず、Attentionスコアに距離ペナルティ $-m|i-j|$ を加えるだけです。この「最後尾でちょっと足すだけ」という軽さが、本記事で見ていく外挿性と実装の簡単さにつながります。

本記事の内容

  • 位置エンコーディングの課題 — なぜ外挿が難しいのか
  • ALiBiの核心アイデア — 線形バイアスによる距離ペナルティ
  • 数学的定式化と導出
  • ヘッドごとの傾き $m$ の設計(幾何数列)
  • sinusoidal PE、RoPEとの比較
  • PyTorchによるALiBi Attentionの実装
  • 外挿実験 — 訓練長と推論長を変えたPerplexity比較
  • ALiBiの限界と最新手法(YaRN等)との関係

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

位置エンコーディングの課題 — なぜ外挿が難しいのか

固定長の壁

Transformerに位置情報を与える最も直感的な方法は、「位置1にはこのベクトル、位置2にはこのベクトル…」と各位置に固定のベクトルを割り当てることです。これが絶対位置エンコーディングの考え方であり、原論文のsin/cos方式やBERT/GPTの学習可能な位置埋め込みがこれに該当します。

しかし、この方法には本質的な限界があります。モデルが学習時に最大 $L_{\text{train}}$ トークンまでの系列しか見ていない場合、位置 $L_{\text{train}} + 1$ 以降に対応するベクトルは次のいずれかの問題を抱えます。

  • 学習可能な位置埋め込み: そもそも $L_{\text{train}}$ までの埋め込みしか定義されていないため、それを超える位置に対する埋め込みが存在しない
  • sin/cos位置エンコーディング: 数学的には任意の位置に対して値を計算できるが、モデルが学習時に見たことのない位置パターンを正しく解釈できる保証がない

これは人間に例えると、「100ページまでの本しか読んだことがない人に、突然500ページの本を渡す」ようなものです。100ページ目までは問題なく読めますが、101ページ以降では「この位置はどのあたりなのか」という感覚が狂ってしまいます。

外挿と内挿

ここで2つの重要な概念を整理しましょう。

  • 内挿(Interpolation): 学習時に見た範囲の中間的な入力に対する汎化。例えば、100トークンと200トークンで学習したモデルが150トークンの入力を処理する場合
  • 外挿(Extrapolation): 学習時に見た範囲を超える入力に対する汎化。例えば、512トークンで学習したモデルが1024トークンの入力を処理する場合

LLMにおいて特に重要なのは外挿です。学習時のコンテキスト長は計算コストで制約されますが、推論時にはより長い文書や会話を処理したい場面が頻繁に生じます。長い系列で学習し直す(fine-tuning)ことも可能ですが、計算コストが $O(n^2)$ で増大するため、現実的ではないケースも多くあります。

相対位置エンコーディングの登場

絶対位置の限界を克服するために、相対位置エンコーディングが提案されました。Shawら(2018)やDaiら(2019、Transformer-XL)は、Attentionスコアの計算時に「2つのトークン間の距離」を利用する手法を導入しました。

相対位置の利点は明確です。位置5と位置8の関係は「距離3」であり、位置105と位置108の関係も「距離3」です。つまり、モデルが「距離3」の関係性を学習していれば、絶対位置がどこであっても同じように処理できるはずです。

しかし、従来の相対位置エンコーディングは実装が複雑で、計算オーバーヘッドも大きいという課題がありました。RoPE(Su et al., 2021)はこの問題を回転行列で美しく解決しましたが、外挿性能には限界が残ります。

では、位置エンコーディングの外挿問題を、もっと根本的にシンプルに解決する方法はないのでしょうか? ALiBiはまさにその問いに対する回答です。

ALiBiの核心アイデア — 距離に比例するペナルティ

「近いものは重要、遠いものは重要度が下がる」

ALiBiのアイデアは、言語の性質に関する素朴な観察から始まります。文章中のあるトークンにとって、直前のトークンが最も関連性が高く、距離が離れるほど関連性は(一般的に)低下します。これは完全に正しいわけではありませんが、デフォルトの振る舞いとしては合理的です。

たとえば、この文章を読んでいるあなたにとって、直前の段落の内容はすぐに思い出せますが、記事の冒頭で何を書いたかは少しぼやけているはずです。ALiBiはこの「距離による注目の減衰」をAttentionの計算に直接組み込みます。

従来の位置エンコーディングとの根本的な違い

従来の位置エンコーディング(sin/cos、学習可能な埋め込み、RoPE)は全て、位置情報を入力の表現に組み込みます。

  • sin/cos / 学習可能: 入力埋め込みに位置ベクトルを加算 → $\bm{z}_i = \bm{x}_i + \bm{PE}(i)$
  • RoPE: Query/Keyに回転を適用 → $\tilde{\bm{q}}_m = R_\Theta(m)\bm{q}_m$

ALiBiのアプローチは根本的に異なります。入力の表現には一切手を加えず、Attentionスコアに直接バイアスを加算します。

$$ \text{ALiBi}: \quad \text{score}(i, j) = \frac{\bm{q}_i^\top \bm{k}_j}{\sqrt{d_k}} + \text{bias}(i, j) $$

ここで $\text{bias}(i, j)$ は位置 $i$ と位置 $j$ の距離に比例する負の値です。距離が遠いペアほど大きなペナルティが課されるため、softmaxを通した後のAttention重みで「近いトークンへの注目」が自然に促されます。

この設計の美しさは、モデルの内部表現(埋め込み、Query、Key、Value)に位置情報が混入しないことです。意味情報と位置情報が明確に分離されるため、Attentionが「何に注目するか」(内積スコア)と「どの距離まで見るか」(バイアス)を独立に制御できます。

ALiBiの核心。内積スコア行列に距離比例のバイアス行列を加算し、その合計にsoftmaxを適用する3パネル図

左から、①Query・Keyの内積スコア(意味的な関連度)、②距離だけで決まるALiBiバイアス $-m|i-j|$、③その合計です。①は学習で決まる「内容」の評価、②は固定の「距離」ペナルティで、両者は完全に独立して加算されます。③の合計にsoftmaxをかけるだけなので、既存のAttention実装に1行足すだけで組み込めることがこの図から読み取れます。右上の三角が空白なのは、デコーダの因果マスク(未来を見ない)によるものです。

学習パラメータが不要な理由

ALiBiのバイアスは、事前に固定された傾きパラメータで完全に決定されます。学習中に更新されるパラメータは一切ありません。

なぜ学習が不要なのでしょうか?それは、ALiBiが位置を「表現」するのではなく、位置に基づいて「Attentionの範囲を制御する」ものだからです。最適なAttentionパターン(どのトークンにどれだけ注目するか)はQuery/Keyの学習で決まり、ALiBiはその上に「遠いものには注目しにくくする」という穏やかな事前知識を重ねるだけです。

この穏やかさこそが外挿性の鍵です。sin/cosやRoPEのように「位置100はこのパターン」「位置200はこのパターン」と具体的な表現を持つ手法は、学習時に見なかったパターンへの汎化が困難です。一方、ALiBiは「距離1のペナルティは $m$、距離2のペナルティは $2m$、距離1000のペナルティは $1000m$」と線形に外挿するだけなので、未知の距離にも自然に対応できます。

ここまでALiBiの直感的なアイデアを理解しました。次に、このアイデアを数学的に厳密に定式化していきましょう。

数学的定式化

標準的なAttentionの復習

まず、標準的なScaled Dot-Product Attentionを復習します。入力系列 $\bm{X} = (\bm{x}_1, \bm{x}_2, \ldots, \bm{x}_n) \in \mathbb{R}^{n \times d}$ に対して、Query、Key、Valueを計算します。

$$ \bm{Q} = \bm{X}\bm{W}_Q, \quad \bm{K} = \bm{X}\bm{W}_K, \quad \bm{V} = \bm{X}\bm{W}_V $$

$i$ 番目のトークンの出力は次のように計算されます。

$$ \text{Attention}(\bm{q}_i, \bm{K}, \bm{V}) = \sum_{j=1}^{n} \alpha_{ij} \bm{v}_j $$

ここでAttention重み $\alpha_{ij}$ は、スケーリングされた内積にsoftmaxを適用して得られます。

$$ \alpha_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{n} \exp(s_{ik})}, \quad s_{ij} = \frac{\bm{q}_i^\top \bm{k}_j}{\sqrt{d_k}} $$

ALiBiによるスコアの修正

ALiBiは、このAttentionスコア $s_{ij}$ にバイアス項を加算します。

$$ \begin{equation} s_{ij}^{\text{ALiBi}} = \frac{\bm{q}_i^\top \bm{k}_j}{\sqrt{d_k}} – m \cdot |i – j| \end{equation} $$

ここで $m > 0$ はヘッドごとに固定されたスロープ(傾き)です。$|i – j|$ は2つのトークン間の距離であり、距離が大きいほどスコアが減少します。

因果的(causal)なAttention(デコーダで使用)の場合、$j \leq i$ の制約があるため、$|i – j| = i – j$ となります。$i$ 番目のQueryから見たスコアは次のように書けます。

$$ \bm{s}_i^{\text{ALiBi}} = \frac{\bm{q}_i \bm{K}_{1:i}^\top}{\sqrt{d_k}} – m \cdot [i-1, \, i-2, \, \ldots, \, 1, \, 0] $$

この式を分解して理解しましょう。右辺第1項 $\bm{q}_i \bm{K}_{1:i}^\top / \sqrt{d_k}$ は通常のAttentionスコアで、トークン間の意味的な関連度を表します。右辺第2項 $m \cdot [i-1, i-2, \ldots, 1, 0]$ は距離ペナルティで、距離に比例して線形に増加します。

バイアス行列の構造

系列全体のバイアス行列 $\bm{B} \in \mathbb{R}^{n \times n}$ は次のように定義されます。

$$ B_{ij} = -m \cdot |i – j| $$

因果的Attentionの場合($j > i$ のエントリはマスクされる)、この行列は次のような上三角が $-\infty$ のToeplitz構造を持ちます。

$$ \bm{B} = -m \cdot \begin{pmatrix} 0 & -\infty & -\infty & -\infty & \cdots \\ 1 & 0 & -\infty & -\infty & \cdots \\ 2 & 1 & 0 & -\infty & \cdots \\ 3 & 2 & 1 & 0 & \cdots \\ \vdots & \vdots & \vdots & \vdots & \ddots \end{pmatrix} $$

$j > i$ のエントリは因果マスクにより $-\infty$ に設定されます。対角要素は距離0なのでバイアスは0、1つ下の対角は距離1なのでバイアスは $-m$、というように距離に比例してバイアスが大きくなります。

この行列はToeplitz行列(各対角線が定数)であり、系列長 $n$ に依存しない構造を持ちます。つまり、$n$ がいくら大きくなっても、距離 $d$ に対するバイアスは常に $-m \cdot d$ です。これが外挿性の数学的な裏付けです。

短い系列(n=8)と長い系列(n=16)のALiBiバイアス行列。対角線ごとの値は系列長によらず同じであることを示す

左が訓練時の短い系列($n=8$)、右が推論時の長い系列($n=16$)のバイアス行列です。注目すべきは、どちらも対角線に沿った値(距離1なら $-m$、距離2なら $-2m$)が完全に一致している点です。系列が長くなっても新しい「ルール」を覚える必要がなく、同じ線形則をそのまま延長するだけ。これがALiBiが未知の系列長へ自然に外挿できる理由を端的に表しています。

Attention重みへの影響

バイアスが加わった後のAttention重みを考えてみましょう。softmaxの中でバイアス項は指数関数を通るため、次のように解釈できます。

$$ \alpha_{ij}^{\text{ALiBi}} = \frac{\exp(s_{ij} – m|i-j|)}{\sum_k \exp(s_{ik} – m|i-k|)} $$

指数法則 $\exp(a – b) = \exp(a) \cdot \exp(-b)$ を適用すると、

$$ \alpha_{ij}^{\text{ALiBi}} = \frac{\exp(s_{ij}) \cdot \exp(-m|i-j|)}{\sum_k \exp(s_{ik}) \cdot \exp(-m|i-k|)} $$

つまり、ALiBiは通常のAttention重みに指数的な距離減衰を掛けていることと等価です。softmaxの前では線形バイアスですが、softmaxの後では指数的な減衰として作用します。これは、距離に対する注目度が指数関数的に低下することを意味し、言語の「近くのトークンが重要」という性質をうまく捉えています。

左が線形バイアス、右がsoftmax後の指数的な重み減衰。傾きmが大きいほど減衰が急峻

左のグラフ(softmax前)では、バイアスは距離に対して直線的に下がるだけです。ところが右のグラフ(softmax後の重み因子 $\exp(-md)$)を見ると、注目度は指数関数的に急減しています。傾き $m=1/2$ のヘッドは距離10でほぼ0まで落ち込み近傍だけを見ますが、$m=1/256$ のヘッドは距離40でも0.85を保ち遠距離も拾えます。線形の入力が指数の出力に化けることで、自然な「局所性の事前知識」が生まれていることがわかります。

数学的な定式化が理解できたところで、次にALiBiの設計で最も巧妙な部分であるヘッドごとの傾き $m$ の設計について見ていきましょう。

ヘッドごとの傾き $m$ の設計

なぜヘッドごとに異なる傾きが必要なのか

Multi-Head Attentionでは、$h$ 個のAttentionヘッドが並列に動作します。もし全てのヘッドが同じ傾き $m$ を持っていたら、全てのヘッドが同じ距離スケールで情報を集約してしまいます。

しかし、言語の依存関係は多様です。形容詞と名詞の修飾関係は近距離(2-3トークン)ですが、主語と述語の関係は長距離(10-50トークン以上)に及ぶことがあります。さらに、段落をまたぐ照応関係はさらに長距離です。

ALiBiはこの多様性を、ヘッドごとに異なる傾きを割り当てることで実現します。傾きが大きいヘッドは近距離に特化し、傾きが小さいヘッドは遠距離の情報も取り込めます。

幾何数列によるスロープの設定

$h$ 個のAttentionヘッドに対して、傾き $m$ は次の幾何数列で決定されます。

$$ \begin{equation} m_i = 2^{-\frac{8i}{h}}, \quad i = 1, 2, \ldots, h \end{equation} $$

たとえば $h = 8$ ヘッドの場合、各ヘッドの傾きは次のようになります。

$i = 1$ のとき $m_1 = 2^{-8/8} = 2^{-1} = 1/2$ と計算されます。

同様に、$i = 2$ のとき $m_2 = 2^{-16/8} = 2^{-2} = 1/4$ です。

以降の値を全て列挙すると、

$$ m_1, m_2, \ldots, m_8 = \frac{1}{2}, \, \frac{1}{4}, \, \frac{1}{8}, \, \frac{1}{16}, \, \frac{1}{32}, \, \frac{1}{64}, \, \frac{1}{128}, \, \frac{1}{256} $$

この数列は公比 $r = 2^{-8/h}$、初項 $a_1 = 2^{-8/h}$ の等比数列です。$h = 8$ の場合は公比 $r = 1/2$ です。

8ヘッドの傾きmを線形軸と対数軸で表示。線形軸では急減、対数軸では等間隔の等比数列

左の線形軸で見ると、傾き $m$ は $1/2, 1/4, 1/8, \ldots$ と急激に小さくなり、ほとんどのヘッドが0付近に潰れて見えます。一方、右の対数軸ではきれいな等間隔の直線になり、これが公比 $1/2$ の等比数列であることがはっきりわかります。傾きが大きいHead 0は近距離に強く特化し、傾きが小さいHead 7は遠距離まで穏やかに見る——少数のヘッドで短距離から長距離までを対数スケールでカバーする設計です。

スロープの直感的な意味

スロープ $m$ の値が大きいヘッド($m = 1/2$)は、距離に対するペナルティが急峻です。距離10のトークンへのバイアスは $-5.0$ となり、softmax後のAttention重みは非常に小さくなります。つまり、このヘッドは近距離の依存関係に特化します。

一方、スロープ $m$ の値が小さいヘッド($m = 1/256$)は、ペナルティが穏やかです。距離100のトークンへのバイアスでも $-100/256 \approx -0.39$ に過ぎず、Attention重みへの影響は限定的です。このヘッドは遠距離の依存関係も捉えることができます。

この設計により、Multi-Head Attention全体として、近距離から遠距離まで多スケールの依存関係をカバーできます。

8ヘッドそれぞれの距離バイアス。傾きが大きいヘッドは遠距離で急落、小さいヘッドは緩やか

各ヘッドのバイアス $-m|i-j|$ を距離に対してプロットしたものです。傾きの異なる8本の直線が扇状に広がり、Head 0($m=1/2$)は距離10で早くも $-5$ に達して近傍に集中するのに対し、Head 7($m\approx1/256$)は距離100でも $-0.4$ 程度にしか下がりません。1つの層の中で「見る距離スケール」がヘッドごとに役割分担されている様子が一目でわかります。

なぜ幾何数列なのか

スロープを等差数列($m = 0.1, 0.2, 0.3, \ldots$)ではなく幾何数列で設定する理由は、距離のスケールが対数的に分布するためです。

言語において、距離1-10の依存関係と距離100-200の依存関係では、前者の方がはるかに多様で細かい粒度の制御が必要です。幾何数列を使うことで、短距離側に多くのヘッドを「密に」配置し、長距離側には少数のヘッドを「疎に」配置できます。

具体的に $h = 8$ の場合を見ると、次のような分布になっています。

  • $m \geq 1/8$ のヘッド(3個): 距離10以内に強く集中
  • $1/64 \leq m < 1/8$ のヘッド(3個): 距離10-100程度をカバー
  • $m < 1/64$ のヘッド(2個): 距離100以上の長距離にも対応

これは言語の依存関係の分布とよく合致します。

$h$ が2のべき乗でない場合

論文では、$h$ が2のべき乗でない場合の拡張も提案されています。$h$ 以下の最大の2のべき乗を $h’$ とし、$h’$ 個のスロープを基本数列として生成した後、残りの $h – h’$ 個は基本数列の隣接要素の幾何平均(つまり基本数列の間に挿入する形)で補完します。

たとえば $h = 6$ の場合、$h’ = 4$ として $m = 2^{-2}, 2^{-4}, 2^{-6}, 2^{-8}$ の4個を生成し、さらに2個を $m = 2^{-1}, 2^{-3}$ のように挿入します。

ここまでスロープの設計について理解しました。次に、ALiBiが他の位置エンコーディング手法と比較してどのような特性を持つかを整理しましょう。

sinusoidal PE、RoPEとの比較

3手法の位置づけ

位置エンコーディングの主要な3手法を、複数の観点から比較します。

観点 sinusoidal PE RoPE ALiBi
位置情報の注入場所 入力埋め込みに加算 Q, Kに回転を適用 Attentionスコアにバイアス加算
位置の種類 絶対位置 相対位置 相対位置(距離ベース)
追加パラメータ 0 0 0
外挿性 低い 中程度 高い
計算コスト 低い 中程度 低い
表現力 中程度 高い 低い(線形制約)
Value表現への影響 あり(入力に加算) なし(Q, Kのみ) なし(スコアのみ)

位置エンコーディング3手法のレーダーチャート。外挿性・表現力・計算の軽さ・実装の簡単さ・意味位置分離を5段階で比較

表の内容をレーダーチャートで俯瞰すると、3手法の性格の違いが立体的に見えてきます。ALiBi(緑)は外挿性・実装の簡単さ・意味/位置の分離で外側に大きく張り出す一方、表現力では内側にへこみます。RoPE(青)は表現力が突出していますが外挿性は中程度です。sinusoidal PE(赤)は計算は軽いものの、外挿性と意味/位置の分離が弱いことがわかります。「万能な1手法」は存在せず、何を重視するかで選択が変わることが読み取れます。

sinusoidal PE(Vaswani et al., 2017)は、sin/cos関数で生成した固定ベクトルを入力埋め込みに加算します。

$$ \bm{z}_i = \bm{x}_i + \text{PE}(i) $$

この方式では、位置情報がトークンの意味表現に混入します。モデルは学習を通じて意味情報と位置情報を分離する必要があり、表現空間の一部が位置情報の表現に「消費」されます。

また、sinusoidal PEのAttentionスコアは次のように展開できます。

$$ (\bm{x}_i + \text{PE}(i))^\top \bm{W}_Q^\top \bm{W}_K (\bm{x}_j + \text{PE}(j)) $$

この展開を行うと、内容-内容、内容-位置、位置-内容、位置-位置の4つの交差項が生じます。ALiBiでは位置情報がスコアに独立に加算されるため、この複雑な交差が発生しません。

外挿性の観点では、sinusoidal PEは数学的には任意の位置 $i$ に対して $\text{PE}(i)$ を計算できますが、モデルが学習時に見たことのない位置パターンを正しく解釈できるかは別問題です。実験的に、sinusoidal PEの外挿性能はALiBiに大きく劣ることが報告されています。

RoPEとの比較

RoPE(Rotary Position Embedding)は、Query/Keyに位置依存の回転を適用し、Attentionスコアが相対位置のみに依存するようにします。

$$ \text{score}(i, j) = (\bm{R}_i \bm{q}_i)^\top (\bm{R}_j \bm{k}_j) = \bm{q}_i^\top \bm{R}_{j-i} \bm{k}_j $$

RoPEとALiBiの最も本質的な違いは、位置情報と内容情報の結合度です。

RoPEでは、位置と内容が回転を通じて乗法的に結合します。つまり、トークンの意味表現と位置情報が不可分に絡み合います。これにより高い表現力が得られますが、学習時に見なかった距離パターンへの汎化が困難になります。

ALiBiでは、位置情報はスコアに加法的に追加されるだけです。内容のスコア $\bm{q}_i^\top \bm{k}_j / \sqrt{d_k}$ と位置バイアス $-m|i-j|$ は完全に独立しており、モデルは内容の重要度と位置の近接性を別々に評価できます。

外挿性の観点では、RoPEは未知の距離に対して回転角度が学習範囲外になるため、性能劣化が生じます。具体的には、RoPEの高周波成分(低次元ペア)が訓練範囲外で急速に振動し、Attentionスコアが不安定になります。一方、ALiBiのバイアスは距離の線形関数であり、未知の距離に対しても関数形が変わりません。

ただし、RoPEにはNTK-awareスケーリングやYaRN(Yet another RoPE extensioN)といった外挿性を改善する拡張手法が多数提案されており、適切なスケーリングを施せばRoPEでも長系列への汎化が可能です。

ALiBiの表現力の制約

ALiBiの弱点も理解しておく必要があります。ALiBiのバイアスは $-m \cdot |i – j|$ という距離の線形関数に固定されているため、距離に対する非線形な位置関係を表現できません。

例えば、「距離5のトークンと距離50のトークンに同程度に注目し、距離10-40のトークンは無視する」といったパターンは、ALiBiのバイアスだけでは直接表現できません(もちろん、Query/Keyの内積スコアでこれを補うことは可能です)。

RoPEは回転による乗法的結合があるため、位置に依存した内容の選択性がより高く、複雑な位置パターンの学習に適しています。これが、LLaMA系列のモデルがRoPEを採用している理由の一つです。

各手法の特性を理論的に比較したところで、次に実際にALiBiをPyTorchで実装してみましょう。

ALiBiの実装(PyTorch)

スロープの計算

まずは、ALiBiのスロープ $m$ を幾何数列で生成する関数と、バイアス行列の可視化を実装します。

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def get_alibi_slopes(n_heads):
    """ALiBiのスロープを幾何数列で生成する"""
    # 2のべき乗ヘッド数の場合
    def get_slopes_power_of_2(n):
        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
        ratio = start
        return [start * (ratio ** i) for i in range(n)]

    if math.log2(n_heads).is_integer():
        return get_slopes_power_of_2(n_heads)

    # 2のべき乗でない場合
    closest_power = 2 ** math.floor(math.log2(n_heads))
    base_slopes = get_slopes_power_of_2(closest_power)
    # 残りのスロープを補間
    extra_slopes = get_slopes_power_of_2(2 * closest_power)
    extra_slopes = extra_slopes[0::2][:n_heads - closest_power]
    return base_slopes + extra_slopes

# 8ヘッドの場合のスロープを確認
slopes_8 = get_alibi_slopes(8)
print("8ヘッドのスロープ:")
for i, s in enumerate(slopes_8):
    print(f"  Head {i}: m = {s:.6f} = 1/{1/s:.0f}")

このコードを実行すると、8ヘッドの場合のスロープが $1/2, 1/4, 1/8, \ldots, 1/256$ の幾何数列になっていることが確認できます。各ヘッドのスロープは2倍ずつ小さくなっており、先ほどの理論で述べた通り、Head 0が最も近距離に集中し、Head 7が最も遠距離まで対応します。

バイアス行列の構築と可視化

次に、ALiBiのバイアス行列を構築して可視化します。

def build_alibi_bias(n_heads, seq_len, causal=True):
    """ALiBiのバイアス行列を構築する"""
    slopes = get_alibi_slopes(n_heads)
    slopes = torch.tensor(slopes, dtype=torch.float32)

    # 距離行列の構築
    positions = torch.arange(seq_len)
    # |i - j| の行列
    distance = torch.abs(positions.unsqueeze(0) - positions.unsqueeze(1))

    # 各ヘッドのバイアス: (n_heads, seq_len, seq_len)
    bias = -slopes.unsqueeze(1).unsqueeze(1) * distance.unsqueeze(0).float()

    if causal:
        # 因果マスク: j > i のエントリを -inf に設定
        causal_mask = torch.triu(
            torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1
        )
        bias = bias.masked_fill(causal_mask.unsqueeze(0), float('-inf'))

    return bias

# バイアス行列の可視化
n_heads = 8
seq_len = 32
bias = build_alibi_bias(n_heads, seq_len, causal=True)
slopes = get_alibi_slopes(n_heads)

fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for h in range(n_heads):
    ax = axes[h // 4, h % 4]
    b = bias[h].clone()
    b[b == float('-inf')] = float('nan')
    im = ax.imshow(b.numpy(), cmap='RdBu_r', aspect='equal',
                   vmin=-8, vmax=0)
    ax.set_title(f'Head {h} (m={slopes[h]:.4f})', fontsize=10)
    ax.set_xlabel('Key position')
    ax.set_ylabel('Query position')
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

plt.suptitle('ALiBi Bias Matrices for 8 Attention Heads', fontsize=14)
plt.tight_layout()
plt.savefig('alibi_bias_matrices.png', dpi=150, bbox_inches='tight')
plt.show()

代表4ヘッドのALiBiバイアス行列ヒートマップ。傾きが大きいヘッドほど対角から離れると急速に濃い青になる

この可視化から、ALiBiのバイアス行列の構造が直感的に理解できます。Head 0($m = 0.5$)では対角線から離れると急速に暗い青(大きな負のバイアス)になっており、非常に狭い範囲のトークンにしか注目できないことがわかります。Head 7($m \approx 0.0039$)ではバイアスの変化が緩やかで、遠距離のトークンにもある程度の重みが残ります。また、因果マスクにより右上三角が全てマスクされている(未来のトークンを見ない)ことも確認できます。各対角線が一定値を取るToeplitz構造になっている点も見て取れます。

ALiBi付きSelf-Attentionの実装

ここからは、ALiBiを組み込んだMulti-Head Self-Attentionモジュールを実装します。

class ALiBiMultiHeadAttention(nn.Module):
    """ALiBiを使用したMulti-Head Self-Attention"""

    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Query, Key, Valueの線形射影
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        # ALiBiのスロープ(固定パラメータ、学習しない)
        slopes = torch.tensor(get_alibi_slopes(n_heads),
                              dtype=torch.float32)
        self.register_buffer('slopes', slopes)

    def _build_alibi_bias(self, seq_len, device):
        """指定された系列長のALiBiバイアスを動的に構築"""
        positions = torch.arange(seq_len, device=device)
        distance = torch.abs(
            positions.unsqueeze(0) - positions.unsqueeze(1)
        ).float()
        # (n_heads, seq_len, seq_len)
        bias = -self.slopes.unsqueeze(1).unsqueeze(1) * distance.unsqueeze(0)
        # 因果マスク
        causal_mask = torch.triu(
            torch.ones(seq_len, seq_len, device=device, dtype=torch.bool),
            diagonal=1
        )
        bias = bias.masked_fill(causal_mask.unsqueeze(0), float('-inf'))
        return bias

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # Q, K, V の計算
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        # ヘッドに分割: (batch, n_heads, seq_len, d_k)
        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Attentionスコア: (batch, n_heads, seq_len, seq_len)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        # ALiBiバイアスを加算
        alibi_bias = self._build_alibi_bias(seq_len, x.device)
        scores = scores + alibi_bias.unsqueeze(0)  # ブロードキャスト

        # softmax と重み付き和
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        # ヘッドの結合
        output = output.transpose(1, 2).contiguous().view(
            batch_size, seq_len, self.d_model
        )
        output = self.W_o(output)

        return output, attn_weights

この実装のポイントは3つあります。第一に、ALiBiのスロープは register_buffer で登録されており、モデルの学習パラメータには含まれません。optimizer.step() で更新されることはなく、モデルの保存・読み込み時には自動的に保持されます。第二に、バイアス行列は _build_alibi_bias メソッドで動的に構築されるため、推論時に訓練時と異なる系列長を入力しても問題なく動作します。第三に、通常のMulti-Head Attentionとの違いは scores = scores + alibi_bias の1行だけであり、既存のコードベースへの統合が非常に容易です。

動作確認とAttentionパターンの可視化

実装したALiBi Attentionの動作を確認し、Attentionパターンを可視化しましょう。

# モデルの初期化と動作確認
d_model = 64
n_heads = 8
seq_len = 32
batch_size = 1

torch.manual_seed(42)
model = ALiBiMultiHeadAttention(d_model, n_heads)

# ランダムな入力
x = torch.randn(batch_size, seq_len, d_model)
output, attn_weights = model(x)

print(f"入力形状: {x.shape}")
print(f"出力形状: {output.shape}")
print(f"Attention重み形状: {attn_weights.shape}")

# Attentionパターンの可視化
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for h in range(n_heads):
    ax = axes[h // 4, h % 4]
    attn = attn_weights[0, h].detach().numpy()
    im = ax.imshow(attn, cmap='viridis', aspect='equal')
    slope = get_alibi_slopes(n_heads)[h]
    ax.set_title(f'Head {h} (m={slope:.4f})', fontsize=10)
    ax.set_xlabel('Key position')
    ax.set_ylabel('Query position')
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

plt.suptitle('ALiBi Attention Patterns (Random Input)', fontsize=14)
plt.tight_layout()
plt.savefig('alibi_attention_patterns.png', dpi=150, bbox_inches='tight')
plt.show()

代表4ヘッドのALiBi Attention重み。傾きが大きいHead 0は対角付近に集中、小さいHead 7は広く分散

Attentionパターンの可視化から、ALiBiの効果が明確に見て取れます。Head 0($m = 0.5$)では対角線付近にAttention重みが集中しており、各トークンは直近の数トークンにのみ強く注目しています。一方、Head 7($m \approx 0.0039$)ではAttention重みがより広く分散しており、遠距離のトークンにも一定の重みが残っています。全てのヘッドにおいて、因果マスクにより右上三角が0になっていること(未来のトークンへのAttentionがないこと)も確認できます。この多スケールなAttentionパターンにより、モデルは局所的な構文関係から大域的な意味関係まで、さまざまなスケールの依存関係を並列に捉えられます。

実装が完成したところで、次にALiBiの最大の強みである外挿性能を実験的に検証してみましょう。

外挿実験 — 訓練長と推論長を変えたPerplexity比較

実験の設計

ALiBiの外挿性能を定量的に評価するために、小規模な言語モデルを構築し、特定の系列長で訓練した後に、訓練時より長い系列でのPerplexityを測定します。

Perplexity(困惑度)は言語モデルの性能を測る標準的な指標で、次の式で定義されます。

$$ \text{PPL} = \exp\left(-\frac{1}{N}\sum_{i=1}^{N}\log p(x_i | x_{

Perplexityが低いほどモデルの予測精度が高いことを意味します。外挿性能が高い手法は、推論時の系列長が訓練時を超えてもPerplexityの上昇が小さいはずです。

以下の実験では、合成データを使って3つの位置エンコーディング手法(sin/cos、RoPE、ALiBi)の外挿性能を比較します。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import math

# === 位置エンコーディング手法の定義 ===

class SinusoidalPE(nn.Module):
    """sin/cos絶対位置エンコーディング"""
    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float()
            * -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]
class RoPEAttention(nn.Module):
    """RoPE付きSingle-Head Attention(簡略化版)"""
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)

    def _apply_rope(self, x, positions):
        d = x.shape[-1]
        theta = 10000.0 ** (-2 * torch.arange(d // 2,
                            device=x.device).float() / d)
        angles = positions.unsqueeze(-1) * theta.unsqueeze(0)
        cos_a = torch.cos(angles)
        sin_a = torch.sin(angles)
        x_even = x[..., 0::2]
        x_odd = x[..., 1::2]
        x_rot = torch.zeros_like(x)
        x_rot[..., 0::2] = x_even * cos_a - x_odd * sin_a
        x_rot[..., 1::2] = x_even * sin_a + x_odd * cos_a
        return x_rot

    def forward(self, x):
        B, L, D = x.shape
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        positions = torch.arange(L, device=x.device).float()
        Q = self._apply_rope(Q, positions)
        K = self._apply_rope(K, positions)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(D)
        mask = torch.triu(
            torch.ones(L, L, device=x.device, dtype=torch.bool), diagonal=1
        )
        scores = scores.masked_fill(mask, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        return torch.matmul(attn, V)
class ALiBiAttention(nn.Module):
    """ALiBi付きSingle-Head Attention"""
    def __init__(self, d_model, n_heads=1):
        super().__init__()
        self.d_model = d_model
        self.slope = get_alibi_slopes(max(n_heads, 2))[0]
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, L, D = x.shape
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(D)
        # ALiBiバイアス
        positions = torch.arange(L, device=x.device)
        dist = torch.abs(positions.unsqueeze(0) - positions.unsqueeze(1)).float()
        bias = -self.slope * dist
        mask = torch.triu(
            torch.ones(L, L, device=x.device, dtype=torch.bool), diagonal=1
        )
        bias = bias.masked_fill(mask, float('-inf'))
        scores = scores + bias
        attn = F.softmax(scores, dim=-1)
        return torch.matmul(attn, V)
class SimpleLM(nn.Module):
    """簡易言語モデル(位置エンコーディング手法を差し替え可能)"""
    def __init__(self, vocab_size, d_model, pe_type='alibi', max_len=512):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pe_type = pe_type

        if pe_type == 'sinusoidal':
            self.pe = SinusoidalPE(d_model, max_len)
            self.attn = self._simple_attn(d_model)
        elif pe_type == 'rope':
            self.attn = RoPEAttention(d_model)
        elif pe_type == 'alibi':
            self.attn = ALiBiAttention(d_model)

        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)

    def _simple_attn(self, d_model):
        """sin/cos PE用の通常のAttention"""
        class SimpleAttn(nn.Module):
            def __init__(self, d):
                super().__init__()
                self.W_q = nn.Linear(d, d, bias=False)
                self.W_k = nn.Linear(d, d, bias=False)
                self.W_v = nn.Linear(d, d, bias=False)
                self.d = d
            def forward(self, x):
                B, L, D = x.shape
                Q, K, V = self.W_q(x), self.W_k(x), self.W_v(x)
                scores = Q @ K.transpose(-2, -1) / math.sqrt(D)
                mask = torch.triu(torch.ones(L, L, device=x.device,
                                             dtype=torch.bool), diagonal=1)
                scores = scores.masked_fill(mask, float('-inf'))
                return F.softmax(scores, dim=-1) @ V
        return SimpleAttn(d_model)

    def forward(self, x):
        h = self.embedding(x)
        if self.pe_type == 'sinusoidal':
            h = self.pe(h)
        h = h + self.attn(h)
        h = self.norm(h)
        return self.head(h)
def train_and_evaluate(pe_type, train_len=64, eval_lengths=None,
                       vocab_size=256, d_model=64, epochs=100,
                       n_samples=200, seed=42):
    """学習と外挿評価を行う"""
    torch.manual_seed(seed)
    np.random.seed(seed)

    if eval_lengths is None:
        eval_lengths = [32, 64, 96, 128, 192, 256]

    max_eval = max(eval_lengths)
    model = SimpleLM(vocab_size, d_model, pe_type, max_len=max_eval)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    # 合成データ: 単純なパターン(繰り返しシーケンス + ノイズ)
    def gen_data(length, n):
        data = []
        for _ in range(n):
            # 周期的パターン + ランダム摂動
            period = np.random.randint(3, 10)
            base = np.random.randint(0, vocab_size, size=period)
            seq = np.tile(base, length // period + 1)[:length]
            noise_mask = np.random.random(length) < 0.1
            seq[noise_mask] = np.random.randint(0, vocab_size,
                                                 size=noise_mask.sum())
            data.append(seq)
        return torch.tensor(np.array(data), dtype=torch.long)

    # 訓練
    train_data = gen_data(train_len, n_samples)
    model.train()
    for epoch in range(epochs):
        idx = np.random.permutation(n_samples)[:32]
        batch = train_data[idx]
        logits = model(batch)
        # 次トークン予測の交差エントロピー
        loss = F.cross_entropy(
            logits[:, :-1].reshape(-1, vocab_size),
            batch[:, 1:].reshape(-1)
        )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 評価(各系列長でのPerplexity)
    model.eval()
    ppls = {}
    with torch.no_grad():
        for L in eval_lengths:
            eval_data = gen_data(L, 50)
            logits = model(eval_data)
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, vocab_size),
                eval_data[:, 1:].reshape(-1)
            )
            ppls[L] = math.exp(loss.item())

    return ppls

# 3手法の比較実験
train_len = 64
eval_lengths = [32, 48, 64, 80, 96, 128, 160, 192, 256]

results = {}
for pe_type in ['sinusoidal', 'rope', 'alibi']:
    print(f"\n--- {pe_type.upper()} ---")
    ppls = train_and_evaluate(
        pe_type, train_len=train_len,
        eval_lengths=eval_lengths, epochs=150
    )
    results[pe_type] = ppls
    for L, ppl in ppls.items():
        marker = " <-- train length" if L == train_len else ""
        print(f"  Length {L:>4}: PPL = {ppl:.2f}{marker}")
# 結果の可視化
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

colors = {'sinusoidal': '#e74c3c', 'rope': '#3498db', 'alibi': '#2ecc71'}
labels = {'sinusoidal': 'Sinusoidal PE', 'rope': 'RoPE', 'alibi': 'ALiBi'}

# (a) Perplexity vs. 系列長
for pe_type, ppls in results.items():
    lengths = sorted(ppls.keys())
    ppl_vals = [ppls[L] for L in lengths]
    ax1.plot(lengths, ppl_vals, 'o-', color=colors[pe_type],
             linewidth=2, markersize=6, label=labels[pe_type])

ax1.axvline(x=train_len, color='gray', linestyle='--',
            alpha=0.7, label=f'Train length ({train_len})')
ax1.set_xlabel('Sequence Length', fontsize=12)
ax1.set_ylabel('Perplexity', fontsize=12)
ax1.set_title('Perplexity vs Sequence Length', fontsize=14)
ax1.legend(fontsize=10)
ax1.grid(True, alpha=0.3)

# (b) 訓練長に対するPerplexity比率
for pe_type, ppls in results.items():
    base_ppl = ppls[train_len]
    lengths = sorted(ppls.keys())
    ratios = [ppls[L] / base_ppl for L in lengths]
    ax2.plot(lengths, ratios, 's-', color=colors[pe_type],
             linewidth=2, markersize=6, label=labels[pe_type])

ax2.axvline(x=train_len, color='gray', linestyle='--', alpha=0.7)
ax2.axhline(y=1.0, color='gray', linestyle=':', alpha=0.5)
ax2.set_xlabel('Sequence Length', fontsize=12)
ax2.set_ylabel('PPL Ratio (relative to train length)', fontsize=12)
ax2.set_title('Extrapolation Performance (PPL Ratio)', fontsize=14)
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('alibi_extrapolation.png', dpi=150, bbox_inches='tight')
plt.show()

系列長ごとのPerplexityとPPL比。ALiBiは訓練長64を超えても安定、sinusoidalとRoPEは長系列で悪化

上の図は、直近2トークンで次トークンが決まる位置不変な局所パターンを訓練長64で学習させ、評価系列長を変えて測定したPerplexityです(学習可能で再現性のあるタスクを選んでいます)。左がPerplexityの絶対値、右が訓練長を基準にしたPPL比です。この外挿実験の結果から、3つの重要な知見が読み取れます。

第一に、内挿領域(系列長が訓練長64以下)では、3手法とも同程度のPerplexityを達成しています。位置エンコーディングの手法の違いは、学習時に見た範囲内ではほとんど性能差を生じません。

第二に、外挿領域(系列長が訓練長64を超えた範囲)では、手法間の差が顕著になります。Sinusoidal PEは系列長が訓練長の2倍(128)を超えるとPerplexityが急激に悪化する傾向が見られます。RoPEはSinusoidal PEよりも安定していますが、系列長が訓練長の3-4倍になると性能低下が顕在化します。ALiBiは系列長が訓練長の4倍(256)になっても比較的安定したPerplexityを維持しています。

第三に、右のグラフ(PPL比率)が特に示唆的です。ALiBiの曲線は1.0に近い水平に近い線を描くのに対し、Sinusoidal PEの曲線は急激に上昇します。これは、ALiBiの線形バイアスが未知の距離に対しても自然に一般化できることの実証です。

ただし、この実験は小規模な合成データによるものであり、実際のLLMでの性能差はモデルサイズ、データ、訓練設定に大きく依存します。Press et al.(2021)の原論文では、WikiText-103での大規模実験で同様の傾向がより顕著に確認されています。

外挿実験でALiBiの優位性を確認しましたが、ALiBiも万能ではありません。次に、ALiBiの限界と、それを克服する最新手法について見ていきましょう。

ALiBiの限界と最新手法(YaRN等)との関係

ALiBiの限界

ALiBiは外挿性能において優れていますが、以下の限界が知られています。

限界1: 表現力の制約

ALiBiのバイアスは距離の線形関数 $-m \cdot |i – j|$ に固定されています。これは「距離が離れるほど単調にペナルティが増加する」という仮定を暗黙的に置いています。しかし、実際の言語では「直前のトークンよりも、10トークン前のトークンが重要」というパターンも存在します(例: 長距離の照応関係、段落の主題文への参照など)。ALiBiではこのような非単調な距離パターンをバイアスで直接表現できず、Query/Keyの内積に頼る必要があります。

限界2: 双方向Attentionへの適用

ALiBiは元々、因果的(causal)なAttention、つまり自己回帰モデルのデコーダ向けに設計されています。BERTのような双方向のAttentionに適用する場合、バイアスの設計を修正する必要があります。絶対値距離 $|i – j|$ をそのまま使えば双方向に適用可能ですが、原論文での実験は主にcausal Attentionに対して行われており、双方向Attentionでの効果は限定的に検証されています。

限界3: 非常に長い系列での飽和

ALiBiは外挿性能が高いとはいえ、無限に外挿できるわけではありません。訓練長の10倍以上の系列では、Attention重みが近距離のトークンに極端に偏り、長距離の情報が事実上失われるケースが報告されています。特にスロープが大きいヘッドでは、長距離のバイアスが非常に大きな負の値になり、softmaxを通すとAttention重みが実質的にゼロになります。

限界4: 実際のLLMでの採用状況

ALiBiはBLOOM(BigScience, 2022)やMPT(MosaicML, 2023)で採用されましたが、LLaMA、Mistral、GPT-4といった現在最も広く使われているモデルはRoPEを採用しています。これは、RoPEの方が表現力が高く、適切なスケーリング手法と組み合わせることで外挿性も確保できるためと考えられます。

RoPEの外挿性を改善する手法

RoPEの外挿性の問題に対して、多くの改善手法が提案されています。

Position Interpolation(Chen et al., 2023) は、推論時の位置インデックスを訓練時の範囲に圧縮(内挿)するシンプルなアイデアです。

$$ \text{位置} \, i \mapsto \frac{i \cdot L_{\text{train}}}{L_{\text{test}}} $$

例えば、512トークンで訓練したモデルに1024トークンを入力する場合、位置1024を位置512にマッピングします。これにより、全ての角度が訓練時の範囲内に収まり、外挿ではなく内挿の問題に帰着させます。ただし、位置の解像度が半分になるため、近距離の区別能力が低下するトレードオフがあります。

NTK-awareスケーリング は、RoPEの角速度 $\theta_i = 10000^{-2i/d}$ の基底 $10000$ をスケーリングする手法です。

$$ \theta_i’ = \left(\alpha \cdot 10000\right)^{-2i/d} $$

ここで $\alpha$ はスケーリング係数です。高周波成分($i$ が小さい)は外挿への影響が大きいため、NTK-awareスケーリングでは低周波成分はそのまま、高周波成分を選択的に補間します。

YaRN(Yet another RoPE extensioN, Peng et al., 2023) は、NTK-awareスケーリングをさらに発展させた手法です。YaRNの核心は、RoPEの各周波数成分を3つのカテゴリに分類し、それぞれに適切な処理を行うことです。

  • 低周波成分(長い波長): 外挿しても安定 → スケーリングなし
  • 高周波成分(短い波長): 外挿で最も影響を受ける → 線形内挿
  • 中間成分: 上記2つの間をスムーズに遷移

さらに、YaRNはAttentionスコアに温度パラメータ $\sqrt{s}$ を導入し、分布のシャープネスを調整します。この温度補正は、実はALiBiのスロープの役割と類似しています。

ALiBiと最新手法の関係

ALiBiの思想は、最新の手法にも影響を与えています。

Local Attentionとの併用: Sliding Window Attention(Mistral、Gemmaなどで採用)は、ALiBiの「近距離を重視する」という思想を極端にした手法です。一定の窓幅を超えるトークンへのAttentionを完全にカットオフします。ALiBiはこれをソフトに行うものと解釈できます。

Attention Sinkとの関連: StreamingLLM(Xiao et al., 2023)は、自己回帰モデルのAttentionが最初の数トークン(attention sink)に集中する現象を発見しました。ALiBiのバイアスは最初のトークンへの距離ペナルティが最も大きいため、この現象との相互作用が指摘されています。

ハイブリッドアプローチ: 一部の研究では、RoPEとALiBiの利点を組み合わせた手法(例えば、RoPEで位置情報を埋め込みつつ、追加のバイアス項で外挿性を改善する)が検討されています。

手法選択のガイドライン

現時点での実践的なガイドラインをまとめると、以下のようになります。

  • 外挿性を最重視する場合: ALiBi。追加パラメータなし、実装が簡単、外挿性が最も安定
  • 表現力と外挿性のバランスを重視する場合: RoPE + YaRN。現在のLLMで最も広く採用されており、エコシステムが充実
  • 短い固定長の系列を扱う場合: sinusoidal PEまたは学習可能な位置埋め込み。外挿が不要なら十分な性能
  • Attention機構のカスタマイズが必要な場合: ALiBi。スコアへの加法的なバイアスであるため、他の手法との組み合わせが容易

まとめ

本記事では、ALiBi(Attention with Linear Biases)の理論から実装、外挿実験、最新手法との関係まで体系的に解説しました。

  • ALiBiの核心アイデア: 入力表現に位置情報を加えるのではなく、Attentionスコアに距離比例の線形ペナルティ $-m \cdot |i – j|$ を加算する。学習パラメータは一切不要
  • ヘッドごとのスロープ設計: 幾何数列 $m_i = 2^{-8i/h}$ でスロープを設定し、近距離から遠距離までの多スケールな依存関係をカバー
  • 外挿性のメカニズム: 線形バイアスは未知の距離に対しても自然に外挿される。softmax後は指数的な距離減衰として作用し、言語の局所性を自然に表現
  • 実験結果: sinusoidal PE、RoPEと比較して、訓練長を超える系列でのPerplexity上昇が最も小さい
  • 限界: 距離に対する表現力が線形に制約される。非常に長い系列での飽和がある。現在の主要LLMではRoPE + YaRN系列の手法が主流

ALiBiは「シンプルさこそが汎化性能の鍵」という設計哲学を体現した手法です。複雑な数学的構造(回転行列や学習パラメータ)を使わず、線形バイアスという最小限の構造で外挿という困難な問題に対処しました。

次のステップとして、以下の記事も参考にしてください。