Inference of Neural Dynamics Using Switching Recurrent Neural Networks

Authors: yongxu  zhang
shreya  saxena
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:06:08

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.

Main paper: 
 introduction: 
 How does complex neural activity lead to dynamic behavior? A foundational principle in theoretical neuroscience suggests that the computations within the nervous system can be explained through the dynamics of the underlying non-linear systems (Breakspear [2017] , Pandarinath et al. [2018] , Durstewitz et al. [2023] ). As our ability to record from larger areas of the brain with unprecedented spatial and temporal resolution increases, we find neural population activity often exhibits distinct dynamical features across time. These features may correspond to distinct internal processes or behavior (Saxena and Cunningham [2019] , Churchland et al. [2012] ). Identifying these dynamical features may help us understand how cognitive functions are implemented in the brain. Previous works have revealed a number of computational strategies through analysis of neural dynamics, particularly across tasks, and during orchestration of precise behavioral states such as during movement generation (Kaufman et al. [2014] , Miri et al. [2017] ).State space modeling is a promising analytical approach for characterizing dynamics of time series; such models can be highly effective and have received substantial attention over many decades on analyzing temporal neural dynamics (He et al. [2023] , Durstewitz et al. [2023] , Sani et al. [2021] ). Well-established methods on inference and learning algorithms have contributed towards learning the parameters of state space models (Blei et al. [2017] , Salimans et al. [2015] , Kingma and Welling [2013] , Linderman et al. [2017] , Fox et al. [2008] ). Neural dynamics can be time-varying due to internal fluctuations of physiological states as well as the external effect of environment. Previous works such as (Kaufman et al. [2014] , Miri et al. [2017] ) have found differences in neural subspaces between distinct behavioral states. Furthermore, correlation-based analysis methods, such as functional connectivity, also reveal that the brain-wide covariation in neural activity changes during movement. For example, West et al. [2022] highlight the changes in functional connectivity in the cerebral cortex, representing a series of changes in the cortical state from rest to locomotion and on return to rest. In practice, researchers typically model switching temporal dynamics by fitting different model parameters to consecutive temporal windows of neural activity (He et al. [2023] , Mitelut et al. [2022] , Song et al. [2022] ).In the past decade, we have seen that neural dynamics can be broadly considered to be non-linear and in a lower dimensional subspace of the recorded neural activity (McKenna et al. [1994] , Rigotti et al. [2013] , Hernandez et al. [2018] , Saxena and Cunningham [2019] , Cunningham and Yu [2014] ). Accordingly, Recurrent Neural Networks (RNNs) have shown efficiency in modeling dynamics due to their non-linearity (Mante et al. [2013] ). While goal-driven models have been very useful to build in normative function in these networks, data-driven RNNs are becoming more popular over recent years (Durstewitz et al. [2023] , Perich et al. [2020] , Duncker and Sahani [2021] , Valente et al. [2022] ). Specifically, these RNNs can be trained to reproduce temporal neural activity from large-scale neural recordings across a set of trials. They are also able to capture a very high amount of explained variance of the neural activity because of their static non-linearities (Durstewitz et al. [2023] , Perich and Rajan [2020] ). Once trained, the internal mechanisms of the RNNs can be analyzed, thus extracting the structure of the neural dynamics. Prior works have studied RNNs that perform a wide range of tasks, e.g., cognitive, sensory, motor tasks, and so on (Dubreuil et al. [2022] , Mastrogiuseppe and Ostojic [2018] ). However, with a few notable exceptions, models do not typically express neural activity as a set of switching nonlinear functions and directly reproduce recorded neural data. Here, we develop Switching Recurrent Neural Networks (SRNNs) which, through switches between different time-varying weights as determined by Markov transitions, are able to directly model neural activity as emanating from a discretely changing set of low-dimensional dynamical models. We perform end-to-end training of SRNNs to reconstruct the neural observations through Variational Inference. We validate SRNNs on a simulated dataset, and then analyze the performance of SRNNs on three different experimental data with distinct recording modalities and behavioral tasks with different animals: (1) electrophysiological recordings of single-unit MC activity from a non-human primate performing a reaching task (Churchland et al. [2012] ), (2) cortex-wide widefield calcium imaging (WFCI) from mice performing a complex self-initiated decision-making task (Musall et al. [2019] ), and (3) WFCI from mice performing a simple self-initiated lever-pull task (Mitelut et al. [2022] ). We are able to not only reconstruct the multidimensional neural activity but accurately predict forward in time and recover behaviorally-relevant switches of neural dynamics. Finally, we visualize the neural dynamics in flow fields plotting and find they are distinct among different states. 
related work: 
 Data-Driven Neural Networks An artificial neural network trained to directly reconstruct observed neural activity is called a data-driven neural network. Prior studies have concentrated on employing various models to extract the underlying neural dynamics and characterize the connectivity between different components of the brain (Perich and Rajan [2020] , Mastrogiuseppe and Ostojic [2018] ). In Perich and Rajan [2020] , the authors introduce Current-Based Decomposition (CURBD), an approach for inferring brain-wide interactions using data driven RNNs that directly reproduce experimentallyobtained neural data. In Mastrogiuseppe and Ostojic [2018] , the authors used low-rank RNNs trained on high-dimensional neural activity obtained from different tasks and successfully characterized the neural dynamics and connectivity between neurons. Importantly, previous work does not consider switches in the underlying low-dimensional neural dynamics; in our work, we focus on the extraction of low-dimensional neural dynamics based on switching data-driven neural networks.Switching State Space Models Hidden Markov Model (HMM) is a commonly-used statistical model based on a Markov process: it dictates the evolution of observations where the evolution is determined by internal factors, which are not directly observable. HMMs are widely used in computational neuroscience to identify different neural events (Baldassano et al. [2017] , Masaracchia et al. [2023] , Baldassano et al. [2016] ). Previous studies have also explored switching dynamical systems, with the observations emanating from a dynamical system that discretely switches according to a Markov process. Auto-regressive HMMs and Switching Linear Dynamical System (SLDS) are commonly used models, where the discrete state transitions and observation dynamics are both linear and dependent solely on the previous timepoints. Various extensions of these also applied changes to these models. For example, in Linderman et al. [2017] , the authors make the current discrete state dependent on the previous continuous state, termed recurrent-SLDS (rSLDS). Researchers have also changed the dynamics process to be nonlinear, such as in a Structured Variational Autoencoder (SVAE) (Johnson et al. [2016] ). Moreover, both nonlinear transitions and nonlinear dynamics processes were also proposed, such as in a Switching Nonlinear Dynamical System (SNLDS) (Dong et al. [2020] ). Out of these, SRNNs have the most comparable structure to SNLDS; the main differences are: (1) we specify the dynamical system as an RNN instead of passing the data to neural networks with nonlinear activation functions; (2) our transition process is the same as rSLDS, which is dependent on the previous hidden state instead of the observation in SNLDS; (3) we constrain the structure of the models in order to predict future observations. Moreover, here we analyze the ability of the model to reconstruct and elucidate the structure of recorded neural activity, while examining prediction as well as region-to-region interactions (Appendix). A relevant recent study also explored the communication between brain regions via multi-regional switching dynamical systems (MR-SDS) (Karniol-Tambour et al. [2023] ), which reveals the relationship between brain regions but does not focus on analyzing behaviorally-relevant neural dynamical states and characterizing distinct neural dynamics. Variational Inference Researchers have proposed numerous powerful algorithms for optimizing the parameters of Switching State Space Models. A well-known approach is the Expectation-Maxmization (EM) algorithm, which is to find maximum likelihood estimates of parameters in State Space Models, e.g., Baum-Welch Algorithm for HMMs (Bishop and Nasrabadi [2006] ). Another common approach to estimating the parameters of Switching State Space Model is Variational Inference (VI), which turns the problem of computing conditional distributions of latent variables into an optimization problem that can be efficiently solved by gradient descent. Essentially, this consists of maximizing the Evidence Lower Bound (ELBO) to learn the parameters (Kingma and Welling [2013] , Linderman et al. [2017] ). The reparameterization trick is a crucial technique in VI, which enables more efficient gradient estimation during optimization. Instead of sampling from the variational distribution itself, reparameterization models the parameters of a distribution, usually Gaussian. This allows the gradient of the objective function with respect to the variational distribution parameters to be computed directly, and backpropagation to be implemented easily. To enable the model to generalize effectively to unseen data, we use Amortized VI. Instead of optimizing for single trials of neural activity independently as in traditional VI, Amortized VI allows us to train the models on multiple trials simultaneously, thus creating a more stable mapping between the observations and the latent variables (Dong et al. [2020] , Ganguly et al. [2023] , Margossian and Blei [2023] ). 
