INFERRING DYNAMICAL SYSTEMS WITH LONG-RANGE DEPENDENCIES THROUGH LINE ATTRACTOR REGU- LARIZATION
Year: 2019

Abstract:
Vanilla RNN with ReLU activation have a simple structure that is amenable to systematic dynamical systems analysis and interpretation, but they suffer from the exploding vs. vanishing gradients problem. Recent attempts to retain this simplicity while alleviating the gradient problem are based on proper initialization schemes or orthogonality/unitary constraints on the RNN’s recurrence matrix, which, however, comes with limitations to its expressive power with regards to dynamical systems phenomena like chaos or multi-stability. Here, we instead suggest a regularization scheme that pushes part of the RNN’s latent subspace toward a line attractor configuration that enables long short-term memory and arbitrarily slow time scales. We show that our approach excels on a number of benchmarks like the sequential MNIST or multiplication problems, and enables reconstruction of dynamical systems which harbor widely different time scales.

heading: 1 INTRODUCTION
text: Theories of complex systems in biology and physics are often formulated in terms of sets of stochastic differential or difference equations, i.e. as stochastic dynamical systems (DS). A long-standing desire is to retrieve these generating dynamical equations directly from observed time series data (Kantz & Schreiber, 2004). A variety of machine and deep learning methodologies toward this goal have been introduced in recent years (Chen et al., 2017; Champion et al., 2019; Jordan et al., 2019; Duncker et al., 2019; Ayed et al., 2019; Durstewitz, 2017; Koppe et al., 2019), many of them based on recurrent neural networks (RNN) which can universally approximate any DS (i.e., its flow field) under some mild conditions (Funahashi & Nakamura, 1993; Kimura & Nakano, 1998). However, vanilla RNN as often used in this context are well known for their problems in capturing long-term dependencies and slow time scales in the data (Hochreiter & Schmidhuber, 1997; Bengio et al., 1994). In DS terms, this is generally due to the fact that flexible information maintenance over long periods requires precise fine-tuning of model parameters toward ’line attractor’ configurations (Fig. 1), a concept first propagated in computational neuroscience for addressing animal performance in parametric working memory tasks (Seung, 1996; Seung et al., 2000; Durstewitz, 2003). Line attractors introduce directions of zero-flow into the model’s state space that enable long-term maintenance of arbitrary values (Fig. 1). Specially designed RNN architectures equipped with gating mechanisms and (linear) memory cells have been suggested for solving this issue (Hochreiter & Schmidhuber, 1997; Cho et al., 2014). However, from a DS perspective, simpler models that can more easily be analyzed and interpreted in DS terms, and for which more efficient inference algorithms exist that emphasize approximation of the true underlying DS would be preferable.
Recent solutions to the vanishing vs. exploding gradient problem attempt to retain the simplicity of vanilla RNN by initializing or constraining the recurrent weight matrix to be the identity (Le et al., 2015), orthogonal (Henaff et al., 2016; Helfrich et al., 2018) or unitary (Arjovsky et al., 2016). In this way, in a system including piecewise linear (PL) components like rectified-linear units (ReLU), line attractor dimensions are established from the start by construction or ensured throughout training by a specifically parameterized matrix decomposition. However, for many DS problems, line attractors instantiated by mere initialization procedures may be unstable and quickly dissolve during training. On the other hand, orthogonal or unitary constraints are too restrictive for reconstructing DS, and more generally from a computational perspective as well (Kerg et al., 2019): For instance, neither
chaotic behavior (that requires diverging directions) nor settings with multiple isolated fixed point or limit cycle attractors are possible.
Here we therefore suggest a different solution to the problem, by pushing (but not strictly enforcing) ReLU-based, piecewise-linear RNN (PLRNN) toward line attractor configurations along some (but not all) directions in state space. We achieve this by adding special regularization terms for a subset of RNN units to the loss function that promote such a configuration. We demonstrate that our approach outperforms, or is en par with, LSTM and other, initialization-based, methods on a number of ’classical’ machine learning benchmarks (Hochreiter & Schmidhuber, 1997). More importantly, we demonstrate that while with previous methods it was difficult to capture slow behavior in a DS that exhibits widely different time scales, our new regularization-supported inference efficiently captures all relevant time scales.

