Inferring stochastic low-rank recurrent neural networks from neural data

Summary

The authors propose fitting stochastic low-rank RNNs to neural recordings using variational sequential Monte Carlo methods. Such techniques permit modeling of noisy sequences (i.e., trial-to-trial variability), identification of a low-dimensional latent dynamical system, generative sampling of neural trajectories with realistic variability, and interpretation via fixed point analysis. The technique is applied successfully to recover the ground truth dynamics in two synthetic systems, and then to model EEG recordings, hippocampal spiking data, and motor cortical spiking data.

Strengths

- Originality: The technique presented appears original in its combination of existing ideas from low-rank RNNs, variational inference, and sequential Monte Carlo. - Quality: Thoughtful comparisons were made to existing approaches. The results improve upon state-of-the-art techniques (e.g., Generalised Teacher Forcing) in certain settings. - Clarity: The figures and tables are clearly presented and quite interpretable. Much of the writing is clear, although see Weaknesses and Questions for suggestions here.

Weaknesses

- Section 2.2 could benefit from being made more accessible to readers who are not experts in the subdomains of variational inference and sequential Monte Carlo methods. - The authors should make clear how to explicitly implement the technique. - The approach does not outperform state-of-the-art techniques in the Neural Latents Benchmark (NLB). While the authors mention that the "NLB metrics center around evaluating the quality of smooth rates inferred from spikes, which is not the central focus of our method. Rather, we aim to fit an RNN, from which -- by design -- we can sample noisy latent trajectories that reproduce variability in the data." But doesn't LFADS (and NDT?) allow generation of noisy latent trajectories (in the LFADS "factors") that reproduce variability in the data?

Questions

- What is meant exactly by "tractable dynamics"? Is there a distinction between "tractable" and "interpretable"? - Why is the proposed technique compared to Generalised Teacher Forcing in the EEG experiments, but not in the synthetic setups (Fig 3) or the spiking data (Figs 5-7)?

Rating

6

Confidence

2

Soundness

4

Presentation

3

Contribution

3

Limitations

- The authors note just one assumption / limitation of their approach: an assumption of correlated Gaussian noise in the recurrent dynamics. How does the technique fair when observations from the true underlying system reflect private noise processes (e.g., measurement noise, variability in neurotransmitter release, etc)? - One limitation that was not discussed or addressed is that the approach does not model the effects of unobserved inputs (beyond noting that correlated noise in the dynamics may arise due to unobserved inputs). Does this imply that the technique as presented is only appropriate for settings where the true dynamics can be reasonably modeled as an autonomous system?

Summary

This paper proposes a low-dimensional, nonlinear dynamics model of neural data based on low-rank RNNs. In particular, the model dynamics are a discretization of the dynamics of a low-rank RNN in the space spanned by the column factors of the low-rank RNN matrix. The model is fit using variational SMC. The approach is validated in simulation and applied to three neural datasets of EEG, hippocampal, and motor cortex recordings. Additionally, for low-rank RNNs the authors introduce an exact fixed point finding procedure that runs in polynomial time.

Strengths

The modeling approach and fitting methods are generally clearly described. The model is applied in multiple synthetic and real applications, where it recovers known ground truth or produces sensible results. Additionally, the proposed analytical technique for finding fixed points of low-rank RNNs appears to be generally useful for anyone working with such models.

Weaknesses

It appears that the formulation of the model in the low-dimensional space loses some generality relative to the original low-rank RNNs. As the authors point out, for non-constant inputs the activity can lie outside the low-d subspace. While the proposed approximation in appendix C.2 may work for the settings in this paper, it is not necessarily clear how the approach would perform if the inputs to the original low-rank RNN are significantly non stationary. All of the models fit in this paper have small latent dimensionalities (typically 2-3 latent dimensions). It is not clear how well the approach scales to larger dimensionalities than those considered here.

Questions

- The authors point out in the discussion that the proposed model can be viewed as a neuralODE, when discussing the FINDR method. This paper could be improved by more clearly and directly discussing how the model and inference approach relate to other models and fitting methods in the neuralODE literature (another neuro example is ODIN from Versteeg et al). Since this model can be viewed as a neuralSDE, should it be using techniques from the neuralSDE literature for fitting? - Generally, low-rank RNNs can have transient dynamics outside the column space of the weight matrix factors, e.g. via inputs as the paper mentions. The current approach may be limited to handling only certain types of inputs. I think it would be helpful for the authors to describe this difference and the limitations more clearly in the main text.

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes

Summary

The authors describe an elegant method to infer a low-dimensional description of stochastic neural dynamics using variational Sequential Monte Carlo and Low Rank RNNs. They apply this method to simulated data and to three experimental datasets (EEG, hippocampus, and motor cortex). They also describe an elegant analytic approach to determine the fixed points of the learned dynamical system.

Strengths

The paper is very well-written. The motivating problem is well-established as an important question in computational neuroscience. The discussion and summary of prior work on this topic is executed very well. In addition to showing strong empirical results in simulation and real experimental analysis, the authors provide an illuminating mathematical analysis of low-rank RNNs and a procedure to efficient numerical procedure to find fixed points. A key modeling advance is the ability to fit a model with stochastic and nonlinear dynamics. This is in contrast to most RNNs, which currently use a deterministic transition.

Weaknesses

I enjoyed this paper and found very few weaknesses. However, I would ideally like to see the variational SMC method benchmarked against simpler baselines -- mostly I would be interested in SMC methods with simpler proposals (e.g. bootstrap particle filtering). Additionally, I would like to see more details about performance as a function of (a) number of particles, and (b) amount of observed data.

Questions

n/a

Rating

8

Confidence

5

Soundness

4

Presentation

4

Contribution

4

Limitations

The limitations are well discussed.

Summary

This work focused on inferring stochastic low-rank structure from neural data. It developed a low-rank RNNs as state space models, and using sequential monte Carlo (SMC) to learn the model's parameters. The proposed model is efficient in finding all fixed points in a polynomial cost instead of exponential cost. It is evaluated on multiple settings, one from student-teacher setups, given ground truth teacher RNN, and infer the structure and statistics, and demonstrated on extracting latents from EEG data, and recover interpretable latent dynamics from spike recordings in hippocampus, and allow position decoding.

Strengths

**Motivation** This work focused on important question in neuroscience, given partial observation and initial states, building stochastic models is essential to model complex neural dynamics. **Method** The proposed method is efficient in terms of the low-rank structure, and as well as low complexity in finding fixed points. **Evaluation** The method is demonstrated on extensive neural datasets, as well as synthetic networks. Evaluations are compared with deterministic model GTF.

Weaknesses

**Novelty** Low-rank RNN has been demonstrated on many applications in recent works, the novelty is limited. **Experiments** Adding simulation with partial observations to demonstrate the effectiveness of stochastic low-rank RNN. **Results** Limited performance in recovering latent structure and predicting dynamics, how does GTF perform in this task? 1. Fig 3 a and b. 2. Fig 4, the generated trace is not similar to EEG signal ground truth in the qualitative evaluation. 3. Fig 5c, pairwise correlation is lower than 0.10. 4. Fig 7e, the mismatch between true and generated in mean ISI dissimilarity.

Questions

1. Explain the limited performance of model as mentioned above, how does GTF perform in the same task? 2. Ablation studies of stochasticity and low-rank constraints.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

No potential negative societal impact.

Reviewer Response

Thanks for the additional experiments. I retain my positive score.