不確実性を意識した連想記憶「Kalman Delta Networks」を提案

線形アテンションは、効率的な長文コンテキスト推論と定数メモリでのデコーディングのために、最先端の言語モデルでますます広く使われています。しかし、その固定サイズの再帰的メモリでは、各トークンごとにオンライン判断が必要です。つまり、「何を記録するか」「既存の関連付けをどの程度上書きするか」を、将来のクエリがどのような情報を必要とするかを知る前に決定しなければならないのです。

デルタ則モデルは、この更新強度を現在のトークン埋め込みから学習しますが、メモリ推定に対する信頼度を追跡しません。その結果、各書き込みが蓄積された証拠に応じて適応することができなくなります。この不確実性を明示的に表現するために、再帰的連想記憶を線形ガウス状態空間モデルとして再定式化しました。カルマンフィルタはこのモデルに対して最適な逐次推定器であり、ここに新しいモデルファミリーである「カルマンデルタネットワーク（KDN）」を導入します。

KDN 内では、遷移がメモリ状態とその不確実性の両方を伝播させます。これにより、カルマンゲインが各残差書き込みを、蓄積された証拠と観測の信頼度に基づいて重み付けできるようになります。この定式化の下では、デルタスタイルの更新は、予測共分散に対してトークンごとの等方性近似を代入し、共分散追跡を省略した特別なケースとして現れます。

しかし、正確な追跡には、状態依存のリカッティ再帰を含む密な計算が必要となり、これは GPU 並列化された線形アテンションのスキャン処理には適していません。この課題に対処するため、私たちは 2 つの KDN 近似手法を導入しました。これらはスキャン処理に互換性があります。

対角線 KDN は、オンライン平均場変分推論を通じて各ステップの事後分布を対角ガウス分布族に射影します。一方、等方性 KDN は、ヘッドごとに不確実性のスカラー値 1 つを用いる等方性近似を採用しています。両者の不確実性の再帰計算はモビウスの写像であり、対数並列深さで連想走査を可能にします。750M パラメータと 1.3B パラメータの制御された事前学習において、KDN の各バリアントは最先端の線形アテンションモデルと比較して、パープレキシティと下流タスクの平均精度を一貫して向上させました。