switching recurrent neural network (srnn): 
 A common switching state space model is:EQUATIONwhere z t ∈ {1, ..., K} is the discrete hidden states or labels, controlling the switch of the dynamics, h t is the continuous latent state with dimensionality P , and y t is the observation or output; in this work, the observations are directly the neural activity. We build a data-driven model with neural time-series data y 1:T ∈ R N ×T ×R from N sequences, T time points and R different neural dimensions as the outputs. Here, we implement a transition network, a set of dynamical networks with recurrent layers and a nonlinear activation function, and finally a linear emissions network to reconstruct the observed neural activity. We show the structure of SRNN in Figure 1A ; following are the equations of the generative network:p θz (z t |z t-1 , h t-1 ) = Cat(Sof tmax(f (θ z h t-1 ))) (2) p θ h (h t |h t-1 , z t = k) = N (h t |f (θ h k h t-1 ), σ h ) (3) p θy (y t |h t ) = N (y t |θ y h t , σ y )(4) where Cat represents Categorical distribution, Sof tmax represents Softmax activation function, N represents Gaussian distribution, and σ h and σ y are the covariance matrices of the relevant Gaussian distributions. θ z , θ h , and θ y are parameters of transition networks, dynamical networks, and emission networks respectively, which are learned through gradient descent. Here, we consider the nonlinearity f (•) = tanh(•) as commonly done in neuroscience in order to model the saturation effect of firing rates Pandarinath et al. [2018] , Saxena et al. [2022] .We use VI to learn the parameters of our model. Here, we maximize the posterior: q θ,ϕ (z 1:T , h 1:T |y 1:T ) = q ϕ (h 1:T |y 1:T )p θ (z 1:T |h 1:T , y 1:T ) (5) where q ϕ (h|y) is the posterior of continuous hidden states given the observation, ϕ is the parameters of our inference network. p θ (z|h, y) is the Bayesian posterior of discrete hidden states given the continuous hidden states and the observations, which we compute using the generative network via Bayes' rule.Inference Network and q ϕ (h 1:T |y 1:T ) We show the structure of our inference network in Figure 1B ; we build an inference network to model q ϕ (h|y) as the following:EQUATIONwhere µ and σ are the mean and variance of a Gaussian distribution. Here, we use the reparameterization trick to infer the continuous hidden states Kingma and Welling [2013] :EQUATION) where ϵ ∈ N (0, 1) is a Gaussian noise parameter. The mean µ and variance σ are typically learned via training neural networks Kingma and Welling [2013] . However, here, we do not find apparent improvement with learning the σ in our work. Therefore, as done in previous studies, we set the variance σ to be constant to reduce the complexity of the model and only optimize the µ Dong et al.[2020], Ganguly and Earp [2021] ; we use 0.0001 for σ. Since our generative model is based on RNNs, we make h t depend on h t-1 to match the recurrence property. Thus, we have µ t = f (y 1:T , h t-1 ) (8) We pass y 1:T to a bidirectional recurrent neural network followed by forwarding the output to a standard recurrent neural network to model Equation 8.Evidence Lower Bound and p θ (z 1:T |h 1:T , y 1:T ) Typically, to learn the parameters via VI, we maximize the evidence lower bound (ELBO). Here, the ELBO is defined as:EQUATION) With Equation 1and Bayes' rule, we getEQUATIONwhere the discrete hidden states z 1:T have been marginalized out given and Russell [2001] , Dong et al. [2020] ).K k=1 p(z t = k) = 1 (MurphyInitialization Previous studies have found that a generative model trained by VI may become stuck in a single discrete state (Alemi et al. [2018] , Dong et al. [2020] ). Researchers have explored various methods to solve this problem, such as initialization of the transitions parameters using the parameters of auto-regressive HMM in Linderman et al. [2017] and entropy regularization in Dong et al. [2020] . In our work, we address this by initializing our model using states from an HMM trained on the same training data, which encourages the model to utilize all states and optimize the weights corresponding to each of the states. Practically, during a first phase of training, we add a term to the loss that penalizes the posterior if it does not utilize all the discrete hidden states of the HMM. In the second phase of training, we remove this penalization and train the model with our loss as described in Equation 15.Optimization We use gradient descent to optimize the parameters θ and ϕ. Given the fact that we use reparameterization to model q ϕ (h 1:T |y 1:T ), and q ϕ (h 1:T |y 1:T ) is independent to θ, the gradient can be approximated as following:EQUATIONwhere h 1:T (ϕ) is the output of inference network, therefore, is dependent on ϕ. Since z p(z)=1, therefore, log p θ (y 1:T , h 1:T (ϕ)) = z p(z) log p θ (y, h|z) = E p(z) log p θ (y, h|z), and with Bayes' rule p θ (y, h|z) = p θ (y,h,z) p(z) , Equation 12is equivalent to:EQUATIONIn Equation 13, the first term can be rewritten by the Markov Chain.E p(z 1:T ) ∇ θ,ϕ log p θ (y 1:T , h 1:T , z 1:T ) = K k=1 p(z 1 = k)∇ log p(z 1 = k)p(h 1 |z 1 = k)p(y 1 |h 1 ) + T t=2 K k=1 J j=1 p(z t = k, z t-1 = j)∇ log p(z t = k|z t-1 = j)p(h t |h t-1 , z t = k)p(y t |h t ) (14)where we use the posterior probability of the discrete hidden states, e.g., p(z 1 = k|h 1:T , y 1:T ) and p(z t = k, z t-1 = j|h 1:T , y 1:T ), which can be obtained easily through Baum-Welch algorithm. The transition, dynamics, and emission are defined in Equation 2, Equation 3, and Equation 4 respectively. The second term of Equation 13 can be optimized via backpropagation through time (BPTT) of the inference network. In summary, with Equation 11, Equation 13, and Equation 14, our loss function is:EQUATIONwhere we use the fact that -E q ϕ (h 1:T |y 1:T ) log q ϕ (h 1:T |y 1:T ) is the entropy H of q ϕ (h 1:T |y 1:T ). We minimize the loss function via Adam optimizer. Our implementation is based on Pytorch 2.2.1 and we train our models using NVIDIA A100 GPUs. We provide the analysis and the results in this paper; the original code for the entire SRNNs framework on Pytorch has been made public https://github.com/saxenalab-neuro/SRNN.Dynamical Features: Fixed points and other dynamical features are crucial for understanding how nonlinear dynamical systems process sequences and maintain information (Sussillo and Barak [2013] , Golub and Sussillo [2018] ). We visualize the flow fields of RNNs, which provides a graphical representation of the evolution of network states over time. To compute these flow fields, we train our model on the entire dataset (without separating a training set from a test set) in order to capture neural dynamics across the entire data. We use 10, 000 initial hidden states uniformly sampled in a grid to compute the flow fields. Neural Activity Prediction: Given the fact that SRNNs are generative models, we test all models on prediction of future neural activity using past observations. Specifically, the models are provided a time-series of previously unseen neural activity y 1:t , with which we can infer discrete hidden states z 1:t and continuous hidden states h 1:t . Next, we sample z t:t+K and h t:t+K through transitions of discrete states and dynamics of continuous states, i.e., using Equations 2 and 3, respectively. Finally, we predict the neural activity in the future time points y t:t+K using the emissions network, i.e., Equation 4. Here, we change the bidirectional layers of our inference network to standard forward recurrent layers to match the prediction tasks (see Figure 1C ). We explore the prediction performance by giving the model an input in different lengths (t ∈ [t 0 , T ], here we use t 0 = 10 timepoints), and we also explore the prediction performance by predicting different lengths of neural activity (K ∈ [10, 20, 30, 40] timepoints). 
results: 
 We detail the results of SRNNs on inferred neural dynamics of a simulated dataset and three diverse experimental datasets ranging from monkey electrophysiology to mouse calcium imaging.On each dataset, we do N -fold cross-validation, where N equals to the number of conditions, sessions, or subjects in the dataset. All the results in this section are reported on the test set. Additionally, we compare our method with other switching dynamical systems: Switching Linear Dynamical Systems (SLDS), recurrent Switching Linear Dynamical Systems (rSLDS) (Linderman et al. [2017] ), Switching Non-linear Dynamical Systems (SNLDS) (Dong et al. [2020] ), and multiregional Switching Dynamical Systems (mrSDS) (Karniol-Tambour et al. [2023] ). Additionally, we provide comparisons with Latent Factor Analysis with Dynamical Systems (LFADS) in Appendix We first apply SRNN on a well-known simulated dynamical system, the Lorenz attractor, which is famous for its "butterfly" shape (see Appendix A for relevant equations; the dimensionality of the data is R = 3). The dynamics can be expressed as a 2-state switching system. In our work, we generate the Lorenz data using the method and code in (Linderman et al. [2017] , Dong et al. [2020] ).In Figure 2 , we show that SRNNs are able to reconstruct the Lorenz attractor precisely, and recover the switching states correctly, as well as identify the underlying dynamics in each discrete state.We explore switching models on the firing rates of single units recorded from monkey motor and premotor cortex (R = 180), while the monkey performs a reaching task in 27 reaching conditions given targets at different locations. More details can be found in Appendix B and Churchland et al. [2012] . We have 5 different behaviorally-relevant states: (1) resting;(2) delay before the 'go' cue;(3) reaction time (from receiving 'go' cue to the beginning of movement); (4) movement execution; (5) hold at the target. Using SRNNs, we recover these states purely using the neural activity.Firstly, we train SRNNs and competing methods on the two different types of reaching conditions (curved reaches and direct reaches) separately. We use N -fold cross-validation, where N equals to the number of conditions in the data; N = 18 for curved reaches and N = 9 for direct reaches.We show an example of comparison between neural dynamical states recovered by different models and the behavioral states, in Figure 3A by SRNNs in both curved and direct reaching match their corresponding behavioral states. We use P = 16 for the latent state (h) dimensionality for all three models in Figure 3A . We quantify the difference between the recovered neural dynamical states and the behavioral states by considering them an imbalanced multi-class classification, i.e., one state as one class, the error of the classification measures the difference between the neural and behavioral states. Specifically, it is defined as the average of '1-recall' obtained on each class (Mosley [2013] ). We claim that 0 error between the recovered neural dynamical states and behavioral states does not necessarily mean that the recovered neural dynamical states are perfectly identified, because the neural dynamical switches may happen consistently leading to or following the behavioral state switches. However, these states recovered by the models with small errors can still be considered as behaviorally-relevant states, which can help us understand the cognitive function implemented in the brain. We show the error for different hidden states in Figure 3B , which illustrates that SRNNs outperform SLDS, rSLDS, SNLDS, and mrSDS on recovering the behavior-relevant neural dynamical states for 4 different values for the latent state dimensionality P . We determine the number of discrete latent states via a hyperparameter sweep for SRNNs. We consider the following metrics to compare the performance of SRNNs with different K: (1) Convergence to lower number of discrete states and reuse of discrete states:We tested our model by changing the number of hidden states K while keeping the number of continuous latent states P constant. We show a comparison in Figure D.3A. We found that 61% of SRNNs with a higher number of discrete hidden states (e.g., K = 6) finally converge to the optimal number of discrete hidden states, i.e., K = 5, and 94% of SRNNs with a lower number of discrete hidden states (e.g., K = 4) had at least one hidden state reused after other states; in other words, SRNNs are not able to perform well with 4 unique discrete hidden states without reusing one of them.(2) Reconstruction performance plateau: While keeping other hyperparameters constant, the reconstruction accuracy plateaus at the same number of discrete states as in the behavior. Thus, we can set K as the minimum number of discrete states as it takes for the model to perform well. We show that SRNNs with K = 5 have lower reconstruction error in Figure D.3B. Moreover, we also implemented a 'co-smoothing' analysis (Yu et al. [2008] Karniol-Tambour et al. [2023] ); we show the results in Figure D.3C, where we found that K = 5 also does well in reconstructing the data with a 'co-smoothing' neuron drop-out analysis.(3) Variability across conditions: In stereotyped tasks or experiments, such as reaching, there may not be a significant amount of variability in the timing of behavior across conditions, and this variability can thus be used as a metric for determining the optimal number of discrete states. Here, we found that SRNNs with K = 5 have much lower variability on recovered behaviorally-relevant states than K = 4 and K = 6 (i.e., 0.098 for K = 5, 0.384 for K = 4, and 0.282for K = 6 in Figure D.3D).We show the performance of all models on the reconstruction of the neural activity in Figure 3C ; the mean squared error (MSE) between desired neural activity and the reconstruction by models indicate that SRNNs also outperform other models on reconstruction of neural activity. We note that as P increases, the MSE difference between SRNNs and other models becomes smaller. Next, we show the prediction capability of future neural activity by SRNNs and competing methods in Figure 3D . We found that the prediction capabilities of SRNNs outperform competing models with linear dynamics. In addition, we show the comparison of behavioral states recovering and neural activity reconstructing performance between SRNNs with bidirectional recurrent inference networks and SRNNs with forward recurrent networks in We visualize the neural dynamics by plotting the flow field of the latent states h in each discrete SRNN state in Figure 3E , where we used P = 16 based on Figures 3B, C . In order to visualize the flow fields in 2D, we apply principal component analysis (PCA) on the latent states of SRNNs to decrease the dimensionality of the flow fields to 2D. We see that the neural dynamics in different discrete states are distinct in Figure 3E : the plots reveal that the neural dynamics change from relatively stable in the first two states, to unstable after the 'go cue', and finally oscillatory during the execution and end of movement. We also show the flow fields for direct reaches in Figure D .1, and we found that the neural dynamics between curved and direct reaches are comparable during the first three discrete states, but their dynamics become different in 'movement' states, which matches the difference between the reaching type.We show the performance of SRNNs on a decision-making task, with neural activity recorded using widefield calcium imaging (WFCI). The number of neural components is R = 210. The details of this dataset and preprocessing can be found in Appendix B of this paper and Musall et al. [2019] . In this data, mice have 5 behaviorally-relevant states, i.e., (1) from baseline to holding of the handle;(2) holding the handle; (3) stimulus presentation; (4) delay from the end of stimulus to the start of licking; (5) licking the spout.Here, we use N = 8 pseudo-sessions, and perform N -fold cross-validation. In Figure 4A , we show an example of comparison between neural dynamical states recovered by different models, along with the behaviorally-relevant states described above. Here, we use a latent dimensionality of P = 8 for all models. We also quantify the discrete state recovery using the same approach as described in the previous section in Figure 4B , where we found that SRNNs are able to identify 5 behaviorally-relevant neural dynamical states, especially as the latent state dimensionality increases. Here, we use the same approach as in reaching task to determine K. We show the performance of SRNNs with different K in 4C, with the SRNN slightly outperforming other models. Analogous to the reaching datasets, a higher latent dimensionality results in better reconstruction. Notably, the prediction capabilities of SRNNs far outperform competing models with linear dynamics (Figure 4D ). In addition, we show the reconstruction of neural activity for an example component with the corresponding ground truth in Figure 4F , with more results shown in Figure D.2C in the Appendix. Furthermore, we visualize the neural dynamics in Figure 4E with P = 8. We again apply PCA to visualize the flow fields, where we see that the flow fields of different RNNs are also distinct. However, the difference is smaller as compared with the reaching dataset, which may reveal that WFCI may have fewer change in dynamics over the course of a trial, potentially due to the slow timescales of calcium indicators.Finally, we explore the utility of SRNNs on another WFCI data, where large-scale neural activity from the mouse dorsal cortex is recorded while mice are engaged in a spontaneous lever-pull behavior for water reward (R = 16). More details can be found in Appendix B and Mitelut et al. [2022] . Here, while the observed behavior happens in a very short time because the mice pull the lever very fast, Mitelut et al. [2022] found that the neural activity has obvious inhibitory dynamics starting from around 5 seconds before the self-initiated lever pull behavior. Thus, the mice have 3 behaviorallyrelevant states: (1) from resting to the starting of inhibitory dynamics (as found in Mitelut et al. [2022] ); (2) from the starting of inhibitory dynamics to behavior onset; (3) lever-pull execution.We use data from 6 different mice, and perform 6-fold cross-validation. In Figure 5 , we show the results with the same quantification methods as in the datasets above. SRNNs also outperform in recovering the behaviorally-relevant neural dynamical states; an example is shown in Figure 5A . While SRNNs can recover the inhibitory dynamics of neural activity from around 3 seconds before the behavior onset, SLDS and rSLDS can only identify a short state before behavior onset. We use P = 2 hidden states for this example since more hidden states do not improve the recovery of these discrete states (Figure 5B ). We use the same approach as above to determine K. We show the performance of SRNNs with different K in Figure 5 : WFCI Self-Initiated Lever Pull: (A-F) Legend similar to that of Figure 4 . Moreover, the reconstruction of neural activity using SRNNs performs slightly worse than other methods (Figure 5C ); however, the prediction of neural activity even 1 or 1.33 seconds ahead is better than competing methods. We show the reconstruction of neural activity from one example brain region in Figure 5F , with more results shown in Figure D.2D. Here, the color bars represent the performance of neural prediction. Indeed, all three methods have acceptable reconstruction and SRNNs have better prediction capabilities. Lastly, we visualize the neural dynamics in Figure 5E using P = 2. We find that the latent dynamics are only different from the starting of the inhibitory dynamics to behavior onset, which may reflect the mechanism through with internal dynamics in the brain achieve movements, with primarily changes in preparatory neural dynamics (here the inhibitory state) leading to upcoming movements (Churchland et al. [2012] ). Finally, we explore modular SRNNs constrained to the activity of different regions in order to elucidate region-to-region communication in Figure D.9. 
limitations and conclusions: 
 We developed a novel model for recovering neural dynamics and neural state changes, termed Switching Recurrent Neural Networks (SRNNs). The results show that SRNNs outperform competing models in both recovering behaviorally-relevant neural states, and predicting future neural activity. However, certain limitations of this method still exist. Firstly, we set the maximum number of discrete states manually based on our observations of behavior, but we note that the reconstruction accuracy does not change drastically if we change this hyperparameter. Secondly, analogous to other generative models trained via Variational Inference, the model needs a good initialization to be trained efficiently; without this, the model may get stuck in one state and not adequately capture the switching nature of neural dynamics. Lastly, the complexity of the model increases the training time, for example, one experiment for SRNNs trained on reaching datasets may take several hours (e.g., around 6.5 hours on one session of curved reaching on NVIDIA A100 GPUs), which is more than the amount of time taken by SLDS and rSLDS (typically less than one hour). In conclusion, we use SRNNs to identify behaviorally-relevant neural dynamical states. We find that the neural dynamics have behaviorally-relevant switches across time and we are able to use SRNNs to capture these switches as well as the corresponding dynamical features. Straightforward extensions of this method include adding inputs or control signals to SRNNs to explore the identification of behaviorally-relevant neural dynamics while providing external stimuli.While this paper focuses on neuroscience impact, methods such as SRNNs have the potential to be applied in health-related fields for positive societal impact towards design of intervention based on accurate prediction of neural dynamics. No negative societal impact is noted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We provide details in the Methods and Results section.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We provide details in the Results section.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: This is addressed in the Conclusions section. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We have cited all the source of our datasets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: The equations of Lorenz Systems are:EQUATIONEQUATIONEQUATIONwhere x, y, z represent 3 dimensions of the system, and σ, ρ, β are three parameters. Here we use commonly used values for the parameters, i.e., σ = 10, ρ = 28, β = 2.667.Electrophysiology Recordings during a Monkey Reaching Task In the reaching dataset, the recording starts with the monkey in a resting state, and a target shows up after 400ms. The monkey then gets a 'go' cue after a delay. Thereafter, the monkey starts moving from the origin to the targets and hold at the targets upon reaching. More details can be found in Churchland et al. [2012] . We explore the relationship between the change of neural dynamics and the change of behavioral states.Here, we use the trial-average neural activity which is averaged across multiple trials with the same conditions, details of the trial-average data can be found in Churchland et al. [2012] . The data is recorded for 180 neurons and has 236 time points representing 2.36 seconds.In self-initiated decision making dataset, mice initiated trials by touching either of two handles and hold a handle for 1s followed by sensory stimuli. The sensory stimulus was presented for 600ms, and after a 500ms pause with no stimulus and then the stimulus was repeated for another 600ms. After the second stimulus, a 1000ms delay was imposed, then the mice are required to lick one of two spouts, the mice were rewarded with a drop of water if they licked the spout twice on the same side as the stimulus. More details can be found in Musall et al. [2019] . Here, we average across multiple trials recorded from the same mouse to create pseudo-sessions. Specifically, we averaged around 80 trials of 516 trials in one mice to produce N = 8 pseudo-sessions, and fit the trial-averaged data using our models. The details of the data can be found in Musall et al. [2019] . The data is also processed via LocaNMF Saxena et al. [2020] , and the 210 temporal components returned by LocaNMF are used. The data has 189 time points representing 6.3 seconds.Widefield Calcium Imaging during Mouse Self-initiated Lever Pull Behavior In self-initiated lever pull data, mice were trained to pull a lever and hold it at an angle (for > 100ms) in order to receive a water supplement. We then apply LocaNMF Saxena et al. [2020] while spatially aligning the imaged neural activity with the Allen mouse brain coordinate framework Wang et al. [2020] using affine transformations, as previously performed in Musall et al. [2019] , and take 16 components as identified by LocaNMF, which form our input signals, with each input dimension from one brain region. In this work, we focus on the signals around the lever pull, i.e., from 10 seconds before lever pull to 1 second after lever pull (overall 330 time points). More details can be found in Mitelut et al. [2022] . Likewise, we use trial-averaged data which is averaged across all trials of the same mouse. We use the data across N = 6 mice; details of this data can be found in Mitelut et al. [2022] .We also use our model to uncover the relationship between different brain regions using self-initiated lever pull dataset recorded by WFCI, because the data has only one component for each brain region. We add a constraint to our SRNNs, specifically, we force 4 hidden states of SRNNs to reconstruct neural activity from 4 brain region independently via constraining the output weights of SRNNs. We focus on communication between 'Motor', 'Visual', 'Somatosensory', and 'Other' regions except these 3. Finally, we visualize the recurrent weights of each RNN of SRNNs in Figure D.9A, we find that the region-to-region communication is highly different between different states. Moreover, the self-communication during rest state (diagonal of SRNN1) is at a low amplitude. Nevertheless, this self-communication becomes higher during the 'inhibitory dynamics' (SRNN2), then decreases in 'after lever pull' (SRNN3). Additionally, the 'Visual' region has more communication with other regions during the states prior to the behavior onset (SRNN1 and SRNN2), the 'Somatosensory' region has little communication at early states, but it then increases close to behavior onset. The 'Motor' region actively communicates with other regions near behavior. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: We cite all the source of our datasets which are all in the public domain; we provide our full code for reproduction.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 


Decision: This paper describes a novel approach to modeling neural dynamics using switching RNNs.  The paper generated considerable interest and discussion among the reviewers. Ultimately, I am persuaded by the comments of the more positive reviewers, and feel that it should be accepted to the meeting. Please thoroughly revise the manuscript to address all reviewer comments and discussion points.

Reviews: review: paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: - Lack of comparison with other methods:
The 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.
In 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.

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

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

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

- Experiments:
The 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].

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