heading: 2 RELATED WORK
text: Long-range dependency problems in RNN. Error gradients in vanilla RNN trained by some form of gradient descent, like back-propagation through time (BPTT, Rumelhart et al. (1986)), tend to either explode or vanish due to the large product of derivative terms that results from recursive application of the chain rule over time steps (Hochreiter, 1991; Bengio et al., 1994; Hochreiter & Schmidhuber, 1997). Formally, RNN zt = Fθ (zt−1, st) are discrete time dynamical systems that tend to either converge, e.g. to fixed point or limit cycle attractors, or diverge (to infinity or as in chaotic systems) over time, unless parameters of the system are precisely tuned to create directions of zero-flow in the system’s state space (Fig. 1), called line attractors (Seung, 1996; Seung et al., 2000; Durstewitz, 2003). Convergence of the RNN in general implies vanishing and global divergence exploding gradients. To address this issue, RNN with gated memory cells have been specifically designed (Hochreiter & Schmidhuber, 1997; Cho et al., 2014), but these are complicated and tedious to analyze from a DS perspective. Recently, Le et al. (2015) observed that initialization of the recurrent weight matrix W to the identity in ReLU-based RNN may yield performance en par with LSTMs on standard machine learning benchmarks. For a ReLU with activity zt ≥ 0, zero bias and unit slope, this results in the identity mapping, hence a line attractor configuration. Talathi
& Vartak (2016) expanded on this idea by initializing the recurrence matrix such that its largest absolute eigenvalue is 1, arguing that this would leave other directions in the system’s state space free for computations other than memory maintenance. Later work enforced orthogonal (Henaff et al., 2016; Helfrich et al., 2018; Jing et al., 2019) or unitary (Arjovsky et al., 2016) constraints on the recurrent weight matrix during training. While this appears to yield long-term memory performance superior to that of LSTMs, these networks are limited in their computational power (Kerg et al., 2019). This may be a consequence of the fact that RNN with orthogonal recurrence matrix are quite restricted in the range of dynamical phenomena they can produce, e.g. chaotic attractors are not possible since diverging eigen-directions are disabled. Our approach therefore is to establish line attractors only along some but not all directions in state space, and to only push the RNN toward these configurations but not strictly enforce them, such that convergence or divergence of RNN dynamics is still possible. We furthermore implement these concepts through regularization terms in the loss functions, such that they are encouraged throughout training unlike when only established through initialization.
Dynamical systems reconstruction. From a natural science perspective, the goal of reconstructing the underlying DS fundamentally differs from building a system that ’merely’ yields good ahead predictions, as in DS reconstruction we require that the inferred model can freely reproduce (when no longer guided by the data) the underlying attractor geometries and state space properties (see section 3.5, Fig. S2; Kantz & Schreiber (2004)). Earlier work using RNN for DS identification (Roweis & Ghahramani, 2002; Yu et al., 2006) mainly focused on inferring the posterior over latent trajectories Z = {z1, . . . ,zT } given time series data X = {x1, . . . ,xT }, p(Z|X), and on ahead predictions (Lu et al., 2017), hence did not show that inferred models can generate the underlying attractor geometries on their own. Others (Trischler & D’Eleuterio, 2016; Brunton et al., 2016) attempt to approximate the flow field, obtained e.g. by numerical differentiation, directly through basis expansions or neural networks, but numerical derivatives are problematic for their high variance and other numerical issues (Raissi, 2018; Baydin et al., 2018; Chen et al., 2017). Some approaches assume the form of the DS equations basically to be given (Raissi, 2018; Gorbach et al., 2017) and focus on estimating the system’s latent states and parameters, rather than approximating an unknown DS based on the observed time series information alone. In many biological systems like the brain the intrinsic dynamics are highly stochastic with many noise sources, like probabilistic synaptic release (Stevens, 2003), such that models that do not explicitly account for dynamical process noise (Champion et al., 2019; Rudy et al., 2019) may be less suitable. Finally, some fully probabilistic models for DS reconstruction based on GRU (Fraccaro et al. (2016), cf. Jordan et al. (2019)), LSTM (Zheng et al., 2017), or radial basis function (Zhao & Park, 2017) networks are not easily interpretable and amenable to DS analysis. Most importantly, none of these previous approaches considers the long-range dependency problem within more easily tractable RNN for DS reconstruction.

heading: 3 MODEL FORMULATION AND OPTIMIZATION APPROACHES

