🤖 machine learning

Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking

本論文は、情報形式のカルマンフィルタを用いて厳密なベイズフィルタリングを再定式化することで、明示的な不確実性を伴う非線形かつスキャン並列な状態更新を実現し、MambaやGLAといった既存の線形計算量モデルの表現力と状態追跡能力を凌駕しつつ、計算効率を維持する並列化可能なシーケンス混合層であるKalman Linear Attention (KLA) を導入するものである。

原著者： Vaisakh Shaj, Cameron Barker, Aidan Scannell, Andras Szecsenyi, Elliot J. Crowley, Amos Storkey

公開日 2026-06-11

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

CC BY 4.0

原著者： Vaisakh Shaj, Cameron Barker, Aidan Scannell, Andras Szecsenyi, Elliot J. Crowley, Amos Storkey

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

大きな絵：AIの「思考」における新しい方法

あなたは非常に長い本を読もうとしているところだと想像してください。

旧世代のAI（Transformer）: 文を理解するために、AIは次の単語を決めるために、これまでの本のすべての単語を振り返ります。それは、一冊の本を見つけるために図書館全体を引き出す司書のようなものです。非常に正確ですが、本が長くなるにつれて信じられないほど遅くなり、コストもかかります。
現在の効率的なAI（Mamba, GLA）: これらは、小さな固定サイズのノートを持っている司書のようなものです。彼らは最も重要なことだけを書き留め、残りのことは忘れてしまいます。高速ですが、単に古いノートに新しい情報を「加算」するだけなので、微妙なつながりを見逃したり、物語が複雑になりすぎると混乱したりすることがあります。

この論文は、新しい司書を紹介しています。それが「Kalman Linear Attention (KLA)」です。

KLAはノートを持っていますが、単に事実を書き留めるのではなく、その事実に対してどれくらい自信を持っているかを記録します。彼らは自問自答します。「私はこれを覚えているが、どの程度確信しているだろうか？この新しい情報はゲームチェンジャー（状況を一変させるもの）なのか、それとも単なる些細な詳細なのか？」

コアとなるアイデア：「自信」のノート

著者たちは、現在の効率的なAIモデルが、記憶を単純な数学の方程式（数値を足し合わせるだけ）として扱っていることに気づきました。しかし、本当の思考には不確実性が伴います。

「信念状態（Belief State）」: KLAは単に事実を保存するのではなく、事実と信頼スコア（例：「降水確率80%」という天気予報のようなもの）を一緒に保存します。
「門番（Gatekeeper）」: 新しい情報が入ってきたとき、KLAは単にそれを追加するわけではありません。まず、その自信を確認します。
- AIが現在の記憶に非常に自信がある場合、新しいノイズ混じりの情報を無視します。
- AIが確信が持てない場合、新しい情報に注意を払い、記憶を更新します。
魔法のトリック（並列化）: 通常、自信を確認しながらステップごとに記憶を更新することは、時間がかかります（スタンプを押すために列を作って待っている人々の列のようなものです）。この論文の大きな突破口は、この「自信のチェック」を、コンベアベルトのように一度にまとめて行うことができることを示した点です。これにより、現在の最も高速なモデルと同じ速度を実現しました。

創造的な比喩

1. 「懐疑的な探偵」 vs 「メモ係」

現在のモデル（メモ係）: 目撃者が言ったことをすべてノートに書き留める探偵を想像してください。目撃者が「車は赤でした」と言えば、探偵は「赤」と書きます。もし後で「実は、たぶんマルーン色だったかも」と言ったら、探偵は単にその横に「マルーン」と書き加えます。ノートは散らかり、探偵は何が起きたのかについて混乱してしまうかもしれません。
KLA（懐疑的な探偵）: この探偵には「自信メーター」があります。
- 目撃者: 「車は赤でした。」
- 探偵: 「わかりました。赤であることに90%の自信があります。そう書き留めておきましょう。」
- 目撃者: 「待ってください、青い車も見えた気がします。」
- 探偵: 「ふむ、私の『赤』に対する自信は高いですが、この新しい手がかりは怪しいですね。赤への自信を少し下げつつ、青い車のこともメモしておきますが、まだ『赤』を消去はしません。」
- 結果: 探偵は、情報に圧倒されることなく、より明確で正確な犯罪現場のイメージを維持できます。

2. 「信号機」システム

AIの記憶をハイウェイ（高速道路）と考えてください。

線形モデル（Linear Models）: すべての車（新しい単語）がただ交通の流れに合流します。それはスムーズで直線的な流れです。
KLA: すべての車に信号機が付いています。
- ハイウェイが空いているとき（自信が低いとき）、信号は「青」になります。新しい車は容易に合流できます。
- ハイウェイが激しい渋滞に巻き込まれているとき（自信が高いとき）、信号は「赤」になります。新しい車は待機するか、非常に慎重に合流しなければなりません。
- イノベーション: この論文は、1,000マイルのハイウェイにあるすべての信号機を、一つ一つのマイルマーカーで止まって確認するのではなく、同時並行で計算する方法を編み出しました。