[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.
questions: questions: - 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).

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

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

- 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.
limitations: limitations: As 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.
paper_summary: paper_summary: The authors propose to model time series neural population activity using switching recurrent neural networks. The generative model includes discrete latent states
strength_weakness: strength_weakness: High-level:
- 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.
- 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).

Details:
- 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.
- 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.
questions: questions: - Are the predictions in Figure 2 cross validated (eg., using the technique described in section 3.3?
- 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).
- Why does SRNN perform worse than SLDS and rSLDS in Figure 5CD?
limitations: limitations: The 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.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: * 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.

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

* 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.
questions: questions: * 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?

* Equation (2) also seems to suggest that the transition probability only depends on $h_{t-1}$ and not $z_{t-1}$. Is this correct?

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

* 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)

* 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?
limitations: limitations: The discussion adequately acknowledges limitations.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.

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

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

4. 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.
questions: questions: 1. 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?

2. 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:
i. 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.
ii. 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.
limitations: limitations: The authors have addressed limitations in the last section of the paper, and I do not envision any societal impact of this work.

Comments: title: Response to rebuttal
A Potential Glitch
Reviewer Response
comment: I have read the authors' responses as well as the comments from the other reviewers. I stand by my rating of 3 due to the limited originality and innovation of the contribution. I do not see the impact of this work meeting the high standard of NeurIPS. While the authors did describe the differences between their approach and previous related approaches, the novelty appears quite limited, and the authors have not demonstrated which of these differences (if any) lead to improved performance.
We thank you for your response. We include a detailed feature comparison between SNLDS, mrSDS, and SRNNs in the following table. Specifically, SRNNs differ fundamentally from SNLDS and mrSDS in their **generative model, inference network**, and **initialization**, as detailed below. Moreover, we examine whether these models, trained directly on neural data, can reconstruct _**behaviorally-relevant discrete states**_: a question that has not been investigated in either of the previous studies. We found that, without fail, **SRNNs have better performance** in this goal than competing methods, including SNLDS and mrSDS.