heading: 3.1 MODEL AND PRELIMINARIES
text: Assume we are given two multivariate time series S = {st} and X = {xt}, one we will denote as ’inputs’ (S) and the other as ’outputs’ (X). We will first consider the ’classical’ (supervised) machine learning setting where we wish to map S on X through a RNN with latent state equation zt = Fθ (zt−1, st), as for instance in the ’addition problem’ (Hochreiter & Schmidhuber, 1997). In DS reconstruction, in contrast, we usually have a dense time series X from which we wish to infer (unsupervised) the underlying DS, where S may provide an additional forcing function or sparse experimental inputs or perturbations.
The latent RNN we consider here takes the specific form
zt = Azt−1 +Wφ(zt−1) +Cst + h+ εt, εt ∼ N (0,Σ), (1)
where zt ∈ RM×1 is the hidden state (column) vector of dimensionM ,A ∈ RM×M a diagonal and W ∈ RM×M an off-diagonal matrix, st ∈ RK×1 the external input of dimension K, C ∈ RM×K the input mapping, h ∈ RM×1 a bias, and εt a Gaussian noise term with diagonal covariance matrix diag(Σ) ∈ RM+ . The nonlinearity φ(z) is a ReLU, φ(z)i = max(0, zi), i ∈ {1, . . . ,M}. This specific formulation is originally motivated by firing rate (population) models in computational neuroscience (Song et al., 2016; Durstewitz, 2017), where latent states zt may represent
membrane voltages or currents, A the neurons’ passive time constants, W the synaptic coupling among neurons, and φ(·) the voltage-to-rate transfer function. However, for a RNN in the form zt = Wφ (zt−1) + h, note that the simple change of variables yt →W−1(zt − h) will yield the more familiar form yt = φ (Wyt−1 + h) (Beer, 2006).
Besides its neuroscience motivation, note that by letting A = I , W = 0, h = 0, we get a strict line attractor system across the variables’ whole support which we conjecture will be of advantage for establishing long short-term memory properties. Also we can solve for all of the system’s fixed points analytically by solving the equations z∗ = (I −A−WDΩ)−1 h, with DΩ as defined in Suppl. 7.1.2, and can determine their stability from the eigenvalues of matrixA+WDΩ. We could do the same for limit cycles, in principle, which are fixed points of the r-times iterated map F rθ , although practically the number of configurations to consider increases exponentially as 2M ·r. Finally, we remark that a discrete piecewise-linear system can, under certain conditions, be transformed into an equivalent continuous-time (ODE) piecewise-linear system ζ̇ = GΩ(ζ(t), s(t)) (Suppl. 7.1.2, Ozaki (2012)), in the sense that if ζ(t) = zt, then ζ(t + ∆t) = zt+1 after a defined time step ∆t. These are among the properties that make PLRNNs more amenable to rigorous DS analysis than other RNN formulations.
We will assume that the latent RNN states zt are coupled to the actual observations xt through a simple observation model of the form
xt = Bg(zt) + ηt, ηt ∼ N (0,Γ) (2)
in the case of real-valued observations xt ∈ RN×1, where B ∈ RN×M is a factor loading matrix and diag(Γ) ∈ RN+ the diagonal covariance matrix of the Gaussian observation noise, or
p̂i,t := p̂t (xi,t = 1) = ( eBi,:zt ) N∑ j=1 eBj,:zt −1 , (3) in the case of multi-categorical observations xi,t ∈ {0, 1}, ∑ i xi,t = 1.

heading: 3.2 REGULARIZATION APPROACH
text: We start from a similar idea as Le et al. (2015), who initialized RNN parameters such that it performs an identity mapping for zi,t ≥ 0. However, 1) we use a neuroscientifically motivated network architecture (eq. 1) that enables the identity mapping across the variables whole support, zi,t ∈ [−∞,+∞], 2) we encourage this mapping only for a subset Mreg ≤ M of units (Fig. S1), leaving others free to perform arbitrary computations, and 3) we stabilize this configuration throughout training by introducing a specific L2 regularization for parametersA,W , and h in eq. 1.
That way, we divide the units into two types, where the regularized units serve as a memory that tends to decay very slowly (depending on the size of the regularization term), while the remaining units maintain the flexibility to approximate any underlying DS, yet retaining the simplicity of the original RNN model (eq. 1). Specifically, the following penalty is added to the loss function (Fig. S1):
Lreg = τA Mreg∑ i=1 (Ai,i − 1)2 + τW Mreg∑ i=1 M∑ j=1 j 6=i W 2i,j + τh Mreg∑ i=1 h2i (4)
While this formulation allows us to trade off, for instance, the tendency toward a line attractor (A→ I , h→ 0) vs. the sensitivity to other units’ inputs (W → 0), for all experiments performed here a common value, τA = τW = τh = τ , was assumed for the three regularization factors.

heading: 3.3 OPTIMIZATION PROCEDURE FOR MACHINE LEARNING BENCHMARKS
text: For comparability with other approaches like LSTMs (Hochreiter & Schmidhuber, 1997) or iRNN (Le et al., 2015), we will assume that the latent state dynamics eq. 1 are deterministic (i.e., Σ = 0), will take g(zt) = zt and Γ = IN in eq. 2 (leading to an implicit Gaussian assumption with identity covariance matrix), and will use stochastic gradient descent (SGD) for training
to minimize the squared-error loss across R samples, L = ∑R n=1 ( x̂ (n) T − x (n) T )2 , between estimated and actual outputs for the addition and multiplication problems, and the cross entropy loss L = ∑R n=1 ( − ∑10 i=1 x (n) i,T log(p̂ (n) i,T ) )
for sequential MNIST, to which penalty eq. 4 was added for the regularized PLRNN (rPLRNN). We used the Adam algorithm (Kingma & Ba, 2014) from the PyTorch package (Paszke et al., 2017) with a learning rate of 0.001, a gradient clip parameter of 10, and batch size of 16. In all cases, SGD is stopped after 100 epochs and the fit with the lowest loss across all epochs is chosen.

