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

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

cloud distributedcomputing parallelcomputing

distributed computing, parallel computing, cloud

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

Florian Hess1,2

Florian Götz1,3

arXiv:2605.12683v1 [cs.LG] 12 May 2026

Daniel Durstewitz1,2,4 1

Dept. of Theoretical Neuroscience, Central Institute of Mental Health, Mannheim, Germany 2 Faculty of Physics and Astronomy, Heidelberg University, Germany 3 Faculty of Mathematics and Computer Science, Heidelberg University, Germany 4 Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University, Germany

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, O(log T ). Since sequence lengths have been practically limited due to the linear runtime complexity 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 > 104 ) 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.

Introduction

Understanding and predicting the behavior of complex nonlinear systems from neural circuits and climate dynamics to fluid flows and ecological networks is a central aim across the natural and engineering sciences [6, 36, 74, 35]. A particularly ambitious goal is to reconstruct the underlying dynamical system (DS) directly from observed time series, a problem known as dynamical systems reconstruction (DSR). Beyond short-term forecasting, DSR requires that the learned model faithfully reproduces the long-term statistical and geometric properties of the true system, such as attractor geometry, power spectra, and Lyapunov exponents. Preprint.

Central to the definition of DS is the flow operator, which provides a recursive rule of how the DS evolves in time [71, 38]. DSR methods that approximate it are inherently recursive [30, 52], such as RNNs, commonly trained by backpropagation through time (BPTT; [81]). While several welldocumented pathologies of BPTT – most notably exploding gradients under chaotic dynamics [52] – can be successfully mitigated by control-theoretic training algorithms such as sparse and generalized teacher forcing (STF/GTF) [52, 33], the computational cost of BPTT remains fundamentally linear in the sequence length T , rendering training on problems with long intrinsic timescales prohibitively expensive. Indeed, DSR applications have historically been confined to modest sequence lengths to keep training tractable [11, 33, 77], leaving open the question of whether learning from substantially longer sequences offers any benefit for reconstruction quality. Recent advances in parallel-in-time sequence modeling offer a path forward, and we examine two common paradigms with respect to their performance in DSR. Linear training-time recurrences with a nonlinear readout, as instantiated by modern SSMs, admit trivial parallelization via linear scans and avoid chaos-induced exploding gradients by construction [52, 55, 82]. By exposing a duality between these linear SSMs and nonlinear RNNs, we show, however, that this convenience comes at a cost: common diagonal parameterization of the linear recurrence used during training imposes structural limitations that often prevent the model from learning accurate nonlinear dynamics. Moreover, its training suffers from exposure bias, degrading autoregressive roll-outs at test-time [4, 62]. The second paradigm – general nonlinear RNNs parallelized via DEER [43] – is in principle better suited to DSR, but naive application fails on chaotic data because the Jacobian products driving DEER’s Newton updates diverge whenever the underlying dynamics exhibits positive Lyapunov exponents [52, 24]. Our main methodological contribution is to resolve this tension by combining DEER with Generalized Teacher Forcing (GTF) [33]. The resulting algorithm, GTF-DEER, inherits the long-sequence scalability of DEER while retaining GTF’s chaos-taming properties: GTF turns the model into a stable DS during training such that GTF-DEER enjoys the average-case O((log T )2 ) scaling established in [24], even when the underlying system is chaotic. Empirically, GTF-DEER delivers speedups of up to 870× over sequential training while matching or improving reconstruction quality. Equipped with GTF-DEER, we then revisit the central empirical question: does training on much longer sequences actually improve DSR? Leveraging the ability to train stably on trajectories of length T > 104 , we find that access to long trajectories allows the DSR model to efficiently capture long time scales if present in the data, yielding substantial gains in long-term statistics, which an equivalent linear SSM based model cannot match. Taken together, our results establish GTF-DEER as a direct replacement of sequential training while highlighting long-sequence training as a largely untapped lever for DSR.

Related work

Dynamical systems reconstruction Data-driven methods for DSR fall into two broad categories: The first approximates the vector field that underlies the observed dynamics, assuming the process is governed by differential equations. Sparse Identification of Nonlinear Dynamics (SINDy) and its variants [14, 13, 15, 51, 17] are particularly popular in the physical sciences and enjoy quick training through least-squares regression, but rely on pre-defined function libraries and struggle with noisy, non-stationary and partially observed empirical data. Neural ODE/PDE methods [16, 37, 2, 39, 1] offer universal approximation capabilities and can be augmented with physical priors [61, 42], but are difficult to train in practice. The second and larger class of models directly approximates the flow operator through black-box universal approximators based on neural networks, including neural operators [41, 47], Koopman operators [48, 56, 12, 53, 3, 80] and hybrids thereof, reservoir computers [59, 58, 75, 60, 57, 23], and RNNs trained through BPTT [76, 77, 11, 33, 63, 9], often accompanied by specialized control techniques that ease optimization and address exploding gradients under chaos [52, 11, 33, 64]. The latter approach in particular achieves SOTA performance on a wide-range of benchmark systems and performs well even on challenging empirical systems [33, 79, 10], which has led to its adoption as a backbone in foundation models for DSR [31]. Efficient sequence modeling A central challenge in sequence modeling is the stable capture and retrieval of long-term dependencies [5, 34, 29, 28, 55, 82]. One practical bottleneck when using autoregressive models such as SSMs or RNNs is the O(T ) runtime of BPTT [81]. Recent developments in the field of sequence models address this issue by parallelizing the inherently 2

sequential operation on parallel accelerators, such as GPUs or TPUs, using parallel associative scans [7, 49, 70]. Parallel scans evaluate linear recurrences, often used as the core block of modern SSMs [55, 70, 28], in O(log T ) time. Furthermore, [43] introduced the DEER framework which parallelizes general nonlinear sequence models by reformulating the forward pass as a fixed-point iteration problem, where each iteration can be solved using a parallel scan. While the worst-case runtime complexity is O(T log T ) [25], the average-case observed in practice for models exhibiting contracting dynamics scales as O((log T )2 ) [24], sparking a renaissance for nonlinear RNNs in long-term sequence modeling. Although DEER demonstrates promising performance in various ML problems, its applicability to the field of DSR remains impractical due to guaranteed worst-case scaling when evaluating chaotic RNNs.

Theoretical background

3.1

Dynamical systems reconstruction with autoregressive models

Given observed time series data X ∈ RTobs ×N originating from some underlying physical process, DSR seeks to learn a generative model that is able to both perform accurate short-term predictions and reproduce the long-term behavior of the observed system. To tackle this task, we consider parameterized state space models of the form zt = Fθ (zt−1 , xt−1 , st ),

x̂t = Gψ (zt ),

(1)

where θ and ψ denote parameter vectors, zt is an M -dimensional state vector, x̂t are N -dimensional predicted observations, xt−1 is an optional teacher signal used during training, and st are Kdimensional optional external inputs. Fθ is a discrete-time universal DS modeling a latent process which is coupled to the observations (data) through Gψ . The aim of training is to learn {θ, ψ} such that the SSM approximates the flow operator of the underlying DS, and after training we have xt ≈ Gψ (Fθ (. . . Fθ (Fθ (z0 , s1 ), s2 ) . . . , st )) =: Gψ (Fθ◦t (z0 , s1:t )).

(2)

In this work, we will investigate two general parameterizations of Eq. (1). 3.1.1

Linear training-time recurrences

In one setting, Fθ is strictly linear, which shifts the burden of capturing nonlinearities in the data to a nonlinear observation function Gψ which feeds back into the system at test time, zt = Azt−1 + U xt−1 + Cst + h, M ×M

x̂t = Bϕ(V zt + b), M ×N

M ×K

(3) M

N ×L

where A ∈ R (often chosen to be diagonal), U ∈ R ,C ∈R ,h∈R ,B ∈R , V ∈ RL×M , b ∈ RL and ϕ is a nonlinear function such as the ReLU(z) = max(0, z). This is essentially the architectural setup of modern SSMs, where linear recurrences are followed by non-linear, point-wise sequence transformations [55, 54, 70, 28], often implemented by MLPs. In our case, the MLP is a simple one-hidden-layer neural network with hidden layer size L. In the following, we will refer to models defined by Eq. (3) as ‘LSSM’. A crucial insight is that while the recurrence in Eq. (3) constitutes a linear, non-autonomous DS during training, the model can produce nonlinear dynamics during evaluation by replacing the teacher signals xt−1 with predictions x̂t−1 . Indeed, during trajectory generation after training, the recursion of Eq. (3) turns into zt = Azt−1 + W ϕ(V zt−1 + b) + Cst + h

(4)

where U B =: W ∈ RM ×L with rank(W ) ≤ min(N, M, L). This allows SSMs to exhibit inherently nonlinear traits such as chaotic dynamics and multistability during autoregressive generation. While the major feature of Eq. (3) is the fact that the forward pass can be calculated in logarithmic time using parallel scan, another insight is that linear recurrences do not suffer from chaos-induced exploding gradients [52], and hence remedy associated training instabilities by design [52, 55, 82]; see Appx. A for details. Training of LSSMs makes use of a variant of teacher forcing originating from the sequence modeling field, where ground-truth data from the previous time step xt−1 is fed through a dedicated input layer (U ) [4, 26]; a generic training algorithm is provided in Alg. 1. This type of training suffers from exposure bias, where the absence of forcing signals during evaluation (Eq. (4)) significantly degrades autoregressive roll-outs [62]. While methods such as scheduled sampling address this issue [4, 77], they destroy linearity and therefore efficient parallelization of the recurrence in Eq. (3) during training; see Appx. C for details. 3

3.1.2

Non-linear training-time recurrences

To compare training algorithms, we will consider a second parameterization of Eq. (1), which uses a nonlinear RNN as the latent model, coupled to the observations through a simple linear mapping: zt = Azt−1 + W ϕ(V zt−1 + b) + Cst + h,

x̂t = Bzt .

(5)