(1) **Generative Model**: SRNNs utilize RNN as a generative model, while SNLDS and mrSDS employ a MLP with nonlinear activation functions. RNNs may offer greater interpretability and accuracy compared to feedforward models, and are very relevant to the neuroscience community for insights on computational mechanisms [1][2].

(2) **Inference Network**: mrSDS uses a non-causal transformer encoder as an inference network and SNLDS uses a non-causal combination of RNNs as an inference network. SRNNs are able to **predict** neural activity many time-steps ahead since we have implemented both **causal and non-causal** RNNs as possible inference networks. This important capability brings us one step closer to real-time and closed-loop neuroscience applications.

(3) **Initialization**: SRNN has a different initialization method compared to SNLDS and mrSDS (*unknown*). In training SRNNs, we rely on HMMs, whereas SNLDS uses uniform entropy regularization. We would like to clarify that the initialization of mrSDS is unknown, therefore, we used both initialization approaches when we implemented mrSDS and report the better performance for mrSDS.

**In summary, our primary contributions are**: (1) characterizing and interpreting low-dimensional switching nonlinear neural dynamics using SRNNs, (2) enabling causal prediction of neural activity, and (3) demonstrating that SRNNs outperform SNLDS and mrSDS in _**recovering behaviorally-relevant states and reconstructing and predicting neural activity**_, making them more reliable for interpreting behaviorally-relevant neural dynamics.