heading: 3.4 OPTIMIZATION PROCEDURE FOR DYNAMICAL SYSTEMS RECONSTRUCTION
text: For DS reconstruction we request that the latent RNN approximates the true generating system of equations, which is a taller order than learning the mapping S →X or predicting future values in a time series (cf. sect. 3.5). This point has important implications for the design of models, inference algorithms and performance metrics if the primary goal is DS reconstruction rather than ’mere’ time series forecasting. In this context we consider the fully probabilistic, generative RNN eq. 1.
Together with eq. 2 (where we take g(zt) = φ(zt)) this gives the typical form of a nonlinear state space model (Durbin & Koopman, 2012) with observation and process noise. We solve for the parameters θ = {A,W ,C,h,Σ,B,Γ} by maximum likelihood, for which an efficient ExpectationMaximization (EM) algorithm has recently been suggested (Durstewitz, 2017; Koppe et al., 2019), which we will briefly summarize here. Since the involved integrals are not tractable, we start off from the evidence-lower bound (ELBO) to the log-likelihood which can be rewritten in various useful ways:
log p(X|θ) ≥ EZ∼q[log pθ(X,Z)] +H (q(Z|X)) = log p(X|θ)−DKL (q(Z|X)‖pθ(Z|X)) =: L (θ, q) (5)
In the E-step, given a current estimate θ∗ for the parameters, we seek to determine the posterior pθ (Z|X) which we approximate by a global Gaussian q(Z|X) instantiated by the maximizer (mode) Z∗ of pθ(Z|X) as an estimator of the mean, and the negative inverse Hessian around this maximizer as an estimator of the state covariance, i.e.
E[Z|X] ≈ Z∗ = arg max Z log pθ(Z|X) = arg max Z [log pθ(X|Z) + log pθ(Z)− log pθ(X)]
= arg max Z
[log pθ(X|Z) + log pθ(Z)] , (6)
since Z integrates out in pθ(X) (equivalently, this result can be derived from a Laplace approximation to the log-likelihood, log p(X|θ) ≈ log pθ(X|Z∗)+log pθ(Z∗)− 12 log |−L
∗|+const, where L∗ is the Hessian evaluated at the maximizer). We solve this optimization problem by a fixed-point iteration scheme that efficiently exploits the model’s piecewise linear structure (see Suppl. 7.1.3, Durstewitz (2017); Koppe et al. (2019)).
Using this approximate posterior for pθ(Z|X), based on the model’s piecewise-linear structure most of the expectation values Ez∼q [φ(z)], Ez∼q [φ(z)zᵀ], and Ez∼q [φ(z)φ(z)ᵀ], could be solved for (semi-)analytically (where z is the concatenated vector form of Z, as in Suppl. 7.1.3). In the Mstep, we seek θ∗ := arg max
θ L(θ, q∗), assuming proposal density q∗ to be given from the E-step,
which for a Gaussian observation model amounts to a simple linear regression problem (see Suppl. eq. 23). To force the PLRNN to really capture the underlying DS in its governing equations, we use a previously suggested (Koppe et al. 2019) stepwise annealing protocol that gradually shifts the burden of fitting the observations X from the observation model eq. 2 to the latent RNN model eq. 1 during training, the idea of which is to establish a mapping from latent statesZ to observations X first, fixing this, and then enforcing the temporal consistency constraints implied by eq. 1 while accounting for the actual observations.

heading: 3.5 PERFORMANCE MEASURES
text: Measures of prediction error. For the machine learning benchmarks we employed the same criteria as used for optimization (MSE or cross-entropy, sect. 3.3) as performance metrics, evaluated across left-out test sets. In addition, we report the relative frequency Pcorrect of correctly predicted trials
across the test set. A correct trial in the addition and multiplication task is defined as an absolute prediction error smaller than 0.04 (analogous to Talathi & Vartak (2016)), while a correct trial in the sequential MNIST data set is defined as one for which the largest probability p̂i∗ = max
i p̂i,T
indicated the correct class xi∗,T = 1.
Agreement in attractor geometries. From a DS perspective, it is not sufficient or even sensible to judge a method’s ability to infer the underlying DS purely based on some form of (ahead-)prediction error like the MSE defined on the time series itself (Ch.12 in Kantz & Schreiber (2004)). Rather, we require that the inferred model can freely reproduce (when no longer guided by the data) the underlying attractor geometries and state space properties. This is not automatically guaranteed for a model that yields agreeable ahead predictions on a time series. Vice versa, if the underlying attractor is chaotic, with a tiny bit of noise even trajectories starting from the same initial condition will quickly diverge and ahead-prediction errors are not even meaningful as a performance metric (Fig. S2A). To quantify how well an inferred PLRNN captured the underlying dynamics we therefore followed Koppe et al. (2019) and used the Kullback-Leibler divergence between the true and reproduced probability distributions across states in state space, thus assessing the agreement in attractor geometries (cf. Takens (1981); Sauer et al. (1991)) rather than in precise matching of time series,
DKL (ptrue(x)‖pgen(x|z)) ≈ K∑ k=1 p̂ (k) true(x) log
( p̂
(k) true(x)
p̂ (k) gen(x|z)
) , (7)
where ptrue(x) is the true distribution of observations across state space (not time!), pgen(x|z) is the distribution of observations generated by running the inferred PLRNN, and the sum indicates a spatial discretization (binning) of the observed state space (see Suppl. 7.1.4 for more details). We emphasize that p̂(k)gen(x|z) is obtained from freely simulated trajectories, i.e. drawn from the prior p̂(z), not from the inferred posteriors p̂(z|xtrain). (The form of p̂(z) is given by the dynamical model eq. 1 and has a ’mixture of piecewise-Gaussians’ structure, see Koppe et al. (2019).) In addition, to assess reproduction of time scales by the inferred PLRNN, we computed the average correlation between the power spectra of the true and generated time series.

