A New Look at Low-Rank Recurrent Neural Networks

We thank the reviewer for their time & valuable feedback on our manuscript. We especially appreciate their positive remarks regarding the clarity of our presentation, the efficiency of our method and its significance in providing novel insights on the expressivity of low-rank RNNs.

Weaknesses:

1.Training RNNs: We thank the reviewer for highlighting this, and apologize for not discussing our comparison appropriately. In Section 6, both our framework and the networks trained with BPTT are trained from a set of teacher trajectories, which are simulated from the underlying ODE (eg. for binary decision making - they start somewhere on the grid and eventually converge to one of the two fixed points). However, unlike previous sections, our method here can be broken into two main steps.

a: We use Euler integration to compute the value of dz/dt along each point of the trajectory. This depicts our target vector for regression.

b: Our basis matrix is evaluated at grid points along these trajectories.

Finally, both our method and the networks trained with BPTT must reproduce a set of target trajectories (eg. blue lines in Fig 5-A) from an input pulse train (eg. representing bit value for each channel over the time interval). We thus believe our comparison is equivalent in that both methods are uncovering underlying dynamics (from data points) that solve the task. We’ve added additional details on training in SI-F.

2. Novelty: We apologize for the lack of discussion regarding similarity to previous methods. Specifc details are provided here (however we’re working on adding a section in our supplement discussing them further) -

2.1: Neural ODEs: Previous NEF literature details this equivalence & we’re working on adding this into the supplement.

2.2: Beiran et al. 2021: Our work is inspired by the methods presented in this paper. However, they use networks with constant input (text under eqn 3.3) for majority of their analysis. While universal approximation is in the appendix, there is no mention of static or time-varying basis functions (Section 3, 5 of our paper)/specific function classes (added SI: A.1 & Fig. SI-2) that can be approximated, which we believe we’re the first to do. Additionally there is no previous work on finding smallest networks (Section 4 of our paper) using this regression. We've updated Sec 1,7 to discuss these points.

3. Limitation: known dynamic g(z): We apologize for the lack of clarity as our method can be used without explicitly knowing g(z) (i.e as long as it can be approximated – see 1. Training RNNs). We’ve updated our discussion section to highlight this. Additionally, from a neuroscience perspective interpreting the recurrent connectivity & linking it to the dynamics remains an important question. Specifically, our work builds on the theoretical foundations laid out in task training RNN literature (e.g. Mante & Sussillo 2013, Sussillo & Barak 2013, Langdon & Engel, biorXiv 2022, Luo, Kim, et al.,biorXiv 2023) & low-rank RNNs (e.g Mastrogiuseppe & Ostojic, 2018; Beiran et al., 2021; Dubreuil et al., 2022; Valente et al., 2022) which focus on training RNNs to perform exactly the tasks in Sec 6 of our paper. Thus while the dynamics are known, our perspective provides novel representational & interpretational insights.

Questions:

NEF: We thank the reviewer for bringing this to our attention. We’re working on including this in our supplement.

Eqn 11: Thanks for pointing this out, Eqn 11 is missing terms! We’ve updated it to reflect the necessary changes.

Training v/s analysis: We apologize for the confusion & hope our discussion above (1. Training RNNs, 3. Limitation:known dynamic g(z)), clarifies the applicability of our method. Additionally, we note alternate training strategies such as FORCE/Full-FORCE, BPTT all require a teacher signal (like our method) & learn either a readout weight/the network connectivity. Our method also provides a way to learn the network connectivity (i.e an appropriate weight vector/matrix n/N), without needing a readout weight. Taken together our method provides an alternate training strategy (inspired by NEF), with the additional advantage of improved interpretability & analysis.

We would humbly ask the reviewer to consider increasing their score if they are persuaded by our clarifications and additional analyses. Thank you again for the helpful comments and suggestions.

We thank the reviewer for their detailed and thoughtful response & accurate summarization regarding our training clarification.