We thank you again for your questions. We hope this comparison is compelling and we look forward to your response for further discussion if needed. 

| Models (on Reaching Data) | Generative Model | Inference Network  |Initialization | Behavioral States Recovery (**Error**)  |Reconstruction (**MSE**)| Multi-regional setting | Neural Prediction | Code |
| :----:       |    :----:   |        :----: |    :----:   |        :----: |    :----:   |        :----: |    :----:   |     :----: |
| SNLDS      | MLP       | NNs   |Entropy regularization| 0.46 ($\pm$0.125)      | 0.00404 ($\pm$0.00170)       | No  |Not implemented (non-causal inference network)|Tensorflow|
| mrSDS   | MLP        | Transformer      |*Unknown*| 0.32 ($\pm$0.139)      | 0.00806 ($\pm$0.00273)       | **Yes**   |Not implemented (non-causal inference network)|*Unknown*|
| SRNN (Ours)   | RNN       | NNs    |HMM| **0.27 ($\pm$0.093)**      | **0.00230 ($\pm$0.00041)**      | **Yes**   | **Yes** (causal inference network)|Pytorch|

[1] Durstewitz, Daniel, Georgia Koppe, and Max Ingo Thurm. "Reconstructing computational system dynamics from neural data with recurrent neural networks." Nature Reviews Neuroscience 24.11 (2023): 693-710.

