💬 NLP

Bayesian Attention Mechanism: A Probabilistic Framework for Positional Encoding and Context Length Extrapolation

本論文は、既存の位置符号化手法を統合する確率的枠組みであるベイズ型注意機構（BAM）を導入し、最先端の長文脈汎化を実現するための一般化ガウス事前分布を提案し、これにより最小限のパラメータオーバーヘッドでトレーニング文脈長の500倍の長さにおいても正確な情報検索を可能にする。

原著者： Arthur S. Bianchessi, Yasmin C. Aguirre, Rodrigo C. Barros, Lucas S. Kupssinskü

公開日 2026-05-11

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Arthur S. Bianchessi, Yasmin C. Aguirre, Rodrigo C. Barros, Lucas S. Kupssinskü

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

現代の AI チャットボットの頭脳であるトランスフォーマー言語モデルを、巨大な本を読もうとする超賢い図書館司書だと想像してみてください。この司書は単語の意味を理解するのが得意ですが、奇妙な問題を抱えています。順序感覚がないのです。「The cat sat on the mat（猫はマットの上に座った）」という文を渡しても、「cat」と「sat」のどちらが先でどちらが後か区別できません。なぜなら、単語がどこに位置しようとも、彼らにとって単語は同じように見えるからです。

これを解決するため、通常は図書館司書に**位置エンコーディング（PE）**を与えます。これは、すべての単語に貼られた「私は 1 番目の単語です」「私は 2 番目の単語です」などと書かれた、小さくて目に見えない名札や色付きのシールのようなものです。これにより、司書は物語の流れを理解できるようになります。

しかし、これらの「シールシステム」のほとんどには欠点があります。短い物語ではよく機能しますが、本が非常に長くなると（10 万ページの長編小説のように）、機能しなくなります。シールが混乱しすぎたり、薄れてしまったりするため、司書は本の冒頭部分を無視し始めるのです。

新しいアイデア：「ベイズ的注意機構（BAM）」

この論文の著者たちは、これらのシールを捉える新しい方法を提案しており、それをBAMと呼んでいます。単語に固定されたシールを貼り付けるのではなく、BAM は単語の位置を確率推測として扱います。

以下がその比喩です：
図書館司書が長い文書の中に隠された特定の情報（「パスキー」）を見つけようとしていると想像してください。

旧式の方法（ALiBi など）： 司書は、答えが現在の単語のすぐ近くにある可能性が最も高いと仮定し、離れるほど答えが見つかる可能性は低くなると考えます。これは家の中で紛失した鍵を探すようなものです。まずポケットを確認し、次にテーブルを確認し、あまりに「遠い」ためガレージでは探さなくなります。
BAM 方法： 司書は、答えがどこにあるかもしれないという「事前信念（確率マップ）」を使用します。このマップは固定されたものではなく、調整可能な柔軟なルールです。

秘密のソース：「一般化ガウス」マップ

この論文は、一般化ガウス事前分布と呼ばれる特定の種類の確率マップを導入しています。このマップを、司書が特定の単語を見る可能性の高さを表す山や谷のある地形だと考えてみてください。

「局所的」な山： マップには、現在の単語のすぐ隣に急峻な山を持つことができます。これにより、司書は即座の文脈（文法や文構造など）を理解できます。
「長距離」の尾： ここが魔法の場所です。著者たちは、特定のつまみ（ $\beta$ と呼ばれます）を調整することで、マップの形状を変えられることを発見しました。
- つまみを一方の方向に回すと、司書は本の遠く離れた部分を無視します（旧式の方法のように）。
- つまみをもう一方の方向に回す（具体的には負の値に設定する）と、司書は直近の隣接単語を見るのをやめ、非常に遠く離れた単語に集中し始めます。

彼らは実際に何を達成したのか？

この論文は、これが病気を治したり、小説を書いたりするものだと主張しているわけではありません。彼らは非常に具体的で制御されたタスクでテストを行いました。