heading: 4 NUMERICAL EXPERIMENTS

heading: 4.1 MACHINE LEARNING BENCHMARKS
text: We compared the performance of our rPLRNN to other models on the following three benchmarks requiring long short-term maintenance of information (as in Talathi & Vartak (2016) and Hochreiter & Schmidhuber (1997)): 1) The addition problem of time length T consists of 100 000 training and 10 000 test samples of 2 × T input series S = {s1, . . . , sT }, where entries s1,: ∈ [0, 1] are drawn from a uniform random distribution and s2,: ∈ {0, 1} contains zeros except for two indicator bits placed randomly at times t1 < 10 and t2 < T/2. Constraints on t1 and t2 are chosen such that every trial requires a long memory of at least T/2 time steps. At the last time step T , the target output of the network is the sum of the two inputs in s1,: indicated by the 1-entries in s2,:, x target T = s1,t1 + s1,t2 . 2) The multiplication problem is the same as the addition problem, only that the product instead of the sum has to be produced by the RNN as an output at time T , xtargetT = s1,t1 ·s1,t2 . 3) The MNIST dataset (LeCun & Cortes, 2010) consists of 60 000 training and 10 000 28× 28 test images of hand written digits. To make this a time series problem, in sequential MNIST the images are presented sequentially, pixel-by-pixel, scanning lines from upper left to bottom-right, resulting in time series of fixed length T = 784.
On all three benchmarks we compare the performance of the rPLRNN (eq. 1) to several other models summarized in Table 1. To achieve a meaningful comparison, all models have the same number of hidden states M , except for the LSTM, which requires three additional parameters for each hidden state and hence has only M/4 hidden states, yielding the overall same number of trainable parameters as for the other models. In all cases, M = 40, which initial numerical exploration suggested to be a good compromise between model complexity (bias) and data fit (variance) (Fig. S3).
Fig. 2 summarizes the results for the machine learning benchmarks. As can be seen, on the addition and multiplication tasks, and in terms of either the MSE or percentage correct, our rPLRNN outperforms all other tested methods, including LSTMs. Indeed, the LSTM performs even significantly worse than the iRNN and the iPLRNN. The large error bars in Fig. 2 result from the fact that the networks mostly learn these tasks in an all-or-none fashion, i.e. either learn the task and succeed in almost 100 percent of the cases or fail completely. The results for the sequential MNIST problem are summarized in Fig. 2C. While in this case the LSTM outperforms all other methods, the rPLRNN is almost en par with it. In addition, the iPLRNN outperforms the iRNN. Similar results were obtained for M = 100 units (M = 25, respectively, for LSTM; Fig. S6). While the rPLRNN in general outperformed the pure initialization-based models (iRNN, npRNN, iPLRNN), confirming that a line attractor subspace present at initialization may be lost throughout training, we conjecture that this difference in performance will become even more pronounced as noise levels or task complexity increase.

heading: 4.2 NUMERICAL EXPERIMENTS ON A DYNAMICAL SYSTEM WITH DIFFERENT TIME SCALES
text: Here our goal was to examine whether our regularization approach would also help with the identification of DS that harbor widely different time scales. By tuning systems in the vicinity of line attractors, multiple arbitrary time scales can be realized in theory (Durstewitz, 2003). To test this, we used a biophysically motivated (highly nonlinear) bursting cortical neuron model with one voltage and two conductance recovery variables (see Durstewitz (2009)), one slow and one fast (Suppl. 7.1.5). Reproduction of this DS is challenging since it produces very fast spikes on top of a slow nonlinear oscillation (Fig. 3D). Time series of standardized variables of length T = 1500 were generated from this model and provided as observations to the rPLRNN inference algorithm. rPLRNNs with M = {8 . . . 18} states were estimated, with the regularization factor varied within τ ∈ {0, 101, 102, 103, 104, 105}/1500. Fig. 3A confirms our intuition that stronger regularization leads to better DS reconstruction as assessed by the KL divergence between true and generated state distributions. This decrease in DKL is accompanied by a likewise decrease in the MSE between the power spectra of true (Suppl. eq. 27) and generated (rPLRNN) voltage traces as τ increased (Fig. 3B). Fig. 3D gives an example of voltage traces and gating variables freely simulated (i.e., sampled) from the generative rPLRNN trained with τ = 23 , illustrating that our model is in principle capable of capturing both the stiff spike dynamics and the slower oscillations in the second gating variable at the same time. Fig. 3C provides more insight into how the regularization worked: While the high frequency components (> 50 Hz) related to the repetitive spiking activity hardly benefitted from increasing τ , there was a strong reduction in the MSE computed on the power spectrum for the lower frequency range (≤ 50 Hz), suggesting that increased regularization helps to map slowly evolving components of the dynamics.