[2] Barak, Omri. "Recurrent neural networks as versatile tools of neuroscience research." Current opinion in neurobiology 46 (2017): 1-6.
We thank you for your response and for your consideration in raising the score. We didn’t find this discrepancy in recovering behavioral vs. reconstructing neural data in our model. We definitely agree that this discrepancy in existing models is interesting to explore.
My primary concern remains only partially addressed: The contribution beyond other switching nonlinear dynamical systems models is not clear. I appreciate the attempts to implement and compare against SNLDS and mrSDS. And I appreciate that mrSDS was very recently accepted (though the paper has been on arxiv for many months before the ICLR acceptance) and that code may not have been publicly available. However, I expect a compelling description of the novelty of your approach relative to these existing approaches in the literature. What is the innovation beyond those approaches? And can you demonstrate that those specific innovations underlie the stated empirical performance benefits over SNLDS and mrSDS?
Thank you for your detailed response. I appreciate the new results: comparisons with SNLDS and mrSDS, as well as the detailed experiments on discrete state switching. I still find myself a bit perplexed by the discrepancy in behavior and neural data results and I wish I understood the advantages / flaws of each of these models better. Overall, since the authors have addressed some of my concerns, I will raise my score to beyond the accept threshold.
Dear Reviewers, 

Due to a potential glitch, the reviewer comments might not be visible to all reviewers. Therefore, we have included the full reviewer comments in our response. Thank you!
>*I would like to thank the authors for their response to my comments and clarifications. I would like to raise my score from 5 (Borderline accept) to 6 (Weak accept).*

>*However, I still believe that there needs to be comparisons with simpler models such as ARHMMs and their extensions [1, 2]. In particular, I do not agree with "...SLDS and rSLDS, where all models are initialized with ARHMM, and thus expect that the log-likelihoods are comparable or higher in these baseline models...". SLDS and rSLDS rely on approximate inference which may frequently lead to bad optima. In contrast, ARHMMs sometimes tend to perform better, thanks to their exact M-step update.*

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

>*If accepted, I highly suggest including these additional comparisons with ARHMMs and their variants.*

We thank you for your response and raising the score from 5 to 6. We also thank you for pointing out that “ARHMMs sometimes tend to perform better, thanks to their exact M-step update.” Yes, you are right, we applied ARHMM on the reaching data and found that ARHMM outperform SLDS and rSLDS in the reconstruction of neural activity. Therefore, we included a comparison between SRNNs vs ARHMMs on reaching data are as following:

|    | SRNN(P=16, K=5) | ARHMMs (R=180, K=5)    |
| :---        |    :----:   |          ---: |
| Reconstruction (**MSE**)      |   **0.00230($\pm$0.00041)**   |          0.00232($\pm$0.00085) |
| Behavioral States Recovery (**Error**)      | **0.27($\pm$0.093)**       |  0.78($\pm$0.22)   |
| Training on Single Condition (**MSE**)   | **0.01117($\pm$0.00099)**        | 0.01216($\pm$0.00104)     |