彼らは実際に何を証明したのか？

この論文は、これが病気を治したり株価を予測したりすると主張しているわけではありません。焦点は言語モデリング（AIの読解力や執筆力を賢くすること）にあります。ここで彼らが示したことは以下の通りです。

より速く、より賢い: KLAは、現在の最も高速なモデル（Mambaなど）と同じ速度でありながら、より難しい論理パズルを解くことができます。
「置換（Permutation）」テスト: 彼らはAIに「A5」と呼ばれるタスクを与えました。これは、アイテムを特定の順序でシャッフルしなければならない複雑なパズルのようなものです。
- 現在の高速なモデル（Mambaや標準的なTransformer）は、AIを巨大かつ深く設計しない限り、このパズルに失敗しました。
- KLAは、非常に小さく浅いモデルでこれを解きました。 これは、KLAが競合モデルよりも複雑で変化する状態を追跡できることを証明しています。
「ロングコンテキスト」への対応: AIが2,000語前の物語を覚えておく必要がある場合、KLAは他の高速モデルよりも優れた精度で正しい詳細を見つけ出すことができました。
スケールへの適応: 彼らは数十億語のデータを用いてモデルを訓練しました。モデルがクラッシュすることはありませんでした。これは、この「不確実性に基づいた」アプローチが、大規模で現実世界のAIシステムに使用できることを証明しています。

「秘伝のソース」：OUプロセス

論文では「オルンシュタイン＝ウーレンベック（OU）過程」というテクニカルな用語に触れています。

比喩: ボールに取り付けられたゴムバンドを想像してください。ボールを引っ張ると、ゴムバンドがボールを中心へと引き戻します。
AIにおける役割: このゴムバンドは、AIの「自信」が暴走（無限大に膨れ上がったり、ゼロに墜落したりすること）するのを防ぎます。これにより、AIの記憶を安定させ、モデルを壊すことなく何層にも深く積み重ねることを可能にします。この「ゴムバンド」がなければ、モデルを大きくしたときに不安定になってしまいます。

まとめ

Kalman Linear Attentionは、自信に満ちた、懐疑的な探偵のように機能する、新しいタイプのAIメモリです。単に暗記するのではなく、自分が知っていることに対してどれほど確信を持っているかを追跡します。これにより、ノイズをフィルタリングし、現在の高速なAIモデルよりもはるかにうまく重要な詳細に集中することができます。著者たちは、これが難しい論理パズルに対して有効であり、膨大な量のデータで訓練可能であることを証明しました。

技術要約：Kalman Linear Attention (KLA)

問題提起

MambaやGated Linear Attention (GLA) といった現在の状態空間モデル (SSM) は、Transformerの二次的なスケーリングのボトルネックに対処する、線形計算量かつ並列化可能な代替案を提供している。しかし、これらのモデルは線形またはアフィンな隠れ状態の更新に依存しており、ソフトマックス正規化によって誘発される非線形な相互作用と比較すると、表現力が制限されている。さらに、既存の線形RNNやSSMは、隠れ状態を決定論的なものとして扱うため、状態の不確実性を表現するための明示的なメカニズムを欠いている。本論文は、線形更新の効率性と非線形フィルタリングの表現力・堅牢性の間のギャップに対処し、正確なカルマンフィルタの更新を効率的なシーケンスミキサーとして実装できるかという問いを投げかける。

手法：Kalman Linear Attention (KLA)

著者らは、シーケンスモデリングを正確なベイズフィルタリングとして定式化する確率的シーケンスミキサー、Kalman Linear Attention (KLA) を提案する。核心となる革新は、並列計算を可能にするために、カルマンフィルタを情報形式（精度および自然パラメータを使用）へと再パラメータ化することにある。

1. 確率的定式化

KLAは、シーケンスミキシングを線形ガウス状態空間モデルにおける推論としてモデル化する：

ダイナミクス事前分布: 潜在状態の進化は、正確に離散化されたオルンシュタイン＝ウーレンベック（OU）過程としてモデル化される。これにより、標準的なSSMの決定論的なダイナミクスとは異なり、状態遷移における不確実性を捉えるプロセスノイズ項 ( $\bar{p}_t$ ) が導入される。
観測モデル: 入力トークンは、観測演算子 ( $k_t$ ) と、トークンの証拠に対する信頼性を表す値の精度 ( $\Lambda^v_t$ ) によってパラメータ化された、潜在状態のノイズを含む観測として扱われる。
リードアウト: 出力は、クエリ演算子 ( $q_t$ ) を介して投影された、潜在状態の事後平均から導出される。

2. メビウス変換による並列化

古典的なカルマンフィルタは本質的に逐次的である。KLAは、フィルタの情報形式を利用することでこれを克服する：

精度再帰: 事後精度 ( $\lambda_t$ ) の更新は、前の精度 ( $\lambda_{t-1}$ ) のメビウス（分数線形）変換であることが示されている。
$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$
決定的なことに、メビウス変換は結合的に合成可能である。これにより、精度更新を、現代のMambaのようなSSMと同等の並列効率を持つ並列プレフィックススキャン（結合スキャン）を介して計算できる。
平均再帰: 精度パスが並列に計算された後、事後平均（情報平均）の更新は、同様に並列化可能なアフィン変換に従う。

