Kalman Delta Networks: Uncertainty-aware Associative Memory

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.

한 줄 요약

KDN은 불확실성을 추적하는 Kalman 필터 기반의 새로운 연관 기억 모델로, 기존 DeltaNet 대비 성능을 개선한다.

핵심 기여도

DeltaNet, Gated DeltaNet, KDA를 Kalman 필터 기반의 선형-가우시안 상태공간 모델로 통합.
Diagonal KDN과 Isotropic KDN이라는 두 가지 GPU-적합한 KDN 근사 모델을 제안.
750M 및 1.3B 파라미터 규모에서 WikiText과 LAMBADA 퍼플렉시티가 기존 모델 대비 낮아짐.
6개 태스크의 평균 제로샷 정확도가 Mamba-3, KDA, GDN-2 대비 높음.

핵심 아이디어

기존 Delta-rule 모델은 토큰 임베딩에서 업데이트 강도를 학습하지만, 저장된 연관성에 대한 신뢰도를 추적하지 못한다. 이는 반복적으로 확인된 정보가 덮어쓸 위험을 초래한다. 본 연구는 Kalman 필터를 사용하여 연관 기억을 선형-가우시안 상태공간 모델로 재구성함으로써, 불확실성을 명시적으로 추적하는 Kalman Delta Networks(KDN)를 제안한다. Kalman 필터는 기억 상태와 그 불확실성을 동시에 전파하며, Kalman 이득을 통해 관측 신뢰도와 누적 증거에 따라 잔차 업데이트를 가중한다. DeltaNet 업데이트는 예측 공분산을 토큰 단위의 등방성 근사로 대체한 특별한 경우로 해석된다.

기술적 접근법

**모델 구조**: KDN은 KDA의 β_t 게이트를 과정 노이즈 ω_t와 관측 노이즈 r_t로 대체.
**Diagonal KDN**: 각 단계의 사후 분포를 대각 가우시안으로 투영하는 온라인 mean-field 변분 추론 사용.
**Isotropic KDN**: 헤드당 단일 불확실성 스칼라로 등방성 근사 적용.
**불확실성 업데이트**: Möbius 맵을 사용하여 로그 병렬 깊이를 유지.
**파라미터 규모**: 750M 및 1.3B 파라미터 규모에서 실험.
**디자인**: 16 또는 18개의 반복 헤드, 128차원, d_model=2048 또는 2304.

주요 결과

**WikiText**: 750M/1.3B 파라미터에서 KDN은 기존 최고 성능 모델 대비 퍼플렉시티가 낮음.
**LAMBADA**: KDN은 Mamba-3, KDA, GDN-2 대비 퍼플렉시티 개선.
**제로샷 정확도**: 6개 태스크 평균 정확도에서 KDN이 기존 모델 대비 +1.6% 개선.
**RULER 점수**: 14셀 평균에서 Diagonal KDN이 750M/1.3B 규모 모두 최고 성능.

의의 및 한계

KDN은 Delta-rule 모델과 Mamba 계열 모델을 통합하는 새로운 관점을 제시하며, 불확실성 추적을 통해 기억 업데이트의 신뢰도를 개선한다. 특히, Diagonal KDN과 Isotropic KDN은 GPU 병렬 처리에 적합한 구조를 가지며, 기존 모델 대비 퍼플렉시티와 제로샷 정확도를 향상시킨다. 그러나 정확한 Kalman 필터는 GPU 병렬 처리에 부적합한 Riccati 재귀를 포함하므로, 이는 근사 모델로만 실용화 가능하다. 또한, 저장된 연관성의 회전을 허용하는 더 복잡한 전이 모델의 확장 가능성은 여전히 개방적이다.

실용적 활용

KDN은 대규모 언어 모델에서 효율적인 장문 추론과 상수 메모리 디코딩이 필요한 상황에 적용 가능하다. 특히, 제로샷 학습 및 다중 태스크 처리에서 뛰어난 성능을 보이므로, 챗봇, 문서 요약, 코드 생성 등 다양한 NLP 분야에 활용될 수 있다.