With ϕ(·) = ReLU(·), the latent model is known as a shallow piecewise-linear RNN (shPLRNN; [33]), an established RNN architecture designed for DSR. Indeed, up to the constraint of a low-rank connectivity matrix, which can be manually enforced in the latter parameterization, the recurrences of Eqs. (4) and (5) are mathematically equivalent. This is important as it allows us to directly compare training algorithms and parameterizations, unconfounded by any architectural differences at test time, studying optimized models with parameters θRNN ∪ ψRNN = {A, W , V , b, C, h, B} and θLSSM ∪ ψLSSM = {A, W , V , b, C, h, B}. All results established below hold in the presence of external inputs s, but we will drop them from notation for brevity. Due to the nonlinearity of the latent model in Eq. (5), the forward pass is evaluated sequentially and cannot be naively parallelized. Hence, a common method to train for DSR is to generate trajectories of length T and then compare them to the data x1:T . To avoid training instabilities caused by, for example, exploding gradients, training is accompanied by control-theoretic forcing methods such as Generalized Teacher Forcing (GTF; [19, 33]). GTF alters the forward pass of the model during training by linearly interpolating between the latent state and a teacher signal before application of the sequence model: zt = Fθ ◦ δα (zt−1 , z t−1 ) = Fθ ((1 − α)zt−1 + αz t−1 ) = Fθ (z̃t−1 ),

(6)

where 0 ≤ α ≤ 1 is the forcing strength and z denotes the teacher signal which is generally computed by inversion of the observation model or output layer, i.e. z t = G−1 ψ (xt ), and z̃ is the forced state. During training, this leads to decomposition of the Jacobian as JF ◦δ (zt−1 ) = (1 − α)JF (z̃t−1 ),

(7)

where JF (z) := ∂F∂z(z) . The forcing parameter α ∈ [0, 1] controls the norm of the Jacobians during BPTT and can be optimally and adaptively adjusted in training to mitigate exploding gradients [33]. For linear observation models (Eq. (5)) and the common case where the RNN has more units than there are observed dynamical variables (M > N ), [64] introduced a correction to GTF (6) which includes the row-space projector of B in the forcing equation: δα,B (zt , xt ) = (I − αB + B)zt + αB + xt ,

(8)

where B denotes the pseudo-inverse of B, such that Eq. (7) becomes JF ◦δα,B (zt−1 ) = JF (z̃t−1 )Pα ,

(9)

where Pα := I − αB + B. While this correction does not preserve the mitigation of exploding gradients in general, it dampens Jacobian singular values along directions corresponding to expansion in observation space. Indeed, we find empirically that Jacobian damping still holds in practice in this case, which we attribute to mixing of directions through the Jacobian product. Note that for M ≤ N and assuming full column-rank, B + B = I and vanilla GTF is recovered. For a detailed overview of the training protocol under GTF, see Alg. 2. 3.1.3

Parallel-in-time algorithms for nonlinear sequence models

Training nonlinear flow operators on long sequences sequentially is prohibitive due to the linear O(T ) scaling of BPTT in sequence length. To make training feasible, we will make use of the recently proposed DEER algorithm [43, 25]. Given an initial condition z0 , let z1:T be a series of candidate latent states and define the residual vector r(z1:T ) := [z1 − F (z0 ), . . . , zT − F (zT −1 )] ∈ RM T . ∗ ∗ ∗ The fixed point z1:T = F (z0:T −1 ) is the only solution with zero residual, i.e. r(z1:T ) = 0. Conventionally, the forward pass is generated sequentially by iterating F starting from z0 . DEER turns the roll-out of the map F into a root-finding problem for the residual, which is solved using (0) Newton’s method. Starting from an initial guess z1:T , where the superscript indicates the current ∗ Newton iteration, the true trace z1:T is approximated by iteratively solving the update equation (i+1) (i) (i) −1 z1:T = z1:T − (Jr (z1:T )) r z1:T , (10) 4

where Jr := ∂z∂r . Inverting the Jacobian matrix explicitly is infeasible for large M T . Instead, it is 1:T more practical to multiply both sides of the equation by the Jacobian and exploit its block-bidiagonal (i+1) (i+1) (i) structure (cf. Eq. (40) in Appx. D) to obtain a recursive formula for ∆zt := zt − zt : h i (i+1) (i) (i+1) (i) ∆zt = JF zt−1 ∆zt−1 − r(z1:T ) (11) t

Eq. (11) is linear in ∆z and hence each Newton iteration can be solved in O(log T ) time by a parallel associative scan. Due to its usage of the full Jacobian and hence full matrix products, increasing the dimensionality of the state of the sequence model can become a bottleneck. To address this, [25] suggested the use of quasi-Newton methods which replace the full Jacobian in Eq. (11) by its diagonal diag(JF ) ∈ RM , reducing the work W of each Newton iteration from Wfull = O(M 3 T ) to Wdiag = O(M T ). While DEER and its variants can be directly used to train sequence models on data from nonlinear DS, the methods suffer from Jacobian divergence when tasked to generate orbits from systems with (on average) unstable dynamics, i.e. where ∥JF (zt )∥ > 1 for most t [24]. This divergence is inevitable when training on chaotic systems [52]. Since most if not all natural complex DS are chaotic [69, 21, 27, 73, 22], this needs to be amended.

Methods

4.1

DSR from long sequences

To enable stable training from long sequences of chaotic dynamics, we propose to combine GTF and DEER, which we coin GTF-DEER. By simply replacing the RNN map Fθ with the teacher-forced variant (6), the residual becomes the vector with entries [rα (z1:T , x1:T )]t = zt − Fθ (δα,B (zt−1 , xt−1 ))

(12)

for t = 2 . . . T and [rα (z1:T , x1:T )]1 = z1 − Fθ (z0 ) since z0 is given. Hence the linear recursion uses the decomposed Jacobians (9) h i h i (i+1) (i) (i+1) (i) ∆zt = JF (z̃t−1 )Pα ∆zt−1 − rα z1:T , x1:T (13) t

Similar to sequential GTF, Jacobian divergence caused by chaotic dynamics is tamed by Pα which for M ≤ N strictly turns a DS with diverging state space directions into a contracting one, improving DEER convergence [24]. To demonstrate this formally, we need the following definitions. Definition 1. The largest Lyapunov exponent (LLE) of an orbit of the iterated map F starting in z0 is λ(z0 ) = limT →∞ T1 log ∥JT JT −1 . . . J1 ∥2 ,

(14)

t−1 ) where Jt := J (zt−1 ) := ∂F∂z(zt−1 .

We say the iterated map F is divergent in z0 if λ(z0 ) > 0. When λ(z0 ) < 0, F is contracting. Proposition 1. Let an SSM be given by a latent model F : RM → RM and a linear observation model G : RM → RN , z 7→ Bz where M ≤ N and B is assumed to have full rank. Let the supremal Jacobian norm satisfy σ̃max := sup ∥J (z)∥2 = σmax (J (z)) | z ∈ RM > 1. Choose −1 the GTF forcing strength α ∈ (α∗ , 1], where α∗ := 1 − σ̃max [33], and define the effective contraction rate ρ := (1 − α) σ̃max < 1. (15) Then the forced system F ◦ δα,B is globally contracting with rate ρ. Its LLE satisfies λGTF ≤ log ρ < 0.

(16)

Proof. See Appx. B. Proposition 1 says that given a suitable α, training a nonlinear sequence model using GTF-DEER leads to guaranteed convergence of the forward pass regardless of the dynamics that underlie the data, enabling stable and efficient parallel-in-time training for DSR. We validate this theoretical finding in Sect. 5.1 and investigate the effect of M > N on GTF-DEER convergence empirically. 5

Moreover, since we directly want to match forced model trajectories to targets given by the data, (0) we can initialize the Newton iterations of DEER using the forcing signals, i.e. z1:T = z 1:T , which speeds up convergence over the typical initialization by zeros [43, 25]. We provide the training routine using GTF-DEER in Alg. 3. 4.2

Capturing long-term dependencies

Regularization To effectively capture long-term dependencies, sequence models must maintain stable flow of information over extended temporal horizons [5, 34, 29, 82]. Crucially, for maintaining connections among states across time we must not only avoid chaotic divergence, but also prevent error signals propagated backward from decaying too fast. In this work we therefore make use of Manifold Attractor Regularization (MAR; [67]) to ameliorate diminishing gradients by equipping the sequence model with a latent subspace that adaptively captures slow timescales in the data. In Appx. E, we review the MAR contribution to the loss function and show how this regularization leads to stable error propagation by imposing a block-diagonal identity structure in the model Jacobian JF . Latent state warm-up To provide the RNN with sufficient dynamical context during training, we employ latent state warm-up [78, 66]. During training we sample sequences of length T , where the first Tw steps are used for warm-up. The warm-up consists of forcing the sequence model with data x1:Tw , either by feeding them through U for the LSSM (Eq. (3)) or by GTF in the case of the nonlinear RNN (Eqs. (6) and (8)) with α = 1. This yields a latent state zTw which holds a compressed history of the warm-up signal, which is then used to predict the remaining T − Tw time steps that contribute to the loss function given by the mean-squared-error MSE(xTw +1:T , x̂Tw +1:T ). Crucially, the forward pass of the entire sequence of length T including the warm-up can be computed by a single call to GTF-DEER, enabling long contexts during training (cf. Alg. 3). Performance measures To evaluate DSR performance we use an established long-term measure which computes the KL divergence between multivariate state distributions in observation space, Dstsp [40, 33]. For partially observed systems, we first perform a delay embedding [72, 65] of ground-truth and generated trajectories and compute Dstsp on these embeddings. We will make DE delay embeddings explicit in notation by writing the measure as Dstsp . We also report the 128-step root-mean-square-error (RMSE) as a measure of short- to medium-length prediction accuracy. For mathematical and experimental details on evaluation measures, see Appx. G.

Results

We will first evaluate the computational efficiency of GTF-DEER by performing runtime analyses. We will then investigate the benefit of long-sequence training for DSR through a set of ablation experiments and comparisons to state-of-the-art sequence models. Specifically, we 1) ask whether long sequences improve DSR and 2) reveal and explain the shortcomings of SSMs for DSR. 5.1

Runtime convergence and performance of GTF-DEER

