KOSS: Kalman-Optimal Selective State Spaces for Long-Term Sequence Modeling

TL;DR

KOSS introduces a Kalman-based selective state space model that enhances long-term sequence modeling by providing a theoretically grounded, context-aware, and scalable approach, outperforming existing models in accuracy and stability.

Contribution

KOSS is the first to formulate selective state space modeling as latent uncertainty minimization using Kalman filtering, enabling stable, context-aware, and scalable long-term sequence modeling.

Findings

Achieves over 79% accuracy on a selective copying task with distractors.

Reduces MSE by up to 36.23% across nine long-term forecasting benchmarks.

Demonstrates robustness in real-world radar tracking applications.

Abstract

Recent selective state space models (SSMs), such as Mamba and Mamba-2, have demonstrated strong performance in sequence modeling owing to input-dependent selection mechanisms. However, these mechanisms lack theoretical grounding and cannot support context-aware selection from latent state dynamics. To address these limitations, we propose KOSS, a Kalman-optimal Selective State Space model that formulates selection as latent state uncertainty minimization. Derived from estimation theory, KOSS adopts a continuous-time latent update driven by a Kalman gain that dynamically modulates information propagation based on content and context, enabling a closed-loop, context-aware selectivity mechanism. To ensure stable computation and near-linear scalability, KOSS employs global spectral differentiation for frequency-domain derivative estimation, along with a segment-wise scan for hardware-efficient processing. On a selective copying task with distractors, KOSS achieves over 79\% accuracy while baselines drop below 20\%, demonstrating robust context-aware selection. Furthermore, across nine long-term forecasting benchmarks, KOSS reduces MSE by 2.92--36.23\% and consistently outperforms state-of-the-art models in both accuracy and stability. To assess real-world applicability, a case study on secondary surveillance radar (SSR) tracking confirms KOSS's robustness under irregular intervals and noisy conditions and demonstrates its effectiveness in real-world applications. Finally, supplementary experiments verify Kalman gain convergence and the frequency response of spectral differentiation, providing theoretical support for the proposed closed-loop design.