Tractable Dendritic RNNs for Identifying Unknown Nonlinear Dynamical Systems

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 ) . 3 MODEL FORMULATION AND THEORETICAL CONSIDERATIONS . 3.1 PIECEWISE LINEAR RECURRENT NEURAL NETWORK ( PLRNN ) . 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 . 3.2 DENDRITIC COMPUTATION AND SPLINE BASIS EXPANSION . 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 .

Summary: 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.