卡爾曼 Delta 網路：不確定性感知的聯想記憶

摘要

線性注意力在前沿語言模型中越來越常用於高效的長上下文推理和固定記憶體解碼。然而，其固定大小的遞迴記憶需要在每個詞元上進行線上決策：在不知道未來查詢需要哪些資訊的情況下，決定要寫入什麼以及以多大強度覆寫現有關聯。Delta 規則模型從當前詞元嵌入中學習這種強度，但不追蹤記憶估計的信心，使得每次寫入無法適應累積的證據。為了明確表示這種不確定性，我們將遞迴聯想記憶重新表述為線性-高斯狀態空間模型，卡爾曼濾波器是該模型的最佳遞迴估計器，並引入了一個新的模型家族：卡爾曼 Delta 網路（KDNs）。在 KDN 中，轉移過程同時傳播記憶狀態及其不確定性，使卡爾曼增益能夠根據累積證據和觀測可靠性對每次殘差寫入進行加權。在這種表述下，Delta 風格的更新成為一種特殊情況，它用逐詞元的各向同性替代量代替預測協方差，並省略了協方差追蹤。然而，精確追蹤需要一個密集的、依賴狀態的 Riccati 遞迴，這不適合 GPU 並行的線性注意力掃描。為了解決這個問題，我們引入了兩種與掃描兼容的 KDN 近似。對角 KDN 通過線上平均場變分推論將每個單步後驗投影到對角高斯族上，而各向同性 KDN 則使用每個頭具有單一不確定性純量的各向同性近似。它們的不確定性遞迴是 Mobius 映射，使得具有對數並行深度的聯想掃描成為可能。在 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.