Title: Inference of Neural Dynamics Using Switching Recurrent Neural Networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: Neural Dynamics

Variational Inference

Recurrent Neural Networks
Abstract: Neural population activity often exhibits distinct dynamical features across time, which may correspond to distinct internal processes or behavior. Linear methods and variations thereof, such as Hidden Markov Model (HMM) and Switching Linear Dynamical System (SLDS), are often employed to identify discrete states with evolving neural dynamics. However, these techniques may not be able to capture the underlying nonlinear dynamics associated with neural propagation. Recurrent Neural Networks (RNNs) are commonly used to model neural dynamics thanks to their nonlinear characteristics. In our work, we develop Switching Recurrent Neural Networks (SRNN), RNNs with weights that switch across time, to reconstruct switching dynamics of neural time-series data. We apply these models to simulated data as well as cortical neural activity across mice and monkeys, which allows us to automatically detect discrete states that lead to the identification of varying neural dynamics. In a monkey reaching dataset with electrophysiology recordings, a mouse self-initiated lever pull dataset with widefield calcium recordings, and a mouse self-initiated decision making dataset with widefield calcium recording, SRNNs are able to automatically identify discrete states with distinct nonlinear neural dynamics. The inferred switches are aligned with the behavior, and the reconstructions show that the recovered neural dynamics are distinct across different stages of the behavior. We show that the neural dynamics have behaviorally-relevant switches across time and we are able to use SRNNs to successfully capture these switches and the corresponding dynamical features.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper develops a switching RNN (SRNN) framework to model neural activity. It builds up on switching linear dynamical system models that are used in neuroscience to segment and extract underlying dynamics of observed neural activity. The different segments corresponding to unique dynamics often reflect distinct behavioral states. The crucial novelty of this work is that they allow the dynamics to be non-linear, unlike SLDS and rSLDS, making the model more expressive. They fit these models using VI using an inference network. Finally, they apply SRNN to synthetic data, as well as 3 distinct neural datasets and show that it outperforms SLDS and rSLDS on segmenting activity into behavioral modules where each module corresponds to distinct dynamics. They visualize these underlying dynamics, and also evaluate their fitted model on predicting future neural activity.\n\nSTRENGTHS:\n1. As we move towards large-scale neural datasets, it is crucial to scale model complexity in order to fully harness these datasets. This paper makes a step in that direction by allowing for non-linear dynamics, while also providing an appropriate fitting approach.\n2. The experiment section is extensive, and I appreciate the application to multiple neural datasets. I particularly found the results on the decision-making dataset to be most impressive. \n3. The literature review is thorough, and the authors do a good job of situating their work in the context of other related studies.\n\nWEAKNESSES:\n1. The authors mention switching nonlinear dynamical systems (Dong et al. 2020), and discuss how their work differs from Dong et al. I think it is important to either provide an experimental comparison to SNLDS or a justification for why these existing models are insufficient to explain neural datasets, as the main novelty/motivation for SRNN and SNLDS is very much related (also noted by the authors in the paper). More on this in the  question section.\n\n2. Behavioral segmentations are somewhat subjective in nature, and while I can see that in the experiments shown here they make sense, in a real world setup we may want to infer the number of such segmentations from the data. Here the authors set the number of discrete states to the # of true behavioral states, however this might not be known in practice. Furthermore, there might be distinct sets of dynamics within one behavioral state due to other reasons not totally explicit from behavior. From the current set of results, it is not clear if SRNN is capable of inferring the # of underlying states. I will elaborate more in the questions section on this as well. \n\n3. I also think the paper will benefit from some editing by the authors. The references are not formatted properly, commas are missing. The referencing to supplementary figures doesn't seem to be working, it links back to figures in the main text. I also think the authors can trim some of the background, such as the section on VI, in favor of explaining some of the experiments such as the Lorenz attractor setup in more detail.\n\n4. While I appreciate the extensive experiments, I find it hard to reconcile some of the results. It seems like in some of the plots (Fig 3C/D, Fig 5C/D) prediction + reconstruction performance across all models is similar. However, the discrete states being inferred look hugely inaccurate for SLDS and rSLDS. I wonder if the authors have thoughts on why this happens.\n\nQUESTIONS:\n1. I feel it is crucial to understand the pros and cons of SNLDS vs SRNN, and understand if the differences between the two models are consequential in modeling neural data. For example, I understand some of the modeling differences such as the dependence of discrete state transitions on the data in SNLDS vs the previous continuous state in this paper, but I am not sure if one of the two is a better assumption. Additionally, I am also struggling to understand why can SNLDS not be used for prediction?\n\n2. In the real-world experiments, were the number of discrete states across all models being compared set to the true number of behavioral states? I would be curious to see how results vary across different # of discrete states across these models, perhaps MSE on prediction or ELBO vs # of discrete states is a possible way to show this. This is for two mains reasons:\ni. In a new dataset we might not know ground-truth behavioral segmentations, and will want to be able to infer the number of such segmentations from the data. Hence, it would be interesting to see if SRNN can be used to do so.\nii. The fact that rSLDS does fine it predicting data but infers discrete states inaccurately makes me wonder if it is clustering data differently, perhaps collapsing 2 states into one or further segmenting one behavioral state into multiple slightly different states.\n\nLIMITATIONS:\nThe authors have addressed limitations in the last section of the paper, and I do not envision any societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors develop a new class of probabilistic nonlinear state space models called switching RNNs. In essence, this extends the well-known switching linear dynamical system (SLDS) model to switch between nonlinear dynamics governed by a stochastic RNN.\n\nSTRENGTHS:\n* The results shown in panels A of Figs 3, 4, and 5 are nice and convincing.\n\nWEAKNESSES:\n* Like many other deep learning based approaches, the model is not particularly interpretable. For example, panel F in Figs 3, 4, and 5 shows 2D flow fields for the different hidden states, but the RNN hidden state is 16-dimensional. Here the authors have used PCA to attempt to find a reasonable 2D flow field, but I know from experience that this has the potential to very poorly capture the true dynamics of the system. Intuitively, even small variance dimensions can matter a lot if the flow field changes rapidly along that dimension.\n\n* There are many tunable parameters in this model (e.g. number of continuous and number of discrete states). It is unclear how to choose these on datasets without ground truth, or at least good educated guesses.\n\n* Related to above, I worry a lot about the identifiability of this model. A nonlinear RNN without discrete switching can already model any flow field if given enough units. Thus a model with many continuous states (e.g. $P=128$) but zero discrete states may perform equally well to a model with few continuous states (e.g. $P=16$ or $P=8$) but a handful of discrete states. How would one then go about choosing between these models? Adding discussion or ideally some sort of mathematical analysis regarding the statistical identifiability of the model would be very helpful.\n\nQUESTIONS:\n* Equation (2) seems wrong to me. The nonlinearity $f(\\cdot) = \\tanh(\\cdot)$ doesn't make sense here since $p(z_t \\mid z_{t-1}, h_{t-1})$ should be a positive number. Perhaps you meant to use a softmax nonlinearity here?\n\n* Equation (2) also seems to suggest that the transition probability only depends on $h_{t-1}$ and not $z_{t-1}$. Is this correct?\n\n* Related to the above, it wasn't obvious to me whether you are allowing the continuous hidden state to impact the transition for the discrete state. Essentially I am wondering if your model is analogous to the switching LDS (where the continuous hidden state doesn't impact the transition statistics) or instead the recurrent switching LDS (where the continuous state does impact the discrete transition probabilities). In Figure 1B, what is the meaning of the red dashed arrow? Does that carry any difference to the black arrows?\n\n* Are the good results shown in panel A of Figs 3, 4, and 5 due solely to differences in the initialization procedure across models? (see top of page 5)\n\n* Regarding identifiability, what happens if you run the model multiple times from different random seeds? Do you recover the same flow fields and fixed point structure?\n\nLIMITATIONS:\nThe discussion adequately acknowledges limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to model time series neural population activity using switching recurrent neural networks. The generative model includes discrete latent states\n\nSTRENGTHS:\nThe proposed method does appear to outperform related switching linear dynamical systems approaches in certain contexts.\n\nWEAKNESSES:\nHigh-level:\n- The contribution beyond other switching nonlinear dynamical systems models is not clear. Such models include the cited Dong et al., 2020, as well as Karniol-Tambour et al., ICLR 2024. If there is a contribution beyond these works, the authors should compare against those existing related methods.\n- The authors do not demonstrate an ability to automatically determine the appropriate number of discrete states. One approach to this might be \"co-smoothing\" (see Yu et al., Gaussian Process Factor Analysis, 2009).\n\nDetails:\n- The mathematical details and notation are often unclear. For example, equation 2 does not appear to be a valid probability distribution, given the description that f(.) = tanh(.). Shouldn't this instead be a categorical distribution or similar? Relatedly, f is also used in equation 8, but from the context it appears to denote something entirely different.\n- The authors should more clearly describe the cross-validation techniques for used for each dataset. The blanket statement in the intro to Section 4 (\"On each dataset, we do N-fold cross-validation, where N equals to the number of conditions, sessions, or subjects in the dataset\") obscures how cross-validation was actually applied in each instance.\n\nQUESTIONS:\n- Are the predictions in Figure 2 cross validated (eg., using the technique described in section 3.3?\n- In Fig 3, are the authors modeling single-trials or condition averages (ie PSTHs)? This should be addressed. It looks like they are predicting condition averages since the \"true\" neural activity in 3E takes on continuous values (rather than indicating spike times or binned spike counts).\n- Why does SRNN perform worse than SLDS and rSLDS in Figure 5CD?\n\nLIMITATIONS:\nThe authors address several limitations, including their need to manually set the number of discrete states, their need for good parameter initializations, and the heavy computational requirements for fitting their models.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes switching recurrent neural networks (SRNN), which allow the RNN weights to switch across time. The RNN weights switch based on a latent Markovian process of discrete states. The authors apply SRNN to a simulated dataset following the Lorenz attractor and three real-world neural recordings.\n\nSTRENGTHS:\n- Clarity: The authors clearly explain the problem, related work, and methodology with well-written equations and easy-to-understand figures. \n\n- Extensive use of datasets: The paper applies SRNN to numerous real-world neural datasets, illustrating the effectiveness of SRNN in accurately segmenting different datasets in an unsupervised fashion.\n\nWEAKNESSES:\n- Lack of comparison with other methods:\nThe paper compares SRNN to (r)SLDS models. However, there exist many other models for unsupervised segmentation. For example, ARHMMs and their extensions are simple yet powerful and interpretable models for segmentation [1, 2]. The authors should cite and consider comparisons with multiple model classes.\nIn addition, the paper notes in line 103 that SRNNs have the most comparable structure to SNLDS, but the authors do not make comparisons. The authors should also cite and compare with [3], which has switching nonlinear dynamics.\n\n[1] Wiltschko, A. B., Johnson, M. J., Iurilli, G., Peterson, R. E., Katon, J. M., Pashkovski, S. L., ... & Datta, S. R. (2015). Mapping sub-second structure in mouse behavior. Neuron, 88(6), 1121-1135.\n\n[2] Lee, H. D., Warrington, A., Glaser, J., & Linderman, S. (2023). Switching autoregressive low-rank tensor models. Advances in Neural Information Processing Systems, 36, 57976-58010.\n\n[3] Karniol-Tambour, O., Zoltowski, D. M., Diamanti, E. M., Pinto, L., Tank, D. W., Brody, C. D., & Pillow, J. W. (2022). Modeling communication and switching nonlinear dynamics in multi-region neural activity. bioRxiv, 2022-09.\n\n- Experiments:\nThe simulated experiment with the Lorenz attractor shows that SRNN does well when it has access to noiseless observations with known state dimensions. In order to have a more convincing simulated experiment, the authors could consider the following. First project the Lorenz attractor to a higher dimensional space and add additive Gaussian noise. Then fit SRNN (and other compared models) to the dataset to see if it can recover the Lorenz attractor and true latent state dimension (using some metric on held-out data). Another simulated experiment could be done with a dataset that simulates the NASCAR track [1,2].\n\n[1] Linderman, S. W., Miller, A. C., Adams, R. P., Blei, D. M., Paninski, L., & Johnson, M. J. (2016). Recurrent switching linear dynamical systems. arXiv preprint arXiv:1610.08466.\n\n[2] Lee, H. D., Warrington, A., Glaser, J., & Linderman, S. (2023). Switching autoregressive low-rank tensor models. Advances in Neural Information Processing Systems, 36, 57976-58010.\n\nQUESTIONS:\n- What are some failure modes of the model? Does extra flexibility mean that SRNNs need more data than simpler models such as ARHMMs or SLDSs? I'm curious how the SRNNs would do in a low-data regime (e.g., sample a small amount of dataset from an SLDS).\n\n- Have you tried fitting the model to datasets other than neural data? Based on how well SRNNs do in segmenting the neural datasets, I'm curious how SRNNs would do on other types of datasets, such as mouse behavioral dataset [1].\n\n[1] Wiltschko, A. B., Johnson, M. J., Iurilli, G., Peterson, R. E., Katon, J. M., Pashkovski, S. L., ... & Datta, S. R. (2015). Mapping sub-second structure in mouse behavior. Neuron, 88(6), 1121-1135.\n\n- How are the hyperparameters selected? Based on how long it takes to fit each SRNN to the datasets, I wonder if it is feasible to sweep over the hyperparameter space.\n\nLIMITATIONS:\nAs the authors noted, some limitations of the model are that the model needs good initialization and that the model takes considerably more amount of time to train than simpler models such as SLDSs.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}