Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems
Authors: Jimmy T.H. Smith
Scott Linderman
David Sussillo
Venue: NeurIPS 2021 Poster (2021)
Keywords: recurrent neural networks
RNNs
switching linear dynamical systems
SLDS
interpretability
dynamical systems
reverse engineering
fixed points

Abstract: Recurrent neural networks (RNNs) are powerful models for processing time-series data, but it remains challenging to understand how they function. Improving this understanding is of substantial interest to both the machine learning and neuroscience communities. The framework of reverse engineering a trained RNN by linearizing around its fixed points has provided insight, but the approach has significant challenges. These include difficulty choosing which fixed point to expand around when studying RNN dynamics and error accumulation when reconstructing the nonlinear dynamics with the linearized dynamics. We present a new model that overcomes these limitations by co-training an RNN with a novel switching linear dynamical system (SLDS) formulation. A first-order Taylor series expansion of the co-trained RNN and an auxiliary function trained to pick out the RNN's fixed points govern the SLDS dynamics. The results are a trained SLDS variant that closely approximates the RNN, an auxiliary function that can produce a fixed point for each point in state-space, and a trained nonlinear RNN whose dynamics have been regularized such that its first-order terms perform the computation, if possible. This model removes the post-training fixed point optimization and allows us to unambiguously study the learned dynamics of the SLDS at any point in state-space.  It also generalizes SLDS models to continuous manifolds of switching points while sharing parameters across switches. We validate the utility of the model on two synthetic tasks relevant to previous work reverse engineering RNNs. We then show that our model can be used as a drop-in in more complex architectures, such as LFADS, and apply this LFADS hybrid to analyze single-trial spiking activity from the motor system of a non-human primate.

Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems
Jimmy T.H. Smith
Institute for Computational and Mathematical Engineering Stanford University Stanford, CA 94305 jsmi Somithh1@@taaffoot
Scott W. Linderman
David Sussillo
Department of Statistics Stanford University
Department of Electrical Engineering Stanford University Stanford, CA 94305
Stanford, CA 94305
scott. www.liindem@@taaan. edu
sussillo@stanf edu
Abstract
Recurrent neural networks (RNNs) are powerful models for processing time-series data, but it remains challenging to understand how they function. Improving this understanding is of substantial interest to both the machine learning and neuro- science communities. The framework reverse engineering a trained RNN by linearizing around its fixed points has provided insight, but the approach has signif- icant challenges. These include difficulty choosing which fixed point to expand around when studying RNN dynamics and error accumulation when reconstructing the nonlinear dynamics with the linearized dynamics. We present a new model that overcomes these limitations by co-training an RNN with novel switching linear dynamical system (SLDS) formulation. first-order Taylor series expansion the co-trained RNN and an auxiliary function trained pick out the RNN's fixed points govern the SLDS dynamics. The results are a trained SLDS variant that closely approximates the RNN, an auxiliary function that can produce fixed point for each point in state-space, and a trained nonlinear RNN whose dynamics have been regularized such that its first-order terms perform the computation, if possible. This model removes the post- testntrainig fixed point optimization and allows us to unambiguously study the learned dynamies of the SL.DS at any point in state-space. It also generalizes SLDSS models to continuous manifolds of switching points while sharing parameters across switches. We validate the utility of the model synthetic tasks relevant to previous work reverse engineering RNNs. We then show that our model be used as a drop-in more complex architectures, such LFADS, and apply this LFADS hybrid to analyze single-trial spiking activity from the motor system of non-human primate.
Introduction
Recurrent neural networks (RNNs) are 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 I In neuroscience. RNNs have been used for modeling large-scale neural recordings and as generator of scientific
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

hypotheses by studying the network's 's learned representations [0 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 significant challenge.
An important line of work to improve understanding of RNN computations uses dynamical systems theory (1HH1)). In particular, Sussillo and Barak [15] proposed reverse engineering a trained RNN by using numerical optimization 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) (2 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) [17 to address these challenges. Given an RNN we would like to train and analyze, we introduce a separate network consisting of 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 SL.SS 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 solve tasks that are described well by the linearized dynamics around its fixed and slow points [15]. These points are state vectors h E IRD that, given an input u* E RV, do not significantly change when applying the RNN update function, F F R² + E., That is, h' x F(h", :": We can find these points numerically by minimizing a loss function
L(h) = ||h F(h, 0)
(1)
using nuto-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 35). which says that the dynamics of nonlinear system in a domain near a hyperbolic fixed point is qualitatively the same as the dynamics of its linearization near this point. potential theoretical issue arises with non-hopeeb fixed points, though empirically, this has not seemed to pose a major issue for this approach. Finally, we note recent alternative method to finding fixed points presented in (36] that makes use of mathematical objects called directional fibers.
2

---

A. Nonlinear RNN Dynamics
B. Linearized RNN dynamics
c. Switching Linear Dynamical System
D. Jacobian SLDS
Ble
Elac
[(a
Figure 1: A. Example of a nonlinear dynamical system learned by a recurrent neural network (RNN). B. Linearization of the RNN dynamics around the nearest fixed point (dots). C. A switching linear dynamical system (SLDS) approximates the nonlinear dynamics by stochastically jumping between three linear regimes. D. The Jacobian SLDS co-trains an RNN (gray arrows) with expansion network (E(.)), which maps the current JSLDS state (white dots, a,) to an expansion point (purple dots, E(a)) )) near to the expansion point to obtain a linear system (purple arrows) that approximates the local dynamics.
fixed point of the RNN. The JSLDS linearizes the RNN dynamics around
2.2 Switching linear dynamical systems
Models based on 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 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 SLDS to break down complex. nonlinear time series into sequence of simpler local linear dynamics. Explicitly, let y, E R. denote the observation data at time and let h, E R represent the corresponding continuous latent state. A SLDS models the expected observation value as Elya] = g(h,) where g is mapping from TR D to N Given the discrete latent state 2, E {1, K} and input u, R the SLDS continuous latent states follow linear dynamics,
b ~ N (A h V(()) + B(=+), Q(=+)).
(2)
The current discrete state determines the current linear dynamics A E DDD input matrix V(*)) RDt bias term b(a.) E IR?' and noise covariance Q(=i) R DaD In the basic SLDS, Markov transition matrix generally defines the discrete state switching probability. An extension 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
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 balance between interpretability and expressivity for many problems. For example, one can improve the expressivity by increasing the number of discrete states, but this come at the cost of interpretability and 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 achieve automated fixed point finding and accurate SLDS approximations of nonlinear RNNs.
3.1 Motivation
Let F denote a nonlinear RNN with previous state h,- E RD, input u, U and parameters Writing the RNN update equation and its first order Taylor series expansion around points hn and u²
3

