Predictability Enables Parallelization of Nonlinear State Space Models

TL;DR

The paper shows that evaluating nonlinear state-space models in parallel hinges on the conditioning of a residual-based merit function, which is governed by the system's predictability. By connecting the Polyak–Łojasiewicz constant of the merit function to the largest Lyapunov exponent, it proves that predictable (negative $\lambda$) dynamics yield well-conditioned landscapes and enable global linear convergence of DEER, with a sublinear $O((\log T)^2)$ time in long sequences. It also characterizes the basin of quadratic convergence and demonstrates a sharp threshold near $\lambda=0$ in experiments across RNNs, Langevin dynamics, and chaotic observers. The results provide a design principle: ensuring predictability makes merit-function based parallelization practical, and guide when to use parallel evaluation versus sequential rollout. Overall, the work offers a theoretical framework and practical guidance for leveraging parallelism in nonlinear state-space modeling by tying dynamical stability to optimization geometry and convergence.

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 (arXiv:2309.12252) or DeepPCR (arXiv:2309.16318) have shown that evaluating a state space model can be recast as solving a parallelizable optimization problem, and sometimes this approach can yield dramatic speed-ups in evaluation time. However, the factors that govern the difficulty of these optimization problems remain unclear, limiting the larger adoption of the technique. In this work, we establish a precise relationship between the dynamics of a nonlinear system and the conditioning of its corresponding optimization formulation. We show that the predictability of a system, defined as the degree to which small perturbations in state influence future behavior, impacts the number of optimization steps required for evaluation. In predictable systems, the state trajectory can be computed in $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 demonstrates that for predictable systems, the optimization problem is always well-conditioned, whereas for unpredictable systems, the conditioning degrades exponentially as a function of the sequence length. We validate our claims through extensive experiments, providing practical guidance on when nonlinear dynamical systems can be efficiently parallelized, and highlighting predictability as a key design principle for parallelizable models.

Figures (8)

• Figure 1: Predictable nonlinear state space models can be recast as well-conditioned, parallelizable optimization problems.
• Figure 2: Threshold phenomenon in DEER convergence based on system predictability. In a family of RNNs, DEER has fast convergence for predictable systems and prohibitively slow convergence for chaotic systems. Left (Theory): We depict Theorem \ref{['theorem:PL-LLE']}, illustrating how the conditioning of the optimization problem degrades as $T$ and the LLE ($\lambda$) increase. Center (Experiment): We vary $\lambda$ across the family of RNNs, and observe a striking concordance in the number of DEER optimization steps empirically needed for convergence with our theoretical characterization of the conditioning of the optimization problem. Right: For 20 seeds, each with 50 different values of $\lambda$, we plot the relationship between $\lambda$ and the number of DEER steps needed for convergence for the sequence length $T=1000$ (gray line in left and center panels). We observe a sharp increase in the number of optimization steps at precisely the transition between predictability and unpredictability.
• Figure 3: DEER converges quickly for Langevin dynamics in a two-well potential. (Left) An illustration of the two-well potential state space in $D=2$. We superimpose a contour plot of the potential on a color scheme showing the spectral norm of the dynamics Jacobian (blue indicates stability, red instability). (Center) A trace plot for the $y$-coordinate. The LLE of the system is $-1.45$. (Right) We observe that this system, which has negative LLE, enjoys sublinear scaling in the sequence length $T$ in the number of DEER iterations needed to converge. We plot the median number of DEER steps to convergence over 20 random seeds.
• Figure 4: Robust relationship in mean field RNN between variance parameter $g$ and LLE of the system. For 20 seeds, we observe a robust and monotonic relationship between the scalar parameter $g$ and the LLE of the resulting mean-field RNN. The plot above is made for $50$ different values of $g$ from $0.5$ to $2.0$ (linearly spaced).
• Figure 5: Chaotic behavior means numerically zero merit function can still be far from sequential trajectory. For $g=1.85$ and $T=1000$, we show the final DEER vs sequential trajectory. The DEER trajectory has merit function \ref{['eq:residual_and_loss']} numerically equal to zero. However: (Left) the mean absolute deviation (MAD) at each time point $t$ between the final DEER iteration $\mathbf{s}_t^{(T)}$ and the sequential rollout $\mathbf{s_t^*}$ grows exponentially. This exponential growth of error is a signature of chaos: compare, for example, with Figure 9.3.5 of strogatz2018nonlinear. The saturation of the error eventually occurs because of the saturating nonlinearity present in the RNN. (Right) We visualize the first coordinate of both the final DEER iteration and the sequential trajectory, showing that while they initially coincide, they diverge around $t=300$.

Theorems & Definitions (33)

• Definition 1: Predictability and Unpredictability
• Proposition 1
• proof
• Theorem 2
• proof
• Theorem 3
• proof
• Theorem 4
• proof
• Theorem 5