heading: 5 CONCLUSIONS
text: In this work we have introduced a simple solution to the long short-term memory problem in RNN that on the one hand retains the simplicity and tractability of vanilla RNN, yet on the other hand does not curtail the universal computational capabilities of RNN (Koiran et al., 1994; Siegelmann & Sontag, 1995) and their ability to approximate arbitrary DS (Funahashi & Nakamura, 1993; Kimura & Nakano, 1998; Trischler & D’Eleuterio, 2016). We achieved this by adding regularization terms to the loss function that encourage the system to form a ’memory subspace’, that is, line attractor
dimensions (Seung, 1996; Durstewitz, 2003) which would store arbitrary values for, if unperturbed, arbitrarily long periods. At the same time we did not rigorously enforce this constraint which has important implications for capturing slow time scales in the data: It allows the RNN to slightly depart from a perfect line attractor, which has been shown to constitute a general dynamical mechanism for regulating the speed of flow and thus the learning of arbitrary time constants that are not naturally included qua RNN design (Durstewitz, 2003; 2004). This is because as we come infinitesimally close to a line attractor and thus a bifurcation in the system’s parameter space, the flow along this direction becomes arbitrarily slow until it vanishes completely in the line attractor configuration (Fig. 1). Moreover, part of the RNN’s latent space was not regularized at all, leaving the system enough degrees of freedom for realizing arbitrary computations or dynamics. We showed that the rPLRNN is en par with or outperforms initialization-based approaches and LSTMs on a number of classical benchmarks, and, more importantly, that the regularization strongly facilitates the identification of challenging DS with widely different time scales in PLRNN-based algorithms for DS reconstruction. Future work will explore a wider range of DS models and empirical data with diverse temporal and dynamical phenomena. Another future direction may be to replace the EM algorithm by black-box variational inference, using the re-parameterization trick for gradient descent (Kingma & Welling, 2013; Rezende et al., 2014; Chung et al., 2015). While this would come with better scaling in M , the number of latent states (the scaling in T is linear for EM as well, see Paninski et al. (2010)), the EM used here efficiently exploits the model’s piecewise linear structure in finding the posterior over latent states and computing the parameters (see Suppl. 7.1.3). It may thus be more accurate and suitable for smaller-scale problems where high precision is required, as often encountered in neuroscience or physics.

heading: 6 ACKNOWLEDGEMENTS
text: This work was funded by grants from the German Research Foundation (DFG) to DD (Du 354/10-1, Du 354/8-2 within SPP 1665) and to GK (TRR265: A06 & B08). We would like to cordially thank Dr. Zahra Monfared for her careful reading of the manuscript and her thoughtful suggestions.

heading: 7 SUPPLEMENTARY MATERIAL

