Predictability Enables Parallelization of Nonlinear State Space Models
Xavier Gonzalez, Leo Kozachkov, David M. Zoltowski, Kenneth L. Clarkson, Scott W. Linderman
Abstract
The rise of parallel computing hardware has made it increasingly important to understand which nonlinear state space models can be efficiently parallelized. Recent advances like DEER [1] and DeepPCR [2] recast sequential evaluation as a parallelizable optimization problem, sometimes yielding dramatic speedups. However, the factors governing the difficulty of these optimization problems remained unclear, limiting broader adoption. In this work, we establish a precise relationship between a system's dynamics and the conditioning of its corresponding optimization problem, as measured by its Polyak-Łojasiewicz (PL) constant. We show that the predictability of a system, defined as the degree to which small perturbations in state influence future behavior and quantified by the largest Lyapunov exponent (LLE), impacts the number of optimization steps required for evaluation. For predictable systems, the state trajectory can be computed in at worst O((log T ) 2 ) time, where T is the sequence length: a major improvement over the conventional sequential approach. In contrast, chaotic or unpredictable systems exhibit poor conditioning, with the consequence that parallel evaluation converges too slowly to be useful. Importantly, our theoretical analysis shows that predictable systems always yield well-conditioned optimization problems, whereas unpredictable systems lead to severe conditioning degradation. We validate our claims through extensive experiments, providing practical guidance on when nonlinear dynamical systems can be efficiently parallelized. We highlight predictability as a key design principle for parallelizable models.
Recent work addresses this mismatch by reformulating sequential dynamics into parallelizable optimization problems. Notably, the DEER/DeepPCR algorithm [1, 2] evaluates nonlinear state space dynamics by minimizing a residual-based merit function, facilitating efficient parallel computation * Equal contribution.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext e76baf93-3c2a-4e70-abf3-4e6f05faeebfCited by top-tier papers3
- ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language ModelsFederico Danieli, Pau Rodríguez, Miguel Sarabia, Xavier Suau et al.ICLR 2026 · 18 citations
- Parallelizing MCMC Across the Sequence LengthDavid M. Zoltowski, Skyler Wu, Xavier Gonzalez, Leo Kozachkov et al.NeurIPS 2025 · 6 citations
- Why Are Linear RNNs More Parallelizable?William Merrill, Hongjian Jiang, Yanhong Li, Anthony Lin et al.ICML 2026 · 5 citations
Builds on27
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 3,482 citations
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space DualityTri Dao, Albert GuICML 2024 · 1,407 citations
- Resurrecting Recurrent Neural Networks for Long SequencesAntonio Orvieto, Samuel L. Smith, Albert Gu, Anushan Fernando et al.ICML 2023 · 474 citations
- Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth ApproachJonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer et al.NeurIPS 2025 · 431 citations
- Parallelizing Linear Transformers with the Delta Rule over Sequence LengthSonglin Yang, Bailin Wang, Yu Zhang, Yikang Shen et al.NeurIPS 2024 · 412 citations
Related papers
- Towards Scalable and Stable Parallelization of Nonlinear RNNsXavier Gonzalez, Andrew Warrington, Jimmy T. H. Smith, Scott W. LindermanNeurIPS 2024 · 47 citations
- Enhancing Robustness in Deep Reinforcement Learning: A Lyapunov Exponent ApproachRory Young, Nicolas PugeaultNeurIPS 2024 · 5 citations
- DeNOTS: Stable Deep Neural ODEs for Time SeriesIlya Kuleshov, Evgenia Romanenkova, Vladislav Andreevich Zhuzhel, Galina Boeva et al.ICLR 2026 · 2 citations
- Linear Dynamical Systems as a Core Computational PrimitiveShiva KaulNeurIPS 2020 · 8 citations
- Oscillatory State-Space ModelsT. Konstantin Rusch, Daniela RusICLR 2025