We evaluate the efficacy of GTF-DEER by performing runtime comparisons to the sequential baseline as well as performing ablations on several parameters of the GTF-DEER algorithm. Fig. 1A shows median runtimes for combined forward+backward passes for both sequential and parallel evaluation using GTF-DEER while training to reconstruct dynamics of the chaotic Lorenz-63 attractor (N = 3). While the sequential approach exhibits linear scaling throughout, GTF-DEER enables sublinear scaling with sequence length, leading to a speedup of up to 870× over sequential evaluation (see Fig. 1A setting M = 4, T = 32,768). However, for large problems, i.e. large B × T × M , GPU memory bandwidth quickly saturates leading to near linear to constant scaling, a common bottleneck in parallel scan implementations on GPU devices [50, 28]. Nevertheless, for all settings using GTF-DEER outperforms naive sequential training in terms of raw runtime. Fig. 1B shows the convergence of the DEER forward pass under increasing values of α ∈ [0, 1]. For GTF-DEER, α does not only control Jacobian norms during the backward pass, addressing the exploding gradient problem in face of chaos [52, 33], but also controls convergence of the linear recurrence in Eq. (13). If α is too small, the Jacobians will diverge in chaotic settings, and hence 6

Figure 1: A: GTF-DEER scales favorably in sequence length, but the GPU quickly saturates for larger problems. All analyses were performed on an NVIDIA RTX 6000 Blackwell (96GB) GPU. Note the logarithmic scaling of the y-axis. B: The forcing parameter α controls Jacobian norms and hence reduces the number of Newton iterations needed for convergence of the GTF-DEER forward pass. (0) Initializing the model trajectory using estimated latents through z1:T = B + x1:T (‘P-inv’) improves GTF-DEER convergence over naive initializations for N = M ; however, the effect vanishes for M > N.

the forward pass becomes numerically unstable such that a large number of iterations is needed to converge. Increasing α pushes Jacobian norms below 1 such that the forced model constitutes a stable DS, leading to convergence of GTF-DEER [24]. In fact, for sufficiently large α, the forward pass even converges optimally in just 2 Newton iterations for M ≤ N !1 For M > N , only a subspace of the latent space of the RNN is forced such that GTF-DEER needs more iterations to converge to the true latent dynamics. Finally, the different curves in Fig. 1B show the benefit of readily available forcing targets: Initializing the latent trajectory with estimated forcing signals by pseudo-inversion of the observation model drastically reduces the number of Newton iterations for a given α and M ≤ N . The benefit vanishes when moving to the overdetermined case of M > N . This is expected, as the RNN dynamics will veer off the estimated forcing signals as GTF is only applied to the zero-error manifold [64]. For more details on the experimental setup underlying Fig. 1 we refer to Appx. H.1. 5.2

GTF-DEER ablations

An important observation is that using diagonal approximation of the Jacobians in Eq. (13) for the forward pass (quasi-DEER, [25]) affects gradients during the backward pass, as the same diagonal Jacobians will be used (see Appx. D for more details). In Fig. 2, we test how Jacobian diagonalization affects GTF-DEER convergence and reconstruction quality. We show loss curves and Newton iterations of shPLRNNs (M = 5, L = 50) trained on the Lorenz-63 system under two settings: In the ‘FO’ setting, the shPLRNN is tasked to reconstruct the dynamics from the fully observed system (N = 3), while in the more challenging ‘PO’ setting only the x-component is observed (N = 1). We then consider training with diagonalized Jacobians (‘quasi’) as well as full Jacobians, i.e. standard GTF-DEER, while keeping all other training-related hyperparameters the same. We find that while full-Jacobian training handled both settings with minimal Newton iterations and provided good DE reconstruction quality (‘PO’: Dstsp = (7.5 ± 6.2) · 10−3 , ‘FO’: Dstsp = (8.7 ± 3.8) · 10−3 ), diagonal approximations needed ≈ 100× more Newton iterations for the forward pass to converge, failed to yield convergent loss curves in the more challenging ‘PO’ setting, and fell behind in reconstruction DE performance across both settings (‘PO + quasi’: Dstsp = (7 ± 7) where 14/20 runs diverged, ‘FO −2 + quasi’: Dstsp = (4.4 ± 3.5) · 10 ). This result highlights the role of the Jacobians during the backward pass: Since the diagonal approximations of the Jacobians are used to compute the gradient, they lose information on temporal mixing effects. This prevents the DSR model from learning a faithful embedding of the partially observed dynamics, which is also reflected in the erratic loss curve 1 That 2 instead of just 1 Newton iterations are required is due to the fact that one additional iteration is needed to verify

that the previous iteration has indeed converged.

in Fig 2A. Thus, while full Jacobians in theory increase the time demand of each Newton step, this is more than compensated for by the manifold fewer iterations needed and the more stable training.

Figure 2: Training dynamics and reconstruction quality of a shPLRNN (M = 5, L = 50) on the Lorenz-63 system under fully observed (FO, N = 3) and partially observed (PO, N = 1) conditions. A: Training loss curves. B: Newton iterations per training step. C: State-space divergence Dstsp (lower is better); red crosses mark runs where trajectory divergence during generation produced NaNs. D: Wall-clock runtime per training step. All reported values are median ± MAD across 20 runs.

5.3

Reconstructing dynamical systems with long-term dependencies

To demonstrate the advantage of processing long sequences, we generated data from 1) a Lorenz-96 system (N = 6) [45] augmented with a sinusoidal forcing term of period 15,000 discrete time steps, thereby introducing an explicit long time scale into the dynamics, and 2) a bursting neuron biophysical model [21] which exhibits long inter-burst-intervals of > 104 time steps (see Appx. F for further details on data generation). We trained shPLRNNs with M = 10 for the Lorenz96 system and M = 6 for the bursting neuron, where both settings share L = 128, to reconstruct the underlying dynamics while systematically varying the training sequence length. To control for the fact that longer sequences expose the model to more data per gradient update, we held the product B · T = 215 = 32,768 fixed. For full experimental details see Appx. H. DE Figure 3A reports the long-term reconstruction measure Dstsp as a function of sequence length. For both systems, the measure improves with increasing sequence length, confirming that the model benefits substantially from training specifically on longer contexts, while keeping the total amount of training data constant. The RMSE(128), however, already saturates for shorter sequence lengths, as expected by design of this short-term measure, with mean ± SEM over all sequence lengths and runs given by RMSE(128) = 0.27 ± 0.02 for the Lorenz-96 and 0.047 ± 0.003 for the bursting neuron system, respectively (see also Fig. A3). Training at such sequence lengths becomes tractable only through the favorable scaling of GTF-DEER: whereas training on sequences of length 32,768 required on average ≈ 10 minutes for the Lorenz-96 setting on a single NVIDIA RTX 6000 Blackwell GPU, sequential GTF [33] would require approximately 34 hours.

5.4

Linear vs. nonlinear training-time recurrences

