Predictability Enables Parallelization of Nonlinear State Space Models

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 (arXiv:2309.12252) and DeepPCR (arXiv:2309.16318) 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-Lojasiewicz (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.

💡 Research Summary

The paper investigates when nonlinear state‑space models (SSMs) can be evaluated in parallel rather than sequentially. Traditional recurrent neural networks, physics simulators, and many other dynamical systems require a step‑by‑step computation because each state depends on the previous one. Recent works such as DEER and DeepPCR showed that the sequential recursion can be reformulated as a residual‑based sum‑of‑squares optimization problem. In this formulation the residual vector stacks the violations of the dynamics at every time step, and the loss is the squared norm of this residual. The Jacobian of the residual has a bidiagonal block structure, which allows each Gauss‑Newton update to be solved with a parallel scan in (O(\log T)) time, where (T) is the sequence length.

The bottleneck, however, is the number of optimization iterations required for convergence. This number is governed by the conditioning of the loss, which can be captured by the Polyak‑Łojasiewicz (PL) constant (\mu). The authors first prove that the loss satisfies the PL condition with (\mu = \inf_s \sigma_{\min}^2(J(s))), where (J(s)) is the residual Jacobian and (\sigma_{\min}) denotes its smallest singular value. Consequently, a small (\mu) means a very flat region around the optimum, leading to slow convergence.

To connect (\mu) with intrinsic properties of the dynamical system, the paper introduces predictability measured by the largest Lyapunov exponent (LLE) (\lambda). The LLE quantifies how the norm of a product of Jacobians grows with time: (\lambda = \lim_{T\to\infty} \frac{1}{T}\log|J_T\cdots J_1|). If (\lambda<0) the system is contracting (predictable); if (\lambda>0) it is chaotic (unpredictable). Assuming a mild “burn‑in” condition that bounds finite‑time Jacobian products between (b e^{\lambda k}) and (a e^{\lambda k}) (with constants (a\ge1, b\le1)), the authors derive explicit bounds on the PL constant:

\