• 1. The reviewer is correct: our method does depend on how accurate the estimate of the underlying ODE is. It is also true in the case of very noisy/sparse resolution our characterization would have larger errors than those shown in the experiments. However, we have added an extra simulation Fig SI-10 to show our method can easily be extended to the high-dim noisy data setting. Specifically, we simulate high-dimensional data from the low-dimensional trajectories and add gaussian noise to this (details in SI-F.1). This resembles noisy rate based neural data. To apply our framework to this high-dimensional noisy time-series, we project the data on the top PC. This allows us to recover a noisy estimate of the true underlying trajectory. Our framework can now be applied on these recovered trajectories.

• 2. Additionally, we provide two examples (3 bit flip-flop and bi-stable ODE) comparing our framework to BPTT – with the main aim of expressing how our method can be used to train low-rank RNNs for neuroscience tasks (note in such tasks the dynamic is low-dimensional while the observation, i.e neural recordings could still be high dimensional) in significantly lower time, while maintaining high interpretability (through the geometric interpretation of basis functions).

We thank the reviewer for pointing us to “vector field regression” approaches - we have added this into our discussion in Sec 7 and clarified the added advantage of our framework as an alternative training strategy (with the caveat of knowing/being able to effectively estimate the ODE).

Finally, we agree with the comments made by the reviewer, but hope our clarifications/reframing expresses the novelty of our work. Additionally we hope our extended simulations showcase the novelty and insights of our work and the ability to adapt our system to multiple settings. We’d kindly request the reviewer to consider adjusting their score if they’re satisfied with our responses.

We thank the reviewer for their careful consideration & insightful feedback. We appreciate their positive remarks on the beauty & efficiency of our perspective, & clarity of our examples. Below, we address the reviewer's concerns & suggestions.

Weaknesses:

Embed dynamical system in low-rank RNN: The reviewer is correct: our method can be used for tasks with known dynamics. However, in Section 6, our framework & networks trained with BPTT infer dynamics from the same set of training trajectories. The dynamic (regression target) is approximated using finite differencing (euler method). Thus even when the ODE is unknown, its value is estimated at points along the trajectory. We’ve added section SI-F to clarify this.

a. Use-cases: Task-trained RNNs interpret governing dynamics of specific tasks linked to brain computation (e.g., Kanitscheider & Fiete 2017; Turner, E, et al. 2021; Luo, Kim, et al. 2023). However, degeneracies in the network/task solution space often lead to incorrect conclusions. We propose designing ODEs to solve tasks, which our method embeds into the RNN. We’ve added two additional dynamical portraits for evidence accumulation tasks (Fig SI-5).

BPTT Comparison: Our framework & BPTT networks are trained from a set of teacher trajectories simulated from an underlying ODE. (eg.: decision making, trajectories start from a random initial location on \z\ until convergence to a fixed point). Unlike previous sections, we used finite-differencing to compute $dz/dt$ from sample trajectories. We then fit the $n$ weights by regressing these derivatives against a design matrix formed by evaluating the nonlinear basis functions at the points in these trajectories. Thus, both training methods (ours & BPTT training) had access to the same training data. Both our method & BPTT-trained networks reproduce target trajectories (e.g., blue lines in Fig 5-A) from an input pulse train (e.g., representing bit values). Thus, both methods uncover underlying dynamics that solve the task. We’ve added details in SI-F to clarify this.

Universal approximator: Reviewer is correct: low-rank RNNs are a subset of RNNs. However, universal approximation is largely overlooked in the low-rank RNN literature. For eg., Beiran et al., 2021 use networks with constant input (text under eqn 3.3). While universal approximation is mentioned in the appendix, there is no mention of basis functions/specific function classes that can be approximated, which we’re the first to do. To further clarify this we’ve added SI: A.1 & Fig. SI-2.

Fig 4 & contribution: Reviewer is correct: dynamics in Fig 4 can be found via other methods. To clarify – our goal was to validate our results grounded in field-accepted theories (eg. need for perturbations). However, unlike other methods, our framework adds novel interpretations of time-varying dynamics through time-varying basis functions. (i.e contribution of each neuron’s time-varying response in approximating the dynamic).

Non-autonomous to autonomous: Doing this adds an extra dimension to the system which could play a critical role in neuroscientific insight. While not explored here, we’ve added this in our discussion section.

NEF & Regression: Thanks for this suggestion, we’re working on adding this in the supplement.

Typos: Thanks for pointing these out! We’ve corrected them in the revised paper.

Questions:

Smallest RNN using BPTT & comparison: Current methods train networks (different ranks & sizes) with BPTT & pick those with best performance. Dynamics solving tasks are then interpreted via dim reduction (eg. Valente et al. -NeurIPS '22, Sussillo, Kanaka, et al. Neuron 2016,Sussillo, D., & Abbott, L. F., Neuron 2009). We instead directly embed hypothesized ODEs governing tasks (dim=rank) without training multiple networks. We use regression to learn the required weighting & provide theoretical & empirical results (OMP), identifying the smallest RNN approximating a function. We’ve added tables comparing training times for networks (trained with BPTT) in SI-F, Table 1,2.

Algo for main method & OMP: Our primary method (Sec 3 & 3.1) involves:

1. Generating a random basis matrix.
2. Using regression to learn the weight vector/matrix.

Alternatively, Sec 4 provides an algorithm to optimally create a sparse basis matrix, .

Neural ODE equivalence: Previous NEF literature details this equivalence & we’re working on adding this into the supplement. In case of noise in the latent space, our framework would map to neural SDEs.

“(although not shown here) ... fails”: We’ve added a section (A.2, Fig SI-3) on why inputs are needed in the 2D case when the function lacks odd symmetry.

We humbly ask the reviewer to consider raising their score if they feel our revised paper (with added simulations, analyses, clarifications, & discussions) would make a worthwhile contribution to the meeting.

I sincerely thank the authors for their efforts to improve their paper with the suggestions from all reviewers.

BPTT Comparison: Thank you for clarifying this point. I think it would be important to add a sentence in the main paper under the BPTT section, such as "trained from a set of teacher trajectories simulated from an underlying ODE." Including this directly in the main text, rather than explaining it solely in the supplementary material, would improve clarity for readers.

Universal Approximator: I understand and appreciate the authors' contributions, which are clearly valuable, but some of the wording might lead to confusion. For example, the statement "More generally, our perspective clarifies limits on the expressivity of low-rank RNNs, such as the fact that without inputs, a low-rank RNN with sigmoidal nonlinearity can only implement odd-symmetric functions," might make it seem like this limitation is unique to low-rank RNNs.

It could be true that universality hasn’t been discussed in the low-rank RNN literature before, but as we’ve discussed earlier, low-rank RNNs inherit some of the same challenges as general RNNs. In this case, the fact that low-rank RNNs are not universal approximators without inputs or biases also seems to be an issue for RNNs in general. So, when the authors write "limits on the expressivity of low-rank RNNs," it could come across as misleading, as it suggests this is a problem unique to low-rank RNNs. I hope I’m not misunderstanding here. I think it would be better to focus more on the geometric interpretation, which is the authors' main contribution. Alternatively, if the authors want to emphasize this point, I suggest clarifying it in the abstract as well, as a misleading abstract would not be helpful for the paper.

Overall, the main issue seems to be how the contribution is framed. For example, the authors provide a formal connection between neural ODEs and low-rank RNNs. However, the wording, "... and shows that a low-rank RNN corresponds to a neural ODE with a single hidden layer [REFS]," could be revised to better reflect their work. It might be clearer to say something like, "This connection has been noted in previous work [REFS], and here we formalize it." This is just one example, but as other reviewers have also pointed out—and I agree—the contributions and the writing don’t fully align.

I appreciate the author's engagement during the rebuttal process.

Universal Approximator: I believe this change has clarified the paper's contributions. If you feel it is sufficient to include this information in the paper but not in the abstract, please feel free to disregard my earlier suggestion about including it in the abstract. As it stands, the abstract with the parenthesis reads awkwardly unless rewritten for better flow.

Overall, I still believe the feedback was incorporated into the manuscript at the sentence level rather than holistically. The other example is, the sentence, "Additionally, although not explored here, the conversion of non-autonomous dynamics to autonomous dynamics could play a critical role in neuroscientific insight," disrupts the flow and seems to have been added later. I kindly suggest that you feel free to ignore my recommendation if you disagree—I am entirely fine with that if you provide your reasoning. Otherwise, I encourage you to consider incorporating the feedback more holistically.

The figures are great, explanatory, and demonstrate the contributions. I increased my score slightly; however, I still think the paper is not yet in its best form due to a reduced flow caused by sentence-based additions and would benefit a lot if the structural writing changes were made.

Additional Note: Still have some typos: Line 200, ReLu, Line 242 phi(x).

We thank the reviewer for their valuable feedback, suggestions, & positive comments on the clarity of our presentation & novel method for identifying the smallest RNN. We were unaware of some past findings the reviewer highlighted & appreciate the opportunity to refine our claims & appropriately credit previous work.

Weaknesses:

1. Non-linear recurrent dynamics in NEF: The reviewer is correct & we apologize for the confusion. [1] shows applications of non-linear recurrent dynamics. We didn’t intend to claim NEF is limited to feed-forward networks, only that most applications we’d seen fell into that category. However, we realize this was inaccurate & we’ve updated Sec 1,7 to correct this.

2. Fitting low-rank RNNs with regression:

2.1 NEF: The reviewer is correct: NEF can & has been used with rate neurons. Part of our motivation was the (what seemed to us) relative disconnect between low-rank RNN work & the NEF work. As the reviewer noted, networks constructed by NEF are low-rank, but haven’t (to our knowledge) been framed as low-rank RNNs. Similarly, we hadn’t seen previous low-rank RNN work making the explicit connection to NEF (we’re grateful to pointing us to [2])! Overall, we feel this connection isn’t as well known as it should be, evidenced by most papers on low-rank RNNs using BPTT & feel our work provides intuition for the role that individual neurons play in low-rank RNNs (e.g. Fig 1 our paper). We apologize for the oversight & will correct our text to focus more on novel insights (e.g. OMP).

2.2 [2]: Our work is inspired by [2], however, they use networks with constant input (text under eqn 3.3) for majority of their analyses. While universal approximation is in the appendix, there is no mention of basis functions/specific function classes (added SI: A.1 & Fig. SI-2) that can be approximated, which we believe we’re the first to do. Also, there is no previous work on finding smallest networks (OMP) using this regression.

3.New insights on universality: We thank the reviewer for raising this point

• 1. As discussed above, [2] considers constant inputs for majority of their analyses. Second, their exploitation of symmetry arises exactly because of this. Eg: [2] Section 5.1: “if a network generates one nontrivial stable fixed point, additional stable fixed points appear at symmetric points in the collective space”. Our results explicitly clarify why this occurs: namely, when the RNN has an odd-symmetric nonlinearity in the absence of inputs. (However, it would not generally hold for other nonlinearities).

• 2. [2]: no mention of basis functions, which we feel provides nice linear algebraic/geometric insight into how a low-rank RNN constructs a representation of an ODE. Our insight is in: viewing neuron specific inputs as shifted basis functions which is arguably more intuitive. Additionally, we expand on how specific non-linearities ([2] restricts analysis to tanh - text under eqn 2.1) can approximate specific function classes better, which to our knowledge hasn’t been previously reported (added SI: A.1 & Fig. SI-2).

• 3. Previous work doesn't explore how universal approximation can be met with a sparse/finite basis set, which we show through OMP (within a margin of error).

4. New insights input driven dynamics: We apologize for the confusion. Our goal was to validate our results grounded in accepted theories (eg, identifiability in [3]). Our novel insight is in interpreting time-varying dynamics in low-rank RNNs through time-varying basis functions (i.e contribution of each neuron’s response): top row Fig 4. Also, prior work ([3], or others) to our knowledge, doesn’t provide explicit ODEs for such “dynamic attractors” via moving fixed points/no fixed points, which we address - Eqn 13.

5. Neural ODE equivalence: We thank the reviewer for pointing this out. We weren’t aware of [4] which connects low-rank RNNs and neural ODEs ([2] omits inputs in Eqn 2 & discusses it only in the appendix). We’ve updated Sec 3 to cite [4], [5] & clarify the connection. We acknowledge there exist other useful links to the neural ODE literature that we haven’t explored in this work.

6. Non-linear oscillators: The reviewer is correct, we deeply apologize. We meant “nonlinear & non-symmetric” (updated Sec 3.1).

Questions:

1. Convergence of OMP: Thanks for this question. OMP converges to a local optimum. However, we show performance is highly consistent across runs in simulation (added SI-E, Fig SI-9).
2. Line 277; Ref [2]; radially odd symmetric: We thank the reviewer for pointing this out. The governing ODE for our limit cycle is non-symmetric: $dz_1/dt = az_1 - z_2 - 0.35$ & $dz_2/dt = az_2 + z_1 + 0.5$, $a = 1 - r^2r$. Since $0.35$ & $0.5$ don’t depend on $z_1/z_2$, the function is not odd symmetric. We’ve added Section A.2 & Fig. SI-3 & updated Section 3.1 to address this.

1. regarding the NEF

I appreciate the author’s restating in the introduction. However some claims with respect to the NEF are still overstated, e.g.,

This presents an extension to the NEF method

Our novelty lies in characterizing a randomized basis for encoding/embedding, thus deviating from the classical NEF setting which carefully crafts how neurons are expected to respond to stimuli

The NEF/nengo typically also starts by sampling random tuning curves / basis encoding function and then solving for the linear decoding functions with least squares (exactly as you do).

I think it is quite crucial to do at least a high-level explanation of the NEF in the main text, and reformulate said sentences in order to provide the reader clarity about what is new (and what is not).

2. Fitting low-rank RNNs

In this work, we address the issues of RNN training

You only address fitting RNNs to known flow-fields, for which methods without e.g., the problem of exploding gradients do exist [1]. The closest is simple teacher forcing (TF): given a (discretised) RNN/neural ODE $F$ and pairs of $z_t$ and $z_{t+1}$, minimise $\lvert z_{t+1} -F(z_t) \rvert$. This is very similar to what you end up doing (especially when using the finite difference method to estimate $\dot{z}$ from trajectories). The main difference would be that you use the NEF approach of sampling random encoding weights and solving for the decoding weights with regression, instead of finding all parameters with back propagation. It could be good to name TF (and extensions, e.g., [1] also fit one hidden layer ODEs/RNNs) - TF probably would have been a fairer comparison than BPTT, but that likely is out of scope now.

3. On universality

It could be true that no-one named NEF networks low-rank RNNs explicitly, although I am sure people were aware of the low-rank structure of NEF networks (see also e.g., [2]). I do now agree with the authors that the geometrical basis functions explanation / NEF view adds some clarification to [3].

4. Input driven dynamics

Thanks for clarifying the additional insights to previous work, I now agree there is something new here. I do think the insight that small perturbations lead to different trajectories in different models is similar to [4] (where perturbations from the mean, were used to distinguish input versus autonomous dynamics).

5. Neural ODE Thanks, this is fine now!

6/Q2. Odd-symmetry of the non-linear oscillator

It is indeed correct that by adding the offset to $z_1$ and $z_2$, the oscillator becomes non-odd-symmetric. For reference, I do want to point out that the new equations are different from the ones on the original manuscript, which were $\dot{\theta} = 1, \dot{r} = 1-r^2$ - which did imply the odd-symmetric system.

Note that this 2D system is highly nonlinear and non-symmetric — unlike recent work focused on oscillations in nonlinear dynamical systems

I still think this should be rephrased, you are not the first to get RNNs to display limit cycles (even when adding some offset), maybe it is best to omit the unlike… part altogether?

Q1. Smallest network

The fact that you converge to a local minimum should be named explicitly in the main text in my opinion - the method finds a small network, not necessarily the smallest network

————

To avoid being overly negative, I do want to highlight that there is definitive value in this work, e.g., the OMP algorithm and the clarifications on universal approximations in low-rank RNNs. The writing is easy to follow and the figures are also great. The current narratives of “we extend the NEF”, and “we overcome limitations of BPTT”, are in my opinion (as I hopefully clarified above), not the ideal way to present these contributions.

————

[1] Hess et al., 2023. Generalized Teacher Forcing for Learning Chaotic Dynamics

[2] Barak and Romani, 2020. Mapping low-dimensional dynamics to high-dimensional neural activity: A derivation of the ring model from the neural engineering framework

[3] Beiran et al., 2021. Shaping Dynamics With Multiple Populations in Low-Rank Recurrent Networks

[4] Galgali et al., 2023. Residual dynamics resolves recurrent contributions to neural computation

We sincerely thank the reviewer for their valuable feedback & positive remarks on the originality & significance of our work in advancing efficient interpretation & training of low-rank RNNs. To address the reviewer’s comments, we’ve conducted new analyses, detailed below & believe the revised paper is much stronger. We humbly request the reviewer to consider adjusting their score if our response satisfactorily addresses their concerns.

Weaknesses:

1. Method not clear enough

   a. Rank $r$: We thank the reviewer for calling attention to this.

a.1. Benchmark systems: As correctly identified by the reviewer, our approach introduces a new perspective on low-rank RNNs by modeling them as low-dimensional ODEs. Thus given the underlying ODE, the rank of the network equals the dimensionality of the ODE (e.g., Section 3,3.1: 1D, 2D ODE & SI Fig. 3: 3-dim Lorenz system).

a.2. Real world datasets: Our work builds on theoretical foundations laid out in low-rank RNN literature (e.g., Mastrogiuseppe & Ostojic, 2018; Beiran et al., 2021; Dubreuil et al., 2022; Valente et al., 2022) & the “task training” RNN literature (e.g., Mante & Sussillo 2013; Sussillo & Barak 2013; Luo, Kim, et al., bioRxiv 2023), which focus on training RNNs to perform the decision-making tasks we consider in Section 6 of our paper. Thus, while these tasks aren’t high-dimensional, they're challenging for animals to perfect & perform. Additionally, we acknowledge the reviewer's criticism on not providing a method to determine the optimal rank. We agree this is an interesting problem, which we don’t yet know how to solve. As far as we’re aware, none of the previous literature provides a theoretical solution to this problem (e.g., current methods try models with different ranks & pick those with the highest performance & interpretability). The significance of our method is: given the ODE of some dimension, we directly embed it into the low-rank RNN (where rank = ODE dim) without needing to train multiple models using FORCE/backprop.

b. Activation function role: The reviewer is correct: different activation functions do influence performance. We've added a comparative analysis (SI: A.1 & Fig. SI-2) for two different ODEs, one which is better approximated with Tanh & the other with ReLU. Our analysis suggests the following criteria:
- Piecewise linear dynamics can be better approximated by smaller ReLU networks (Similar performance is achieved with larger Tanh networks).
- Tasks with more non-linear dynamics are better approximated by smaller Tanh networks.

c. Probability distributions: We’ve added an ablation study (SI: E & Fig. SI-9) showing similar performance (similar # of neurons with OMP) on the bi-stable ODE, regardless of the distribution (uniform/standard normal) of the basis functions. Thus, our results suggest: 1. the specific distribution of basis functions isn’t significant (if they cover the same domain/support of the function being approximated); 2. Critically, shifted basis functions ( ($I_i$)) reconstruct non odd-sym functions (Fig. 1C–F & SI-1) – irrespective of the distribution from which they’re drawn.

2.Training time BPTT: We’ve added tables comparing training time across models & tasks (SI-F & Table 1,2).

3.Lorenz Attractor: We’ve added a figure (SI-7) showing a sample true & fitted trajectory & its corresponding time-series.

4.Interpretability: We thank the reviewer for this question, as this is a key contribution of our paper. Our work shows: a low-rank RNN implementing a particular dynamical system can be seen as providing a description of an ODE using a linear combination of nonlinear basis functions; the role of each neuron in a low-rank RNN is thus to contribute a single basis function. This perspective removes the need to use dimensionality reduction methods on trained RNNs (e.g., PCA, which other methods use to interpret recurrent connectivity/dynamics), or numerical optimization methods to find fixed points (common approach to interpret full rank RNNs). Thus, to summarize the key features:

• 1. Each neuron represents a specific non-linear basis
• 2. The ODE is a weighted sum of each basis, showing each neurons contribution
• 3. OMP selects optimal neuron profiles for approximating a function, creating smaller networks
• 4. Time-varying dynamics map directly to time-varying basis functions (or neuron responses).

We’ll edit the text (Sec. 3 & 7) to clarify these points.

Real World Applicability: We build on literature on “task training” RNNs (“weaknesses”- a. rank r - 2: real world datasets), which focus on training models to perform these “simple” (i.e rank 1 or 2) neuroscience tasks (eg. decision-making: Section 6 of our paper). However, to the reviewer’s comment, we’ve added two additional dynamical portraits representing evidence accumulation (Fig SI-5).