3. 不確実性によるゲーティング

ゲーティングがしばしば線形またはトークン依存で決定論的である線形SSMとは異なり、KLAのゲーティングは精度の比率から自然に生じる。「忘却ゲート」と「書き込みゲート」は、事前状態と新しい観測の間の相対的な不確実性によって駆動される、履歴依存かつ非線形なものである。これにより、モデルはすべてのトークンに対して明示的な学習済みゲーティングパラメータを必要とすることなく、信頼度に基づいて情報を選択的に取り込むことができる。

4. アーキテクチャ

KLAは、標準的なSSMまたはアテンション層のドロップインリプレイスメントとして設計されている。1次元因果畳み込みとSiLU活性化関数を組み込んだ、Mambaに似た融合MLPアーキテクチャに従う。著者らは、高帯域幅メモリ（HMB）内で拡張された状態を展開することを避けるために、メビウススキャンのための効率的なCUDAおよびTritonカーネルを実装している。

主な貢献

結合再パラメータ化: 著者らは、情報形式における対角線形ガウス・カルマンフィルタがメビウス構造を持ち、ベイズフィルタリングのための正確な並列プレフィックススキャンを可能にすることを実証した。
不確実性駆動の非線形性: KLAは、精度の比率から導かれる非線形ゲーティングを導入し、線形/アフィンなSSMの更新よりも厳密に高い表現力を提供しながら、線形時間のスキャン構造を維持する。
KLAレイヤー: 明示的な信念状態の不確実は保持しつつ、深いネットワーク内でスタック可能な、実用的なドロップインシーケンスミキサーを提供する。
スケーリングと検証: 本論文は、数十億トークン規模で訓練された最初のスタック型ベイズフィルタリング・プリミティブを提示し、その安定性と性能を検証している。

実験結果

著者らは、合成タスク、状態追跡ベンチマーク、および大規模言語モデリングにおいてKLAを評価している。

合成タスク (MAD-Lab): KLAは、特にFuzzy RecallやSelective Copyといった不確実性下での選択を必要とするタスクにおいて、GDN、GLA、Mamba、およびmLSTMを凌駕する。Compressionにおいて最先端の結果を達成しており、これは情報形式のフィルタリングが最適なコンテキスト圧縮器として機能することを示唆している。
状態追跡 (A5 置換): NC1完全な状態追跡のベンチマークであるA5置換タスクにおいて、KLAは定数深度（1〜2ブロック）でこのタスクを解決する。対照的に、線形SSMやアテンションベースのモデルは、シーケンス長とともに増大する深度を必要とする。これは、KLAの分数線形更新が、並列訓練可能性を維持しながら、線形回帰モデルの表現力の限界を打破することを示している。
長文脈連想想起 (MQAR): KLAは、シーケンス長2048において、Mambaを上回り、GLAを大幅に上回る、ほぼ完全な精度（>95%）をMQARタスクで達成する。
大規模事前学習:
- スタンドアロン: 45Mおよび180Mパラメータにおいて、スタンドアロンのKLAは、ゼロショットの常識ベンチマーク（例：HellaSwag、PIQA）において、GPT、Mamba、およびGDNと同等またはそれ以上の性能を示す。
- ハイブリッド: GPTの最終アテンション層をKLAに置き換えた場合（GPT+KLA）は、ハイブリッドの中で最も強力な性能を示し、GPT+MambaおよびGPT+GDNを上回る。
アブレーション: 著者らは、オルンシュタイン＝ウーレンベック（OU）事前分布が深いスタックの安定性に不可リオであることを示している。平均回帰的な確率的ダイナミクスがない場合、KLAの深いスタックの訓練は不安定になる。

意義と主張

本論文は、KLAが線形状態空間モデルの効率性と非線形アテンションメカニズムの表現力の間のギャップをうまく埋めることを主張している。シーケンスミキシングを正確なベイズフィルタリングとして捉えることで、KLAは以下を提供する：

明示的な不確実性: 決定論的なミキサーとは異なり、KLAは情報ゲーティングを直接制御する明示的な信念状態の不確実性を保持する。
厳密に高い表現力: メビウス更新により、KLAは並列訓練能力を犠牲にすることなく、線形SSMが達成できない状態追跡タスク（A5など）を解決できる。
スケーラビリティ: 本研究は、正確なベイズフィルタリングのプリミティブがスタック可能であり、大規模な言語モデリングにおいても数十億トークン規模で訓練可能であることを示し、確率的フィルタリングは本質的に逐次的である、あるいは大規模な言語モデリングにとって計算的に困難であるという概念に異を唱えている。

著者らは、KLAを単なる理論的な好奇心としてではなく、長文脈処理に必要な計算効率を維持しながら、現代の言語モデルの堅牢性と状態追跡能力を強化する、実用的でドロップイン可能なプリミティブとして位置づけている。