heading: 7.1 SUPPLEMENTARY TEXT
text: 7.1.1 Simple exact PLRNN solution for addition problem
The exact PLRNN parameter settings (cf. eq. 1) for solving the addition problem with 2 units (cf. Fig. 1C) are as follows:
A = ( 1 0 0 0 ) ,W = ( 0 1 0 0 ) ,h = ( 0 −1 ) ,C = ( 0 0 1 1 ) ,B = (1 0) (8)
7.1.2 Conversion from discrete to continuous time PLRNN
Under some conditions we can translate the discrete into an equivalent continuous time PLRNN. Using DΩ(t) as defined below (7.1.3) for a single time step t, we can rewrite (ignoring the noise term and inputs) PLRNN eq. 1 in the form
zt+1 = F (zt) = WΩ(t)zt + h, withWΩ(t) := A+WDΩ(t), (9)
where Ω(t) := {m|zm,t > 0} is the set of all unit indices with activation larger 0 at time t. To convert this into an equivalent (in the sense defined in eq. 11) system of (piecewise) ordinary differential equations (ODE), we need to find parameters W̃Ω and h̃,
ζ̇ = G(ζ) = W̃Ωζ(t) + h̃, (10)
such that z0 = ζ(0)⇒ z1 = F (z0) = ζ(∆t), (11)
where ∆t is the time step with which the empirically observed time series X was sampled. From these conditions it follows that for each of the s ∈ {1, . . . , 2M} subregions (orthants) defined by fixed index sets Ωs ⊆ {1, . . . ,M} we must have
(A+WDΩs − I) z0 + h = ∫ ∆t
0
W̃Ωsζ(t) + h̃ dt, (12)
where we assume that DΩs is constant for one time step, i.e. between 0 and ∆t. We approach this by first solving the homogeneous system using the general ansatz for systems of linear ODEs,
(A+WDΩs − I) z0 ! = ∫ ∆t 0 W̃Ωs ∑ k cke λ̃ktṽk dt (13)
= ∑ k ckW̃Ωsvk ∫ ∆t 0 eλ̃kt dt (14)
= ∑ k ckλ̃kvk 1 λ̃k ( eλ̃k∆t − 1 ) (15)
⇒WΩsz0 ! = ∑ k ckvke λ̃k∆t (16)
= V e λ̃1∆t . . . 0 ... . . .
... 0 . . . eλ̃M∆t  ︸ ︷︷ ︸
:=Λ
c (17)
⇒WΩs = V ΛV −1. (18)
where we have used z0 = ∑ k ckvk on lines 15 and 16. Hence we can infer matrix W̃Ωs from the eigendecomposition of matrix WΩs , by letting λ̃k = 1∆t log λk, where λk are the eigenvalues of WΩs , and reassembling
W̃Ωs = V 1
∆t log(Λ)V −1. (19)
We obtain the general solution for the inhomogeneous case by requiring that for all fixed points z∗ = F (z∗) of the map eq. 9 we have G(z∗) = 0. Using this we obtain
h̃ = −W̃Ωs (I −WΩs)−1 h (20)
Assuming inputs st to be constant across time step ∆t, we can apply the same transformation to input matrix C. Fig. S5 illustrates the discrete to continuous PLRNN conversion for a nonlinear oscillator.
Note that in the above derivations we have assumed that matrix WΩs can be diagonalized, and that all its eigenvalues are nonzero (in fact, WΩs should not have any negative real eigenvalues). In general, not every discrete-time PLRNN can be converted into a continuous-time ODE system in the sense defined above. For instance, we can have chaos in a 1d nonlinear map, while we need at least a 3d ODE system to create chaos (Strogatz, 2015).
7.1.3 More details on EM algorithm
Here we briefly outline the fixed-point-iteration algorithm for solving the maximization problem in eq. 6 (for more details see Durstewitz (2017); Koppe et al. (2019)). Given a Gaussian latent PLRNN and a Gaussian observation model, the joint density p(X,Z) will be piecewise Gaussian, hence eq. 6 piecewise quadratic in Z. Let us concatenate all state variables across m and t into one long column vector z = (z1,1, . . . , zM,1, . . . , z1,T , . . . , zM,T )
ᵀ, arrange matrices A, W into large MT ×MT block tri-diagonal matrices, define dΩ := ( 1z1,1>0,1z2,1>0, . . . ,1zM,T>0 )ᵀ as an indicator vector with a 1 for all states zm,t > 0 and zeros otherwise, and DΩ := diag(dΩ) as the diagonal matrix formed from this vector. Collecting all terms quadratic, linear, or constant in z, we can then write down the optimization criterion in the form
Q∗Ω(z) = − 1
2 [zᵀ (U0 +DΩU1 +U ᵀ 1DΩ +DΩU2DΩ) z − zᵀ (v0 +DΩv1)
− (v0 +DΩv1)ᵀ z] + const. (21)
In essence, the algorithm now iterates between the two steps:
1. Given fixedDΩ, solve z∗ = (U0 +DΩU1 +U ᵀ 1DΩ +DΩU2DΩ) −1 (v0 +DΩv1) 2. Given fixed z∗, recomputeDΩ
until either convergence or one of several stopping criteria (partly likelihood-based, partly to avoid loops) is reached. The solution may afterwards be refined by one quadratic programming step. Numerical experiments showed this algorithm to be very fast and efficient (Durstewitz, 2017; Koppe et al., 2019). At z∗, an estimate of the state covariance is then obtained as the inverse negative Hessian,
V = (U0 +DΩU1 +U ᵀ 1DΩ +DΩU2DΩ) −1 . (22)
In the M-step, using the proposal density q∗ from the E-step, the solution to the maximization problem θ∗ := arg max
θ L(θ, q∗), can generally be expressed in the form
θ∗ = (∑ t E [αtβᵀt ] )(∑ t E [βtβᵀt ] )−1 , (23)
where, for the latent model, eq. 1, αt = zt and βt := [ zᵀt−1, φ(zt−1) ᵀ, sᵀt , 1 ]ᵀ ∈ R2M+K+1, and for the observation model, eq. 2, αt = xt and βt = g (zt).
7.1.4 More details on DS performance measure
The measure DKL introduced in the main text for assessing the agreement in attractor geometries only works for situations where the ground truth ptrue(X) is known. Following Koppe et al. (2019), here we would like to briefly indicate how a proxy for DKL may be obtained in empirical situations where no ground truth is available. Reasoning that for a well reconstructed DS the inferred posterior pinf(z|x) given the observations should be a good representative of the prior generative
dynamics pgen(z), one may use the Kullback-Leibler divergence between the distribution over latent states, obtained by sampling from the prior density pgen(z), and the (data-constrained) posterior distribution pinf(z|x) (where z ∈ RM×1 and x ∈ RN×1), taken across the system’s state space:
DKL (pinf(z|x)‖pgen(z)) = ∫ z∈RM×1 pinf(z|x) log pinf(z|x) pgen(z) dz (24)
As evaluating this integral is difficult, one could further approximate pinf(z|x) and pgen(z) by Gaussian mixtures across trajectories, i.e. pinf(z|x) ≈ 1T ∑T t=1 p(zt|x1:T ) and pgen(z) ≈ 1 L ∑L l=1 p(zl|zl−1), where the mean and covariance of p(zt|x1:T ) and p(zl|zl−1) are obtained by marginalizing over the multivariate distributions p(Z|X) and pgen(Z), respectively, yielding E[zt|x1:T ], E[zl|zl−1], and covariance matrices Var(zt|x1:T ) and Var(zl|zl−1). Supplementary eq. 24 may then be numerically approximated through Monte Carlo sampling (Hershey & Olsen, 2007) by
DKL (pinf(z|x)‖pgen(z)) ≈ 1
n n∑ i=1 log pinf(z (i)|x) pgen(z(i)) , z(i) ∼ pinf(z|x) (25)
For high-dimensional state spaces, for which MC sampling becomes challenging, there is luckily a variational approximation of eq. 24 available (Hershey & Olsen, 2007):
DvariationalKL (pinf(z|x)‖pgen(z)) ≈ 1
T T∑ t=1 log
∑T j=1 e
−DKL(p(zt|x1:T )‖p(zj |x1:T ))∑T k=1 e −DKL(p(zt|x1:T )‖p(zk|zk−1)) , (26)
where the KL divergences in the exponentials are among Gaussians for which we have an analytical expression.
7.1.5 More details on single neuron model
The neuron model used in section 4.2 is described by
−CmV̇ = gL(V − EL) + gNam∞(V )(V − ENa) + gKn(V − EK) + gMh(V − EK) + gNMDAσ(V )(V − ENMDA) (27)
ḣ = h∞(V )− h
τh (28)
ṅ = n∞(V )− n
τn (29) σ(V ) = [ 1 + .33e−.0625V ]−1 (30)
whereCm refers to the neuron’s membrane capacitance, the g• to different membrane conductances, E• to the respective reversal potentials, and m, h, and n are gating variables with limiting values given by
{m∞, n∞, h∞} = [ 1 + e({VhNa,VhK ,VhM}−V )/{kNa,kK ,kM} ]−1 (31)
Different parameter settings in this model lead to different dynamical phenomena, including regular spiking, slow bursting or chaos (see Durstewitz (2009) for details). Parameter settings used here were: Cm = 6 µF, gL = 8 mS, EL = −80 mV, gNa = 20 mS, ENa = 60 mV, VhNa = −20 mV, kNa = 15, gK = 10 mS, EK = −90 mV, VhK = −25 mV, kK = 5, τn = 1 ms, gM = 25 mS, VhM = −15 mV, kM = 5, τh = 200 ms, gNMDA = 10.2 mS.