Finally, we compare the DSR performance of LSSMs (Eq. (3) to that of nonlinear RNNs (Eq. (5)), trained by GTF-DEER, on the forced Lorenz-96 system. For comparability, we fix common hyperparameters, i.e. both models use M = 10, L = 128, T = 81,920 and B = 1. To investigate the effect of the low-rank constraint in the LSSM (Eq. (4)), we also train a low-rank version of the shPLRNN, where W := WL · WR with WL ∈ RM ×r and WR ∈ Rr×L . To account for MAR applied to the shPLRNN, we add a similar regularization term to the training of the LSSM, see Appx. E for details. Figure 3B summarizes the results of this ablation quantitatively through DE Dstsp , while qualitative example reconstructions are provided in Fig. A2. The LSSMs as well as the shPLRNN with their rank limited to the number of observed variables, r = N = 6, all fall behind in reconstructing the limiting dynamics of the forced Lorenz-96 system, failing to capture the 8

DE Figure 3: A: Long-term measure Dstsp as a function of sequence length for shPLRNNs trained on the DE forced Lorenz-96 and bursting neuron system. B: Dstsp evaluated on the forced Lorenz-96 system for different models (see Fig. A2 for qualitative comparison). Note that even Mamba-2 cannot match the performance of the r = 7 shPLRNN trained with GTF-DEER even when allowed many more parameters. C: Example reconstructions for long-sequence GTF-DEER (top) vs. a model trained only on standard-length sequences (bottom) for the bursting neuron (left) and forced Lorenz-96 (right).

influence of the sinusoidal forcing pattern. Simply increasing the rank by just 1 (r = 7) equips the shPLRNN with the expressivity needed to capture the latent forcing dynamics, hence significantly improving reconstruction quality. We also trained a Mamba-2-based model [18] with a comparable parameter count on this task (see Appx. H for details). Even though Mamba-2 avoids the low-rank constraint through its selection mechanism, where recurrence parameters are parameterized by the data, A → A(xt ) and U → U (xt ), it performs much worse than the r = 7 shPLRNN trained by GTF-DEER and requires more trainable parameters (see Mamba-2-7k vs. Mamba-2-17k), while at the same time exhibiting much higher variance in training outcome. The latter observation highlights another important feature of the training algorithm: GTF-DEER improves autoregressive roll-outs at test time, since GTF reduces exposure bias by keeping α, i.e. the forcing strength, minimal. Applying similar strategies to address exposure bias to linear training-time recurrences destroys their amenability to efficient parallelization (see Appx. C).

Discussion

We introduced GTF-DEER, a parallel-in-time training algorithm that combines the long-sequence scalability of DEER with the chaos-taming properties of Generalized Teacher Forcing (GTF). Through forcing, GTF-DEER turns the RNN during training into a stable DS, improving empirical runtime scaling from the worst-case of O(T log T ) under unstable dynamics to the average case of O((log T )2 ) for contractive RNNs [24]. Leveraging fast training from arbitrary sequence lengths, we provided what is to our knowledge the first systematic evidence that DSR benefits from training on such long sequences when the data carry slow timescales. Furthermore, we showed that feeding predictions of linear SSMs (LSSMs) with nonlinear read-out back into the recurrence during autoregressive generation imposes a low-rank constraint that hinders the LSSM from inferring unobserved dynamical variables in the data. Even when SSMs ameliorate this problem by introducing gating or selection mechanisms [28, 18], DSR quality suffers due to the inherent exposure bias of conventional teacher forcing [4, 62], which degrades autoregressive roll-outs during testing. For modern SSMs, this problem cannot be resolved without trading the linear recurrence for a (locally) nonlinear one, such that parallelization by parallel scan and similar algorithms breaks down (cf. Appx. C, [7, 29, 70, 28, 18]). 9

Limitations GTF-DEER inherits DEER’s cubic work in the latent dimension M per Newtoniteration, Wfull = O(M 3 T ), and the diagonal (quasi-DEER) approximation that would alleviate this degrades gradients in partially observed settings substantially (Sect. 5.2). This sets practical limits to the latent dimension M (see Fig. 1A). Finally, Proposition 1 strictly holds for M ≤ N , while for M > N we only find empirical evidence that GTF-DEER’s convergence also holds for the settings and problems considered in this work.

Acknowledgments and Disclosure of Funding This work was supported by individual grants Du 354/15-1 (project no. 502196519) and Du 354/18-1 (project no. 567025973) from the German Research Foundation (DFG), by the German Ministry for Research, Astronautics, and Technology (BMFTR) through NAILIt (“Neuro-Inspired AI for Learning & Inference in Non-Stationary Environments”), grant number 01GQ2509A.

References [1] Julius Aka, Johannes Brunnemann, Jörg Eiden, Arne Speerforck, and Lars Mikelsons. Balanced neural ODEs: nonlinear model order reduction and koopman operator approximations. In The Thirteenth International Conference on Learning Representations, 2025. [2] Victor M. Martinez Alvarez, Rareş Roşca, and Cristian G. Fălcuţescu. Dynode: Neural ordinary differential equations for dynamics modeling in continuous control. arXiv preprint arXiv:2009.04278, 2020. [3] Omri Azencot, N. Benjamin Erichson, Vanessa Lin, and Michael W. Mahoney. Forecasting Sequential Data using Consistent Koopman Autoencoders. In Proceedings of the 37th International Conference on Machine Learning, 2020. [4] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015. [5] Y. Bengio, P. Simard, and P. Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157–166, 1994. [6] Upinder S. Bhalla and Ravi Iyengar. Emergent properties of networks of biological signaling pathways. Science, 283(5400):381–387, 1999. [7] Guy E. Blelloch. Prefix sums and their applications. 1990. [8] Jonah Botvinick-Greenhouse. Invariant measures for data-driven dynamical system identification: Analysis and application. arXiv preprint arXiv:2502.05204, 2025. [9] Manuel Brenner, Christoph Jürgen Hemmer, Zahra Monfared, and Daniel Durstewitz. Almostlinear rnns yield highly interpretable symbolic codes in dynamical systems reconstruction. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 36829–36868. Curran Associates, Inc., 2024. [10] Manuel Brenner, Florian Hess, Georgia Koppe, and Daniel Durstewitz. Integrating Multimodal Data for Joint Generative Modeling of Complex Dynamics. In Proceedings of the 41st International Conference on Machine Learning, pages 4482–4516. PMLR, July 2024. ISSN: 2640-3498. [11] Manuel Brenner, Florian Hess, Jonas M. Mikhaeil, Leonard F. Bereska, Zahra Monfared, PoChen Kuo, and Daniel Durstewitz. Tractable Dendritic RNNs for Reconstructing Nonlinear Dynamical Systems. In Proceedings of the 39th International Conference on Machine Learning, pages 2292–2320. PMLR, June 2022. ISSN: 2640-3498. [12] Steven L. Brunton, Marko Budišić, Eurika Kaiser, and J. Nathan Kutz. Modern koopman theory for dynamical systems. SIAM Review, 64(2):229–340, 2022. 10

[13] Steven L. Brunton and J. Nathan Kutz. Data-driven science and engineering: Machine learning, dynamical systems, and control. Cambridge University Press, 2019. [14] Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences USA, 113(15):3932–3937, 2016. [15] Kathleen Champion, Bethany Lusch, J. Nathan Kutz, and Steven L. Brunton. Data-driven discovery of coordinates and governing equations. Proceedings of the National Academy of Sciences USA, 116(45):22445–22451, 2019. [16] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems 31, 2018. [17] Alexandre Cortiella, Kwang-Chun Park, and Alireza Doostan. Sparse identification of nonlinear dynamical systems via reweighted l1-regularized least squares. Computer Methods in Applied Mechanics and Engineering, 376:113620, April 2021. [18] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 10041–10071. PMLR, 21–27 Jul 2024. [19] Kenji Doya. Bifurcations in the learning of recurrent neural networks. In Proceedings of the 1992 IEEE International Symposium on Circuits and Systems, 1992. [20] Daniel Durstewitz. Implications of synaptic biophysics for recurrent network dynamics and active memory. Neural Networks, 22(8):1189–1200, 2009. [21] Daniel Durstewitz and Thomas Gabriel. Dynamical Basis of Irregular Spiking in NMDA-Driven Prefrontal Cortex Neurons. Cerebral Cortex, 17(4):894–908, April 2007. [22] Richard J. Field, Endre Koros, and Richard M. Noyes. Oscillations in chemical systems. ii. thorough analysis of temporal oscillation in the bromate-cerium-malonic acid system. Journal of the American Chemical Society, 94(25):8649–8664, 1972. [23] Daniel J. Gauthier, Erik Bollt, Aaron Griffith, and Wendson A. S. Barbosa. Next generation reservoir computing. Nature Communications, 12(1):5564, September 2021. Number: 1 Publisher: Nature Publishing Group. [24] Xavier Gonzalez, Leo Kozachkov, David Zoltowski, Kenneth Clarkson, and Scott Linderman. Predictability enables parallelization of nonlinear state space models. In Annual Conference on Neural Information Processing Systems, 2025. [25] Xavier Gonzalez, Andrew Warrington, Jimmy T. Smith, and Scott W. Linderman. Towards scalable and stable parallelization of nonlinear rnns. Advances in Neural Information Processing Systems, 37:5817–5849, 2024. [26] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http://www.deeplearningbook.org. [27] R. B. Govindan, K. Narayanan, and M. S. Gopinathan. On the evidence of deterministic chaos in ecg: Surrogate and predictability analysis. Chaos: An Interdisciplinary Journal of Nonlinear Science, 8(2):495–502, 1998. [28] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First conference on language modeling, 2024. [29] Albert Gu, Karan Goel, and Christopher Ré. Efficiently Modeling Long Sequences with Structured State Spaces, August 2022. arXiv:2111.00396 [cs]. 11

[30] Niclas Alexander Göring, Florian Hess, Manuel Brenner, Zahra Monfared, and Daniel Durstewitz. Out-of-Domain Generalization in Dynamical Systems Reconstruction. In Proceedings of the 41st International Conference on Machine Learning, pages 16071–16114. PMLR, July 2024. ISSN: 2640-3498. [31] Christoph Jürgen Hemmer and Daniel Durstewitz. True zero-shot inference of dynamical systems preserving long-term statistics. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [32] John R. Hershey and Peder A. Olsen. Approximating the kullback leibler divergence between gaussian mixture models. 2007 IEEE International Conference on Acoustics, Speech and Signal Processing - ICASSP ’07, 4:IV–317–IV–320, 2007. [33] Florian Hess, Zahra Monfared, Manuel Brenner, and Daniel Durstewitz. Generalized Teacher Forcing for Learning Chaotic Dynamics. In Proceedings of the 40th International Conference on Machine Learning, pages 13017–13049. PMLR, July 2023. ISSN: 2640-3498. [34] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8):1735–1780, nov 1997. [35] Eugene M. Izhikevich. Dynamical systems in neuroscience: the geometry of excitability and bursting. Computational neuroscience. MIT Press, Cambridge, Mass, 2007. OCLC: ocm65400606. [36] Eugenia Kalnay. Atmospheric Modeling, Data Assimilation and Predictability. Cambridge University Press, 2003. [37] Daniel Karlsson and Olle Svanström. Modelling Dynamical Systems Using Neural Ordinary Differential Equations, 2019. [38] Anatole Katok, A. B. Katok, and Boris Hasselblatt. Introduction to the Modern Theory of Dynamical Systems. Cambridge University Press, 1995. Google-Books-ID: 9nL7ZX8Djp4C. [39] Joon-Hyuk Ko, Hankyul Koh, Nojun Park, and Wonho Jhe. Homotopy-based training of neuralodes for accurate dynamics discovery. Advances in Neural Information Processing Systems, 36:64725–64752, 2023. [40] Georgia Koppe, Hazem Toutounji, Peter Kirsch, Stefanie Lis, and Daniel Durstewitz. Identifying nonlinear dynamical systems via generative recurrent neural networks with applications to fMRI. PLOS Computational Biology, 15(8):e1007263, 2019. [41] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar, et al. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2020. [42] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. [43] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. In International Conference on Learning Representations, 2024. [44] Edward N. Lorenz. Deterministic nonperiodic flow. Journal of atmospheric sciences, 20(2):130– 141, 1963. [45] Edward N. Lorenz. Predictability: A problem partly solved. In Proc. Seminar on predictability, volume 1, 1996. [46] Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017. 12

