TRACTABLE DENDRITIC RNNS
Year: 2021

Abstract:
In many scientific disciplines, we are interested in inferring the nonlinear dynamical system underlying a set of observed time series, a challenging task in the face of chaotic behavior and noise. Previous deep learning approaches toward this goal often suffered from a lack of interpretability and tractability. In particular, the highdimensional latent spaces often required for a faithful embedding, even when the underlying dynamics lives on a lower-dimensional manifold, can hamper theoretical analysis. Motivated by the emerging principles of dendritic computation, we augment a dynamically interpretable and mathematically tractable piecewise-linear (PL) recurrent neural network (RNN) by a linear spline basis expansion. We show that this approach retains all the theoretically appealing properties of the simple PLRNN, yet boosts its capacity for approximating arbitrary nonlinear dynamical systems in comparatively low dimensions. We introduce two frameworks for training the system, one based on fast and scalable variational inference, and another combining BPTT with teacher forcing. We show that the dendritically expanded PLRNN achieves better reconstructions with fewer parameters and dimensions on various dynamical systems benchmarks and compares favorably to other methods, while retaining a tractable and interpretable structure.

heading: 1 INTRODUCTION
text: For many complex systems in physics, biology, or the social sciences, we do not know or have only rudimentary knowledge about the dynamical system (DS) that may underlie those quantities that we can empirically observe or measure. Data-driven approaches aimed at automatically inferring the generating DS from time-series observations could therefore strongly support the scientific process, and various such methods have been proposed in recent years (Raissi et al., 2018; Zhu et al., 2021; Yin et al., 2021; Norcliffe et al., 2021; Mohajerin & Waslander, 2018; Karl et al., 2017; Chen et al., 2018; Strauss, 2020). However, due to the often high-dimensional, complex, chaotic, and inherently noisy nature of real-world DS, like the brain, weather-, or ecosystems, this remains a formidable challenge. Moreover, although the true DS may evolve on a lower-dimensional manifold in its state space, the system used for approximation usually needs to be of higher dimensionality to achieve a proper embedding (Takens, 1981; Sauer et al., 1991; Kantz & Schreiber, 2004). This is especially true when the approximating system is of a different functional form than the one that would most naturally describe the data generation process (but is unknown), for instance, when we attempt to approximate a system of exponential or trigonometric functions by polynomials.
In this work we sought to improve the capacity and expressiveness of a specific class of recurrent neural networks (RNNs), achieving agreeable solutions with fewer dimensions and parameters while retaining a set of desirable theoretical properties. Specifically, we build on piecewise-linear RNNs (PLRNNs) based on ReLU activation functions, for which fixed points, periodic orbits, and other dynamical properties can be derived analytically (Schmidt et al., 2021; Koppe et al., 2019), and for which dynamically equivalent continuous-time (ordinary differential equation, ODE) systems can be constructed (Monfared & Durstewitz, 2020b). Inspired by principles of dendritic computation in biological neurons (Fig. 1), each PLRNN unit was endowed with a set of nonlinear pre-processing subunits (“dendritic branches”), such that it effectively takes on the role of an equivalent much larger network. Mathematically, this comes down, in our case, to enhancing each latent unit with a linear spline basis expansion as popular in statistics (Hastie et al., 2009). Through this trick, we achieve
a powerful RNN which provides reconstructions of underlying nonlinear DS in lower-dimensional latent spaces than were needed by conventional PLRNNs. At the same time, model inference can be performed within the scalable framework of sequential variational auto-encoders (SVAE) (Archer et al., 2015; Girin et al., 2020; Krishnan et al., 2017), or with classical Back-Propagation-ThroughTime (BPTT; Rumelhart et al. (1986)) augmented by teacher-forcing (TF; Williams & Zipser (1989); Pearlmutter (1990)). We further prove that these modifications preserve the mathematical and dynamical accessibility of the resulting system, e.g., such that fixed points, cycles, and their stability, can still be computed analytically.
Besides its effectiveness in capturing complex dynamical systems in fewer dimensions within a tractable framework, our approach highlights more generally how principles of dendritic signal processing may be harvested in the design of RNNs. Strongly nonlinear local computations are known for decades to occur within dendritic trees of biological neurons (Mel, 1994; Poirazi et al., 2003), but have hardly been exploited so far for machine learning models.

heading: 2 RELATED WORK
text: One class of DS reconstruction models attempts to discover governing equations from the flow field estimated from data through differencing the time series. Sparse Identification of Nonlinear Dynamics (SINDy), for instance, does so by sparsely regressing on a rich library of basis functions using the least absolute shrinkage and selection operator (LASSO) (Brunton et al., 2016; Rudy et al., 2017; de Silva et al., 2020). Other methods approximate the flow field using graph reconstruction via differential equations (Chen et al., 2017), sparse autoencoders (Heim et al., 2019), shallow multi-layer perceptrons reformulated as RNNs (Trischler & D’Eleuterio, 2016), or deep neural networks (Chen et al., 2018). Some works aimed at directly learning the system’s underlying Hamiltonian (Chen et al., 2020; Greydanus et al., 2019). Generally, numerical derivatives obtained from time series tend to be more noise-prone than the time series observations themselves (Baydin et al., 2018; Chen et al., 2017; Raissi, 2018). This can be a problem particularly if only comparatively short trajectories were empirically observed or when the underlying systems are very high-dimensional, as in these cases the system’s flow field may be (severely) under-sampled. Methods directly based on numerical derivatives also need to be augmented by other techniques, like delay embeddings (Kantz & Schreiber, 2004) or deep auto-encoders (Champion et al., 2019), if not all the system’s dimensions were observed.
Various RNN architectures such as Long-Short-Term-Memory networks (LSTMs) (Zheng et al., 2017), Reservoir Computing (RC) (Pathak et al., 2018), or PLRNNs (Koppe et al., 2019; Schmidt et al., 2021) have been employed to infer DS directly from the observed time series without going through numerical derivatives. More recently, transformers (Shalova & Oseledets, 2020a;b) were used as black box approaches for DS prediction. Except for PLRNNs, however, all these systems, although optimized for DS reconstruction and prediction, rest on relatively complex model formulations that are not easy to tackle and analyze from a DS perspective (Fraccaro et al.). The ability to gain deeper insights into the specific DS properties and mechanisms of the recovered system is, however, often crucial for its applicability to science and engineering problems. Transformers, unlike RNNs, do not even constitute DS themselves (as they explicitly forgo any temporal recursions), and therefore are not directly amenable to DS theory tools. Moreover, most of these models, RC in particular, need very high-dimensional latent spaces, which further adds to their black-box nature.
Better interpretability and tractability is achieved by using PLRNNs (Koppe et al., 2019; Schmidt et al., 2021) or by (locally) linearizing nonlinear systems through ideas from Koopman operator theory (Azencot et al., 2020; Brunton et al., 2017; Yeung et al., 2017). In such systems, certain DS properties can be analytically accessed (Schmidt et al., 2021; Monfared & Durstewitz, 2020a), or the resulting equations can be more easily interpreted by a human reader (Heim et al., 2019). On the downside, usually one needs to move to very high dimensions to represent the DS in question properly. Here we aim to overcome this limitation by augmenting PLRNNs with linear basis expansions without altering their analytical accessibility.
Finally, probabilistic (generative) latent variable models such as state space models have been applied to the problem of posterior inference of latent state paths zt ∼ p(zt|x1:T ) of DS given time series observations {x1:T } (Pandarinath et al., 2018; Ghahramani & Roweis, 1998; Durstewitz, 2017; Krishnan et al., 2017). The advantage here is that they also account for uncertainty in the model formulation or latent process itself and yield the full distribution over latent space variables (Karl
et al., 2017). For DS reconstruction, however, we need to move beyond posterior inference: We require that samples drawn from the model’s prior distribution p(z) after training exhibit the same temporal and geometric structure as those produced by the unknown DS.
Here we embed PLRNNs augmented with a linear spline expansion into a fully probabilistic, variational approach that scales well with system size by employing stochastic gradient variational Bayes (SGVB; (Kingma & Welling, 2014; Rezende et al., 2014)), thereby combining the advantages of the two classes of models reviewed above. On the other hand, we show that the model can also be efficiently trained by BPTT using a specific form of TF (Appx. 6.1).

heading: 3 MODEL FORMULATION AND THEORETICAL CONSIDERATIONS