We found that SRNNs and ARHMMs exhibit very similar reconstruction accuracy, both outperform SLDS and rSLDS. However, SRNNs show a slight advantage in terms of the mean and standard deviation of the mean squared error (MSE) between neural activity and the reconstructions. Despite their performance in reconstruction, ARHMMs are not able to recover behaviorally relevant states, and the states identified by ARHMMs lack interpretability.

We agree that the extensions of ARHMMs, such as Switching autoregressive low-rank tensor models, are interesting to compare. We will include a comprehensive comparison between all these models mentioned above in the final version, if accepted.
>*Thanks for the additional points and clarifications. I retain my score of a "weak accept" as I think the work is novel, technically correct, and could be of interest to the neural modeling community. I still have reservations about impact of this method on the broader neuroscience community, given challenges related to interpretability and somewhat weak reasons to prefer this method over a higher dimensional RNN.*

We thank you for your response and retaining a positive rating. We agree on the trade-off between switching low-dimensional RNNs and non-switching higher-dimensional RNNs. If accuracy is the only metric, we agree that high-dimensional standard RNNs are able to achieve this goal. However, for achieving interpretability in (a) identifying discrete behaviorally-relevant switches in dynamics, and (b) distinctive flow-field visualization, we demonstrate the utility of switching RNNs in the paper. As a concrete example comparing the two models, in order to achieve similar reconstruction accuracy as a P=2, K=3 SRNN (lever-pull data), we need a standard RNN with dimensionality P=5. This higher dimensionality in the standard RNN naturally comes at the cost of accuracy in the 2D flow-field representations, and thus may result in compromised interpretability of the dynamics.

The trade-offs between low-dimensional switching RNNs and higher-dimensional standard RNNs are interesting, and we will add quantitative accuracy tradeoffs with varying values for P and K in the final version, if accepted.
I would like to thank the authors for their response to my comments and clarifications. I would like to raise my score from 5 (Borderline accept) to 6 (Weak accept).

However, I still believe that there needs to be comparisons with simpler models such as ARHMMs and their extensions [1, 2]. In particular, I do not agree with "...SLDS and rSLDS, where all models are initialized with ARHMM, and thus expect that the log-likelihoods are comparable or higher in these baseline models...". SLDS and rSLDS rely on approximate inference which may frequently lead to bad optima. In contrast, ARHMMs sometimes tend to perform better, thanks to their exact M-step update.

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

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

If accepted, I highly suggest including these additional comparisons with ARHMMs and their variants.
Thanks for the additional points and clarifications. I retain my score of a "weak accept" as I think the work is novel, technically correct, and could be of interest to the neural modeling community. I still have reservations about impact of this method on the broader neuroscience community, given challenges related to interpretability and somewhat weak reasons to prefer this method over a higher dimensional RNN.
**Details of Determining relevant hyperparameters (Weakness 2 and Question 2)**

We now add a discussion for determining relevant hyperparameters, **such as the number of discrete and continuous latent states**. Importantly, we include a comprehensive hyperparameter sweep for the number of discrete states for both the reaching and decision-making datasets. We detail out the salient results for the reaching dataset, where there are 5 behavioral states in the task (decision-making results below):

(1) **Convergence to lower number of discrete states**: We tested our model by increasing the number of hidden states K while keeping the number of continuous latent states P constant. We found that **61%** of SRNNs with a higher number of discrete hidden states (e.g., K=6) finally converge to the optimal number of discrete hidden states, i.e., K=5. 

(2) **Reuse of discrete states**: We also test our model by decreasing the number of hidden states. We found that **94%** of SRNNs with a lower number of discrete hidden states (e.g., K=4) had at least one hidden state reused after other states: in other words, SRNNs are not able to perform well with 4 unique discrete hidden states without reusing one of them. 

(3) **Reconstruction performance plateau**: While keeping other hyperparameters constant, the reconstruction accuracy plateaus at the same number of discrete states as in the behavior, thus we can use the minimum number of discrete states as it takes for the model to perform well. We have included a figure detailing this in the attached PDF (Figure 2). Moreover, we also implemented a ‘co-smoothing’ method as suggested by Reviewer hFy4 [Yu et al., 2009 and Karniol-Tambour et al., ICLR 2024], we show the results in the attached PDF (Figure 3), where we found that K=5 also does well in reconstructing the data with a ‘co-smoothing’ neuron drop-out analysis.

(4) **Variability across conditions**: In stereotyped tasks or experiments, such as reaching, there may not be a significant amount of variability in the timing of behavior across conditions, and this variability can thus be used as a metric for determining the number of discrete states. **Here, we found that SRNNs with K=5 have much lower variability on recovered behaviorally-relevant states than K=4 and K=6 (i.e., 0.098 for K=5, 0.384 for K=4, and 0.282 for K=6).**

Furthermore, we did the same analysis for the decision-making data (Figure 4 in the attached PDF) and we found the same situation as in the reaching dataset. **Higher K also converges to a smaller number of discrete states, and K=5 has the smallest variability across different pseudo-sessions (i.e., 0.106 for K=5, 0.287 for K=4, and 0.290 for K=6).** 

These results demonstrate appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance, and (3) variability across conditions or trials. We show comprehensive results for these metrics on two experimental datasets where we recover the appropriate number of behaviorally relevant states using any of these three metrics. We have now provided this in the paper, while adding a discussion for the general case.

Additionally, we show a comparison between values for another important hyperparameter P in the original submission.
**Details of Determining relevant hyperparameters (Weakness 2)**

We now add a discussion for determining relevant hyperparameters, **such as the number of discrete and continuous latent states**. Importantly, we include a comprehensive hyperparameter sweep for the number of discrete states for both the reaching and decision-making datasets. We detail out the salient results for the reaching dataset, where there are 5 behavioral states in the task (decision-making results below):

(1) **Convergence to lower number of discrete states**: We tested our model by increasing the number of hidden states K while keeping the number of continuous latent states P constant. We found that **61%** of SRNNs with a higher number of discrete hidden states (e.g., K=6) finally converge to the optimal number of discrete hidden states, i.e., K=5. 

(2) **Reuse of discrete states**: We also test our model by decreasing the number of hidden states. We found that **94%** of SRNNs with a lower number of discrete hidden states (e.g., K=4) had at least one hidden state reused after other states: in other words, SRNNs are not able to perform well with 4 unique discrete hidden states without reusing one of them. 