[47] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218–229, March 2021. Number: 3 Publisher: Nature Publishing Group. [48] Bethany Lusch, J. Nathan Kutz, and Steven L. Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nat Commun, 9(1):4950, December 2018. arXiv: 1712.09707. [49] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations, 2018. [50] Duane Merrill and Michael Garland. Single-pass parallel prefix scan with decoupled look-back. NVIDIA, Tech. Rep. NVR-2016-002, 2016. [51] Daniel A. Messenger and David M. Bortz. Weak SINDy: Galerkin-Based Data-Driven Model Selection. Multiscale Modeling & Simulation, 19(3):1474–1497, January 2021. Publisher: Society for Industrial and Applied Mathematics. [52] Jonas Mikhaeil, Zahra Monfared, and Daniel Durstewitz. On the difficulty of learning chaotic dynamics with RNNs. Advances in Neural Information Processing Systems, 35:11297–11312, December 2022. [53] Ilan Naiman and Omri Azencot. A Koopman Approach to Understanding Sequence Neural Models. arXiv:2102.07824 [cs, math], October 2021. arXiv: 2102.07824. [54] Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel L Smith. Universality of linear recurrences followed by non-linear projections: Finite-width guarantees and benefits of complex eigenvalues. In International Conference on Machine Learning, pages 38837–38863. PMLR, 2024. [55] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 26670–26698. PMLR, 2023. [56] Samuel E. Otto and Clarence W. Rowley. Linearly recurrent autoencoder networks for learning dynamics. SIAM Journal on Applied Dynamical Systems, 18(1):558–593, 2019. [57] Dhruvit Patel and Edward Ott. Using machine learning to anticipate tipping points and extrapolate to post-tipping dynamics of non-stationary dynamical systems. Chaos (Woodbury, N.Y.), 33(2):023143, February 2023. [58] Jaideep Pathak, Brian Hunt, Michelle Girvan, Zhixin Lu, and Edward Ott. Model-Free Prediction of Large Spatiotemporally Chaotic Systems from Data: A Reservoir Computing Approach. Phys. Rev. Lett., 120(2):024102, 2018. [59] Jaideep Pathak, Zhixin Lu, Brian R. Hunt, Michelle Girvan, and Edward Ott. Using Machine Learning to Replicate Chaotic Attractors and Calculate Lyapunov Exponents from Data. Chaos: An Interdisciplinary Journal of Nonlinear Science, 27(12):121102, December 2017. arXiv: 1710.07313. [60] Jason A. Platt, Stephen G. Penny, Timothy A. Smith, Tse-Chun Chen, and Henry D. I. Abarbanel. Constraining chaos: Enforcing dynamical invariants in the training of reservoir computers. Chaos: An Interdisciplinary Journal of Nonlinear Science, 33(10), 2023. [61] M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686–707, February 2019. [62] Marc’Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. In 4th International Conference on Learning Representations, ICLR 2016, 2016. 13

[63] T Konstantin Rusch, Siddhartha Mishra, N Benjamin Erichson, and Michael W Mahoney. Long expressive memory for sequence modeling. In International Conference on Learning Representations, 2022. [64] A Erdem Sağtekin, Colin Bredenberg, and Cristina Savin. Error forcing in recurrent neural networks. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [65] Tim Sauer, James A. Yorke, and Martin Casdagli. Embedology. Journal of statistical Physics, 65(3):579–616, 1991. [66] Julian D. Schiller, Malte Heinrich, Victor G. Lopez, and Matthias A. Müller. Tuning the burn-in phase in training recurrent neural networks improves their performance. In The Fourteenth International Conference on Learning Representations, 2026. [67] Dominik Schmidt, Georgia Koppe, Zahra Monfared, Max Beutelspacher, and Daniel Durstewitz. Identifying nonlinear dynamical systems with multiple time scales and long-range dependencies. In Proceedings of the 9th International Conference on Learning Representations, 2021. [68] Bernard W Silverman. Density estimation for statistics and data analysis. Routledge, 2018. [69] B. Sivakumar. Chaos theory in geophysics: past, present and future. Chaos, Solitons & Fractals, 19(2):441–462, 2004. [70] Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. International Conference on Learning Representations (ICLR), 2023. [71] Steven H. Strogatz. Nonlinear dynamics and chaos: with applications to physics, biology, chemistry, and engineering. Chapman and Hall/CRC, 2024. [72] Floris Takens. Detecting strange attractors in turbulence. In Dynamical Systems and Turbulence, Warwick 1980, volume 898, pages 366–381. Springer, 1981. [73] Peter Turchin and Andrew D. Taylor. Complex dynamics in ecological time series. Ecology, 73(1):289–305, 1992. [74] Eli Tziperman, Harvey Scher, Stephen E. Zebiak, and Mark A. Cane. Controlling spatiotemporal chaos in a realistic el niño prediction model. Phys. Rev. Lett., 79:1034–1037, Aug 1997. [75] Pietro Verzelli, Cesare Alippi, and Lorenzo Livi. Learn to synchronize, synchronize to learn. Chaos: An Interdisciplinary Journal of Nonlinear Science, 31(8):083119, August 2021. [76] Pantelis R. Vlachas, Wonmin Byeon, Zhong Y. Wan, Themistoklis P. Sapsis, and Petros Koumoutsakos. Data-driven forecasting of high-dimensional chaotic systems with long shortterm memory networks. Proc. R. Soc. A., 474(2213):20170844, 2018. [77] Pantelis R Vlachas and Petros Koumoutsakos. Learning on predictions: Fusing training and autoregressive inference for long-term spatiotemporal forecasts. Physica D: Nonlinear Phenomena, 470:134371, 2024. [78] Pantelis R. Vlachas, Jaideep Pathak, Brian R. Hunt, Themistoklis P. Sapsis, Michelle Girvan, Edward Ott, and Petros Koumoutsakos. Backpropagation Algorithms and Reservoir Computing in Recurrent Neural Networks for the Forecasting of Complex Spatiotemporal Dynamics. arXiv:1910.05266 [physics], February 2020. arXiv: 1910.05266. [79] Eric Volkmann, Alena Brändle, Daniel Durstewitz, and Georgia Koppe. A scalable generative model for dynamical system reconstruction from neuroimaging data. Advances in Neural Information Processing Systems, 37:80328–80362, 2024. [80] Rui Wang, Yihe Dong, Sercan Ö Arik, and Rose Yu. Koopman Neural Forecaster for Time Series with Temporal Distribution Shifts, October 2022. arXiv:2210.03675 [cs, stat]. [81] Paul J. Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10):1550–1560, 1990. 14

[82] Nicolas Zucchet and Antonio Orvieto. Recurrent neural networks: vanishing and exploding gradients are not the end of the story. Advances in Neural Information Processing Systems, 37:139402–139443, 2024.

Appendix A

Absence of chaos-induced exploding gradients in linear SSMs

The gradient propagation properties of linear SSMs are well-established in prior work [52, 55, 82] and we briefly recap them here for completeness. Consider a discrete-time linear SSM with time-invariant dynamics: zt = Azt−1 + vt , (17) where zt ∈ RM is the latent state, A ∈ RM ×M is the state transition matrix, and vt ∈ RM is an extrinsic input. Differentiating the recurrence yields a constant Jacobian Ji = ∂zi /∂zi−1 = A, so the BPTT chain-rule product collapses to t Y ∂zt = Ji = At−r , ∂zr i=r+1

∂zt ∂zr

≤ ∥A∥t−r .

(18)

Gradient propagation is therefore independent of the state trajectory {zi } and inputs {vi }, and is fully determined by the spectrum of A. As a consequence, linear SSMs cannot exhibit chaos-induced gradient explosion: constraining the eigenvalues of A (typically near the unit disk) is sufficient to control gradient magnitudes.

Proof of Proposition 1: GTF-DEER convergence

Proof. Since M ≤ N and the observation matrix B has full column rank, B + B = IM , and by Eq. (8), Pα = (1 − α) IM . The Jacobian decomposition simplifies to Eq. (7): J˜t := J˜(zt−1 ) := JF ◦δα (zt−1 ) = (1 − α) JF (z̃t−1 ).

(19)

Taking spectral norms leads to: (1 − α) ∥JF (z̃t−1 )∥2 ≤ (1 − α) σ̃max = ρ.

(20)

Since α > α∗ = 1 − 1/σ̃max , we have (1 − α) < 1/σ̃max and hence ρ < 1. For any product of k consecutive forced Jacobians: J˜t J˜t−1 · · · J˜t−k+1

≤ 2

k−1 Y

J˜t−i

i=0

≤ ρk 2

(21)

Then by Def. 1, we have TY −1 1 log J˜T −i T →∞ T i=0

≤ lim

λ(z) = lim

T →∞ T

log ρT = log ρ < 0.

(22)

Thus the forced system is contracting.

Scheduled sampling breaks the linearity of LSSMs

Conventional teacher forcing [26] feeds the ground-truth observation xt−1 into the latent recurrence zt = Azt−1 + U xt−1 + Cst + h.

(23)

at every step. While the following holds for any forcing strategy that re-introduces a feedback from model-predicted observations x̂ into the latent recurrence, we will exemplify the problem using the well established method of scheduled sampling [4]. Scheduled sampling switches between teacher forcing and free-running generation during training by replacing xt−1 with the model’s own prediction x̂t−1 with a scheduled probability ϵt ∈ [0, 1] , xt−1 with probability 1 − ϵt , x̃t−1 = (24) x̂t−1 with probability ϵt , 16

so that the latent recurrence used during training reads zt = Azt−1 + U x̃t−1 + Cst + h.

(25)

Whenever the prediction branch is taken, substituting x̂t−1 = Bϕ(V zt−1 + b) from Eq. (3) into Eq. (25) yields (26) zt = Azt−1 + W ϕ(V zt−1 + b) + Cst + h, with W = U B, i.e. exactly the test-time recurrence of Eq. (4). The transition zt−1 7→ zt is therefore no longer linear (affine) in zt−1 but contains the nonlinearity ϕ. Incompatibility with parallel scan Evaluation of Eq. (3) via parallel scan [7, 49, 70] hinges on each transition being representable as an affine map zt = At zt−1 + bt , identified with the pair (At , bt ). Composition of two such maps, (A2 , b2 ) ◦ (A1 , b1 ) = A2 A1 , A2 b1 + b2 , (27) is associative and again affine, such that partial compositions over disjoint chunks of the training sequence can be combined in any order, enabling an O(log T )-depth reduction. In the standard LSSM recurrence the input term U xt−1 depends only on observed data and contributes solely to bt , leaving the map affine in the latent state. Once xt−1 is replaced by x̂t−1 , the input becomes a nonlinear function of zt−1 and the transition is no longer of this affine form. The composition of two such nonlinear transitions cannot in general be represented by a fixed, finite-parameter operator that is independent of the latent state, such that the recurrence can not be reduced to an associative scan. Training must instead unroll Eq. (26) sequentially, recovering the O(T ) bottleneck.

GTF-DEER gradients

We define the per-timestep MSE loss function ℓ : RN × RN → R, 1 2 ∥x − x̂∥2 , N and the total loss, L : (RN )T × (RN )T → R, for the full time series, ℓ(x, x̂) =

