AI・機械学習
上級

ALiBi(Attention with Linear Biases)(アリビ)

Attention計算時にクエリとキーの距離に比例した線形バイアスを減算することで、位置情報を注入する手法。訓練時より長いシーケンスへの外挿能力に優れ、BLOOM-176BやMPT-7B/30Bで採用された。

0 回閲覧
0 いいね

ALiBi(Attention with Linear Biases)とは

ALiBi(Attention with Linear Biases)は、Press et al.が2022年の論文「Train Short, Test Long」で提案した位置エンコーディング手法である。トークン埋め込みへの位置ベクトル加算を完全に廃止し、代わりにAttentionスコアにクエリ・キー間の距離に比例した負のバイアスを加算する。

概要

ALiBiの核心はシンプルである。Attention行列の(i, j)要素に対して -m × |i - j| というバイアスを加算する。ここで m はヘッドごとに異なる固定スロープ(傾き)で、2^(-8/n) の等比数列から取る(n はAttentionヘッド数)。

この線形バイアスにより、遠いトークンほどAttentionスコアが減衰する局所性バイアスが自然に生まれる。結果として:

  • 位置埋め込みパラメータが完全に不要(メモリ削減)
  • 訓練時のシーケンス長より2-4倍長い推論が可能(外挿能力)
  • 実装が極めてシンプル(数行のコード追加)

数学的定式化

n ヘッドのMulti-Head Attentionに対して:

  • ヘッド h のスロープ: m_h = 2^(-8h/n)(h = 1, 2, ..., n)
  • Attention(Q, K, V) = softmax(QK^T / √d + B) × V
  • バイアス行列 B[i][j] = -m_h × |i - j|

例えば8ヘッドの場合、スロープは: m = [1/2, 1/4, 1/8, 1/16, 1/32, 1/64, 1/128, 1/256]

採用モデルと性能比較

モデルパラメータ訓練長実効推論長備考
BLOOM-176B176B2,048~8,192BigScience、46言語対応
MPT-7B7B2,048~8,192MosaicML、商用利用可
MPT-30B30B8,192~32,768MosaicML、8K訓練
BTLM-3B-8K3B8,192~32,768Cerebras、効率重視

一方、2024年以降の主要モデル(LLaMA 3、Mistral、Qwen2.5)はRoPEを採用し、ALiBiは新規採用が減少傾向にある。

ALiBiの長所と短所

長所:

  • 外挿能力: 訓練長の2-4倍まで性能劣化なし。Sinusoidal PEが訓練長を超えると即座に崩壊するのと対照的
  • パラメータフリー: 位置埋め込みの学習パラメータが不要。BLOOM-176Bで約70MB相当のパラメータ削減
  • 実装容易: Attention計算にバイアス行列を加算するだけ。既存コードへの統合が容易
  • 訓練効率: 短いシーケンスで訓練し、長いシーケンスで推論できるため、訓練コストを削減可能

短所:

  • 表現力の限界: 線形バイアスは「近いトークンほど重要」という仮定を強制する。コード生成や数学的推論のように、遠距離のトークンが重要な場合に不利
  • RoPEに対する性能差: 同じ訓練データ・モデルサイズでの比較ではRoPEの方がperplexityが低い傾向(EleutherAIの比較実験)
  • 外挿の限界: 4倍以上の外挿では性能が顕著に劣化。RoPE+YaRNの32倍外挿には及ばない
  • エコシステム: 2025年以降の主要フレームワーク(vLLM、TGI)はRoPE最適化が優先され、ALiBiのGPUカーネル最適化は停滞

実装例

import torch

def alibi_bias(n_heads, max_len):
    slopes = torch.tensor([2 ** (-8 * i / n_heads) for i in range(1, n_heads + 1)])
    positions = torch.arange(max_len)
    distances = positions.unsqueeze(0) - positions.unsqueeze(1)
    bias = -slopes.unsqueeze(1).unsqueeze(1) * distances.abs().unsqueeze(0)
    return bias  # shape: (n_heads, max_len, max_len)

よくある質問(FAQ)

Q1: ALiBiとRoPEはどちらが優れている? A: 2026年時点ではRoPEが主流である。RoPE+YaRNは32倍以上の外挿が可能で、128Kトークン以上のコンテキスト対応に優れる。ALiBiは実装のシンプルさと4倍程度の外挿には有利だが、超長文処理ではRoPEに劣る。

Q2: ALiBiはなぜ外挿できる? A: 線形バイアスは任意の距離 |i-j| に対して定義されるため、訓練時に見ていない距離にも自然に一般化する。ただし距離が大きくなるとバイアスの絶対値も線形に増大するため、実効的なAttention範囲は有限となる。

Q3: BLOOM-176Bが ALiBiを選んだ理由は? A: BigScienceプロジェクトは2021-2022年に設計されており、RoPEの実用性がまだ十分に検証されていなかった。ALiBiの外挿性能とパラメータ効率がその時点で最良の選択肢であった。46言語・2048トークン訓練で8Kトークン推論を実現した。

まとめ

  • ALiBiはAttentionスコアに距離比例の線形バイアスを加算する位置エンコーディング
  • パラメータフリーで訓練長の2-4倍外挿が可能
  • BLOOM-176B、MPT-7B/30Bで採用実績
  • 2024年以降はRoPEが主流だが、シンプルさでは優位が残る