カルマン・デルタネットワーク：不確実性を考慮した連想記憶

要旨

線形アテンションは、効率的な長文脈推論と一定メモリでのデコーディングのために、最先端の言語モデルでますます用いられている。しかし、その固定サイズの再帰メモリは、各トークンにおいてオンライン決定を必要とする。すなわち、将来のクエリがどの情報を必要とするかを知る前に、何を書き込み、既存の連想をどれだけ強く上書きするかである。デルタ則モデルはこの強度を現在のトークン埋め込みから学習するが、メモリ推定の信頼度を追跡しないため、各書き込みが蓄積された証拠に適応することを妨げている。この不確実性を明示的に表現するため、我々は再帰的連想記憶を線形ガウス状態空間モデルとして再定式化する。このモデルではカルマンフィルタが最適な再帰的推定量であり、新たなモデル群であるカルマンデルタネットワーク（KDN）を導入する。KDN内では、遷移がメモリ状態とその不確実性の両方を伝播させ、カルマンゲインが各残差書き込みを蓄積された証拠と観測信頼性によって重み付けすることを可能にする。この定式化の下では、デルタ型更新は、予測共分散の代わりにトークンごとの等方性代理を代入し、共分散追跡を省略する特殊ケースとして現れる。しかし、厳密な追跡は、GPU並列の線形アテンションスキャンに適さない密な状態依存リカッチ再帰を伴う。この問題に対処するため、スキャン互換の2つのKDN近似を導入する。対角KDNは、オンライン平均場変分推論を通じて各1ステップ事後分布を対角ガウス族へ射影する。一方、等方KDNは、ヘッドごとに単一の不確実性スカラーを用いる等方近似を用いる。これらの不確実性再帰はメビウス写像であり、対数並列深さの連想スキャンを可能にする。750Mおよび1.3Bパラメータでの制御された事前学習において、KDN変種は最先端の線形アテンションモデルに対してパープレキシティと下流タスクの平均精度を一貫して改善する。

English

Linear attention is increasingly used in frontier language models for efficient long-context inference and constant-memory decoding. Its fixed-size recurrent memory, however, requires an online decision at each token: what to write and how strongly to overwrite existing associations before knowing which information future queries will require. Delta-rule models learn this strength from the current token embedding but do not track confidence in the memory estimate, preventing each write from adapting to accumulated evidence. To represent this uncertainty explicitly, we reformulate recurrent associative memory as a linear--Gaussian state-space model, for which the Kalman filter is the optimal recursive estimator, and introduce a new family of models, Kalman Delta Networks (KDNs). Within KDNs, the transition propagates both the memory state and its uncertainty, allowing the Kalman gain to weight each residual write by accumulated evidence and observation reliability. Under this formulation, Delta-style updates emerge as a special case that substitutes a token-wise isotropic surrogate for predictive covariance and omits covariance tracking. Exact tracking, however, entails a dense, state-dependent Riccati recursion that is poorly suited to GPU-parallel linear-attention scans. To address this issue, we introduce two scan-compatible KDN approximations. Diagonal KDN projects each one-step posterior onto the diagonal Gaussian family through online mean-field variational inference, whereas Isotropic KDN uses an isotropic approximation with a single uncertainty scalar per head. Their uncertainty recurrences are Mobius maps, enabling associative scans with logarithmic parallel depth. Across controlled pretraining at 750M and 1.3B parameters, KDN variants consistently improve perplexity and mean downstream accuracy over state-of-the-art linear-attention models.