(28)

T T 1X 1 X 2 ℓ(xt , x̂t ) = ∥xt − x̂t ∥2 , T t=1 N T t=1

(29)

L(x1:T , x̂1:T ) =

where in the model (1), x̂t = Gψ (zt ) are the observations associated with the latent states; we consider x1:T ∈ RN T ∼ = (RN )T as a stacked vector; and in practice truncate a warm-up period 0 < Tw < T and only supply vectors x, x̂ of size N T̃ , T̃ := T − Tw , cf. Sect. 4.2 and Eq. (60). To optimize the parameters θ of the RNN F̃θ,α = Fθ ◦ δα,B , we need to compute the gradient of the loss function, ∂L ∂L ∂z = (30) ∂θ ∂z ∂θ for all elements θ (for notational simplicity, we choose not to write any index) in the abstract parameter vector θ which comprises all matrix weights. In the DEER forward pass, we use the residual rθ,α (zt−1 , zt ) = zt − F̃θ,α (zt−1 , xt−1 ) ∈ RM , Rθ,α (z1:T ) = (rα (z0 , z1 ), . . . , rα (zT −1 , zT )) ∈ R

(31) ,

(32)

where z0 is a fixed initial state (which is not predicted), to obtain the predicted time series z1:T ∈ RM T ∼ = (RM )T . At the end of each forward pass, Rα (z1:T , θ) := Rθ,α (z1:T ) = 0 ∈ RM T ,

(33)

so by the implicit function theorem, ∂Rα ∂z1:T ∂Rα + = 0 ∈ RM T . ∂z1:T ∂θ ∂θ 17

(34)

Rearranging this equation gives us ∂z1:T =− ∂θ

∂Rα ∂z1:T

−1

∂Rα . ∂θ

(35)

We insert this into the gradient of the loss: −1 ∂L ∂L ∂Rα ∂Rα =− , ∂θ ∂z1:T ∂z1:T ∂θ | {z }

(36)

=:V ⊤ ∈R1×M T

where we write V = (v1 , . . . , vT ) ∈ RM T ∼ = (RM )T . Rather than explicitly inverting the huge ∂Rα Jacobian matrix ∂z1:T , we instead solve the adjoint problem V⊤

∂L ∂Rα =− . ∂z1:T ∂z1:T

(37)

To make this more explicit, we use ∂rα,θ (zt−1 , zt ) ∂ F̃θ,α (zt−1 ) =: −JF̃ (zt−1 ) =: −J˜t ∈ RM ×M , =− ∂zt−1 ∂z ∂rα,θ (zt−1 , zt ) = IM , ∂zt

(38) (39)

and write 

I  M −J˜  2  ∂Rα  = 0  ∂z1:T  ..  .  0



···

−J˜3 .. .

IM .. .

··· .. . .. . −J˜T

 0     0    0   IM

(40)

and ∂Rα ∂Rα ∂ F̃θ,α ∂ F̃θ,α = =− . ∂θ ∂θ ∂θ ∂ F̃θ,α

(41)

The block-bidiagonal structure of the matrix turns (37) into a (backward) recursion in t: vT⊤ = −

∂L , ∂zT

⊤ vt−1 = vt⊤ JF̃ (zt−1 ) −

∂L , ∂zt−1

(42)

where, substituting in the observation model, x̂ = Gψ (z), ∂L ∂L ∂Gψ 2 ∂Gψ = = (x̂t−1 − xt−1 )⊤ . ∂zt−1 ∂ x̂t−1 ∂zt−1 NT ∂zt−1

(43)

In the same way as in the forward pass, this recursion can be solved in parallel via an associative scan. Note that for quasi-DEER, the Jacobians JF̃ in (42) are replaced by a diagonal approximation and the equation becomes vT⊤ = −

∂L , ∂zT

⊤

⊤ vt−1 = (vt ⊙ diag (JF̃ (zt−1 ))) −

∂L ∂zt−1

(44)

with the element-wise product ⊙. In summary, the loss is therefore calculated simply by solving the recursion (42) using the final latent trajectory and observations computed via the forward pass, and inserting the result into ∂L ∂ F̃θ,α = −V ⊤ . ∂θ ∂θ 18

(45)

Model setup, loss and regularizations

Parameterizations For both the RNN and LSSM we parameterize A = diag(tanh(Ā)) where Ā ∈ RM . The nonlinearity avoids instabilities in the Azt−1 term when ∥A∥2 ≥ 1. Initialization To facilitate the capture of long time scales in the data, we initialize the RNN to exhibit long time scales at initialization, by initializing the RNN near the identity: Ā = artanh(κ) IM Wij ∼ U −(1 − κ)L−1/2 , (1 − κ)L−1/2 Vij ∼ U −(1 − κ)M −1/2 , (1 − κ)M −1/2 Cij ∼ U −(1 − κ)K −1/2 , (1 − κ)K −1/2

(46)

b=0 h = 0, where U is the uniform distribution. In practice we use κ = 0.9995. For M ≥ N , the observation to perform identity read-out of the first N units in the RNN, i.e. B = matrix B is initialized IN 0 N ×(M −N ) . For the LSSM, we follow similar strategy: Ā = artanh(κ) IM Uij ∼ U −(1 − κ)N −1/2 , (1 − κ)N −1/2 Cij ∼ U −(1 − κ)K −1/2 , (1 − κ)K −1/2

(47)

h = 0, and the observation model is initialized as Bij ∼ U −(1 − κ)L−1/2 , (1 − κ)L−1/2 Vij ∼ U −(1 − κ)M −1/2 , (1 − κ)M −1/2

(48)

b = 0. For all experiments, we use ϕ(·) = ReLU = max(0, ·). Manifold attractor regularization (MAR) For the nonlinear RNN introduced in Eq. (5), we can encourage a slow manifold in the last Mr units by adding the term [67]   M L X X λMAR 1 |1 − Aii |p + LMAR (θRNN ) = (|Wij |p + |Vji |p ) + |hi |p  (49) Mr L j=1 i=M −Mr +1

to the loss function, where λMAR is a regularization parameter and p ∈ {1, 2} determines the type of penalty. In practice, we found it beneficial to scale regularization terms of connectivity matrices W and V based on their variance at initialization, i.e.   M L X X λMAR 1 |1 − Aii |p + LMAR (θRNN ) = (γW |Wij |p + γV |Vji |p ) + |hi |p  , Mr L j=1 i=M −Mr +1

(50) 1 1 where γW = 3L , γV = 3M . For λMAR → ∞, this regularization leads to stable error propagation by imposing a block-diagonal identity structure in the model Jacobian JF , as we demonstrate below. Recall the nonlinear RNN from Eq. (5) with A = diag(a1 , . . . , aM ): zt = Azt−1 + W ϕ(V zt−1 + b) + Cst + h. 19

(51)

Its Jacobian is JF (zt−1 ) = A + W diag(ϕ′ (V zt−1 + b)) V .

(52)

Let Ms := M − Mr and partition z = (z (s) , z (r) )⊤ into “fast” (unregularized) and “slow” (regularized) units. The diagonal A and the non-linear coupling term decompose as As 0 Ws A= , W = , V = (Vs Vr ) . (53) 0 Ar Wr As λ → ∞, the MAR penalty, Eq. (49), drives Ar → IMr ,

Wr → 0,

Vr → 0.

(54)

(s)

Writing Ds (zt−1 ) := diag ϕ′ (Vs zt−1 + b) , the limiting Jacobian is block-diagonal: λ→∞

JF (zt−1 ) −−−−→

As + Ws Ds (zt−1 )Vs 0

0 IMr

(55)

The latent space decouples into a nonlinear subsystem acting on z (s) and a manifold-attractor subsystem on z (r) with identity dynamics. Consequently, the Jacobian product for a length t sequence and λ → ∞ satisfies ! Qt−1 t−1 Y 0 k=0 [As + Ws Ds (zt−1−k )Vs ] JF (zt−1−k ) = , (56) 0 IMr k=0 where error signals propagate without decay or amplification along the Mr directions, enabling stable capture of long-term dependencies. Observation model regularization In addition to MAR, we reduce the direct contribution of the Mr MAR units to the read-out and hence discourage them from being forced through Eq. (8) by L1 (B) =

N λ1 X N · Mr i=1

M X

|Bij |p ,

(57)

j=M −Mr

where B is the readout matrix (cf. Eq. (5)). We also follow [33] and regularize B to stay wellconditioned by pulling its singular values towards 1: r

L2 (B) =

λ2 X 2 (σi (B) − 1) , r i=1

(58)

where r = rank(B) and σi (B) denotes the i-th singular value of B. Total loss given by

The overall loss function used for training of the nonlinear RNN defined in Eq. (5) is L(xTw +1:T , x̂Tw +1:T ; θRNN , B) = LMSE (xTw +1:T , x̂Tw +1:T ) + LMAR (θRNN ) + L1 (B) + L2 (B),

(59)

with LMSE (xTw +1:T , x̂Tw +1:T ) =

1 N (T − Tw )

T X

∥xt − x̂t ∥2 .

(60)

t=Tw +1

MAR for the LSSM

Application of MAR to LSSMs (Eq. (3)) is straightforward with   M N X X λMAR 1 |1 − Aii |p + LMAR (θLSSM ) = |Uij |p + |hi |p  . Mr N j=1 i=M −Mr +1

(61)

Datasets

Lorenz-63 Our first benchmark system is the three-dimensional Lorenz-63 system [44]. Its dynamics are described by the differential equations dx = σ(y − x), dt dy = x(ρ − z) − y, dt dz = xy − βz. dt

(62)

