Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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. 
main_review: 1) Their contribution (JSLDS) seems genuinely novel.
2) This work is clear and very understandable.
3) The results of this work seem highly important and very interesting.

However, there are a few comments:

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


- Regarding the sentence " Allows the nonlinear RNN to be approximated by a SLDS, if possible” ( page 2, line 64) :
It is not clear that if their technique fails, this necessarily implies that the nonlinear RNN cannot be approximated by a SLDS.

- 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):  
There 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.

limitations_and_societal_impact: I think so
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
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. 

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

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

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

Altogether, the paper exposes convincing applications of a method with a strong potential. 
limitations_and_societal_impact: yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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. 
SLDS 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).

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

I 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?
The 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).

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

How 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?

Finally, 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). 

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

Minor issues:
- 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’.
- 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?
- 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.
- How exactly are the motion, color, choice axes determined; are these the axes along which motion etc. cause most variance?
- What exactly is shown in Fig. 4B?
- p.9: How can a nonlinear system like the JSLDS indicate how ‘nonlinear’ a task is?
limitations_and_societal_impact: Partly -- for limitations see my review above.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
By 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. 
The authors show the utility of the framework on three tasks previously studied for reverse engineering RNNs.

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.

Quality: 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).

Clarity: The submission is clearly written and well organized. It seems like it is straightforward to apply the method, but I did not try to.

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

The main issues that, in my view, could improve the paper are:
1.	The 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.
2.	How 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.
3.	Are 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.
4.	The 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. 
5.	Is 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.
6.	Is there any way to limit the number of expansion points?

Specific points:

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

8.	L36: proposes -> proposed
9.	Choosing 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.
10.	L62, 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
11.	L99: limit expressivity – perhaps limit interpretability?
12.	There 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).
13.	This 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.

limitations_and_societal_impact: The authors discuss societal impact in an appropriate manner.
Limitations: mentioned above or described by authors.
review_title: null

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)