「干し草の山の中の針」テスト： 彼らは、5 桁の数字（「パスキー」）を数千語のテキストの奥深くに隠しました。
- 結果： 他の手法（RoPE や ALiBi など）はテキストが長くなりすぎると諦めてランダムに推測するのに対し、BAM モデルは、モデルが訓練された長さよりも500 倍長いテキストであっても、その数字を完璧に見つけることができました。
「パープレキシティ」テスト： これはモデルが文の次の単語をどの程度予測できるかを測定します。
- 結果： BAM は、既存の最良の手法と同様に単語の予測が得意でした。つまり、この長距離の超能力を得るために一般的な言語能力を犠牲にしたわけではありません。

「目に見えない」コスト

彼らの発見の最もクールな点の一つは、追加コストが非常に低いことです。

彼らは、すでに 1 億 2000 万のパラメータを持つモデルに、1,000 未満の新しいパラメータ（小さなメモリ断片）を追加しました。
これは砂浜に砂粒を一粒加えただけなのに、突然その砂浜全体が地平線を見渡せるようになったようなものです。
モデルの速度が遅くなったり、より多くのコンピューターパワーを消費したりすることはありませんでした。

平易な英語での要約

著者たちは、AI のための新しい「位置シール」システムを構築しました。答えはいつも近くにあると仮定するのではなく、AI に直近の周囲を無視し、過去に埋もれた情報に集中することを可能にする柔軟なルールブックを与えました。

彼らは数学的にこれが機能することを証明し、実験を通じて、他の AI が数千語で迷い込むのに対し、彼らの AI は 50 万語の干し草の山から針を見つけることができることを示しました。彼らはこれを現実世界の医療文書や法的文書でテストしたわけではありませんが、長距離情報を発見する機構が、今やはるかに強力で信頼性のあるものになったことを示しました。

技術概要：ベイズ的注意機構（BAM）

問題定義

トランスフォーマーベースの言語モデルは、トークンの順序を処理するために位置符号化（PE）に依存しているが、既存の手法は文脈長の外挿（トレーニング中に観測されたものよりも長いシーケンスへの一般化能力）に関して、2 つの決定的な限界に直面している：

理論的不透明性：多くの PE 手法（例：正弦波、RoPE、ALiBi、NoPE）は、限られた理論的基盤に基づいた経験的に動機付けられたものであり、外挿時の挙動は十分に理解されていない。
不十分な評価：主要な指標としてパープレキシティへの依存は不十分である。モデルは、長距離情報を真に統合することなく、局所的な注意パターン（最近のトークンのみに焦点を当てること）を通じて低いパープレキシティを達成できる。これはスライディングウィンドウ評価によって浮き彫りにされた欠陥である。

手法：ベイズ的注意機構（BAM）

著者らは、トークンの内容と位置の結合確率モデルの下で計算された値の期待値として自己注意を再定義する理論的枠組みである**ベイズ的注意機構（BAM）**を提案する。

理論的枠組み

結合確率定式化：BAM において、注意重み $p_{ij}$ （トークン $j$ がトークン $i$ に注意を向ける確率）は、結合確率分布として定義される：
$p_{ij} = p(f_{cont}(q_i, k_j) | g_{pos}(i, j)) \times p(g_{pos}(i, j))$
ここで、 $f_{cont}$ は内容の類似性を表し、 $g_{pos}$ は位置関係を表す。
事前分布としての位置符号化：この枠組みは、位置符号化戦略がトークン位置に対する特定の事前分布 $p(g_{pos}(i, j))$ $p (g_{p os} (i, j))$ に対応すると仮定する。
- NoPE は、因果マスクによって制限された一様事前分布として導出される。
- ALiBi は、一様事前分布とラプラス事前分布（距離に基づく指数関数的減衰）の積として導出される。
正規化スカラー（ $Z$ ）：導出は、内容と位置の確率の積を正規化するスカラー $Z$ を導入する。これは、内容と位置の間の相互作用の自由エネルギーまたは情報結合を表す。

提案される解決策：GGD-BAM

この枠組みを活用し、著者らは位置事前分布として一般化ガウス分布（GGD）を利用するGGD-BAMを導入する。注意スコアに追加されるバイアス行列 $B$ は以下のように定義される：
$b_{ij} = -\left| \frac{j - i - \mu}{\alpha} \right|^\beta$
ここで、 $\mu$ 、 $\alpha$ 、 $\beta$ は学習可能なパラメータ（または特定の挙動のために固定される）である。