heading: 7.2 SUPPLEMENTARY FIGURES
text: A W h Fig. S1: Illustration of the L2-regularization for the PLRNN’s auto-regression matrix A, coupling matrix W , and bias terms h. Regularized values are indicated in red, crosses mark arbitrary values (all other values set to 0 as indicated).
Fig. S2: MSE evaluated between time series is not a good measure for DS reconstruction. A) Time graph (top) and state space (bottom) for the single neuron model (see section 4.2 and Suppl. 7.1.5) with parameters in the chaotic regime (blue curves) and with simple fixed point dynamics in the limit (red line). Although the system has vastly different limiting behaviors (attractor geometries) in these two cases, as visualized in the state space, the agreement in time series initially seems to indicate a perfect fit. B) Same as in A) for two trajectories drawn from exactly the same DS (i.e., same parameters) with slightly different initial conditions. Despite identical dynamics, the trajectories immediately diverge, resulting in a high MSE. Dash-dotted grey lines in top graphs indicate the point from which onward the state space trajectories were depicted.
Fig. S3: Performance of the rPLRNN on the addition problem for different A) numbers of latent states M , B) values of τ and C) proportions Mreg/M . Dashed lines denote the values used for the results reported in section 4.1
Fig. S4: Effect of regularization strength τ on rPLRNN network parameters (cf. eq. 1) (regularized parameters for states m ≤ Mreg, eq. 1, in red). Note that some of the non-regularized network parameters (in blue) appear to systematically change as well as τ is varied.
Fig. S5: Illustration of conversion of discrete into continuous time PLRNN for a PLRNN emulation of the nonlinear van-der-Pol oscillator. Shown are the first two latent dimensions. Red lines: continuous solution; blue circles: discrete solution; black bars: perturbations (external inputs).
Fig. S6: Comparison of rPLRNN (τ = 5, MregM = 0.5, cf. Fig. S3) for M = 40 and M = 100 to iRNN (M = 100) and LSTM (M = 100/4) for A) addition problem, B) multiplication problem and C) sequential MNIST. Top row gives loss as a function of time series length T (error bars = SEM), bottom row shows relative frequency of correct trials. Dashed lines indicate chance level, black dots in C indicate individual repetitions of the experiment. Note that the rPLRNN does not improve for M = 100 vs. M = 40.

Review summaries:
This paper proposes a type of regularization for recurrent networks, with the goal of encouraging particular dynamical structures (in this case, line attractors) in the dynamics of the networks. The proposed regularization penalty is only applied to a subset of the recurrent units; the motivation for this is to allow neurons not contained in the subset to learn different structures. The paper applies this regularization method on three example machine learning sequence tasks: an addition task, a multiplication task, and sequential MNIST classification, as well as on learning a 2-D dynamical system model of a bursting neuron with two different timescales.