---

we have
h, F(b,-1, 4,;0) F(h", 0) + OF ōh
OF
(h", :):)) (h--1 - h²) + du (h" u" (u, -u") u²)
Here, OF (h", "hh"" uu: 0) E RDD is the recurrent Jacobian that determines the recurrent local dynamics OF 2h
and Du (h", 0); E Dzt' is the input Jacobian that determines the system's 's input sensitivity. If h' is fixed point of for given u", then uu"" = hh* and equation (3) yields
OF 2h (h", 00 (h1-1 - h") + du
OF (h", 00) (u, -u).
h h
(4)
Eq. (4) gives a LDS that locally approximates F around . In principle, if one knew how to select the correct fixed point and there was always fixed point nearby, one could use eq. 4 run 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 state z indexes the matrices (**) and V(*e) 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.
manner analogous to how the discrete
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 separate SLDS with its own hidden state a, E IR_ that switches around an expansion point eū E RDD
OF a - e Oh (e, 0² (a.-1 - ei)+ +
OF Tu (e, u² (u; - uu)
(5)
Note that eq. (5) shares its parameters 0 with F (eq. 3. i.e. given e, and u², the nonlinear RNN's parameters 0 determine the update matrices. We will generally take ut to be either zero (the average value for examples) or the value of context-dependent static input. The goal is for e to approximate the RNN's fixed and slow points. To accomplish this, we supplement the SLDS with nonlincar auxiliary function E (the expansion network) with separate learned parameters Φ
ej Eia
(6)
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 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 be forward in time independent of the original nonlinear RNN after training. Figure illlutrates the general idea. Given the previous state, a,-1, the expansion network, E, uses this state to select the next expansion point, e, (eq. (6) The system then updates the state by using ef to compute the recurrent Jacobian and input Jacobian of the original nonlinear RNN F (eq. 5). Assuming the expansion network has learned find fixed/slow points, switching between the points et corresponds to reverse engineering nonlinear RNNs using fixed points. The dependence of the expansion point e, on the previous state of the network a links to the recurrent connection in RSLDS.
Note that eq. (5) is intended to closely follow the dynamics of F, and will enforce this in the training procedure. We will refer to the system comprised of eqs. (516) as the JSLDS. We will refer the combination of the JSLDS and the co-trained RNN as the JSI.DS-RNN system. While limitation of our method that it does not currently lend itself easily to a stochastic formulation like the standard SLDS, it does allow for potential continuum of different switches using constant number of parameters. It also automatically determines the number of switches required to solve the task instead of the hyperparameter tuning required determine this in SLDS.
3.3 JSLDS co-training Procedure
We co-train together the nonlinear RNN (eq. 3) and the JSLDS (eqs. 5 Each network can be run forward and solve the task independently. We pass each of their states through the same

---

output activation function to compute two loss functions, CRNN and LISLDS, for the RNN and JSLDS, respectively. In addition, the expansion points should approximate fixed points of F in order to achieve a good JSLDS approximation of the RNN. We also need to ensure the JSLDS states a, approximate the RNN states h,. We achieve these goals by adding to the total loss function fixed point regularizer Rg and an approximation regularizer Rg defined as
Re(0, Σ
llee -F(e,, u';
Ra Ra(0,0) - Σ ||a, h
(8)
Now we define the total training loss as
C(0, Ø) ARNNCRNN + AUSLOS (.R( + (0.
(9)
where ARNN, AISLDS de and Aa control the strengths of the RNN loss, the JSLDS loss, the fixed point regularizer and the approximation regularizer, respectively. In practice, we have found these hyperparameters straightforward to select (see Section A. in the Appendix for more detailed discussion). For a particular optimization iteration, we compute the loss function in eq. 9 and then update all of the parameters and once using standard backpropagation through time (BPTT) methods for RNNs [4]]. Assuming the optimization goes well, the result will be two independent trained systems, with the JSLDS approximating the nonlinear RNN to first order.
In related work, 42) introduced regularization terms that force part of the subspace of piecewise linear RNNs 43 towards plane attractors to mitigate the exploding/vanishing gradient problem [44 45 within a simple RNN architecture. In another relevant work, [46] proposed learning interpretable nonlinear SDEs by modeling the dynamics function as a Gaussian process conditioned on the learned locations of fixed points and associated local Jacobians. We also note that our co-training procedure shares some similarities the adversarial training GANs (47). However, we stress that our method shares e between the co-trained networks, and 0 and ø are each updated at the same time, i.c., we do not alternate between updating / holding & constant and vice versa.
4 Results
We analyze the JSLDS-RNN system on three examples: a synthetic 3-bit memory task, a syn- thetic context-dependent integration task, and multineuronal population recordings from monkey performing a reaching task [ For the synthetic tasks, we use a relative error metric as in [21 to compare the quality of linearized approximation provided by JSLDS and the standard method of linearizing around the fixed/slow points (found numerically) of a standard trained RNN (without JSLDS co-training). The metric computes the relative error, | [h iNNN - hh" ||2/ |2/l|/h, [I,ANN] I|2 between the nonlinear RNN state, hANN, and the state approximated from the linearization method, h, Note that for the standard linearization method we had to resort to only computing one-step ahead dynamics predictions. This approach was necessary because running the linearized dynamics forward for many timesteps accumulates substantial error and causes the trajectory to diverge. In contrast, one-step ahead predictions were not necessary for the ISLDS.
Concretely, for the standard linearization method: we trained an RNN, numerically found its fixed/slow points, and then computed h! for all of the timesteps of the held-out trials using eq. (4). To use eq. (4), we set h in that equation to k,RNN, i.e., the true previous state and we set h to be the nearest fixed/slow point (in Euclidean distance). We have to find the nearest fixed/slow point because, unlike JSLDS, the standard method does not directly link location in state space to the fixed point one should linearize around. For the JSLDS method: we first co-trained the JSLDS and RNN together. We then simulated the JSLDS forward for the entire trajectory of each held-out trial using eqs. 66 to compute all of the JSLDS states a, We then set h!"" equal to a, to compute the relative error for each timestep. For each method, we computed the mean relative error of all the timesteps of a held-out batch of trials. We repeated this experiment 10 times for each method by starting the training from different random weight initializations. We report the mean and standard deviation of the mean relative error across the 10 trials
OOr implementation for the synthetic tasks is available at http:://jiiibbboco//////mmmmmm// public.
5

---

4.1 3-bit discrete memory
This task highlights how JSLDS can automatically learn to switch about a discrete number of fixed points and significantly reduces the linearized approximation error. We trained the JSLDS-RN system to store and output three discrete binary inputs (iguure24) similar to the experiment described in [5]. For our purposes, the models receive three 2-dimensional input vectors where each input vector corresponds to a different channel Each input vector can take value of { OJ, [0. Oj. [0, 1] corresponding to a state of -1, O, +1, respectively. The models have three outputs, each of which needs to remember the last nonzero state of its corresponding input channel. The RNNs used in this experiment were GRUs 48] with state dimension of 100 and linear readout function. See Section A.3 of the Appendix for additional experiment details.
Projecting the co-trained RNN and JSLDS dynamics into the readout space illustrates the close agreement between the networks for predictions on held-out trials (Fig. 23) A benefit of JSLDS is that we can use its expansion points produced along a trajectory (given the trajectory inputs) to approximate the fixed/slow points the co-trained RNN uses along the same trajectory. As a verification, we numerically found the fixed/slow points of the co-trained RNN and projected both the expansion points and the numerical fixed/slow points into the readout space (Fig. 2 We observed that the co co-trained networks learned fixed point solution consisting of 8 marginally stable fixed points (typically 2-3 eigenvalues within .025 of (1,0) in the complex plane). This solution was robust across different random weight initializations. See Section A.3.2 of the Appendix for detailed analysis and discussion of this solution compared to the fixed point solution found by a standard GRU without co-training (Fig. (C)). Figures p and F compare the mean relative error for the standard method and the JSLDS and show an example PCA trajectory of the reconstructed dynamics for each. This shows that JSLDS can simulate forward entire dynamics trajectories with much less error than the standard method (which relies on one-step ahead dynamics generation). Finally, as an additional experiment, we initialized the JSLDS co-training procedure with the trained weights of the standard GRU. We observed the fixed point solution change from the one presented in Fig. _ to one like that presented in Fig. - and the same improved linearized approximations of the dynamies as presented in Fig.
4.2 Context-dependent Integration
This task illustrates that JSLDS can lear to switch about multiple continuous manifolds of fixed points, improve the linearized approximation of the dynamics, and be used to perform a complex analysis similar to that performed in (20). The experiment consists of training the model to contextu- ally decide which of two white noise input streams to integrate (Fig. 3A). The model receives two static context inputs corresponding to motion and color contexts and two time- varrin white noise input streams. It is trained to output the cumulative sum of the white noise stream specified by the active context input. We used vanilla RNN with state dimension of 128 for the co-trained RNN and a linear readout function. See Section A.4. of the Appendix for more experiment details.
After co- tr-iiriiiig the "SLDS-RNN system, we observed close agreement between the JSLDS and RNN on task performance for held-out trials (Fig. B). Next, we analyzed the dynamics of the ISLDS for held-out trials under both contexts, set to different bias levels for both the color and motion input streams. We observed that for either context, the system integrates the relevant input using a single linear mode with an eigenvalue of 1, while the other eigenvalues decay rapidly (Fig. 3C). Next, we report the mean relative error on held-out trials for both JSLDS and linearizing a standard trained RNN in Fig. 3K. Again, JSLDS substantially improves the linearized approximation of the dynamics.
From the analysis in 20], we expect vanilla RNN to solve this type of task by representing the integration of relevant evidence as movement along an approximate line attractor (the choice axis) determined by the top right eigenvector. The solution consists of two line attractors that never exist together: one exists in the motion context and the other in the color context. For a given context (and therefore a specific line attractor), the top left eigenvector (the selection vector) determines the amount of evidence integrated. We expect the selection vector to project strongly onto the relevant input and be approximately orthogonal to the irrelevant input. See the Mathematical Supplement Section 10 in (20) for more details.
We reparameterized the inputs compared to [5]. This ensures the models do not have to act nonlinearly in the inputs and does not change the basic logic of the experiment since interested in nonlinear dynamics.

---

A
C
E
RNN Outpuc
Stable
Chi Input
Seodie-1
RNN: Marginally Stable
Input
Saddle-2
ISDS Marginally Stable
Chi Output
Readout Bostoon POSt0 Dhn #3
MT A
Input
Ch2 input 2
Output
Readout Fondbaa Dion. EZ
s
Ch3 input
Ch3 Input
Ch3 Output
B
D
F
RNN Output
Mean Relative Error on held out trials:
Mean Relative Error on held-out trials: 0.004
JSLDS
PC P2
FC 22
Resdbnu Din.
Readout Din
RNN
RNN
Linearized
Figure 2: 3-bit memory. A. Two-dimensional inputs (dark blue and light blue) corresponding to input states of -1, 0, or 1 enter at random while the corresponding output (dashed red) has to remember the last non-zero input state. B. JSLDS closely approximates the co-trained RNN readout space for held-out trial data. C. Standard GRU (no JSLDS co-training) outputs and numerical fixed points projected into readout space. D. Example PCA trajectory of standard GRU and linearized dynamics (one-step ahead dynamics generation) using numerically optimized fixed points. We also note the mean relative error for held-out trials. E. Comparison of co-trained RNN fixed points (found numerically) and JSLDS expansion points projected into readout space. The solution consists of 8 marginally stable fixed points. JSLDS has changed the fixed point solution compared to the standard GRU's solution (panel C). F. Example PCA trajectory of co-trained RNN and JSLDS dynamics (fully simulating dynamics forward) and the held-out mean relative error.
To verify this holds for the JSLDS, we produce a figure similar to Figures 5 and 6C from 20 by projecting the JSLDS states and expansion points into the 3-dimensional subspace meant match the axes of choice, motion input and color input (Fig. 3D--). Section 7.6 of the Supplementary Information of [20] describes the construction of this subspace in detail and provide brief description in Section A.4.2 in the Appendix. It is analogous to a regression subspace estimated from neural data in that work. was constructed by orthogonalizing the direction of the first right eigenvectors (averaged over expansion points) and the input weight vectors corresponding to the color and motion input streams. Panels D-F and G-I correspond the motion and color contexts, respectively. In panels D and I, we see that for the relevant context input stream, the states move along the axis of choice and the relevant input axis in proportion the strength the input. Panels F and G show that the strength of the nonrelevant input stream does not affect the direction of choice. Next, we analyze the global arrangement (Fig. 3)) of the motion and color context line attractors and selection vectors. As expected, we see that the selection vectors project strongly onto the relevant input axis but are approximately orthogonal to the irrelevant axis. Figure A.3 the Appendix presents the results of performing this experiment for standard trained vanilla RNN without the JSLDS co-training. It appears the JSLDS co-training did not dramatically change the standard trained RNN's fixed point solution for this task.
7

---

A
Color Mation
cumulative sum
Target
RNN output
JSLDS output
Contet (motton color)
Timestops
Real
D
E
F
Motion context expansion points
Choice
Choice
Choice
G
H
Co context expansion
Motton
Motion
Color
Choice
Choice
Choice
K
Motion Inputs:
Color
Strong Choice
Strong Choice
Mean relative error on held-out trials
Medium Choice
Medium Choice
Weak Choice
Weak Choice
Co
Weak Choice 2
Weak Choice
Standard Method
JSLDS
Choice
Medium Choice
Strong Choice 2
Strong Choice
Selection vector
0.38542003378
0.009110.0004
Motion
Figure 3: Context-dependent integration A. One of two white-noise input streams (motion or color) is selected to be integrated based static context input. The other stream is ignored. B. Sample held-out trial outputs show close agreement between JSLDS and RNN. C. Typical eigenvalues at a sample expansion point for motion (red x a's) and color (blue dots) contexts. D-J. JSLDS has learned to switch between two continuous manifolds of fixed points. JSLDS states (averaged) and expansion points are projected into the subspace spanned by the axes of choice, motion and color. Movement along the choice axis represents integration of evidence and the relevant input stream deflects along the relevant input axis. The input axes of E,F,G.H have been intensified. The trials used in and G are the same trials as D-E and H-I, respectively, but re-sorted and averaged according to the direction and strength of the irrelevant input. The expansion points were computed separately for motion (red x's) and color contexts (orange x's 3)) J. Global arrangement of the selection vectors (green lines) and line attractor expansion points for both contexts projected onto the input axes. Inputs are selected by the selection vector (which is approximately orthogonal to the contextually irrelevant input) and integrated along the line attractor. K. JSLDS improves the dynamics approximation compared to linearizing a standard trained RNN
4.3 Monkey reach task with LFADS-JSLDS
Finally, we illustrate how JSLDS can be dropped in as module to improve our understanding of more complex architectures that use RNNs such as LFADS [.. LFADS is a sequential variational auto-encoder 49 50 used to infer latent dynamics from single-trial neural spiking data. A criticism of L.FADS has been that is hard to interpret the RNN generator that produces the dynamies. Here, use JSLDS to improve this understanding by substituting the combined JSLDS-RNV -NNNNN system for
8

---

the standard GRU used in the LFADS generator (Fig. A.4 in the Appendix). We refer to this system as the LFADS-ISLDS, Once trained, we can use either the JSLDS or the co-trained RNN as the generator to produce the firing rates.
We used the monkey J single-trial maze data from Churchland et al. [1 using the same setup as Pandarinath et al. [5] to train the LFADS-ISLDS model. The data consists of 2296 trials of spiking activity recorded from 202 neurons simultaneously while monkey made reaching movements during a maze task [52] 5 across 108 reaching conditions. We used a GRU for the RNN generator. See Section A.5.1 1 of the Appendix for more details. The jPCA method [5T] has been applied to this data before 24|55, so we make use of it to validate our method. It finds linear combinations of principal components that capture rotational structure in data. See [5T] for the full details on jPCA.
We present the LFADS-JSIDS firing rates generated from the inferred initial condition (using the co-trained RNN generator) for several sample neurons in Fig. 4 It was observed in [5] that the standard LFADS population dynamics on single trials exhibit rotational dynamics when projected onto the first two jPC planes. To confirm LFADS-JSLD also exhibits this behavior, we applied jPCA to the co-trained RNN generator states (Fig. B). Next, the sample trials in Figure 4C show how the JSLDS generator closely approximates the RNN generator. Finally, focusing on the JSLDS generator dynamics, we learned that despite minor variations in the expansion points, the JSLDS generator eigenvalues and eigenvectors were the same at every timestep of every trial. So the ISLDS generator learned to represent the dynamics using a single, condition condition-independeer independent linear system. Figure 4D displays the eigenvalues for this system.
Sussillo et al. [24] noted that the dynamics found by linearizing around the fixed point of their regularized RNN model should roughly agree with the dynamies found by applying jPCA directly to the model. Along these lines, we validate the use of JSLDS by confirming that the JSLDS generator dynamics agree with the dynamics found by fitting jPCA the co-trained RNN generator states (Fig. 4E-F). A subspace angle analysis shows that four of the top five planes in the JSLDS state space (defined by the eigenvectors) were similar the first four jPCA planes (Fig. 4F. Using this correspondence, we see that the eigenvalues reported by jPCA (constrained to be purely imaginary) revealed four frequencies that closely agreed with the top four frequencies found by analyzing the ISLDS generator dynamics (Fig 4E). See Section A.5.2 for more details on this subspace analysis. The correspondence between the jPCA and the JSLDS dynamics validates the use of JSLDS for the LFADS generator. Given this, we conclude that LEADS--SLDD learns to represent the dynamics for this task using a single, aadition-inndeenneen linear system, a fact that was not obvious priori. See Section A.5.3 of the appendix for a discussion of numerically finding the fixed points of trained LFADS model without the JSLDS co-training. We observed that the JSLDS co-training did not dramatically change the standard LFADS model's fixed point solution for this task.
5
Discussion
Inspired by ideas from reverse engineering RNNs 15 20 24 25 10 and SLDS models 37
this work addresses the challenging problem of improving our understanding of how RNNs perform computations. We introduced new model, the JSLDS, improve RNNs. We applied it in various settings and to different architectures: GRUs, Vanilla RNNs, and LFADS models. The JSLDS does not require post-training fixed point optimizations, significantly reduces the approximation error associated with reconstructing the nonlinear dynamics using the locally linearized solutions, and maps each point in state space fixed point. These benefits significantly improve our ability to reverse engineer R RNNs, assuming a state-dependent SLDS can provide a good approximation of the original nonlinear RNN trained on a particular task.
ability to reverse engineer
Furthermore, the JSLDS could be valuable tool to investigate the limits of the general framework of reverse engineering RNNS with fixed points. This is because only expect the JSLDS approximation to break down if there is a system with nonlinear dynamics that are not well-described by switching between linearizations around fixed points. In addition to the above, JSLDS generalizes SLDS models to potential continuum of switches with a constant number of parameters and automatically leams the required number of switches.
An area for refinement the expansion network. We observed in some experiments that the expansion network might produce clusters of slightly varying expansion points that all define single linear

---

A
B
Cand. Avg.
RNN JSLDS
Trial
Exp. Points
Nouron
Neuron
PCA #1
D
E
F
0.075
0.075
IPCA plane number
80
(dcareesi
JSLDS evals
-0.075
eigenvalues
-0.075
JPCA evais
2
5
JSLDS eigenvector plane number
Real
Real
Figure 4: Monkey maze task. A. LFADS-JSIL JSLDS firing rates generated from inferred initial conditions for sample neurons. B. Projection of the co-trained RNN generator hidden states onto the first 2 jPC planes. We see the generator exhibits the expected rotational dynamics. C. Sample trial dynamics show low approximation error between RNN and JSLDS. Note also the JSLDS expansion points and RNN numerical fixed point. D. Eigenvalues of the JSLDS observation is that these eigenvalues are the same for every timestep of every trial, i.e. LFADS-ISLDS has learned organize the movement dynamics with single condition-independenl linear system. E. Top 70 eigenvalues from the same linear system shown in D along with the purely imaginary connected eigenvalues associated with the jPCA analysis (green squares). The jPCA eigenvalues (blue line) to their corresponding JSLDS eigenvalues given by the subspace analysis in F. F. Subspace analysis comparing the jPC planes and planes corresponding the top five complex eigenvalue pairs of the JSLDS generator. Color indicates the minimum subspace angle between the corresponding planes. Angles of 30-40 degrees indicate highly overlapping subspaces.
single timestep of single trial. Our
system instead of just producing a single expansion point. Perhaps an additional loss function penalty more specific architectures for particular tasks could help reduce this variation.
Finally, in the 3-bit memory task, the JSLDS co- co-training changed the fixed point structure the co- trained RNN used to solve the task compared to the standard GRU solution. We additionally observed improved linearized dynamics approximations with this new solution. These observations provide evidence that JSLDS can regularize nonlinear RNN towards solutions better described by switching between linearized dynamics around fixed points. This regularization towards switching linear structure could potentially have beneficial performance effects for robustness and generalization held-out data. However, more in-depth and larger-scale studies are required to quantify these potential effects. We also note that method could potentially suffer from the same theoretical limitation discussed for the previous reverse engineering method in Section 2. when near non-hyperbolic fixed points due the Hartman-Grrobmn theorem. However, the potential for method to bias the nonlinear RNN solutions towards solutions well approximated by switching between linearizations around fixed points could alleviate this concern in practice.
Broader Impact While understanding how RNNs perform complex computations could eventually help bound expected model behavior, identify biases, improve robustness to adversarial inputs and suggest ways to improve performance, we foresee no immediate societal consequences of this work. Overall, this work makes it easier reverse engineer RNNs and continue to gain insight into how these models work to benefit both neuroscience and machine learning.
10

---

Acknowledgements We would like to thank Mark Churchland, Matt Kaufman and Krishna V. Shenoy for access the monkey maze data.
J.T.H.S. received funding support from a Stanford Graduate Fellowhip in Science and Engineering (Mayfield fellowship). S.W.L. was supported by grants from the Simons Collaboration the Global Brain (SCGB 697092) and the NIH BRAIN Initiative (U19NS1132) and ROINS113119) D.S. was supported by grant from the Simons Foundation (SCGB 543049, DCS).
References
[1] Rohit Prabhavalkar, Kanishka Rao, Tara N Sainath, Bo Li, Leif Johnson, and Navdeep Jaitly. comparison of requence-to-seqquence models for speech recognition. In Interspeech, pages 939 0000-943,2222 2017.
[2] Debidatta Dwibedi, Jonathan Tompson, and Pierre Sermanet, Temporal reasoning in videos using convolutional gated recurrent units. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, 2018.
[3] Keunwoo Choi, György Fazekas, Mark Sandler, and Kyunghyun Cho. Convolutional recurrent neural networks for music classification. In 2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 2392-2396, 2017.
[4] Duyu Tang, Bing Qin, and Ting Liu. Document modeling with gated recurrent neural network for sentiment classification. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 1422-1432, Lisbon, Portugal, September 2015. Associa- tion for Computational Linguistics.
[5] Chethan Pandarinath, Daniel J O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew Kaufman, Stephen Ryu, Leigh R Hochberg, al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10):805 --5, 2018
[6] Christopher J. Cueva and Xue-Xin Wei. Emergence of grid-like representations by training recurrent neural networks to perform spatial localization. In International Conference on Learning Representations, 2018.
[7] Ben Sorscher, Gabriel Mel, Surya Ganguli, and Samuel Ocko. A unified theory for the origin of grid cells through the lens of pattern formation. H. Wallach, H. Larochelle, Beygelzimer, F. d'Alché-Buc E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[8] Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin J Chadwick, Thomas Degris, Joseph Modayil, et al. Vector- based navigation using grid-like representations artificial agents. Nature, 557(7705):429-433 2018.
[9] Ingmar Kanitscheider and lla Fiete. Emergence of dynamically reconfigurable hippocampal responses by learning to perform probabilistic spatial reasoning. bioRxiv, 2017, doi: 10.1 101/ 239
[10] Rylan Schaeffer, Mikail Khona, Leenoy Meshulam, Brain Laboratory International, and Ila Fiete. Reverse-engineering recurrent neural network solutions to a hierarchical inference task for mice. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4584-4596 Curran Associates, Inc., 2020.
[11] Fu-Sheng Tsung and Garrison Cottrell. Phase-space learning, pages 481-488. 01 1994.
[12] David Zipser. Recurrent network model of the neural mechanism of short-term active memory. Neural Computation, 179) 1991. doi: 111162/nncco.1991...222
[13] Michael Casey. The dynamics of discrete-time computation, with application to recurrent neural networks and finite state machine extraction. Neural Computation, 8:1135-1178, 1996.

---

[14] Paul Rodriguez, Janet Wiles, and Jeffrey L Elman. A recurrent neural network that learns to count. Connection Science, 11(1):5-40, 1999.
[15] David Sussillo and Omri Barak. Opening the black box: Low Low-dimensional dynamics in high- dimensional recurrent neural networks. Neural Comput. 25(3):626-649, March 2013. ISSN 0899-76677. doi: 10. 162//ECCO 1012/EECO_a0004000 a_00099..
[16] I. Jordan, Piotr A. Sokól, and II Memming Park. Gated recurrent units viewed through the lens of continuous time dynamical systems. ArXiv, has/1906.01005, 2019.
[17] Yiwei Fu, Samer Saab Jr, Asok Ray, and Michael Hauser. A dynamically controlled recurrent neural network for modeling dynamical systems. arXiv preprint arXiv:1911.000 0000090 2019.
[18] Duong Nguyen, Said Ouala, Lucas Drumetz, and Ronan Fablet. Variational deep learning for the identification and reconstruction of chaotic and stochastic dynamical systems from noisy and partial observations. arXiv preprint arXiv:2009.088 200002009,,,9 2020.
[19] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadeneshelli Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference Learning Representations, 2021.
URL httptttt:://peeeriiiieeeee htwp://wwww..eetttoorrmmiieeeee-
[20] Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome, Context-dependent computation by recurrent dynamics in prefrontal cortex. nature, 503(7774):77-844 2013.
[21] Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Reverse engineering recurrent networks for sentiment classification reveals line attractor dy- namics. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[22] Niru Maheswaranathan and David Sussillo. How recurrent networks implement contextual processing in sentiment analysis. ArXiv, abs/2004.08013, 2020.
[23] Federico Carnevale, Victor de Lafuente, Ranulfo Romo, Omri Barak, and Néstor Parga. Dynamic control of response criterion in premotor cortex during perceptual detection under temporal uncertainty. Neuron, 86(4):1077-1077, 2015. ISSN 0896-6273, doi: http::/ddoi.or//10.001/// neuron.2015.04.044
[24] David Sussillo, Mark M Churchland, Matthew Kaufman, and Krishna Shenoy. neural net- work that finds naturalistic solution for the production of muscle activity. Nature neuroscience, 18(7):1025-1033 1033, 2015.
[25] Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Universality and individuality in neural dynamics across large populations recurrent networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[26] Arseny Finkelstein, Lorenzo Fontolan, Michael N. Economo, Nuo Li, Sandro Romani, and Karel Svoboda. Attractor dynamics gate cortical information flow during decision-making. bioRxiv, 2019. doi: 10. 10.110//2199..2.1.88622..
[27] G. Ackerson and K. Fu. On state estimation in switching environments. IEEE Transactions Automatic Control, 15(1):10-17, 1970. doi: 19.11999AAC
[28] C. B. Chang and M. Athans. State estimation for discrete systems with switching parameters. IEEE Transactions on Aerospace and Electronic Systems, AESS-14(3:44 425, 1978. doi: (1009/AAES.198.3000
[29] Kevin P. Murphy. Switching Kalman filters. Technical report, 1998.
[30] Zoubin Ghahramani and Geoffrey E. Hinton. Variational learning for switching state-space models. Neural Computation, 1244),831-844, 2000. doi: 10. 10.162289976600000066669999
12

---

[31] Emily Fox, Erik Sudderth, Michael Jordan, and Alan Willsky. Nonparametric Bayesian learning of switching linear dynamical systems. In D. Koller, D. Schuurmans, Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems, volume 21. Curran Associates, Inc., 2009.
[32] Matthew D. Golub and David Sussillo. Fixedpointfinder: A tensorflow toolbox for identifying and characterizing fixed points in recurrent neural networks. Journal of Open Source Software, 3031):1003, 2018.
doi: 10.21 105/joss.01003.
[33] David M Grobman. Homeomorphism systems of differential equations. Doklady Akademii Nauk SSSR, 128(5):880-
1959.
[34] Philip Hartman. A lemma in the theory of structural stability of differential equations. Proceed- ings of the American Mathematical Society, 11(4):610-220, 1960.
[35] David Arrowsmith and Colin Place. Dynamical systems: differential equations, maps, and chaotic behaviour, volume 5. CRC Press, 1992.
[36] Garrett E Katz and James A Reggia. Using directional fibers to locate fixed points of recurrent neural networks. IEEE transactions on neural networks and learning systems, 633))36636 2017.
[37] Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian Learning and Inference Recurrent Switching Linear Dynamical Sys- tems. In Aarti Singh and Jerry Zhu, editors, Proceedings of the 20th International Confer- ence on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learn- ing Research, pages http://prroceedings mlr. www.eeeeesssseeeeeeeee. http:///////ineemammnmmm 11mmmaan1aa 17a. html
Fort Lauderdale, FL, USA, 20 22 Apr 2017. PMLR. URL
[38] Scott Linderman, Annika Nichols, David Blei, Manuel Zimmer, and Liam Paninski. Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity in C. elegans. bioRxiv, 2019. doi: 101//2222..40.
[39] Josue Nassar, Scott Linderman, Monica Bugallo, and II Memming Park. Tree-structured recur- rent switching linear dynamical systems for multi-scale modeling. In International Conference on Learning Representations, 2019.
[40] Joshua Glaser, Matthew Whiteway, John P Cunningham, Liam Paninski, and Scott Linderman. Recurrent switching dynamical systems models for multiple interacting neural populations. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 14867-14878. Curran Associates, Inc., 2020.
[41] P. Werbos. Backpropagation through time: what does it do and how to do it. In Proceedings of IEEE, volume 78, pages 1550-1560, 1990.
[42] Dominik Schmidt, Georgia Koppe, Zahra Monfared, Max Beutelspacher, and Daniel Durstewitz. Identifying nonlinear dynamical systems with multiple time scales and long-range dependencies. In International Conference on Learning Representations, 2021. URL /tttts:://operreeeee
net/ id== XYzwxP IQu6
[43] Georgia Koppe, Hazem Toutounji, Peter Kirsch, Stefanie Lis, and Daniel Durstewitz. Identifying nonlinear dynamical systems via generative recurrent neural networks with applications to fmri. PLoS computational biology, (58))e1000266 2019.
[44] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1735-1780, 1997.
[45] Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
[46] Lea Duncker, Gergo Bohner, Julien Boussard, and Maneesh Sahani. Learning interpretable continuous-time models of latent stochastic dynamical systems, International Conference on Machine Learning, pages 1726-1734 PMLR, 2019.
13

---

[47]
lan Goodfellow, Jean Pouget-Abadics Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets, Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014, htt.hhttpss/prrceeeeiiiaa pdf
[48] Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. In Proceedings of the 2014 Conference Empirical Methods Natural Language Processing (EMNLP), pages 1724-1734, Doha, Qatar, October
2014. Association for Computational Linguistics. doi:
[49] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 2013.
[50] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In International conference on machine learning, pages 278-1126. PMLR, 2014.
[51] Mark M Churchland, John P Cunningham, Matthew T Kaufman, Justin D Foster, Paul Nuyu- jukian, Stephen Ryu, and Krishna Shenoy. Neural population dynamics during reaching. Nature, 487(7405):51-566 2012.
[52] Mark M Churchland, John P Cunningham, Matthew T Kaufman, Stephen I Ryu, and Krishna V Shenoy. Cortical preparatory activity: representation of movement or first cog in a dynamical machine? Neuron, 68(3):387-400, 2010.
14


Reviews:
[{"review_id": "y0C09m2jL07", "reviewer": "Reviewer_kE2y", "summary": "The authors introduce a framework to facilitate the analysis of recurrent neural networks (RNNs) by reverse engineering. An RNN is co-trained on a task together with a so-called Jacobian switching linear dynamical system (JSLDS). The latter is the RNN’s linearization around certain expansion points. The co-training idea results in both networks solving the task, while ensuring that the JSLDS has both similar dynamics and similar fixed points (expansion points) to the RNN.\nBy this construction, the trained JSLDS is easy to analyze while matching the RNN, thereby eliminating many previously needed steps of analysis. This procedure also implicitly regularizes the RNN to exhibit dynamics that can be described by the JSLDS. \nThe authors show the utility of the framework on three tasks previously studied for reverse engineering RNNs.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The idea of co-training a network and a linearization of it simultaneously is new. Both training and RNN and linearizing around expansion points are well-known techniques, and the authors cite previous work adequately (perhaps the work of Katz and Reggia 2017 should be mentioned as an alternative method). Perhaps GANs can be mentioned as a distant analogy.\n\nQuality: The submission is technically sound: the authors test the new framework on a set of previously studied cases. The claims are supported by the experiments, and differences to well-studied works are described. The effect of implicit regularization on the solutions found could be better explored (detailed below).\n\nClarity: The submission is clearly written and well organized. It seems like it is straightforward to apply the method, but I did not try to.\n\nSignificance: There are two results in this paper. One is the introduction of a new method to analyze trained RNNs. The other is a new kind of bias or regularizer when training RNNs. The first is a welcome addition to the existing toolbox, and will probably find wide applications by many practitioners. The second is also significant, and adds to the literature on biases and the space of solutions for different tasks. This latter aspect is significant, yet underexplored in the current version.\n\nThe main issues that, in my view, could improve the paper are:\n1.\tThe two main contributions of the paper (method and biasing) are entangled. The text currently introduces the method-result as the goal, and the biasing-result as something discovered along the way.\n2.\tHow different are the solutions? Figure 2C,E should show this. But – they use a different color scheme, and the JSLDS in overlayed. It would be helpful to show just the RNNs using the exact same analysis. Perhaps even a 2-bit task would be better suited for this comparison. In particular, the JSLDS-RNN seems to have no saddle points between the attractors. That cannot be. There are distinct basins of attraction. This has to be properly explained.\n3.\tAre the solutions “simpler”, as stated by the authors. Figure 2E has a cluster of expansion points (and corresponding RNN fixed points) instead of the single attractors of Figure 2C. This seems like a more complex solution, and not a simpler one. On a similar note, the marginal nature of the clusters is also not necessarily simpler.\n4.\tThe authors state that the method will provide RNNs that can be approximated by a linear system, if such a system exists. This raises the interesting question – which tasks lend themselves to such an approximation? Is there a negative example (otherwise the statement at L279 remains speculative)? Perhaps connected, is the choice of the authors to change the input scheme in the 3-bit example. \n5.\tIs it actually simpler to use the new method? Are there cases where the “original” (train, then find fixed points) method fails, and this one works? Perhaps the half-sentence at L266 “not obvious a priori” hints at this, but deserves more than half a sentence.\n6.\tIs there any way to limit the number of expansion points?\n\nSpecific points:\n\n7.\tThe result of learning for the expansion points is not very clear. First, how does this depend on the choice of the function $G$? Second, if $F$ has multiple fixed points, will the expansion always choose the closest one?\n\n8.\tL36: proposes -> proposed\n9.\tChoosing the closest fixed points. L45 says it might not be the best choice. Is there an example? Furthermore, the zero, first and second order terms of the candidate fixed points might be used as a selection mechanism. If the question is the validity of a linear approximation, then these are the natural quantities.\n10.\tL62, feature 2. It’s not clear whether this is a good or a bad thing. It seems to complicate matters in the 3-bit example\n11.\tL99: limit expressivity – perhaps limit interpretability?\n12.\tThere is no comparison of the LFADS or context integration task with “normal” RNNs. Was there a bias in the type of solutions? For readers that are not fully versed in the original work, and simply want this as an example for the new algorithm, it would be useful to compare (even in supplement).\n13.\tThis might be a bit too much for 1 week… but: is it possible to have an intermediate algorithm. Train an RNN without other constraints. Then train the JSLDS, as a substitute for the fixed-point finder.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discuss societal impact in an appropriate manner.\nLimitations: mentioned above or described by authors.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "p4DQ3OOh2tX", "reviewer": "Reviewer_DUeH", "summary": "The central idea of this paper is to co-train a switching linear dynamical system (SLDS) together with a nonlinear RNN, such that the SLDS, due to its locally linear dynamics, can be used to gain insights into some aspects of the dynamics of the RNN. The SLDS retains the same Jacobians as the accompanying RNN, which are evaluated at approximate fixed points learned through an auxiliary function. The fixed point condition is imposed via a regularization term, as is the agreement in RNN and SLDS dynamics. The method is evaluated on 2 simple benchmarks and a neuroscience dataset.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First, I’d like to point out that the basic problem addressed here has been solved before using piecewise linear RNNs (see for instance https://arxiv.org/pdf/1910.03471.pdf and references therein). Also the idea to impose fixed point conditions via regularization has been introduced previously in http://proceedings.mlr.press/v97/duncker19a/duncker19a.pdf (some of the ideas in this paper appear quite similar to what the authors have done). In contrast to the authors’ own formulation, these previous models enable to learn nonlinear dynamics and to explicitly compute/ determine fixed points within one and the same architecture. That is, there is no need to train two systems in parallel as in the present approach. \nSLDS are also not continuous at the switching points if I recall correctly (which would make them unsuitable as a general emulator of dynamical systems).  In general, there is a wider recent ML literature on dynamical systems that is not really reviewed in here, although it seems relevant (for instance https://arxiv.org/abs/1911.00089, https://arxiv.org/abs/2009.02296, https://arxiv.org/pdf/2010.08895.pdf).\n\nThis being said, in general I still like the authors’ idea, and it may have the advantage that it’s more generally applicable. The co-trained RNN can be anything and doesn’t need to be of a particular form. But at least a proper discussion against the literature background above seems to be indicated. Although, to me the question remains, why not fit SLDSs on the data directly? Piecewise linear systems are generally expressive enough to approximate most nonlinear dynamical system (for instance Lu et al. 2017).\n\nI was less convinced by the evaluation: No comparisons to any other method were performed (VIND, https://arxiv.org/pdf/1811.02459.pdf, https://arxiv.org/abs/1707.09049, or PLRNN as cited above, seem obvious candidates). Only single examples were presented, which could have been cherry-picked, rather than a thorough statistical evaluation based on many different benchmark simulations and parameter initializations. For instance, how would I know that some of the results like those in Fig. 2E are not completely over-interpreted?\nThe neuroscience examples are probably only comprehensible if one had read [5] and [17] before. Without this background, it remains largely unclear what precisely is shown in Figs. 3 and 4. For example, in Fig 4, eigenvalues of precisely what are shown here (all eigenvalues of all Jacobians?), what does jPCA do, what is a ‘subspace angle’, what exactly is meant with ‘decaying oscillatory modes’ in a linear system (stable spiral points)? Many things I could just guess as they were not explained neither in the text nor in the legend, and some of the subpanels were very unclear to me (e.g. Fig. 4E).\n\nIt was also not fully clear to me why co-training RNN and JSLDA works that much better than finding fixed points numerically posthoc and linearizing around them. In terms of runtime I guess it doesn’t make a lot of a difference whether co-training of a second system with resp. regularization terms is involved, or whether this optimization is performed afterwards? But if the correct fixed points of the original RNN can be found, why does linearizing around them work that much worse? Also in this context, can the fixed points not simply be found numerically by iterating the RNN equations forward in time from different initial conditions? That should be cheap and fast.\n\nHow is it ensured that the fixed points are actually stable? If they are repellers, trajectories would exponentially diverge from them and the SLDS approximation would break down, no?\n\nFinally, does the RNN+SLDS learn the dynamics of the underlying experimental system in Fig. 4? Only posterior inference it seems was shown in Fig. 4A, which is not that impressive given that other models are already around that can learn the full dynamics (I conjecture that a simple LDS would have given similarly good agreements). \n\nBottom line: Although I’m kind of fond of the basic idea, it needs to be worked out much more rigorously and convincingly in my mind.\n\nMinor issues:\n- sect. 2.1: A fixed point is a point for which x*=F(x*) holds exactly. If it’s just approximate, it’s not a fixed point but perhaps a ‘slow point’ or ‘attractor ruin’.\n- sect. 4.1: Why do the fixed points of JSLDS need to be found by running the system forward in time? This could have been done right away with the RNN? And I thought the whole point of JSLDS is that the expansion points e* are approximate fixed points and hence returned directly, let alone that they could simply be explicitly computed in linear systems?\n- If there are just 8 fixed points in Fig. 2E, then what is all the jitter about? Also not sure what Fig. 2B tells me. How was the readout function g defined? There are so many unexplained details in all of the figures.\n- How exactly are the motion, color, choice axes determined; are these the axes along which motion etc. cause most variance?\n- What exactly is shown in Fig. 4B?\n- p.9: How can a nonlinear system like the JSLDS indicate how ‘nonlinear’ a task is?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Partly -- for limitations see my review above.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Pcd-qLqTHr", "reviewer": "Reviewer_DWxE", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work opens a wide range of possibilities by presenting a model with a good balance between interpretability and expressivity, and by showing that it can be applied both to purely in silico studies or fitted to neural recordings. \n\nA weak point of the model in its present state seems to be the expansion network which keeps an unclear division of state-space unlike the idealized version illustrated in figure 1. This seems to lead to an unreasonable number of expansion points according to the other figures. \n\nThe model is rather well described, however it could be made clearer in sections 3.2 and 3.3 whether the update matrices of the JSLDS are trained independently of the RNN or whether the derivation operator is backpropagated through from the LDS loss function. \n\nFinally, it would have been interesting to see the JSLDS model used as a lower dimensional model of high-dimensional state-space dynamics, particularly for the neural recordings application, as it is typically done in most LDS studies. \n\nAltogether, the paper exposes convincing applications of a method with a strong potential. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "64LcI3zOj4e", "reviewer": "Reviewer_cBwm", "summary": "This work concerns the challenging problem of improving the understanding of how RNNs perform computations and was inspired by combining ideas from reverse engineering RNNs and SLDS models. The novel SLDS dynamics are governed based on the first-order Taylor series expansion of the co-trained RNN. It is also equipped with an auxiliary function trained to pick out the fixed points of the RNN. Their results are a trained SLDS variant which closely approximates the RNN, an auxiliary function that can produce a fixed point for each point in state-space, and a trained nonlinear RNN whose dynamics have been regularized such that its first-order terms perform the computation, if possible. Moreover, The presented model removes the post-training fixed point optimization, and allows to study the learned dynamics of the SLDS at any point in state-space. It also generalizes SLDS models to continuous manifolds of switching points while sharing parameters across switches. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1) Their contribution (JSLDS) seems genuinely novel.\n2) This work is clear and very understandable.\n3) The results of this work seem highly important and very interesting.\n\nHowever, there are a few comments:\n\n- An important part of this work is to approximate the nonlinear RNN by a SLDS with linearizations around fixed points. From Dynamical systems point of view, due to Hartman-Grobman theorem (linearization theorem), 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 fixed point. While this is not true in case of non-hyperbolic fixed points. So when there is a non-hyperbolic fixed point, the linearization around the fixed point cannot necessarily describe the dynamics of the nonlinear RNN. But, the authors did not discuss how they could overcome the problem of linearizations around *non-hyperbolic fixed points*, although this needs to be addressed.\n\n\n- Regarding the sentence \" Allows the nonlinear RNN to be approximated by a SLDS, if possible” ( page 2, line 64) :\nIt is not clear that if their technique fails, this necessarily implies that the nonlinear RNN cannot be approximated by a SLDS.\n\n- Regarding the sentence \"The difference between these two solutions shows that the JSLDS co-training procedure can change the fixed points of a nonlinear RNN for a particular task and suggests that it can provide regularization towards simpler solutions for some problems” (page 5, line 183):  \nThere is no justification for the latter have of this claim (the regularization). The authors repeat the claim in their conclusion (283. f) and I’m unsure they actually showed this regularization towards simpler solutions.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think so", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree that the paper proposes an interesting approach to the problem of improving the understanding of RNNs. Although some reviewers have some technical concerns at their first reviews, basically those have been resolved by the authors' responses. Thus, although there are some points that should be modified from the current form, I think we can expect the authors modify the paper in the camera-ready by reflecting the discussion. Based on these, I recommend acceptance (poster) for this paper.
Decision: Accept (Poster)