We chose the classical parameters σ = 10, ρ = 28, β = 8/3, which put the system into a chaotic regime. We used a Runge-Kutta 4/5 scheme with ∆t = 0.01 and integrated a trajectory of length 100,000 time steps from an initial point x0 . The trajectory was standardized per dynamical variable after generation. Forced Lorenz-96 As a benchmark for long time scales, we equipped the vanilla Lorenz-96 equation [45] with a sinusoidal forcing term: dxi = (xi+1 − xi−2 ) xi−1 − xi + F0 + A sin(ωt), i = 1, . . . , N, (63) dt where F0 is a constant forcing offset, A is the amplitude and ω the frequency of the sinusoidal forcing. We set F0 = 14, A = 12, ω = 2π / 75 and N = 6. The autonomous Lorenz-96 system (i.e. A = 0) is chaotic for F ≥ 8, such that the sinusoidal forcing leads to a periodic switching between cyclic and chaotic dynamics (see Fig. A1). For integration we used a Runge-Kutta 4/5 scheme with ∆t = 5 · 10−3 and integrated two trajectories from different initial t and x0 of length 500 internal time units, leading to two trajectories of length 100,000 time steps. One trajectory is contaminated with 5% Gaussian observation noise and used for training, while the other trajectory is kept clean and used for testing. Both trajectories are standardized per dynamical variable after generation. Bursting neuron model We used a three-dimensional simplified Hodgkin-Huxley-type neuron model; its dynamical variables are the membrane potential V as well as two gating variables n and h which control the opening of fast and slow potassium channels, respectively [20, 67]: 1h V̇ = I − gL (V − EL ) − gNa m∞ (V ) (V − ENa ) − gK n (V − EK ) C i − gM h (V − EK ) − gNMDA s∞ (V ) (V − EN M DA ) , (64)

n∞ (V ) − n , τn h∞ (V ) − h , ḣ = τh

ṅ =

with m∞ (V ) =

1 + exp (Vh,Na − V )/kNa 1 , n∞ (V ) = 1 + exp (Vh,K − V )/kK

h∞ (V ) =

, 1 + exp (Vh,M − V )/kM 1 . s∞ (V ) = 1 + 0.33 exp −0.0625 V

(65)

The model parameters we used for our experiments are reported in Table A1. Similar to the Lorenz-96 system, we used a Runge-Kutta 4/5 scheme with ∆t = 2.5 · 10−2 and integrated two trajectories from different initial conditions of length 5000 internal time units from which we cut 1000 as transients, leading to two trajectories of length 160,000 discrete time steps. One trajectory is contaminated with 5% Gaussian observation noise and used for training, while the other trajectory is kept clean and used for testing and both trajectories are again standardized per dynamical variable after generation. Moreover, to make reconstruction more challenging, we throw away the h variable after generation, such that the system is partially observed with N = 2. For an example reconstruction, see Fig. 3. 21

Table A1: Neuron model parameter settings I

gNa

ENa

Vh,Na

kNa

Vh,K

τn

Vh,M

τh

gNMDA

ENMDA

−80

−20

−90

−25

25.2

−18

1000

10.2

Evaluation measures

State space divergence Dstsp The state space divergence Dstsp measures the geometrical disagreement between state distributions of the data p(x) and that of model generated trajectories q(x) by evaluating the Kullback-Leibler (KL) divergence: Z p(x) Dstsp := DKL (p(x) || q(x)) = p(x) log dx. (66) q(x) N x∈R In practice, we estimate p(x) and q(x) by placing Gaussian Mixture Models (GMM) along orbits [40, 11]), i.e. T 1 X p̂(x) = N (x; xt , Σ) T1 t=1 (67) T2 1 X N (x; x̂t , Σ), q̂(x) = T2 t=1 where x1:T1 is the ground-truth data of length T1 and x̂1:T2 is a model-generated trajectory of length T2 , N (x; xt , Σ) is a multivariate Gaussian with mean vector xt and covariance matrix 2 Σ = diag([σ12 , . . . , σN ]). The KL-divergence between two GMMs can be approximated using a Monte-Carlo approach [32] n p̂(x(i) ) 1X log , (68) Dstsp = DKL (p̂(x) || q̂(x)) ≈ n i=1 q̂(x(i) ) with n Monte Carlo samples x(i) randomly drawn from the GMM based on observed orbits, p̂(x(i) ). When the observed system is only partially observed and the true dynamics unfold in d > N dimensions, the attractor first has to be unfolded to apply Dstsp correctly [8]. To this end we first apply a delay embedding (DE) [72, 65] of the observed and model generated orbits ⊤ ξt = xt , xt−τ , xt−2τ , . . . , xt−(m−1)τ ∈ RmN (69) with delay τ and embedding dimension m, where the same construction is performed for ξ̂t . We then fit GMMs and calculate Dstsp in embedding space, i.e. n

DE Dstsp = DKL (p̂(ξ) || q̂(ξ)) ≈

p̂(ξ (i) ) 1X log . n i=1 q̂(ξ (i) )

(70)

For all experiments, we use n = 106 and set the bandwidth Σ of Gaussian compartments using Σ = −1/(d+4) fbw · diag([σ̂12 , . . . , σ̂d2 ]), where fbw = T1 (d+2) is a bandwith factor chosen according to 4 Silverman’s rule-of-thumb [68] and σ̂i is the empirical standard deviation of xi,1:T1 . Furthermore, to probe long-term consistency beyond the available data, we choose to generate model orbits for T2 = 3 · T1 . This avoids low Dstsp when the model is only transiently generating the correct long-term behavior. Since the forced Lorenz-96 system (Appx. F) and bursting neuron data sets are only partially observed, we used settings m = 3, τ = 4,096 for the Lorenz-96 and m = 7, τ = 1,024 for the bursting neuron. We explicitly used a large delays to resolve the long time scales present in the data. For the partially observed Lorenz-63 (see Sect. 5.2), we used m = 3, τ = 10. Short-term RMSE defined as

The n-step-ahead RMSE assesses short- to medium-length forecasts and is v u n u1 X 2 RMSE(n) = t xk − Gψ (Fθ◦k (z0 )) (71) n k=1

where x1:n := xti :ti +n−1 is a window starting at index ti of the total ground-truth data x1:Ttot . The initial condition z0 is obtained by performing a warm-up using a history of length Tw , xti −Tw :ti −1 (see Sect. 4.2 for details on warm-up). In practice we used averaged the RMSE over 100 windows from different initial conditions.

Additional experimental details

This appendix section lists specific experimental details of experiments performed in the main paper. All experiments were performed on a compute server with 2x96-Core AMD EPYC 9655 CPUs, 768 GB RAM, and 6x NVIDIA RTX Pro 6000 Blackwell (96 GB) GPUs; however, individual trainings were always performed on a single GPU at a time. H.1

Figure 1

Fig. 1A We trained shPLRNNs (Eq. (5), L = 50) of varying latent dimensions M ∈ {4, 16, 64, 128} on the Lorenz-63 system using GTF-DEER. The batch size was fixed at B = 1, while the sequence length was varied in powers of 2, i.e. T ∈ {128, 256, . . . , 32,768} and we used no warm-up for this experiment (Tw = 0). We measured the time of combined forward+backward passes throughout training. Training was performed for 10, 000 updates (samples). We choose to explicitly train a model to account for the fact that GTF-DEER’s convergence depends on the dynamics of the model, giving us an estimate of the spread of runtimes throughout training. Indeed, for the chosen α = 0.15, GTF-DEER’s convergence stayed consistent throughout training (see also Fig. 2B). The sequential baseline runtime is independent of explicit dynamics, and hence we only measured the time of trained models using Ns = 500 samples. In code, the only difference between the sequential and parallel implementation is the call to the forward pass solver, which is given by a straight-forward jax.lax.scan for the sequential case, and the seq1d (DEER solver) function for parallel one, respectively. Both methods parallelize over the batch size using jax.vmap. All data was gathered on an NVIDIA RTX 6000 Blackwell (96GB) GPU. Fig. 1B For this figure we used a shPLRNN (varying M , fixed L = 50) that was trained on the full Lorenz-63 system (N = 3). We performed forward passes for different values of forcing strength α ∈ [0, 1]. The maximum number of Newton iterations was capped at niter = 500: In practice, GTF-DEER loses its competitive performance over sequential evaluation even for moderate Newton iterations such that 500 is already a generous upper bound of what can be considered a useful regime. The different curves correspond to different initial guess strategies, where ‘Zero’ corresponds to (0) (0) initializing z1:T = 0T ×M . ‘N (0, 1)’ draws entries from a standard Normal distribution, zi,t ∼ (0)

N (0, 1) and ‘P-inv’ uses z1:T = B + x1:T to initialize the Newton iterations. H.2

Figure 2

For this experiment we trained shPLRNN (M = 5, L = 50) on the Lorenz-63 under different conditions. We trained on the full Lorenz-63 system (N = 3) as well as on the partially observed one (N = 1, x-component). Furthermore, we switched between GTF-DEER with full Jacobians (Eq. (13)) and diagonalized ones (quasi-DEER, JF → diag(JF )). We trained 20 models per setting for 20,000 parameter updates. The forcing strength was α = 0.15. H.3

Figure 3A

For Fig. 3A we trained shPLRNNs with M = 10 and L = 128 on the forced Lorenz-96 system and M = 6 for the bursting neuron model (see Appx. F). To probe the effect of increasing sequence length on this problem, we trained the models on sequence lengths T ∈ {256, 512, 1,024, 2,048, 4,096, 8,192, 16,384, 32,768} where for each setting we used Tw = T /2. To keep data parsed by the model per parameter update the same, we adjusted batch size accordingly, such that the product B·T = 215 = 32,768 stayed fixed. This resulted in settings (B, T ) = (256, 128) and (B, T ) = (1, 32,768) for the lowest and highest sequence length used, respectively. For each data point in the figure, we trained 10 independent models. For the Lorenz-96 we used a forcing strength of α = 0.08, for the bursting neuron model α = 0.4. 23

H.4 H.4.1

Figure 3B Mamba-2

As a state-of-the-art SSM baseline, we compare the LSSM (Eq. (3)) and the GTF-DEER-trained shPLRNN (Eq. (5)) to a model in which the linear recurrence of the LSSM is replaced by a single Mamba-2 block [18], while keeping the rest of the architecture, in particular the one-hidden-layer MLP read-out from Eq. (3), identical. We refer to this configuration as Mamba-2-np , where np is the number of parameters. For the non-recurrence machinery of the Mamba-2 block (input projections, depth-wise causal convolution, gating branch, multi-head setup, output projection) we refer the reader to the original work by [18]. In the following, we clarify the, to us important distinction to the LSSM introduced in Eq. (3). Selective SSM recurrence. The main functional difference of the Mamba recurrence is the introduction of a selection mechanism [28, 18], which renders the model non-autonomous even in the absence of external inputs st . Dropping st for brevity and using the LSSM (Eq. (3)) as a template, the core recurrence inside the Mamba-2 block can be written as zt = Āt (xt−1 ) zt−1 + Ūt (xt−1 ) xt−1