$\beta \in (0, 1)$ ：ラプラス事前分布（ALiBi）よりも遅く減衰する重み尾分布を生み出し、局所性を維持しつつ ALiBi よりもさらに注意を拡張することを可能にする。
$\beta < 0$ ：標準的な GGD 定義の新しい緩和である。これは、確率質量がクエリトークンの近くで抑制され、文脈ウィンドウの遠端に集中する「反転」事前分布を作成する。理論的には、これによりモデルは局所的文脈を無視し、任意に長い距離の依存関係に焦点を当てることが可能になる。
スケーラブルなソフトマックス（SSMax）：著者らは、BAM が「注意の減衰」（文脈の拡大に伴う一様分布）の影響を受けやすいと指摘する。彼らは分布の鮮明さを維持するために、ログを動的に再スケーリングする SSMax と BAM を組み合わせる。

主要な貢献

統合された理論的枠組み：BAM は既存の PE 手法の確率的な統合を提供し、NoPE と ALiBi をそれぞれ一様事前分布とラプラス事前分布に明示的にマッピングする。
一般化ガウス事前分布：GGD-BAM の提案は、柔軟な位置事前分布のファミリーを導入する。形状パラメータ $\beta$ を学習することで、モデルは局所的注意と長距離検索モードの間を適応的に切り替えることができる。
最小限のオーバーヘッド：この手法は、各注意ヘッドに対して $\theta_\alpha$ と $\theta_\beta$ を学習し、 $\theta_\mu = 0$ を固定することで、120M パラメータモデルの場合、384 個の追加学習パラメータ（1,000 未満）のみを導入する。
解釈可能性：この枠組みは、位置に対する明示的な確率分布として注意パターンを可視化し、解釈するための原理的な方法を提供する。

実証結果

120M から 11 億パラメータまでのモデルで実験が行われ、FineWeb 10B 上で 512 トークンの文脈長でトレーニングされた。

パスキー検索：GGD-BAM（特に SSMax を使用した場合）は、トレーニング長さの 64 倍に相当する32,000 トークンまでのシーケンスから特定のトークンを検索する際、完全な精度を達成し、トレーニング長さの 1,000 倍に相当する512,000 トークンまで 80% 以上の精度を維持した。対照的に、RoPE、正弦波、NoPE などのベースラインは、トレーニング範囲を超えるとほぼランダムな精度に低下した。ALiBi はパープレキシティの外挿においてより良好な結果を示したが、極端な長さにおける検索タスクでは大幅に失敗した。
RULER ベンチマーク（NIAH）：GGD-BAM は、「干し草の山の中の針（Needle In A Haystack）」サブセットにおいて、すべてのベースライン（RoPE や ALiBi を含む）を上回った。特に Single 1 タスクでは、6,000 トークンを超えて高い精度を維持した唯一の手法であった。
パープレキシティ：GGD-BAM は、長い文脈において ALiBi や RoPE と同等のパープレキシティを維持し、検索能力の向上が言語モデルの流暢さを犠牲にするものではないことを示した。
下流タスク：MMLU および ARC ベンチマークにおいて、BAM SSMax は RoPE SSMax を上回った。
注意の可視化：可視化により、 $\beta \le 0$ のヘッドが局所的文脈を効果的に抑制し、遠くのトークン（パスキー）に焦点を当てていること、そして $\beta \in (0, 1)$ のヘッドが長い尾を持つ注意を維持していることが確認された。

意義と主張

本論文は、BAM が二重の目的を果たすと主張している：

理論的：位置符号化に対する厳密な確率的基盤を提供し、特定の手法（ALiBi など）がパープレキシティには機能するが、局所的な性質により検索には失敗する理由を明確にする。
実用的：文脈長の外挿のための非常に効率的な手法を提供する。注意ヘッドが負の $\beta$ 値を学習することを可能にすることで、モデルはトレーニング中にそのような長さに直接さらされることなく、任意に長い文脈からの情報を検索するために局所性を明示的に抑制できる。

著者らは、パープレキシティが標準的な指標である一方で、真の長文脈能力を評価するには不十分であると強調している。モデルが単に局所的なパターンに依存しているのではなく、遠くの情報を注意していることを検証するためには、検索ベースのタスクが必要である。この研究は、位置事前分布を明示的にモデル化する能力が、次世代の長文脈トランスフォーマーにとって重要な帰納的バイアスであることを示唆している。