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