heading: 3.1 PIECEWISE LINEAR RECURRENT NEURAL NETWORK (PLRNN)
text: Our approach builds on PLRNNs (Durstewitz, 2017; Koppe et al., 2019) because of their mathematical tractability (see Sec. 3.3). PLRNNs are defined by the M -dimensional latent process equation
zt = Azt−1 +Wϕ(zt−1) + h+Cst + ϵt, (1)
which describes the temporal evolution of M -dimensional latent state vector zt = (z1t . . . zMt)T . The self-connections of the units are represented by diagonal matrix A ∈ RM×M , whereas the connections between units are collected in off-diagonal matrix W ∈ RM×M , with the nonlinear activation function ϕ given by the rectified linear unit (ReLU) applied element-wise:
ϕ(zt−1) = max(0, zt−1). (2)
Additionally, the PLRNN comprises a bias term h ∈ RM , potential external inputs st ∈ RK weighted by C ∈ RM×K , and a Gaussian noise term ϵt ∼ N (0,Σ) with diagonal covariance Σ. The PLRNN can be interpreted as a discrete-time neural rate model (Durstewitz, 2017), where the entries of A stand for the individual neurons’ time constants, W for the synaptic connection strengths between neurons, and ϕ(z) for a (ReLU-shaped) voltage-to-spike-rate transfer function. The probabilistic latent RNN Eq. 1 is linked to the N -dimensional observed time series (xt)t=1...T , xt ∈ RN , drawn from an underlying noisy DS, by an observation function (decoder model) which, in the simplest case, may take the linear Gaussian form
xt = Bzt + ηt, (3)
where B ∈ RN×M represents a factor loading matrix and ηt ∼ N (0,Γ) is Gaussian observation noise with diagonal covariance Γ ∈ RN×N .

heading: 3.2 DENDRITIC COMPUTATION AND SPLINE BASIS EXPANSION
text: Dendrites have long been known to play an active and important part in neural computation (Mel, 1994; 1999; Koch, 2004). Active, fast voltage-gated ion channels endow dendrites with strongly nonlinear behavior, giving rise for instance to dendritic Ca2+ spikes that boost synaptic inputs (Schiller et al., 2000; Häusser et al., 2000). It has been suggested previously that different dendritic branches may constitute rather independent computational sub-units whose outputs are combined at the soma, as in a 2-layer neural network (Poirazi et al., 2003; Mel, 1993; 1994), an idea that received strong empirical support especially in recent years (Poirazi & Papoutsi, 2020). Here we mimic this functional setup by modeling dendritic processing through a linear combination of ReLU-type threshold-nonlinearities (Fig. 1), replacing Eq. 2 by
ϕ(zt−1) = B∑ b=1 αb max(0, zt−1 − hb), (4)
with “dendritic input/output” slopes αb ∈ R and “activation” thresholds hb ∈ RM . As in real dendrites, where both ion channels and morphological structure are subject to learning (Poirazi & Papoutsi, 2020; Stemmler & Koch, 1999), we treat these as trainable parameters. We note that Eq. 4 inserted into model Eq. 1 takes the form of a linear spline basis expansion as popular in statistics (Hastie et al., 2009) for approximating arbitrary functions (Wahba, 1990; Storace & De Feo, 2004) in regression settings. For instance, such concepts have been frequently employed within data-analytical
models in neuroscience (Frank et al.; Huang et al.; Qian et al.), but never within the context of DS reconstruction enabling lower-dimensional solutions in mathematically tractable models.
To emphasize the connection to dendritic computation we call the system Eqs. 1, 3, 4, the dendPLRNN.

heading: 3.3 MATHEMATICAL TRACTABILITY AND DYNAMICAL SYSTEMS INTERPRETATION
text: Sharp threshold-nonlinearities (like a ReLU) are a reasonable choice from a neurobiological perspective, as dendrites naturally give rise to this threshold-type behavior (Mel, 1999; Koch, 2004). Another important consideration in choosing this particular form, however, was that it preserves all the theoretically appealing properties of a PLRNN, as we will formally establish below: For PLRNNs fixed points and cycles can be explicitly computed (Schmidt et al., 2021; Koppe et al., 2019), and they can be translated into dynamically equivalent continuous-time systems (Monfared & Durstewitz, 2020b), properties which profoundly ease the analysis of trained systems from a DS perspective. This is crucial for application in the sciences, where we are specifically interested in understanding the underlying system’s dynamics. For PLRNNs, precise connections between the long-term behavior of the system and that of its gradients have also been established (Schmidt et al., 2021). Finally, PLRNNs belong to the class of continuous piecewise-linear (PWL) maps, for which many important types of bifurcations have been well characterized (Feigin, 1995; Hogan et al., 2007; Patra, 2018) (cf. (Monfared & Durstewitz, 2020a) for an overview). Bifurcations are essential to understand how geometrical and topological properties of the system’s state space depend on its parameters or could be controlled, and hence are also important to characterize or improve the training process itself (Doya, 1992; Pascanu et al., 2013; Saxe et al., 2014) or to understand properties of trained systems (Maheswaranathan et al., 2019b;a) .
Our first proposition, therefore, assures that by the particular form of basis expansion introduced in Eq. 4, the system will remain within the class of continuous PWL maps:
Proposition 1. The model defined through Eq. 1 and Eq. 4 constitutes a continuous PWL map.
The proof essentially straightforwardly follows from the model’s definition as a linear spline basis expansion in each unit, but is formally provided in Appx. 6.5.4.
While Proposition 1 is all we need to ensure we can harvest all previously established results on PLRNNs in particular, and on continuous PWL maps more generally, it is revealing to note that any dendPLRNN (Eqs. 1, 4) can be rewritten as a conventional PLRNN, as stated in the following theorem:
Theorem 1. Any M -dimensional dendPLRNN as defined in Eqs. 1, 4, can always be rewritten as a M ×B-dimensional “conventional” PLRNN of the form
ẑt = Ãẑt−1 + W̃ max(0, ẑt−1) + ĥ0 + C̃st + ϵ̃t. (5)
Proof. Straightforward by construction, see Appx. 6.5.5.
This theorem highlights why the dendPLRNN will allow to reduce the dimensionality of the reconstructed system, as it suggests we may often be able to reformulate a high-dimensional PLRNN in terms of an equally powerful lower-dimensional dendPLRNN. In Appx. 6.5.1 we also spell out the exact computation of fixed points and k-cycles for the dendPLRNN.
Finally, the unboundedness of the PLRNN’s latent states due to the ReLU function can cause divergence problems in training. The dendPLRNN, on the other hand, offers a simple and natural way to contain the latent states without violating the basic model description above, as established in the following theorem: Theorem 2. For each basis {αb,hb} in Eq. 4 of a dendPLRNN let us add another basis {α∗b ,h∗b} with α∗b = −αb and h∗b = 0. Then, for σmax(A) < 1, any orbit of this “clipped” dendPLRNN (Eq. 10) will remain bounded.
Proof. See Appx. 6.5.6.
Appx. 6.5 collects further theoretical results, assuring, for instance, that the manifold attractor regularization employed here (see next section) does not interfere with the results above (Proposition 2).

