Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems

1 Introduction . Recurrent neural networks ( RNNs ) are a powerful and popular tool for modeling complex sequence data . They learn to transform input sequences into output sequences by using an internal state that allows data from the past to influence the current state . RNNs have been utilized in various applications such as speech recognition , sentiment analysis , music , and video [ 1–4 ] . In neuroscience , RNNs have been used for modeling large-scale neural recordings and as a generator of scientific 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . hypotheses by studying the network ’ s learned representations [ 5–10 ] . However , RNNs are generally viewed as black boxes . While there has been progress in understanding their operation on simple tasks , rigorously understanding how they solve complex tasks remains a significant challenge . An important line of work to improve our understanding of RNN computations uses dynamical systems theory [ 11–19 ] . In particular , Sussillo and Barak [ 15 ] proposed reverse engineering a trained RNN by using numerical optimization to find the RNN ’ s fixed and slow points . The RNN is linearized around these points , and the resulting linear approximation dynamics are studied to draw insights into how the RNN solves the task . While there are no guarantees concerning the success of this approach , empirically , linearization around fixed and slow points has led to insights in numerous applications [ 10 , 20–26 ] . There are a few drawbacks of this method . First , it requires a separate numerical optimization routine after training the network . Second , it can be ambiguous which fixed point to linearize around for any given point in state space . The standard fixed point finding numerical optimization routine provides a collection of fixed points but no direct link between these points and locations in state space . Finally , simulating nonlinear RNNs with linearizations around fixed points can slowly accumulate significant error , forcing previous attempts to resort to one-step ahead dynamics generation [ 21 ] . These problems can lead to uncertainty of how well switching between linearizations around fixed and slow points describes the nonlinear dynamics . Here , we combine ideas from reverse engineering RNNs and switching linear dynamical systems ( SLDS ) [ 27–31 ] to address these challenges . Given an RNN we would like to train and analyze , we introduce a separate network consisting of a novel SLDS formulation based on the first-order Taylor series expansion of the RNN equation and a learnable auxiliary function that produces the RNN ’ s fixed/slow points . We then define a loss function that includes regularization terms to force the SLDS to approximate the RNN and switch about the RNN ’ s fixed/slow points . After co-training these three functions with standard RNN training methods , the result is an accurate switching linear approximation of the nonlinear RNN and a trained auxiliary function that provides fixed and slow points of the RNN . This architecture and training procedure : 1 . Eliminates the need for post-training fixed point finding . 2 . Generalizes SLDS to be able to switch about continuous manifolds of fixed points . 3 . Enables parameter sharing between the SLDS switches . 4 . Allows the nonlinear RNN dynamics to be approximated by switching between local linearizations around fixed points , if possible . The combination of these benefits significantly simplifies the process of reverse engineering an RNN using fixed points . We illustrate the method on two synthetic tasks and a neural dataset . 2 Review of reverse engineering RNNs and SLDSs . 2.1 Reverse engineering RNNs with fixed points . The motivation for reverse engineering RNNs with fixed point analysis is the hypothesis that trained RNNs use mechanisms to solve tasks that are described well by the linearized dynamics around its ∗fixed and slow points [ 15 ] . These points are state vectors h∗ ∈ RD that , given an input u ∈ RU , do not significantly change when applying the RNN update function , F : RD × RU → RD . That is , h∗ ≈ F ( h∗ ∗ , u ; θ ) . We can find these points numerically by minimizing a loss function ∗ L ( h ) = kh − F ( h , u ; θ ) k22 , ( 1 ) using auto-differentiation methods [ 32 ] . In practice , one typically initializes the candidate fixed points with the hidden states produced from running forward passes given trial inputs . Once we have found the fixed/slow points , we linearize the system around these points and analyze the dynamics of the linearized system to determine how the system computes . This reverse engineering method is supported in theory by the Hartman-Grobman theorem [ 33–35 ] , which says that the dynamics of a nonlinear system in a domain near a hyperbolic fixed point is qualitatively the same as the dynamics of its linearization near this point . A potential theoretical issue arises with non-hyperbolic fixed points , though empirically , this has not seemed to pose a major issue for this approach . Finally , we note a recent alternative method to finding fixed points presented in [ 36 ] that makes use of mathematical objects called directional fibers . 2.2 Switching linear dynamical systems . Models based on a linear dynamical system ( LDS ) are often used to model multi-dimensional time series and lend themselves well to dynamical systems analyses . The basic LDS models time series data using a latent representation that follows linear dynamics . A switching LDS ( SLDS ) augments the basic model with discrete states that correspond to different linear dynamics . This allows a SLDS to break down complex , nonlinear time series into a sequence of simpler local linear dynamics . Explicitly , let yt ∈ RN denote the observation data at time t and let ht ∈ RD represent the corresponding continuous latent state . A SLDS models the expected observation value as E [ yt ] = g ( ht ) where g is a mapping from RD to RN . Given the discrete latent state zt ∈ { 1 , ... , K } and input ut ∈ RU , the SLDS continuous latent states follow linear dynamics , � ht ∼ N A ( zt ) ht−1 + V ( zt ) ut + b ( zt ) , Q ( zt ) . ( 2 ) ∈ RDxD The current discrete state determines the current linear dynamics A ( zt ) , input matrix V ( zt ) ∈ RDxU , bias term b ( zt ) ∈ RD and noise covariance Q ( zt ) ∈ RDxD . In the basic SLDS , a Markov transition matrix generally defines the discrete state switching probability . An extension to the SLDS model is the recurrent switching linear dynamical system ( RSLDS ) which allows the discrete state transition to depend on the previous continuous latent state [ 37–40 ] . This recurrent connection allows more expressivity in the model and corresponds to the idea that the current discrete state ( and therefore the current dynamical regime ) should depend on the current state space location . SLDS models can often offer a balance between interpretability and expressivity for many problems . For example , one can improve the expressivity by increasing the number of discrete states , but this can come at the cost of interpretability and an increase in learnable parameters . In addition , SLDS generally requires hyperparameter tuning to determine the optimal number of discrete states to use . 3 Jacobian Switching Linear Dynamical System . We now present the Jacobian Switching Linear Dynamical System ( JSLDS ) model and training procedure . It combines ideas from reverse engineering RNNs and SLDS to achieve automated fixed point finding and accurate SLDS approximations of nonlinear RNNs . 3.1 Motivation . Let F denote a nonlinear RNN with previous state ht−1 ∈ RD , input ut ∈ RU , and parameters θ . ∗Writing the RNN update equation and its first order Taylor series expansion around points h∗ and u we have ∂F ∂F∗ ∗ ∗ ht = F ( ht−1 , ut ; θ ) ≈ F ( h ∗ , u ; θ ) + ( h ∗ , u ; θ ) ( ht−1 − h ∗ ) + ( h ∗ , u ; θ ) ( ut − u ∗ ) . ∂h ∂u ( 3 ) ∗Here , ∂F ( h∗ , u ; θ ) ∈ RDxD is the recurrent Jacobian that determines the recurrent local dynamics ∂h ∗and ∂F ( h∗ , u ; θ ) ∈ RDxU is the input Jacobian that determines the system ’ s input sensitivity . If∂u ∗h∗ is a fixed point of F for a given u ∗ , then F ( h∗ , u ; θ ) = h∗ and equation ( 3 ) yields ∂F ∂F∗ ∗ ht − h ∗ ≈ ( h ∗ , u ; θ ) ( ht−1 − h ∗ ) + ( h ∗ , u ; θ ) ( ut − u ∗ ) . ( 4 ) ∂h ∂u Eq . ( 4 ) gives a LDS that locally approximates F around h∗ . In principle , if one knew how to select the correct fixed point and there was always a fixed point nearby , one could use eq . ( 4 ) to run a very accurate SLDS approximation of F by switching between fixed points as needed . Notice the fixed point h∗ indexes the recurrent Jacobian and input Jacobian in a manner analogous to how the discrete state zt indexes the matrices A ( zt ) and V ( zt ) in the SLDS of eq . ( 2 ) . In practice , selecting the correct fixed point in real-time is difficult , as described above , which leads to the approach presented in the next section . 3.2 The JSLDS model . Our approach is to co-train the RNN , F , with a novel SLDS formulation based on the Jacobian of F in the spirit of equation ( 4 ) . Specifically , we introduce a separate SLDS with its own hidden state ∗ at ∈ RD that switches around an expansion point e ∈ RD t ∂F ∂F∗ ∗ ∗ ∗ ∗ ∗ at − e = ( et , u ; θ ) ( at−1 − e ) + ( et , u ; θ ) ( ut − u ∗ ) . ( 5 ) t ∂h t ∂u ∗Note that eq . ( 5 ) shares its parameters θ with F ( eq . 3 ) , i.e . given e and u ∗ , the nonlinear RNN ’ s t ∗parameters θ determine the update matrices . We will generally take u to be either zero ( the average ∗value for examples ) or the value of a context-dependent static input . The goal is for e to approximate t the RNN ’ s fixed and slow points . To accomplish this , we supplement the SLDS with a nonlinear auxiliary function E ( the expansion network ) with separate learned parameters φ ∗ e = E ( at−1 ; φ ) . ( 6 ) t The expansion network E returns the learned expansion points and is co-trained with the nonlinear RNN and the JSLDS . Once trained , the goal is for this function to approximate the RNN ’ s fixed/slow points . For the experiments presented in this paper , we define E as a 2-layer multilayer perceptron ( MLP ) with the same dimension per layer as the state dimension of the RNN . We discuss other potential formulations for this network in Section A.2 in the Appendix . Eqs . ( 5-6 ) define the JSLDS model , which can be run forward in time independent of the original nonlinear RNN after training . Figure 1 illustrates the general idea . Given the previous state , at−1 , the ∗expansion network , E , uses this state to select the next expansion point , e ( eq . 6 ) . The system then t ∗updates the state by using e to compute the recurrent Jacobian and input Jacobian of the original t nonlinear RNN F ( eq . 5 ) . Assuming the expansion network has learned to find fixed/slow points , ∗switching between the points e corresponds to reverse engineering nonlinear RNNs using fixed t ∗points . The dependence of the expansion point e on the previous state of the network at−1 links to t the recurrent connection in RSLDS . Note that eq . ( 5 ) is intended to closely follow the dynamics of F , and we will enforce this in the training procedure . We will refer to the system comprised of eqs . ( 5-6 ) as the JSLDS . We will refer to the combination of the JSLDS and the co-trained RNN as the JSLDS-RNN system . While a limitation of our method is that it does not currently lend itself easily to a stochastic formulation like the standard SLDS , it does allow for a potential continuum of different switches using a constant number of parameters . It also automatically determines the number of switches required to solve the task instead of the hyperparameter tuning required to determine this in SLDS .

Summary: This study describes a new interpretable model of neural activity, consisting of a piecewise linear latent dynamical system coupled to a trained RNN. The manuscript outlines intuitions for the model and first applies it to artificial neural networks trained on two cognitive tasks, a working memory task and a context-dependent evidence integration task, showing that their JSLDS model recovers dynamical mechanisms already described in the literature with sets of linear dynamical systems. The authors also propose using the JSLDS model as a replacement for the complex generator RNN of the LFADS model, and apply this idea to neural recordings of a monkey performing a motor control task, showing that a simple linear dynamical system captures most features of this dataset.