ALiBi 位置エンコーディングの盲点：長距離注意が数値精度で消える

導入

ALiBi（Attention with Linear Biases）は、Transformer に長いコンテキストへの汎化能力を持たせるためによく使われる位置エンコーディング手法です。位置ごとの埋め込みを学習するのではなく、注意スコアに距離に応じた線形バイアスを加えることで、近いトークンを優先しつつ過去の情報も参照できるようにします。

しかし論文 When Attention Goes Blind は、この単純で扱いやすい仕組みに、見落とされてきた数値的な失敗モードがあると指摘します。問題はモデルが遠い情報を「好まない」ことではなく、浮動小数点計算の結果として遠い情報への注意が完全にゼロになる可能性がある点です。

核心ポイント

• 原因は浮動小数点のアンダーフロー：ALiBi の線形バイアスは距離に従って大きくなります。距離が長くなると softmax 後の注意確率が表現可能な精度を下回り、ゼロに丸められることがあります。
• 注意ヘッドが部分的に「盲目」になる：注意重みがゼロであれば、その位置から情報を読み出す経路は実質的に閉じます。これは単に小さい重みを与える場合より深刻です。
• 通常の評価では見えにくい：論文によれば、この失敗は長い文脈から特定トークンを探す retrieval タスクで大きく効きます。一方、標準的なデコーダベンチマークへの影響は比較的小さく、問題が表面化しにくいとされています。
• 実際の事前学習済みモデルでも確認：著者らは、ALiBi に基づく最先端の事前学習済みモデルでもこの現象が起きることを示しています。
• 制御された事前学習実験：148M パラメータのデコーダモデルを用いた実験により、この数値失敗と一般的な長文脈外挿の劣化を切り分けています。
• 有効な緩和策：4種類の学習時対策を比較した結果、距離を対数スケールにする方法が passkey retrieval で最も一貫した改善を示しました。

意義と影響

この研究の重要な点は、長文脈モデルの失敗を単なる能力不足ではなく、数値表現とスケーリングの問題として捉え直したことです。注意確率がすでにゼロになっているなら、モデルはその位置を学習でうまく使う以前に、情報へアクセスする数値的な経路を失っています。

ただし、論文は ALiBi を否定しているわけではありません。デフォルトの ALiBi slopes は依然として強力なベースラインであり、特に needle-in-a-haystack retrieval では優れた競争力を保つと報告されています。

実務上の教訓は明確です。ALiBi を使う長文脈モデルでは、通常の言語モデリング評価だけでなく、passkey retrieval や needle-in-a-haystack のような検索型評価を導入すべきです。また、学習時には距離スケーリングの設計を見直し、対数スケールのような緩和策を検討する価値があります。

出典：Hugging Face Daily Papers