heading: 3.4 TRAINING THE DENDPLRNN
text: To infer the parameters θ = {A,W ,h,C,Σ,B,Γ, {αb,hb}} of the dendPLRNN (Eq. 1, 3, 4) from observed data, we apply two different training strategies: First, a fast and scalable variational inference (VI) algorithm which maximizes the Evidence Lower Bound (ELBO) L(θ, ϕ;x) := Eqϕ [log(pθ(x|z)] − KL[qϕ(z|x)||pθ(z)] using the reparameterization trick (Kingma & Welling, 2014), and convolutional neural networks (CNNs) for parameterizing the encoder model qϕ(z|x) (see Appx. 6.1 for details). Furthermore, as proposed in Schmidt et al. (2021), to efficiently capture DS at multiple time scales, we add a regularization term to the ELBO that encourages the mapping of slow time constants and long-range dependencies (so-called manifold attractor regularization, see Eq. 6, with regularization factor λ). Second, we employ “classical” BPTT with a variant of teacher forcing (TF) (Williams & Zipser, 1989; Pearlmutter, 1990). TF here means that the first N latent states zk,lτ+1, k ≤ N , were replaced by observations xk,lτ+1 at times lτ + 1, l ∈ N0, where τ ≥ 1 is the forcing interval (for details, see Appx. 6.1). All code used in here is made freely available at [placeholder].

heading: 4 EXPERIMENTS

heading: 4.1 PERFORMANCE MEASURES
text: In DS reconstruction, we aim to capture invariant properties of the underlying DS like its geometrical and temporal structure. To evaluate the quality of the reconstructions w.r.t. geometrical properties we employed a Kullback-Leibler divergence (Dstsp) that quantifies agreement in attractor geometries (more details in Appx. 6.2), as first suggested in Koppe et al. (2019) (see also Schmidt et al. (2021)). Specifically, this measure evaluates the overlap between the observed data distribution p(xobs) and the distribution p(xgen|zgen) generated from model simulations (i.e., with zgen ∼ pθ(z) after model training) across state space (not time!). Since this measure as originally defined in Koppe et al. (2019) is expensive to compute, for the high-dimensional benchmark DS we used another approximation, details of which are given in Appx. 6.2. Dstsp is evaluated on a set of 100 trajectories, pulled from the learned distribution over initial conditions, with 1000 time steps each. To assess the agreement in temporal structure, a dimension-wise, Gaussian-kernel-smoothed power spectrum correlation (PSC) between ground truth and model-generated trajectories was used (see Appx. 6.2). Finally, we also computed a 20-step-ahead prediction error along test set trajectories (see Appx. 6.2), although not of primary interest in the context of DS reconstruction.

heading: 4.2 DS BENCHMARKS USED FOR EVALUATION
text: We evaluated our approach and the specific role of the basis expansion on five different types of challenging DS benchmarks.
First, the famous 3d chaotic Lorenz attractor (Lorenz-63) originally proposed by Lorenz (1963) (formally defined in Appx. 6.4) has become a popular benchmark for DS reconstruction algorithms. Fig. 2a (l.h.s.) illustrates true (blue) and reconstructed (orange) time series from this system, while the r.h.s. illustrates the chaotic attractor’s geometry in its 3d state space for both the ground truth (blue) and reconstructed (orange) systems. It is important to note that both the time and state space graphs are not merely ahead predictions from the dendPLRNN but are produced by simulating the trained dendPLRNN from some initial condition. This illustrates that the dendPLRNN has captured the temporal and geometrical structure of the original Lorenz-63 system in its own governing equations. Moreover, computing analytically (see Appx. 6.5.1) the fixed points of the reconstructed system, we see that their positions in state space agree well with those of the true system.
Second, a 3d biophysical model of a bursting neuron (see Eq. 15 in Appx. 6.4; Durstewitz (2009)) highlights another aspect of DS reconstruction: Besides an equation for membrane voltage (V ), the model consists of one very fast (n) and one slow (h) variable that control the gating of the model’s ionic conductances. This produces fast spikes that ride on top of a much slower oscillation, making this system challenging to reconstruct. One such successful dendPLRNN reconstruction is illustrated in Fig. 2b (orange) together with time graphs and state space representations of the true system (blue).
Third, the Lorenz-96 weather model is an example of a higher-dimensional, spatially organized chaotic system with local neighborhood interactions that can be extended to arbitrary dimensionality (Eq. 18 in Appx. 6.4). It has also been used more widely for benchmarking DS reconstruction algorithms. For our experiments we employed a 10-dimensional spatial layout. Fig. 3a illustrates time graphs for selected dimensions (l.h.s.), the full evolving spatio-temporal pattern (center), and examples of power spectra (r.h.s.) for both the ground truth system (blue) and an example reconstruction (orange). The spatio-temporal characteristics of the true and the dendPLRNN-generated time series tightly agree.
Fourth, as another high-dimensional example we used a neural population model with structured connectivity tuned to produce coherent chaos (Landau & Sompolinsky, 2018), from which we produced 50d observations (see Appx. 6.4 for details). Fig. 3b provides example time series (l.h.s.), full spatio-temporal patterns (center), and overlaid power spectra (r.h.s.) for time series drawn from the true system (blue) and those simulated by a trained dendPLRNN (orange). Again there is a tight agreement, and again we emphasize that - like in all the other examples - these are not mere model ahead-predictions but fully simulated from some random initial condition.
Finally, we studied a real-world dataset consisting of electroencephalogram (EEG) recordings from human subjects, described in more detail with results (Fig.***) in Appx. ***.

heading: 4.3 BASIS EXPANSION IMPROVES DS RECONSTRUCTION AND LOWERS PARAMETER COSTS
text: While, in theory, the dendPLRNN is equivalent to a larger PLRNN without basis expansion, in practice the smaller basis-expanded models trained more successfully. Fig. 4a summarizes our observations for the VI algorithm on the impact of the basis expansion using the Lorenz-63 DS as an example (see Fig. S3 for further examples): Both the 20-step ahead prediction error as well
as the geometrical reconstruction quality as assessed by Dstsp profoundly improve with the size B of the basis expansion even for the same total number of trainable model parameters (given by M(M + 1 + B + N) + B, where N is the dimensionality of the observed data). Hence, as conjectured in Sec. 3, the basis expansion yields better reconstructions at no additional costs in terms of numbers of model parameters. Fig. 4b looks at the impact of the basis expansion from the angle of dimensionality reduction by systematically varying the number of bases B and latent states M for the Lorenz-63 DS: Following the curves horizontally, it becomes clear that the basis expansion enables to reduce the model’s overall dimensionality without compromising performance.

heading: 4.4 MODEL COMPARISONS
text: We compared our model to the PLRNN without the dendritic expansion and three other algorithms purpose-tailored for DS reconstruction: First, SINDy (Brunton et al., 2016) aims to reconstruct the governing equations by approximating numerical derivatives (obtained by differencing the time series, and applying a variance regularization to reduce noise) through a large library of polynomial basis functions. Sparse (LASSO) regression is used to pick out the right terms from the library (we used the PySINDy implementation (de Silva et al., 2020) with multinomials up to sixth order). Second, Vlachas et al. (2018) used a hybrid of truncated LSTMs and mean-field stochastic models based on Ornstein-Uhlenbeck processes (LSTM-MSM) to approximate the true system’s flow estimated from observed time series. Third, Pathak et al. (2018) built on reservoir computing (RC) for their approach with reservoir parameters chosen to satisfy the “echo state property” (Jaeger & Haas, 2004). For
higher-dimensional systems, a spatially arranged set of reservoirs with local neighborhood relations is employed. For all these systems, optimized hyper-parameters were used as reported by the authors. For our own system, the dendPLRNN, we also performed a grid search for optimal hyper-parameters λreg, τTF, M , and B (see Appx. 6.1 and Table S1 for details). For all four methods, to the degree possible we tried to ensure roughly the same number of trainable parameters (see Table S2).
Results for all four models on all five DS benchmarks employed here are summarized in the upper part of Table S2, using the temporal and geometrical reconstruction measures introduced in Sec. 4.2 (as well as a 20-step-ahead prediction error for comparison). To produce this table, 100,000 time steps for both training and testing were simulated from each ground truth system, all dimensions were standardized to have zero mean and unit variance, and process noise and observation noise (with 1% of the data variance) were added while simulating the (now stochastic) differential equations, and after drawing the observations, respectively (see Appx. 6.4 for further methodological details). To produce statistics, each method was run from a total of 20 randomly chosen initial conditions for the parameters. We also tested all four methods on the real EEG data and on challenging data situations produced using the Lorenz-63 system (Fig. 2a), with either short time series of just 1000 time steps, only partial observations (just state variable x in Eq.14 in Appx. 6.4), or high process and high observation noise (drawing from a Gaussian with dϵ ∼ N (0, 0.1dt× I) for the process noise as described in Appx. 6.4, and using 10% of the observation variance, respectively). SINDy cannot naturally handle missing observations, as it has no latent variables but formulates the model directly in terms of the observations. Therefore, for the partially observed system, we used delay embedding (Takens, 1981; Sauer et al., 1991) to create a 3d dataset, adding two time-lagged versions of x as coordinates.1
A general observation is that indeed all four models are quite powerful for reconstructing the underlying DS. However, in most comparisons the dendPLRNN had an edge over the other methods,
1We point out that this may already impose a restriction for methods like SINDy as one moves to very high-dimensional systems.
or came out second after SINDy, especially when trained by BPTT+TF. SINDy tends to outperform the dendPLRNN on the Lorenz-63 DS, but it completely fails on the bursting-neuron and population model examples, and on the real EEG data, and generally becomes comparatively slow to train on high-dimensional systems. This can be explained by the fact that SINDy already has the correct functional form for the Lorenz-63 (and also Lorenz-96) DS: Both of these have a strictly polynomial form (see Eq. 14 and Eq. 18 in Appx. 6.4), and SINDy works with a set of polynomial library functions to begin with. Hence, SINDy only needs to pick out the right terms from its expansion to succeed, giving it a clear advantage on these model systems by design. On the other hand, as revealed in Table S2, it completely fails on systems which have a different (non-polynomial) functional form (or when the true form, as in the EEG case, is simply not known). SINDy therefore appears less suitable as a general framework for DS reconstruction if an appropriate library of basis functions cannot be specified a priori, unlike the other methods.
While our conclusion is that essentially all of the three tested models LSTM-MSM, RC, and dendPLRNN are suitable for reconstruction of arbitrary DS even in very challenging data situations (Table S2, bottom), LSTM-MSM and RC performed worse on average and have other profound disadvantages compared to our method: First, they are quite complex in their architectures and hence not easily interpretable, i.e. much harder to track and analyze mathematically.2 In contrast, as summarized in Sec. 3.3, the dendPLRNN is a continuous PWL map and as such comes with a huge bulk of already existing theoretical results (Schmidt et al., 2021; Monfared & Durstewitz, 2020b;a), as well as with mathematical tractability (see Fig. 2a and Appx. 6.5.1). On top, the dendPLRNN achieves reconstruction of all DS in (much) lower dimensions than the RC or LSTM-MSM (see Table S2), further adding to its better interpretability. Second, by embedding the dendPLRNN within a SVAE (Archer et al., 2015) framework we also obtain uncertainty estimates on the state trajectories and can perform posterior inference, features that the other models lack.

heading: 5 CONCLUSIONS
text: In this work we augmented PLRNNs (Durstewitz, 2017; Koppe et al., 2019) by a linear spline basis expansion inspired by principles of dendritic computation. We show mathematically that by doing so we remain within the theoretical framework of continuous PWL maps and hence can harvest a huge bulk of existing DS theory (Sec. 3.3), while at the same time achieving better performance with less parameters and in lower dimensions. Another contribution of this work is transferring the PLRNN into the framework of SVAEs which allow for fast and scalable inference and training. These are two key advantages from both a scientific perspective where mechanistic insight and understanding of the system under study are sought, and a prediction perspective where we are also interested in uncertainty estimates.
We close by pointing out two open issues: First, somewhat surprisingly, the BPTT+TF approach to model training clearly outperformed the more sophisticated VI approach. This could be rooted in suboptimal encoder models or in suboptimal sampling from the approximate posterior: While BPTT+TF assesses longer bits of trajectory during optimization, in VI single time-point samples are drawn and the temporal consistency is ensured only through the Kullback-Leibler term in the ELBO. Other more expressive yet still fast to compute encoder models, e.g., based on normalizing flows (Rezende & Mohamed, 2015), may boost performance. Smart initialization techniques (Talathi & Vartak, 2016) or specific annealing and curriculum training protocols (as used in Koppe et al. (2019)) are other amendments to consider. Second, we felt that quantitative measures for assessing the quality of DS reconstructions in high-dimensional, high-noise situations are an interesting research topic in their own right. It is known that “classical” DS measures like Lyapunov spectra or correlation dimensions (Kantz & Schreiber, 2004) are very hard to robustly assess for higher-dimensional or more noisy systems (Schreiber & Kantz, 1996), and are often not even known for comparatively simple models. Yet, the geometrical and temporal measures employed here come with their own pitfalls, some of them alluded to in Appx. 6.2.

heading: ACKNOWLEDGMENTS
text: 2This is especially true for RC. Moreover, the fact that only the weights of the linear output layer are trainable while the recurrent connections within the reservoirs are static, may raise the question of what precisely is learnt in terms of dynamics if the reservoirs themselves cannot adapt to the DS at hand.
Ethics statement The current work performs theoretical analysis and basic research on a novel type of RNN architecture, which is evaluated exclusively on a set of simulated physical and biological systems (no human or animal subjects involved, no privacy concerns etc.). Although it is conceivable that the current models, training algorithms, and results may be used to develop prediction algorithms in sensitive domains (like medical time series or tracking consumer data), there are no immediate ethical implications from this work as far as we can see.
Reproducibility statement All theoretical results in this paper were carefully and thoroughly proven, with all proofs and detailed derivations available in the Appendix. Likewise, we will make available all code used in the empirical section in a way that will allow others to easily reproduce the results from this paper. This means we will include everything, starting with the code for the benchmark models and simulations, the simulated time series data used for evaluation themselves, the code for our own model and training algorithms, up to the meta-files that produce the figures in this work, on our lab github site. All of this will be clearly documented.

heading: 6 APPENDIX

heading: 6.1 FURTHER METHODOLOGICAL DETAILS
text: Manifold Attractor Regularization As proposed in Schmidt et al. (2021), to encourage the discovery of long-term dependencies and slow time scales in the data, a subset of Mreg ≤M states was regularized by adding the following term to the ELBO for the VI approach:
Lreg = λ Mreg∑ i=1 (Aii − 1)2 + Mreg∑ i=1 M∑ j ̸=i (Wij) 2 + Mreg∑ i=1 h2i  . (6) This regularization pushes the regularized subset of states toward a continuous set of marginally stable fixed points that tends to form an attracting manifold in the full state space, which supports the learning of systems with widely differing time scales, such as the bursting neuron model (cf. Sec. 4). We found that for all datasets this regularization significantly helped to discover the underlying dynamics. To put it on equal grounds with the regularization term, the ELBO was divided by the number of time steps T of a given batch. Regularization settings used are summarized in Table S1 along other hyper-parameter settings.
BPTT-TF To train a deterministic version of the dendPLRNN, we employ BPTT with a scheduled version of TF (Williams & Zipser, 1989; Pearlmutter, 1990). To do so, we choose an “identitymapping” for the observation model x̂t = Izt, where I ∈ RN×M with Ikk = 1 if k ≤ N and zeroes everywhere else. This allows us to regularly replace latent states with observations to “recalibrate” the model and break trajectory divergence in case of chaotic dynamics. Consider a time series {x1,x2, · · · ,xT } generated by a DS we want to reconstruct. At times lτ + 1, l ∈ N0, where τ ≥ 1 is the forcing interval, we replace the first N latent states by observations ẑk,lτ+1 = xk,lτ+1, k ≤ N . The remaining latent states, ẑk,lτ+1 = zk,lτ+1, k > N , remain unaffected by the forcing. This means that we optimize the dendPLRNN such that a subspace of the latent space directly maps to the observed time series variables. The forcing interval τ is a hyperparameter, with optimal settings varying depending on the dataset. The settings we chose are summarized in Table S1. With F = {lτ + 1}l∈N0 , the dendPLRNN updates can then be written as
zt+1 = { dendPLRNN(z̃t) if t ∈ F dendPLRNN(zt) else . (7)
The loss is calculated prior to the forcing, such that Lt = ∥xt − Izt∥22 for every time step. To improve performance we employ a mean-centred dendPLRNN (for details see next paragraph). In the evaluation phase, the trained dendPLRNN is simulated freely without any forcing. As the model is deterministic, the initial condition z1 = [x1,Lx1]T is estimated from the first data point x1 with a matrix L ∈ R(M−N)×N which is jointly learned with the other model parameters.
Mean-Centered dendPLRNN Layer normalization has recently been developed as a way of significantly improving RNN training (Ba et al., 2016). Here we adapt the idea of layer normalization to the piecewise-linear nature of our dendPLRNN. Instead of fully standardizing the latent states at every time step before applying the activation function, we only mean-center them:
zt = Azt−1 +Wϕ ( M(zt−1) ) + h0, (8)
where ϕ(·) is given in Eq. 4 and M(zt−1) = zt−1 − µt−1 = zt−1 − 1 1
M M∑ j=1 zj,t−1, where
1 ∈ RM is a vector of ones. Note that this mean-centering is linear and can be rewritten as a matrix-multiplication
M(zt−1) = zt−1 − µt−1
= 1
M M − 1 −1 · · · −1−1 M − 1 · · · −1· · · · · · · · · · · · −1 −1 · · · M − 1  zt−1 = Mzt−1. (9) As Remark 1 points out, all results about the tractability of the dendPLRNN also hold for the mean-centred dendPLRNN.
State clipping Since the ReLU function used in the dendPLRNN is non-saturating, states may diverge to infinity. As Theorem 2 guarantees, there is a simple and natural way to construct a “clipped” dendPLRNN
zt = Azt−1 +W B∑ b=1 αb [ max(0, zt−1 − hb)−max(0, zt−1) ] + h0. (10)
Note that the results of Theorem 2 also hold true when the manifold attractor regularization is applied. This is detailed in Proposition 2 further below.
Approximate posterior and hyperparameter settings To estimate the true unknown posterior p(z|x), we make a Gaussian assumption for the approximate posterior qϕ(z|x) = N (µϕ(x),Σϕ(x)), where mean and covariance are functions of the observations. Without any simplifying assumptions, the number of parameters in Σϕ(x) ∈ RMT×MT would scale unacceptably with time series length T . We therefore made a mean field assumption and factorized qϕ(z|x) across time. Specifically, a time-dependent mean µt,ϕ and covariance Σt,ϕ were parameterized through stacked convolutional networks which take the observations {xt−w...xt+w} as inputs, with w given by the largest kernel size. The mean is given by a 4-layer CNN with decreasing kernel sizes (41, 31, 21 and 11, respectively), with the last layer of the CNN feeding into a fully connected layer. For the diagonal covariance, the observations are mapped directly onto the logarithms of the covariance through a single convolutional layer (with a kernel size of 41) feeding into a fully connected layer. The classical motivation behind using CNNs rests on the assumption that the data contains translationally invariant patterns, and that this allows the recognition model to embed potentially meaningful temporal context into the latent representation (see e.g. Cui et al. (2016)). We note that while the mean-field approximation is computationally highly efficient, it makes potentially strongly simplifying assumptions (Blei et al., 2017; Bayer et al.) that may limit the ability of the encoder model to approximate the true posterior.
To train the dendPLRNN in the VI framework, Adam (Kingma & Ba, 2015), with a batch size of 1000 and learning rate of 1e− 3 was used as the optimizer. For the training with BPTT, we used the Adam optimizer with an initial learning rate of 1e− 3 that was iteratively reduced during training down to 1e− 5. For each epoch we randomly sampled sequences of length Tseq = 500 (except for the Lorenz-63 runs, where Tseq = 200 time steps were sufficient) from the total training data pool of each dataset, which are then fed into the reconstruction method in batches of size 16. Initial weights were drawn from a uniform distribution.
To find optimal hyper-parameters we performed a grid search within λreg ∈ {0, 0.01, 0.1, 1, 10} (VI), τTF ∈ {1, 5, 10, 25, 50, 100} (BPTT-TF), M ∈ {5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 75, 100}, and B ∈ {0, 1, 2, 5, 10, 20, 35, 50}. Hyper-parameters chosen for the benchmarks in Sec. 4 are reported in Table S1 below (note that these may not fully agree with the ranges initially scanned, as given above, since we attempted to adjust them further in order to approximately match the number of parameters among models in Table S2).
Table S1: Hyperparameter settings for dendPLRNN VI/TF for the four different data sets from Sec. 4.
Dataset M B Mreg/M λreg τTF Lorenz-63 22 20 1.0/− 1.0/− −/25 Lorenz-96 42/50 50/30 1.0/− 1.0/− −/10
Bursting Neuron 26 50/47 0.5/− 1.5/− −/5 Neural Population Model 12/75 5/40 0.2/− 1.0/− −/5
EEG 117/128 50/50 0.8/0.5 1.0/1.0 −/10

heading: 6.2 PERFORMANCE MEASURES
text: Geometrical measure Dstsp used for evaluating attractor geometries (Fig. 4) measures the match between the ground truth distribution ptrue(x) and the generated distribution pgen(x | z) through the
discrete binning approximation (Koppe et al., 2019)
Dstsp (ptrue(x), pgen(x | z)) ≈ K∑
k=1
p̂ (k) true(x) log
( p̂ (k) true(x)
p̂ (k) gen(x | z)
) , (11)
where K is the total number of bins, and p̂(k)true (x) and p̂ (k) gen (x | z) are estimated as relative frequencies through sampling trajectories from the benchmark DS and the trained reconstruction method, respectively. A range of 2× the data standard deviation on each dimension was partioned into m bins, yielding a total of K = mN bins, where N is the dimension of the ground truth system. Initial transients are removed from sampled trajectories to ensure that the system has reached a limiting set. If the bin size is chosen too large, important geometrical details may be lost, while if it is chosen too small, noise and (low) sampling artifacts with many empty bins may misguide the approximation above. Here we chose a bin number of m = 30 per dimension as an optimal compromise that distinguished well between successful and poor reconstructions.
Since the number of bins needed to cover the relevant (populated) region of state space scales exponentially with the number of dimensions, for high-dimensional systems evaluating Dstsp as outlined above is not feasible. We therefore resorted to an approximation of the densities based on Gaussian Mixture Models (GMMs), similar to a strategy outlined in (Koppe et al., 2019). Specifically, we approximate the true data distribution by a GMM ptrue(x) ≈ 1T ∑T t=1N (xt,Σ) with Gaussians centered on the observed data points {xt} and covariance Σ, which we treat as a hyper-parameter that determines the granularity of the spatial resolution (similar to the bin size in Eq. 11). We can generate a likewise distribution by sampling trajectories from the trained models (or one very long trajectory) and placing Gaussians on the sampled data points, pgen(x|z) ≈ 1L ∑L l=1N (x̂l | zl,Σ) (in the case of VI, rather than sampling, one could also use the model’s distributional assumptions to build this posterior across the observations). For the Kullback-Leibler divergence between two GMMs efficient approximations are at hand (Hershey & Olsen, 2007). Here we employ a Monte Carlo approximation
D̃stsp ( ptrue(x), pgen(x|z) ) ≈ 1
n n∑ i=1 log 1/T ∑T t=1N (x(i);xt,Σ) 1/L ∑L l=1N (x(i); x̂l,Σ) , (12)
where n Monte-Carlo samples x(i) are drawn from the GMM representing ptrue. In practice, we set the covariance Σ = σ2I equal to a scaled identity matrix, with a single hyperparamter σ2. Scanning the range σ2 ∈ {0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5}, we found that values for σ2 = 0.1 − 1.0 to differentiate best between good and bad reconstructions. We chose σ2 = 1.0 for numerical stability. For this setting, Dstsp as derived with the binning method and D̃stsp computed through the GMMs also correlated highly on the low-dimensional benchmark systems (see Figure S1).
Power Spectrum Correlation The power spectrum correlations (PSC) were obtained by first sampling time series of 100,000 time steps, standardizing these, and computing dimension-wise Fast Fourier Transforms (using scipy.fft) for both the ground truth systems and model-simulated time series. Individual power spectra were then slightly smoothed with a Gaussian kernel, normalized, and the long, high-frequency tails of the spectra, mainly representing noise, were cut off. Smoothing width σ and cutoff values scale linearily with the length of the time series used to compute the spectrum, and were chosen by visual inspection of the individual spectra. Dimension-wise correlations between smoothed power spectra were then averaged to obtain the reported PSC scores.
Mean Squared Prediction Error A mean squared prediction error (PE) was computed across test sets of length T = 1000 by initializing the reconstruction model with the benchmark time series up to some time point t, from where it was then iterated forward by n time steps to yield a prediction at time step t+ n. The n-step ahead prediction error (PE) is then defined as the MSE between predicted and true observations:
PE(n) = 1
(T − n) T−n∑ t=1 ||xt+n − x̂t+n||22. (13)
Note that for a chaotic system initially close trajectories will exponentially diverge, such that PEs for too large prediction steps n are not meaningful anymore (in a chaotic system with noise, for large n
Figure S1: Correlation between the binning approximation (m = 30) and the logarithm of the GMM approximation (σ2 = 1) to Dstsp on the Lorenz-63 system for different noise realizations and variances. Similar as reported for the KLZ approximation in Koppe et al. (2019) we observed a logarithmic relation between the GMM and binning based measures.
Figure S2: Example power spectrum for different values of the smoothing factor σ2.
the PE may be high even when estimated from two different runs of the same ground truth model from the same initial condition; see (Koppe et al., 2019)). How quickly this happens depends on the rate of exponential divergence as quantified through the system’s maximal Lyapunov exponent (Kantz & Schreiber, 2004).

heading: 6.3 FURTHER RESULTS
text: Figure S3: Effect of basis expansion for dendPLRNN trained by VI. (a) Agreement in attractor geometries (left) and 20-step ahead prediction error (right) for the bursting neuron system as a function of the number of bases (B) for fixed numbers of total parameters. (b) Agreement in power spectrum correlation (left) and 20-step ahead prediction error (right) for the Lorenz-96 system as a function of the number of bases (B) for different numbers of latent states (M ).
Figure S4: We observed that the dendPLRNN frequently underwent bifurcations between fixed point or various cyclic solutions until it reaches the chaotic behavior of the ground truth system.
Table S2: Comparison of dendPLRNN (Ours) trained by VI or BPTT+TF, and a standard PLRNN (Schmidt et al., 2021), trained by VI or BPTT+TF on four DS benchmarks (top) and three challenging data situations (bottom). Values are mean ± SEM.
Dataset / Setting Method PSC Dstsp 20-step PE Dynamical variables Parameters
Lorenz
dendPLRNN VI 0.997 ± 0.001 0.80 ± 0.25 2.1e−3 ± 0.2e−3 22 1032 dendPLRNN TF 0.997 ± 0.002 0.13 ± 0.18 9.2e−5 ± 2.8e−5 22 1032 PLRNN VI 0.94 ± 0.004 16.6 ± 0.4 1.8e−1 ± 0.1e−1 22 1032 PLRNN TF 0.994 ± 0.001 0.4 ± 0.09 4.3e−3 ± 0.2e−3 30 1011
Lorenz96
dendPLRNN VI 0.987 ± 0.001 0.10 ± 0.01 3.1e−1 ± 0.9e−1 42 4384 dendPLRNN TF 0.998 ± 0.0001 0.04 ± 0.01 4.1e−2 ± 0.8e−2 50 4480 PLRNN VI 0.93 ± 0.002 1.68 ± 0.03 2.1e−3 ± 0.2e−3 60 4260 PLRNN TF 0.996 ± 0.0003 0.05 ± 0.01 2.2e−1 ± 0.2e−1 64 4700
Bursting Neuron
dendPLRNN VI 0.55 ± 0.03 7.5 ± 0.4 6.1e−1 ± 0.1e−1 26 2052 dendPLRNN TF 0.76 ± 0.04 2.9 ± 1.3 6.1e−2 ± 2.2e−2 26 2040 PLRNN VI 0.54 ± 0.01 17.5 ± 0.5 1.17 ± 0.14 42 2021 PLRNN TF 0.72 ± 0.07 3.2 ± 2.0 7.0e−2 ± 2.7e−2 43 2021
Neural Population Model
dendPLRNN VI 0.45 ± 0.05 0.56 ± 0.05 0.82 ± 0.09 12 821 dendPLRNN TF 0.51 ± 0.01 0.19 ± 0.02 1.53 ± 0.03 75 9990 PLRNN VI 0.48 ± 0.01 11.65 ± 1.32 0.68 ± 0.09 13 832 PLRNN TF 0.38 ± 0.15 9.6 ± 7.3 19 ± 23 98 12102
Low amount of data
dendPLRNN VI 0.967 ± 0.007 4.36 ± 0.10 2.8e−2 ± 0.2e−2 22 1032 dendPLRNN TF 0.97 ± 0.04 6.9 ± 5.3 1.5e−2 ± 0.9e−2 22 1032 PLRNN VI 0.96 ± 0.01 18.1 ± 0.10 1.08 ± 0.02 30 1020 PLRNN TF 0.96 ± 0.04 9.0 ± 5.4 1.8e−2 ± 0.5e−2 30 1011
Partially observed
dendPLRNN VI 0.940 ± 0.006 12.6 ± 1.0 6.5e−2 ± 1.4e−2 22 1032 dendPLRNN TF 0.993 ± 0.003 0.54 ± 0.16 5.3e−3 ± 0.2e−3 22 1032 PLRNN VI 0.944 ± 0.002 17.2 ± 0.2 2.7e−1 ± 0.03e−1 30 1020 PLRNN TF 0.994 ± 0.003 0.56 ± 0.34 5.0e−3 ± 0.2e−3 30 1011
High noise
dendPLRNN VI 0.973 ± 0.006 4.9 ± 0.75 3.5e−2 ± 0.1e−2 22 1032 dendPLRNN TF 0.995 ± 0.002 0.4 ± 0.13 4.6e−3 ± 0.4e−3 22 1032 PLRNN VI 0.94 ± 0.004 18.2 ± 0.04 6.4e−1 ± 0.1e−1 30 1020 PLRNN TF 0.994 ± 0.002 0.5 ± 0.08 4.3e−3 ± 0.2e−3 22 1032
EEG Dataset Electroencephalogram (EEG) data were taken from a study by (Schalk et al., 2000) available at https://physionet.org/content/eegmmidb/1.0.0/. These are 64-channel EEG data obtained from human subjects during different motor and imagery tasks. We trained the dendPLRNN using BPTT+TF on the ”eyes open” baseline time series from subject 0, which had a total of 9760 time steps. The signal was standardized and smoothed with a Hann function, using numpy.hanning and a window length of 15. Results for ground-truth and freely generated EEG signals from several brain regions are shown in figure S5.

heading: 6.4 DETAILS ON DYNAMICAL SYSTEMS BENCHMARKS
text: Lorenz-63 system The famous 3d chaotic Lorenz attractor with the butterfly wing shape, originally proposed in (Lorenz, 1963), is defined by
Figure S5: EEG recordings from frontal, occipital, parietal and temporal lobe vs. freely generated trajectories, sampled from the dendPLRNN, trained with BPTT (M = 128, B = 50, τ = 10,Mreg/M = 0.5, λ = 1.0).
dx dt = σ(y − x) + dϵ1(t) dt ,
dy dt = x(ρ− z)− y + dϵ2(t) dt , (14)
dz dt = xy − βz + dϵ3(t) dt .
Parameters used for producing ground truth data in the chaotic regime were σ = 10, ρ = 28, and β = 8/3. Process noise was injected into the system by drawing from a Gaussian term dϵ ∼ N (0, 0.01dt× I).
Bursting neuron model The 3d biophysical bursting neuron model was introduced in (Durstewitz, 2009) and is defined by one voltage and two ion channel gating variables (one slow and one fast):
−CmV̇ = gL (V − EL) + gNam∞(V ) (V − ENa) + gKn (V − EK) + gMh (V − EK)
+ gNMDA [ 1 + .33e−.0625V ]−1 (V − ENMDA)
(15)
ḣ = h∞(V )− h
τh
ṅ = n∞(V )− n
τn
(16)
The limiting values of the ionic gates are given by
{m∞, n∞, h∞} = [ 1 + e({VhNa,VhK ,VhM}−V )/{kNa,kK ,kM} ]−1 . (17)
We borrowed parameter settings from Schmidt et al. (2021) to place the system into the burst-firing regime:
Cm = 6µF, gL = 8mS, EL = −80mV, gNa = 20mS, ENa = 60mV, VhNa = −20mV, kNa = 15, gK = 10mS, EK = −90mV, VhK = −25mV, kK = 5, τn = 1 ms, gM = 25mS
VhM = −15mV, kM = 5, τh = 200 ms, gNMDA = 10.2mS
Lorenz-96 system The Lorenz-96 is a high-dimensional, spatially extended weather model, also introduced by Edward Lorenz (Lorenz, 1996):
dxi dt = (xi+1 − xi−2)xi−1 − xi + F ++ dϵi(t) dt , i = 1 . . . N, (18)
with (constant) forcing term F . F = 8 is a common choice that leads to chaotic behavior. Process noise was added as for the Lorenz-63 system, dϵ ∼ N (0, 0.01dt× I). In our simulations we used N = 10, but in principle the system allows for arbitrary dimensionality.
Neural population model A larger-scale neural population model was recently introduced in Landau & Sompolinsky (2018) to examine the effect of structured connectivity on top of a randomly initialized network matrix. Specifically, an independently Gaussian distributed weight structure was combined with a rank-1 component with coupling strength J1. The dynamics of the single unit currents were defined as
dh dt = −h+ Jϕ(h) + J1√ N ξvTϕ(h), (19)
where ϕ(h) = tanh(h(t)). We produced a 50-dimensional chaotic network model based on the code provided in Landau & Sompolinsky (2018) using J1 = 0.09 and seeding the random number generator with 35.
The Lorenz-63 and Lorenz-96 systems were simulated using scipy.integrate, while for the bursting neuron and neural population model we used the code provided in Schmidt et al. (2021) and Landau & Sompolinsky (2018), respectively.

heading: 6.5 THEORETICAL ANALYSIS
text: Consider the PLRNN with linear spline basis expansion as defined by Eq. 1, Eq. 4, reproduced here for convenience:
zt =Azt−1 +W B∑ b=1 αb max(0, zt−1 − hb) + h0 +Cst + ϵt, (20)
where ϵt ∼ N(0,Σ), E[ϵt, ϵTt′ ] = 0 for t ̸= t′, αb ∈ R are scalar weighting factors and hb ∈ RM different ReLU “activation thresholds”, and all other parameters are as in conventional PLRNNs (Koppe et al., 2019).
Defining
D (b) Ω(t−1)(zt−1 − hb) := max(0, zt−1 − hb), (21)
Eq. 20 can be rewritten as
zt =
( A+W B∑ b=1 αb D (b) Ω(t−1) ) zt−1
+ W B∑ b=1 αb D (b) Ω(t−1) (−hb) + h0 + Cst + ϵt, (22)
where D(b)Ω(t−1) = diag ( d (b) 1,t−1, d (b) 2,t−1, · · · , d (b) M,t−1 ) are diagonal binary indicator matrices with d (b) m,t−1 = 1 if zm,t−1 > hm,b and 0 otherwise.
6.5.1 FIXED POINTS AND n-CYCLES OF SYSTEM EQ. 22
Defining
DBΩ(t−1) := B∑ b=1 αb D (b) Ω(t−1),
hBΩ(t−1) := B∑ b=1 αb D (b) Ω(t−1)(−hb), (23)
WBΩ(t−1) := A+W D B Ω(t−1),
and considering the autonomous system (i.e., without external inputs or noise terms), Eq. 22 can be rewritten as
zt = W B Ω(t−1) zt−1 +W h B Ω(t−1) + h0. (24)
Fixed points and cycles of Eq. 20, and their eigenvalue spectra, can now be computed in a way analogous to standard PLRNNs. Specifically, solving the equation F (z∗1) = z∗1, fixed points of the dendPLRNN are given by
z∗1 = ( I −WBΩ(t∗1) )−1[ W hBΩ(t∗1) + h0 ] , (25)
where z∗1 = zt∗1 = zt∗1−1, and det(I −WBΩ(t∗1)) = PWB Ω(t∗1) (1) ̸= 0, i.e. WBΩ(t∗1) has no eigenvalue equal to 1.
For n > 1, an n-cycle with periodic points {z∗n, F (z∗n), F 2(z∗n), · · · , Fn−1(z∗n)} of map F can be obtained by solving Fn(z∗n) = z∗n. Therefore, in order to find the periodic points, we first compute Fn in the following way:
zt = F (zt−1) = W B Ω(t−1) zt−1 +W h B Ω(t−1) + h0,
zt+1 = F 2(zt−1) = F (zt) = W B Ω(t) W B Ω(t−1) zt−1 + ( WBΩ(t) W h B Ω(t−1) + W h B Ω(t) ) + ( WBΩ(t) + I ) h0,
zt+2 = F 3(zt−1) = F (zt+1) = W B Ω(t+1)W B Ω(t)W B Ω(t−1) zt−1 + ( WBΩ(t+1)W B Ω(t)Wh B Ω(t−1)
+ WBΩ(t+1)Wh B Ω(t) +Wh B Ω(t+1) ) + ( WBΩ(t+1)W B Ω(t) + W B Ω(t+1) + I ) h0,
...
zt+(n−1) = F n(zt−1) = n+1∏ i=2 WBΩ(t+n−i) zt−1 + n∑ j=2 [ n−j+2∏ i=2 WBΩ(t+n−i) W h B Ω(t+j−3) ]
+ WhBΩ(t+n−2) + ( n∑ j=2 n−j+2∏ i=2 WBΩ(t+n−i) + I ) h0, (26)
where n+1∏ i=2 WBΩ(t+n−i) = W B Ω(t+n−2)W B Ω(t+n−3) · · ·W B Ω(t−1). Defining t+ n− 1 =: t∗n, the periodic point z∗n of the n-cycle of F can now be obtained as the fixed point of the n-times iterated map Fn as
z∗n = ( I − n∏ i=1 WBΩ(t∗n−i) )−1( n∑ j=2 [ n−j+1∏ i=1 WBΩ(t∗n−i)Wh B Ω(t∗n−n+j−2) ] + WhBΩ(t∗n−1)
+ ( n∑
j=2 n−j+1∏ i=1 WBΩ(t∗n−i) + I ) h0 ) , (27)
where z∗n = zt∗n = zt∗n−n, if (I − ∏n i=1 W B Ω(t∗n−i)) is invertible, i.e.
det ( I − n∏ i=1 WBΩ(t∗n−i) ) = P∏n i=1 W B Ω(t∗n−i) (1) ̸= 0,
which implies WΩ∗n := ∏n i=1 W B Ω(t∗n−i) has no eigenvalue equal to 1.
Remark 1. These results about fixed points and n-cycles also hold for the mean-centred dendPLRNN. This can easily be seen by defining WBΩ(t−1) := A+W D B Ω(t−1) M and noting that the elements of D(b)Ω(t−1) are now determined by the mean-centred latent states. That is d (b) m,t−1 = 1 if zm,t−1 −
1 M ∑M j=1 zj,t−1 > hm,b and 0 otherwise. The rest of the calculations then proceeds as above.

heading: 6.5.2 SUB-REGIONS AND DISCONTINUITY BOUNDARIES CORRESPONDING TO SYSTEM EQ. 22
text: Consider system Eq. 22 without external input and noise terms. Denoting hb = (h1,b, h2,b, · · · , hM,b)T in Eq. 22, for b = 1, 2, · · · , B, we can order the elements hj,1, hj,2, · · · , hj,B for every j ∈ {1, 2, · · · ,M}. Without loss of generality, let
hj,1 < hj,2 < · · · < hj,B , j = 1, 2, · · · ,M. (28)
Then, for every j, we define the intervals Ij,b as follows:
Ij,1 := (−∞, hj,1],
Ij,b := (hj,b−1, hj,b], b = 2, 3, · · · , B, (29)
Ij,B+1 := (hj,B ,+∞).
By definition of D(i)Ω(t−1) in Eq. 22, the phase space is separated into (B + 1) M sub-regions by MB(B+1)M−1 hyper-surfaces as discontinuity boundaries. Every sub-region can be defined by the thresholds hb as Cartesian product of suitable intervals in Eq. 29 for j ∈ {1, 2, · · · ,M}. (Note that if in Eq. 28 we had ” ≤ ” instead of strict inequalities ” < ”, obviously the number of intervals, hence sub-regions, would decrease.) In each sub-region the matrices D(b)Ω(t−1), b = 1, 2, · · · , B, have a different configuration. Therefore, in Eq. 24 there are (B + 1)M different forms for DBΩ(t−1), and so for WBΩ(t−1) and h B Ω(t−1) as well. Hence, indexing D B Ω(t−1), W B Ω(t−1) and h B Ω(t−1) as D B (r), W B (r) and hB(r) for r ∈ {1, 2, · · · , (B + 1) M}, Eq. 22 can be written as
zt = W B (r) zt−1 +W h B (r) + h0. (30)
To visualize the sub-regions and their borders, let for example M = 2 and B = 2. In this case there are 9 sub-regions divided by 12 borders. As illustrated in Fig. 6.5.2, there are different matrices D
(b) Ω(t−1), b = 1, 2, and D B (r) = D 2 (r), r = 1, 2, · · · , 9, for each sub-region.

heading: 6.5.3 BOUNDED ORBITS ARE COMPATIBLE WITH THE MANIFOLD ATTRACTOR REGULARIZATION
text: Proposition 2. The results of Theorem 2 are also true when the manifold-attractor regularization, Eq. 6, is strictly enforced for the dendPLRNN, Eq. 10.
Proof. Assume A, W , ϕ̃(zt−1) (see proof of Theorem 2 in Appx. 6.5.6 for the definition) and h0 have the partitioned forms
A =  Ireg OT O Anreg  , W =  Oreg OT S Wnreg  ,
Figure S6: Example of different sub-regions and related matrices D(b)Ω(t−1), b = 1, 2, and D B (r), r = 1, 2, · · · , 9, for M = 2 and B = 2. Here, it is assumed that the components of h1 = (h1,1, h2,1)T and h2 = (h1,2, h2,2)T satisfy Eq. 28 with ” < ”.
h0 =  hreg0 hnreg0  , ϕ̃(zt−1) =  ϕ̃reg(zt−1) ϕ̃nreg(zt−1)  , (31) where IMreg×Mreg =: Ireg ∈ RMreg×Mreg ,OMreg×Mreg =: Oreg ∈ RMreg×Mreg , O,S ∈ R(M−Mreg)×Mreg , the sub-matrices A{Mreg+1:M,Mreg+1:M} =: Anreg ∈ R(M−Mreg)×(M−Mreg) and W{Mreg+1:M,Mreg+1:M} =: Wnreg ∈ R(M−Mreg)×(M−Mreg) are diagonal and off-diagonal respectively. Furthermore, hreg0 , ϕ̃reg(zt−1) ∈ RMreg and h {Mreg+1:M,Mreg+1:M} 0 =: h nreg 0 , ϕ̃{Mreg+1:M,Mreg+1:M}(zt−1) =: ϕ̃nreg(zt−1) ∈ RM−Mreg .
In this case ∥A∥ = σmax(A) = max{1, σmax(Anreg)} and∥∥∥Aj W ϕ̃(zT−1−j)∥∥∥ = ∥∥∥∥∥∥  O
Ajneg S ϕ̃nreg(zt−1) +A j neg Wneg ϕ̃nreg(zt−1) ∥∥∥∥∥∥ = ∥∥∥Ajneg S ϕ̃nreg(zt−1) +Ajneg Wneg ϕ̃nreg(zt−1)∥∥∥ ,
∥∥Aj W h0∥∥ = ∥∥∥∥∥∥  O
Ajneg S h nreg 0 +A j neg Wneg h nreg 0 ∥∥∥∥∥∥ = ∥∥Ajneg S hnreg0 +Ajneg Wneg hnreg0 ∥∥ . (32)
Thus, for σmax(Anreg) < 1
∥zT ∥ ≤ ∥A∥T−1 ∥z1∥ + T−2∑ j=0 ∥∥∥Aj W ϕ̃(zT−1−j)∥∥∥+ T−2∑ j=0 ∥∥Aj h0∥∥
≤ ∥z1∥ + ( c̃+ ∥h0∥ )( ∥S∥+ ∥Wneg∥ ) T−2∑ j=0 ∥Aneg∥j
= ( c̃+ ∥h0∥ )( ∥S∥+ ∥Wneg∥ ) 1− ∥Aneg∥ < ∞. (33)

heading: 6.5.4 PROOF OF PROPOSITION 1
text: Proof. For A = (aij) ∈ RM×M , W = (wij) ∈ RM×M , ϵt = (ϵ1,t, ϵ2,t, · · · , ϵM,t)T, st = (s1,t, s2,t, · · · , sM,t)T and C = (cij) ∈ RM×M , writing Eq. 22 in scalar form yields
zl,t = M∑ j=1 aljzj,t−1 + M∑ j=1 wlj B∑ b=1 αb d (b) j,t−1[zj,t−1 − hj,b] + hl,0 + M∑ j=1 clj sj,t + ϵl,t
= M∑ j=1 ( aljzj,t−1 + wlj B∑ b=1 αb d (b) j,t−1[zj,t−1 − hj,b] ) + hl,0 + M∑ j=1 clj sj,t + ϵl,t
=: M∑ j=1 fl,j(zj,t−1) + hl,0 + M∑ j=1 clj sj,t + ϵl,t =: Fl(zt−1), l = 1, 2, · · · ,M. (34)
Using this, we can write Eq. 22 in the vector form zt = ( F1(zt−1), F2(zt−1), · · · , FM (zt−1) )T . (35)
We show that every Fl is continuous and so Eq. 22 is a continuous PWL map. For this purpose, by Eq. 34, it suffices to prove that every fl,j(zj,t−1) is continuous. According to the definition of the intervals Ij,b, Eq. 29, for any j ∈ {1, 2, · · · ,M} we have
zj,t−1 ∈ Ij,1 ⇒ d(b)j,t−1 = 0 ∀ b = 1, 2, · · · , B,
zj,t−1 ∈ Ij,s ⇒  d (b) j,t−1 = 1, b = 1, 2, · · · , s− 1 d (b) j,t−1 = 0, b = s, s+ 1, · · · , B
s = 2, 3, · · · , B,
zj,t−1 ∈ Ij,B+1 ⇒ d(b)j,t−1 = 1 ∀ b = 1, 2, · · · , B. (36)
Hence, for l, j = 1, 2, · · · ,M , each function fl,j(zj,t−1) can be stated as
fl,j(zj,t−1) =  f (1) l,j = alj zj,t−1; zj,t−1 ∈ Ij,1 f (2) l,j = (alj + α1 wlj) zj,t−1 − α1 wljhj,1; zj,t−1 ∈ Ij,2 ... f (B) l,j = (alj + wlj ∑B−1 b=1 αb) zj,t−1 − wlj ∑B−1 b=1 αb hj,b; zj,t−1 ∈ Ij,B
f (B+1) l,j = (alj + wlj ∑B i=1 αb) zj,t−1 − wlj ∑B b=1 αb hj,b; zj,t−1 ∈ Ij,B+1
(37)
Since for every b = 1, 2, · · · , B,
lim zj,t−1→hj,b
f (b) l,j (zj,t−1) = lim zj,t−1→hj,b f (b+1) l,j (zj,t−1) = f (b) l,j (hj,b), (38)
each function fl,j(zj,t−1) is continuous. Hence, Eq. 22 is a continuous PWL map in z (but has discontinuities in its Jacobian matrix across the borders). Because of these properties, all the results established for standard PLRNNs in Monfared & Durstewitz (2020a;b); Schmidt et al. (2021) apply to the dendPLRNN as well, only that the sub-regions and discontinuity boundaries are different.

heading: 6.5.5 PROOF OF PROPOSITION 1
text: Proof. Defining z̃t as B identical copies of zt,
z̃t =  z̃1,t z̃2,t ... z̃M,t z̃M+1,t ...
z̃BM,t
 :=  zt zt ... zt  BM×1
(39)
and likewise
h̃ =  h̃1 h̃2 ... h̃M h̃M+1
... h̃BM
 =  h1 h2 ... hB  BM×1 , h̃0 =  h̃0,1 h̃0,2 ... h̃0,M h̃0,M+1 ...
h̃0,BM
 =  h0 h0 ... h0  BM×1
ÃBM×BM = diag ( AM×M ,AM×M , · · · ,AM×M︸ ︷︷ ︸
B times
) ,
W̃BM×BM =  α1WM×M α2WM×M . . . αBWM×M α1WM×M α2WM×M . . . αBWM×M ... ... . . . ...
α1WM×M α2WM×M . . . αBWM×M
 ,
C̃st =  c̃s1,t c̃s2,t ... c̃sM,t c̃sM+1,t ...
c̃sBM,t
 =  Cst Cst ... Cst  BM×1 , ϵ̃t =  ϵ̃1,t ϵ̃2,t ... ϵ̃M,t ϵ̃M+1,t ...
ϵ̃BM,t
 =  ϵt ϵt ... ϵt  BM×1
(40)
one can rewrite the dendPLRNNfrom Eq. 20 as
z̃t = Ãz̃t−1 + W̃ max(0, z̃t−1 − h̃) + h̃0 + C̃st + ϵ̃t. (41)
Now performing the substitution
∀ t ẑt ← z̃t − h̃, (42)
Eq. 41 can be rewritten as the M ×B-dimensional “conventional” PLRNN Eq. 5 with
ĥ0 = ( Ã− I ) h̃ + h̃0. (43)

heading: 6.5.6 PROOF OF THEOREM 2
text: Proof. It can easily be shown that for every i ∈ {1, 2, · · · ,M}
αb [ max(max(0, zi,t−1 − hi,b)−max(0, zi,t−1) ] ∈
{ [−αbhib , 0] if sgn(αb) = sgn(hi,b)
[0, αbhi,b] else .
(44)
By defining
B∑ b=1 αb [ max(0, zt−1 − hb)−max(0, zt−1) ] := ϕ̃(zt−1) = ( ϕ̃1(zt−1), · · · , ϕ̃M (zt−1) )T ,
(45)
and
cupi,b =
{ 0 if sgn(αb) = sgn(hi,b)
αbhi,b else , clowi,b =
{ −αbhi,b if sgn(αb) = sgn(hi,b)
0 else , (46)
we can conclude that clowi ≤ ϕ̃i(zt−1) ≤ c up i , where clow/upi = ∑B b=1 c low/up i,b . For ci = max{|clowi |, |c up i |} we have
ϕ̃i(zt−1) 2 ≤ c2i ,
and so letting c = max{c1, c2, · · · , cM} yields ∥∥∥ϕ̃(zt−1)∥∥∥ = √√√√ M∑
i=1
( ϕ̃i(zt−1) )2 ≤ √√√√ M∑
i=1
c2 := c̃. (47)
Since
zt = Azt−1 + W ϕ̃(zt−1) + h0, (48)
for T ∈ N and t = 2, · · · , T , computing z2, z3, · · · , zT recursively leads to
z2 = Az1 + W ϕ̃(z1) + h0
z3 = A 2 z1 + AW ϕ̃(z1) +W ϕ̃(z2) + [ A+ I ] h0
...
zT = A T−1 z1 + T−2∑ j=0 Aj W ϕ̃(zT−1−j) + T−2∑ j=0 Aj h0. (49)
Therefore, by Eq. 47, for every T ≥ 2, we have
∥zT ∥ ≤ ∥A∥T−1 ∥z1∥ + c̃ ∥W ∥ T−2∑ j=0 ∥A∥j + T−2∑ j=0 ∥A∥j ∥h0∥ . (50)
If σmax(A) < 1, then lim T→∞
∥A∥T−1 = 0 and
lim T→∞ ∥zT ∥ ≤ c̃ ∥W ∥ ∞∑ j=0 ∥A∥j + ∞∑ j=0 ∥A∥j ∥h0∥ = c̃ ∥W ∥+ ∥h0∥ 1− ∥A∥ < ∞. (51)

Review summaries:
The paper proposes to modify the nonlinearity of neuron couplings in PLRNN with a linear spline basis expansion. This modification is claimed to boost PLRNN's capacity of capturing arbitrary nonlinear dynamics in low-dimensions. The paper also introduces two training frameworks, variational inference and BPTT. The proposed method is evaluated on various nonlinear dynamical systems.

This paper focuses on the nonlinear dynamical system on time series data. It proposed a piecewise-linear (PL) recurrent neural network (RNN) and attempted to couple it with dendritic computation, a computing paradigm from the neural computation discipline. The training method of the proposed model is through variational inference, and BPTT with teacher forcing. 

The authors proposed an extension of PLRNN with a dendrite-like formulation of the nonlinear activation function. They showed that the proposed dendPLRNN maintains the mathematical tractability and dynamical systems interpretation as PLRNN. They compared with alternatives for DS reconstruction with an overall better performance.