Riccati state space models speed up nonlinear sequence processing
Riccati State Space Models: Non-iterative Parallelization for Nonlinear Sequence Modeling
Machine LearningArtificial Intelligence
Summary
Processing long sequences of data quickly is important for many tasks like speech or sensor analysis. Some models update internal states in ways that can be combined efficiently, but when these updates depend on the state itself in complex ways, this efficiency is lost. The authors introduce a new type of model called RiccatiSSM where the nonlinear updates still combine perfectly, allowing fast, exact parallel computation. This speeds up processing without losing prediction quality.
What this means in practice
- •For machine learning engineers: Accelerate nonlinear recurrent models for long sequence tasks by enabling exact parallel state computations without iterative approximations.
- •For signal processing developers: Implement efficient nonlinear state updates in real-time systems where fast and scalable sequence computations are critical.
Authors
Mónika Farsang, Ramin Hasani, Daniela Rus, Radu Grosu
Abstract
State space models (SSMs) achieve efficient sequence processing because their affine state updates are closed under composition and can therefore be evaluated with an associative parallel scan. Nonlinear recurrent models can provide richer, state-dependent dynamics, but generally lose this compositional structure: parallel evaluation then requires iterative methods that repeatedly linearize and scan the recurrence. We ask, what state-dependent nonlinear dynamics can be designed to remain exactly composable? We answer by introducing RiccatiSSM, a nonlinear SSM, in which each state dimension follows an input-conditioned Riccati differential equation. Its quadratic state dependence makes the local Jacobian explicitly state-dependent, while its exact per-step flow under piecewise-constant inputs is a Möbius transformation. Since Möbius maps are closed under composition and compose through $2\times 2$ matrix multiplication, the complete nonlinear state trajectory can be evaluated exactly with a single associative parallel scan, without iterative linearization. We further derive a constrained parameterization that ensures bounded, contractive dynamics, and avoids poles in the fractional-linear state update. Across long-sequence classification, regression, and forecasting tasks, RiccatiSSM achieves competitive predictive performance while reducing runtime by $22{-}33\%$ compared to the nonlinear LrcSSM under matched architectures. These results demonstrate that state-dependent nonlinear dynamics can retain exact composability and be evaluated efficiently within a single parallel scan.