(3) **Reconstruction performance plateau**: While keeping other hyperparameters constant, the reconstruction accuracy plateaus at the same number of discrete states as in the behavior, thus we can use the minimum number of discrete states as it takes for the model to perform well. We have included a figure detailing this in the attached PDF (Figure 2). Moreover, we also implemented a ‘co-smoothing’ method as suggested by Reviewer hFy4 [Yu et al., 2009 and Karniol-Tambour et al., ICLR 2024], we show the results in the attached PDF (Figure 3), where we found that K=5 also does well in reconstructing the data with a ‘co-smoothing’ neuron drop-out analysis.

(4) **Variability across conditions**: In stereotyped tasks or experiments, such as reaching, there may not be a significant amount of variability in the timing of behavior across conditions, and this variability can thus be used as a metric for determining the number of discrete states. **Here, we found that SRNNs with K=5 have much lower variability on recovered behaviorally-relevant states than K=4 and K=6 (i.e., 0.098 for K=5, 0.384 for K=4, and 0.282 for K=6).**

Furthermore, we did the same analysis for the decision-making data (Figure 4 in the attached PDF) and we found the same situation as in the reaching dataset. **Higher K also converges to a smaller number of discrete states, and K=5 has the smallest variability across different pseudo-sessions (i.e., 0.106 for K=5, 0.287 for K=4, and 0.290 for K=6).** 

These results demonstrate appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance, and (3) variability across conditions or trials. We show comprehensive results for these metrics on two experimental datasets where we recover the appropriate number of behaviorally relevant states using any of these three metrics. We have now provided this in the paper, while adding a discussion for the general case.

Additionally, we show a comparison between values for another important hyperparameter P in the original submission.
**Details of Determining relevant hyperparameters (Weakness 2)**

We now add a discussion for determining relevant hyperparameters, **such as the number of discrete and continuous latent states**. Importantly, we include a comprehensive hyperparameter sweep for the number of discrete states for both the reaching and decision-making datasets. We detail out the salient results for the reaching dataset, where there are 5 behavioral states in the task (decision-making results below):

(1) **Convergence to lower number of discrete states**: We tested our model by increasing the number of hidden states K while keeping the number of continuous latent states P constant. We found that **61%** of SRNNs with a higher number of discrete hidden states (e.g., K=6) finally converge to the optimal number of discrete hidden states, i.e., K=5. 

(2) **Reuse of discrete states**: We also test our model by decreasing the number of hidden states. We found that **94%** of SRNNs with a lower number of discrete hidden states (e.g., K=4) had at least one hidden state reused after other states: in other words, SRNNs are not able to perform well with 4 unique discrete hidden states without reusing one of them. 

(3) **Reconstruction performance plateau**: While keeping other hyperparameters constant, the reconstruction accuracy plateaus at the same number of discrete states as in the behavior, thus we can use the minimum number of discrete states as it takes for the model to perform well. We have included a figure detailing this in the attached PDF (Figure 2). Moreover, we also implemented a ‘co-smoothing’ method, we show the results in the attached PDF (Figure 3), where we found that K=5 also does well in reconstructing the data with a ‘co-smoothing’ neuron drop-out analysis.

(4) **Variability across conditions**: In stereotyped tasks or experiments, such as reaching, there may not be a significant amount of variability in the timing of behavior across conditions, and this variability can thus be used as a metric for determining the number of discrete states. **Here, we found that SRNNs with K=5 have much lower variability on recovered behaviorally-relevant states than K=4 and K=6 (i.e., 0.098 for K=5, 0.384 for K=4, and 0.282 for K=6).**

Furthermore, we did the same analysis for the decision-making data (Figure 4 in the attached PDF) and we found the same situation as in the reaching dataset. **Higher K also converges to a smaller number of discrete states, and K=5 has the smallest variability across different pseudo-sessions (i.e., 0.106 for K=5, 0.287 for K=4, and 0.290 for K=6).** 

These results demonstrate appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance, and (3) variability across conditions or trials. We show comprehensive results for these metrics on two experimental datasets where we recover the appropriate number of behaviorally relevant states using any of these three metrics. We have now provided this in the paper, while adding a discussion for the general case.

Additionally, we show a comparison between values for another important hyperparameter P in the original submission.
**Details of Determining relevant hyperparameters (Q3)**

We now add a discussion for determining relevant hyperparameters, **such as the number of discrete and continuous latent states**. Importantly, we include a comprehensive hyperparameter sweep for the number of discrete states for both the reaching and decision-making datasets. We detail out the salient results for the reaching dataset, where there are 5 behavioral states in the task (decision-making results below):

(1) **Convergence to lower number of discrete states**: We tested our model by increasing the number of hidden states K while keeping the number of continuous latent states P constant. We found that **61%** of SRNNs with a higher number of discrete hidden states (e.g., K=6) finally converge to the optimal number of discrete hidden states, i.e., K=5. 

(2) **Reuse of discrete states**: We also test our model by decreasing the number of hidden states. We found that **94%** of SRNNs with a lower number of discrete hidden states (e.g., K=4) had at least one hidden state reused after other states: in other words, SRNNs are not able to perform well with 4 unique discrete hidden states without reusing one of them. 

(3) **Reconstruction performance plateau**: While keeping other hyperparameters constant, the reconstruction accuracy plateaus at the same number of discrete states as in the behavior, thus we can use the minimum number of discrete states as it takes for the model to perform well. We have included a figure detailing this in the attached PDF (Figure 2). Moreover, we also implemented a ‘co-smoothing’ method as suggested by Reviewer hFy4 [Yu et al., 2009 and Karniol-Tambour et al., ICLR 2024], we show the results in the attached PDF (Figure 3), where we found that K=5 also does well in reconstructing the data with a ‘co-smoothing’ neuron drop-out analysis.

(4) **Variability across conditions**: In stereotyped tasks or experiments, such as reaching, there may not be a significant amount of variability in the timing of behavior across conditions, and this variability can thus be used as a metric for determining the number of discrete states. **Here, we found that SRNNs with K=5 have much lower variability on recovered behaviorally-relevant states than K=4 and K=6 (i.e., 0.098 for K=5, 0.384 for K=4, and 0.282 for K=6).**

Furthermore, we did the same analysis for the decision-making data (Figure 4 in the attached PDF) and we found the same situation as in the reaching dataset. **Higher K also converges to a smaller number of discrete states, and K=5 has the smallest variability across different pseudo-sessions (i.e., 0.106 for K=5, 0.287 for K=4, and 0.290 for K=6).** 

These results demonstrate appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance, and (3) variability across conditions or trials. We show comprehensive results for these metrics on two experimental datasets where we recover the appropriate number of behaviorally relevant states using any of these three metrics. We have now provided this in the paper, while adding a discussion for the general case.

Additionally, we show a comparison between values for another important hyperparameter P in the original submission.