(72)

i.e., structurally identical to the LSSM but with Āt and Ūt now depending on the input xt−1 . Specifically, parameterizations of Āt and Ūt are given by ∆t = softplus(W∆ xt−1 + b∆ ) ∈ R Ut = WU xt−1 Āt = exp(∆t A)

(73)

Ūt = ∆t Ut , where A is a fixed, learnable parameter which for Mamba-2 is restricted to A = a I for a single learned scalar a < 0 per head. The recurrence parameters thus become functions of the input through the data-dependent scalar ∆t and the data-dependent input vector Ut , while A itself remains time-invariant. Setting ∆t = ∆ and Ut = U recovers a fixed-parameter linear recurrence equivalent to the LSSM (Eq. (3)).2 Relation to the LSSM Crucially, the recurrence in Eq. (72) is still linear in the latent state zt−1 for any fixed value of xt−1 , which is what permits parallelization via parallel scan during training [18]. At the same time, the selection mechanism circumvents the rank ≤ min(N, M, L) bottleneck on the effective test-time recurrence (Eq. (4)) discussed in Sect. 3.1: because Ut varies with t, the model is no longer constrained to a single fixed product W = U B of fixed rank, but can in principle realize a different effective input pathway at every step. H.4.2

Hyperparameters

Optimization settings For experiments in Sect. 5.4, all models were trained using the same following settings: batch size B = 1, sequence length T = 81,920 with warm-up Tw = 16,384 and a learning rate decay from ηs = 5 · 10−5 to 10−6 using a cosine decay schedule [46]. Training was performed for 150,000 parameter updates. Mamba-2 For the Mamba-2 models, we evaluated two parameter budgets, denoted Mamba-2-7k with 7,038 and Mamba-2-17k with 16,830 parameters, respectively. The former is matched as closely as possible to the shPLRNN/LSSM budget (≈ 2,200–2,900 parameters) for direct comparison at fixed capacity without overly limiting Mamba-2’s expressiveness; the latter probes whether additional capacity closes the gap. Both use a single Mamba-2 block with hidden dimension D = 16 (d_model) for Mamba-2-7k and D = 32 for Mamba-2-17k. All other parameters stay the same: State size 2 For clarity, Eqs. (72) & (73) describe the Mamba-2 recurrence at the level of abstraction relevant for comparison with the LSSM. In practice, a Mamba-2 block does not operate directly on xt−1 : the input is first lifted to an expanded representation of dimension D = E · N via a linear projection, where E is the expansion factor, then split into H heads, and each head runs an independent SSM with its own data-dependent parameters (∆t , Ut ) shared across the channels within that head, while A is a learned scalar per head. Channel mixing across the observation dimensions therefore occurs in the surrounding input/output projections rather than inside the recurrence itself. We refer the reader to [18] for the full block specification.

N = 16 (d_state), d_conv = 4 and expand = 2. Following the Mamba-2 block is the same MLP read-out as used in the LSSM (Eq. (3)) with L = 128. We used the Mamba-2 implementation from the original code repository https://github.com/state-spaces/mamba, licensed under the Apache License, Version 2.0. RNN The RNN used M = 10 and L = 128 and Mr = 4. We used λMAR = 1, λ1 = 1 and λ2 = 10−4 , however, we found the regularization of the singular values of B did not have a huge impact on training convergence and reconstruction performance. The forcing was α = 0.08. LSSM For the LSSM we kept core settings equal to the ones used in the RNN, i.e. M = 10, L = 128. Additionally, we apply MAR to the LSSM (see Appx. E) where we scanned λMAR ∈ {10−4 , 10−3 , 10−2 , 10−1 , 1, 10} but found that while λMAR = 10 degraded performance, all other settings produced similar results. Hence we reported λMAR = 1 to make clear that MAR does not improve reconstruction for the LSSM. For the latter settings we also used Mr = 4, mirroring the setting of the RNN.

Additional figures

Figure A1: Generated dynamics of a shPLRNN (M = 10, L = 128) trained on the forced Lorenz-96 system. A: Access to long sequences during training enables the model to capture the latent longperiod sinusoidal forcing F (t) in the MAR regularized units, and hence producing accurate long-term rollouts. B: A model trained on short sequence lengths w.r.t. the intrinsic time scales of the data fails to capture and reproduce the latent forcing during autoregressive roll-outs.

Figure A2: Generated dynamics of 3× the length of available ground truth data from models evaluated in Sect. 5.4. LSSMs and shPLRNNs with their rank limited to the number of observed variables r = N = 6 fail to learn the latent sinusoidal forcing. Increasing the rank to 7 in the shPLRNN enables the model to learn the underlying forcing and produce convincing limiting behavior, however, the model still struggles to infer the correct phase from the context. Even though the Mamba-2 based model can circumvent the low-rank constraint due to its data-dependent recurrence weights, it does not match the performance of the shPLRNNs trained with GTF-DEER. Each generated trace was produced using the best model (out of 10) per configuration.

Figure A3: RMSE(128) as a function of sequence length for shPLRNNs trained on the forced Lorenz-96 and bursting neuron systems, complementary to the Dstsp plot in Fig. 3. Median ± MAD, 10 runs per setting.

Training algorithms

Algorithm 1 Linear SSM training (using parallel associative scan) Require: Data X ∈ RTobs ×N , optional external inputs S ∈ RTobs ×K , linear recurrence Fθ (zt−1 , xt−1 , st ), decoder Gψ (zt ), batch size B, sequence length T , warm-up length Tw Ensure: Trained parameters θ, ψ 1: repeat 2:

Sample sequences x0:T ∈ RB×(T +1)×N from X, s1:T ∈ RB×T ×K from S

▷ Implicit vectorized batch processing below, e.g. via jax.vmap

Initialize z0 = 0

z1:T ← associative_scan(Fθ , z0 , x0:T −1 , s1:T )

x̂1:T = Gψ (z1:T )

L ← MSE(xTw +1:T , x̂Tw +1:T )

(gθ , gψ ) ← grad(L; θ, ψ)

(θ, ψ) ← optimize(θ, ψ, gθ , gψ )

◁ ▷ Solve in parallel on GPU

▷ Exclude warm-up steps from loss

10: until convergence

Algorithm 2 Initial value / trajectory matching training with generalized teacher forcing Require: Data X ∈ RTobs ×N , optional external inputs S ∈ RTobs ×K , recurrent model Fθ , observation matrix B ∈ RN ×M , forcing strength α ∈ [0, 1], batch size B, sequence length T Ensure: Trained parameters θ, B 1: repeat 2:

Compute B + ← pinv(B),

P ← IM − αB + B

Sample sequences x0:T ∈ RB×(T +1)×N from X, s1:T ∈ RB×T ×K from S

▷ Implicit vectorized batch processing below, e.g. via jax.vmap

Compute teacher signals z̄0:T = B (x0:T )

Initialize z0 ← z̄0

for t = 1, . . . , T do

z̃t−1 ← P zt−1 + α z̄t−1

zt ← Fθ (z̃t−1 , st )

10:

◁

▷ Build weighted forced state

x̂t ← Bzt

11:

L ← MSE(x1:T , x̂1:T )

12:

(gθ , gB ) ← grad(L; θ, B)

13:

(θ, B) ← optimize(θ, B, gθ , gB )

▷ Backpropagation through time

14: until convergence

Algorithm 3 GTF-DEER: Generalized teacher forcing with DEER parallel solver Require: Data X ∈ RTobs ×N , optional external inputs S ∈ RTobs ×K , recurrent model Fθ , observation matrix B ∈ RN ×M , forcing strength α ∈ [0, 1], batch size B, sequence length T , tolerance ε Ensure: Trained parameters θ, B 1: repeat 2:

Compute B + ← pinv(B),

Pα ← IM − αB + B,

Sample sequences x0:T ∈ RB×(T +1)×N from X, s1:T ∈ RB×T ×K from S

▷ Implicit vectorized batch processing below, e.g. via jax.vmap Compute forcing targets z̄t ← B xt for t = 0, . . . , T

Initialize z0 ← z̄0

Initial guess z1:T ← z̄1:T

for k = 0, 1, . . . until ∥∆z1:T

10: 11: 12: 13: 14: 15: 16:

◁

P1 ← IM − B + B

(k)

▷ Below, fix z0

(0)

= z0 ∀k

▷ Warm start from data (k+1)

∥∞ < ε do ▷ DEER iteration (k) (k) z̃t ← P1 zt + z̄t for t = 1, . . . , Tw ▷ Warm-up (k) (k) z̃t ← Pα zt + αz̄t for t = Tw + 1, . . . , T ▷ Forced inputs (k) ∂Fθ Jt ← ∂z z̃(k) P1 for t = 1, . . . , Tw ▷ Warm-up Jacobians t (k) ∂Fθ Jt ← ∂z z̃(k) Pα for t = Tw + 1, . . . , T ▷ Forced Jacobians t (k) (k) (k) rt ← zt − Fθ (z̃t−1 ) for t = 1, . . . , T ▷ Residuals (k+1) (k) (k+1) (k) Solve ∆zt = Jt−1 ∆zt−1 − rt for t = 1, . . . , T via associative scan (k+1) (k) (k+1) z1:T ← z1:T + ∆z1:T (final) x̂Tw +1:T ← BzTw +1:T

17:

L ← MSE(xTw +1:T , x̂Tw +1:T )

▷ Loss excluding warm-up states

18:

(gθ , gB ) ← grad(L; θ, B)

▷ Single associative scan via IFT

19:

(θ, B) ← optimize(θ, B, gθ , gB )

20: until convergence

Record · ID 180635 · SHA-256 2917b792013a76f3
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.