Parallel Broyden methods for efficiently evaluating nonlinear state space models
Abstract
Nonlinear state space models are typically evaluated sequentially, limiting their ability to exploit modern parallel hardware. However, recent work has shown that state space models can be evaluated in parallel by reformulating evaluation as a root-finding problem, which can be solved with a parallel form of Newton's method. However, exact Newton methods require expensive multiplication of dense Jacobian matrices, and quasi-Newton methods based on diagonal approximations, while more scalable, are often slow to converge because they neglect interactions across state dimensions. We introduce a novel quasi-Newton approach based on Broyden’s method, which captures coupling terms with diagonal-plus-low-rank approximations to blocks of the Jacobian. We construct the approximations from trajectory secants, avoiding automatic differentiation while retaining the favorable scaling of diagonal approximations and remaining compatible with parallel scan through rank compression. We extend previous theoretical results to show that convergence degrades monotonically with approximation error and prove finite-step recovery of the exact trajectory. Empirically, we show that our method is advantageous when interactions are approximately low rank or when Jacobian evaluation is prohibitively expensive. It converges in fewer iterations than other quasi-Newton methods, yielding substantially faster run times for simulating low-rank RNNs, parallel MCMC algorithms, and denoising diffusion models.