Parallel-in-Time Training of Recurrent Neural Networks for Dynamical Systems Reconstruction

TL;DR

This paper introduces GTF-DEER, a parallel-in-time training method for RNNs that enables efficient learning of complex nonlinear dynamical systems from very long sequences, improving data-driven system reconstruction.

Contribution

It develops a novel parallel training framework combining DEER with Generalized Teacher Forcing, allowing stable learning of nonlinear dynamics from extremely long sequences.

Findings

GTF-DEER effectively trains on sequences longer than 10,000 time steps.

Training on long sequences enhances the accuracy of dynamical systems reconstruction.

The method is robust across different types of nonlinear models.

Abstract

Reconstructing nonlinear dynamical systems (DS) from data (DSR) is a fundamental challenge in science and engineering, but it inherently relies on sequential models. Recent breakthroughs for sequential models have produced algorithms that parallelize computation along sequence length $T$, achieving logarithmic time complexity, $\mathcal{O}(\log T)$. Since sequence lengths have been practically limited due to the linear runtime complexity $\mathcal{O}(T)$ of classical backpropagation through time, this opens new avenues for DSR. This paper studies two prominent classes of parallel-in-time algorithms for this task, both of which leverage parallel associative scans as their core computational primitive. The first class comprises models with linear yet non-autonomous dynamics and a nonlinear readout, such as modern State Space Models (SSMs), while the second consists of general nonlinear models which can be parallelized using the DEER framework. We find that the linear training-time recurrence of the first class of models imposes limitations that often hinder learning of accurate nonlinear dynamics. To address this, we augment DEER with Generalized Teacher Forcing (GTF), a novel variant within the more general nonlinear framework that ensures stable and effective learning of nonlinear dynamics across arbitrary sequence lengths. Using GTF-DEER, we investigate the benefits of training on extremely long sequences ($T>10^4$) for DSR. Our results show that access to such long trajectories significantly improves DSR if the data features long time scales. This work establishes GTF-DEER as a robust tool for data-driven discovery and underscores the largely untapped potential of long-sequence learning in modeling complex DS.