Inference of Neural Dynamics Using Switching Recurrent Neural Networks
Authors: Yongxu Zhang
Shreya Saxena
Venue: NeurIPS 2024 poster (2024)
Keywords: Neural Dynamics
Variational Inference
Recurrent Neural Networks
Primary area: neuroscience_and_cognitive_science

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.

Inference of Neural Dynamies Using Switching Recurrent Neural Networks
Yongxu Zhang Shreya Saxena
Yale University
{yongxu. zhang, shreya. sareya.saxenaaarrp... saxena)@yalee..e...aa edu
Abstract
Neural population activity often exhibits distinct dynamical features across time, which may correspond to distinct internal processes 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, 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 simulated data as well as cortical neural activity across mice and monkeys, which allows to automatically detect discrete states that lead the identification of varying neural dynamics. In a monkey reaching dataset with electrophysiology recordings, mouse self-initiated lever pull dataset with widefield calcium recordings, and a mouse self celf-initiated decision making dataset with widefield calcium recording, SRNNs are able 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 Sahaaiorrally-eeevvae switches across time and we are able to use SRNNs to successfully capture these switches and the corresponding dynamical features.
Introduction
How does complex neural activity lead to dynamic behavior? 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 al. (223) ). As ability record from larger areas of the brain with unprecedented spatial and temporal resolution increases, we find neural population activity often cxhibits distinct dynamical features across time. These features may correspond to distinct internal processes or behavior (Saxena and Cunningham [2019), Churchland et al. (2022). Identifying these dynamical features may help understand how cognitive functions imp 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 during movement generation (Kaufman et al. 2014). Miri et al. 2017)).
State space modeling is promising analytical approach for characterizing dynamics of time series; such models can be highly effective and have received substantial attention over many decades analyzing temporal neural dynamics (He et al. [2023]. Durstewitz et al. (2023). (Sani et al.
(202T). Well-established methods on inference and learning algorithms have contributed towards learning the parameters of state space models (Blei al. 2017 Salimans et al. 2015 Kingma and
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Welling [2013), Linderman et al. (2017]. Fox et al. (2008) ) Neural dynamics can be time-varying due to internal iluctuations of phystological states as well as the external effect of environment. Previous works such as (Kaufman et al. 2014), Miri et al. (017)) have found differences in neural subspaces between distinct behavioral states. Furthermore, correlation-basel analysis methods, such as functional connectivity, also reveal that the brain-wide covariation in neural activity changes during movement. For example, WWest 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. [022], Song et al. (2022)) )
In the past decade, we have seen that neural dynamics can be broadly considered to be non non-linee linear and in a lower dimensional subspace of the recorded neural activity (McKenna et al. [1994]. , Rigotti and Yu (014)). Accordingly, Recurrent Neural Networks (RNNs) have shown efficiency modeling dynamies due to their non-linearity (Mante et al. 2003)). 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 [021]. Valente et al. (022]) )
et al. [2013). Hernandez et al. (2018), [Saaeaaaaaaaaa and Cunningham
Specifically, these RNNs can be trained to reproduce temporal neural activity from large-scale neural recordings across set of trials. They are also able to capture very high amount of explained variance of the neural activity because of their static non- lonearicssss (Durstewitz et al. [2023), [,eerih and Rajan (020).. Once trained, the internal mechanisms of the RŇNs 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- -iwdiimnsional dynamical models. We perform end-to-end training of SRNNs to reconstruct the neural observations through Variational Inference. We validate SRNNs on 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 non-human primate performing a reaching task (Churchland et al.) [2022). (2) cortex-wide widefield calcium imaging (WFCI) from mice performing 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 vehaviorally-releaat switches of neural dynamics. Finally, we visualize the neural dynamics in flow fields plotting and find they are distinct among different states.
2 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 (018)). ). 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 experimentally- obtained 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 oowdiimesioonal 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) a commonly-use statistical model based on Markov process: it dictates the evolution of observations where the evolution determined by internal factors, which are not directly observable. HMMs are widely used in computational neuroscience to identify different neural events (Baldassano et al. (017) Masaracchia
jet al. [203),, [Baldassano et al.] (0166). Previous studies have also explored switching dynamical systems, with the observations emanating from a dynamical system that discretely switches according

---

A
B
8
0 YE1
0
SSutch
Figure 1: Switching Recurrent Neural Network (SRNN) (A) Schematic showing an SRNN with neural outputs. (B) Structure of the generative model for SRNN. (C) Structure of Inference Neural Network for SRNN. The hf and h? represent the states of a bidirectional RNN, followed by a hj, which model the inferred h, from the given observation yt-
to 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 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 the previous continuous state, termed recurrent-SI SLDD (rSLDS). Researchers have also changed the dynamics process to be nonlinear, such as in a Structured Variational Autoencoder (SVAB) (Johnson et al. (0166). Moreover, both nonlinear transitions and nonlinear dynamics processes were also proposed, such as in a Switching Nonlinear Dynamical System (SNLDS) (Dong) al. (2000). Out of these, SRNNs have the most comparable structure 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 predict future observations. Moreover, here we analyze the ability of the model to reconstruct and clucidate the structure of recorded neural activity, while examining prediction as well as region-to-regio 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 lo-lorally-eleeaatt 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 to find maximum likelihood estimates of parameters in State Space Models, e.g.. Baum- Algorithm for HMMs (Bishop and Nasrabadi (2006)). Another common approach 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). [inderman et al.) (017)). The reparameterization trick is a crucial technique VI, which enables more efficient gradient estimation during optimization. Instead of sampling from the variational distribution itself, reparameterization models the parameters Gaussian. This allows the gradient of the objective function with respect to the variational distribution parameters to computed directly, and backpropagation to be implemented easily. To enable the model to generalize effectively unseen data, trials of neural activity independently as in traditional VI, Amortized VI allows us to train the models on multiple trials simultaneously, thus creating more stable mapping between the observations and the latent variables (Dong et al. 22020), Ganguly et al. 22233] Margossian and Blei 2023
distribution, usually
Amortized VI. Instead of optimizing for single
3 Methods
3.1 Switching Recurrent Neural Network (SRNN)
common switching state space model is:
T
Poly, h, = Pa,, (11 (p, (vilhr) II
Po. (2x 22- 1111-1/Pp,, (hr 1-,- 2,) (0.,,,uull (vrlhe)
(1)
2
3

---

where 2, E [I....,KK } is the discrete hidden states or labels, controlling the switch of the dynamics, he the continuous latent state with dimensionality P. and Yr 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 y1:r RxNxXxR 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 nonlinear activation function, and finally a linear emissions network to reconstruct the observed neural activity. We show the structure of SRNN in Figure A following are the equations of the generative network:
pe_ (2o121- h2-1) Cat(Softmax(((0hh--)))
(2)
pa, (helhe-1. 2 = =k) (
(3)
pa, (vilbb)) NNvvl0,,e
where Cat represents Categorical distribution, Softmaz represents Softmax activation function, NN represents Gaussian distribution, and σソ and Ty are the covariance matrices of the relevant Gaussian distributions. 8z. 0n. and fy are parameters of transition networks, dynamical networks, and emission networks respectively, which are learned through gradient descent. Here, consider the nonlinearity f()) tanh(.) as commonly done neuroscience in order to model the saturation effect of firing
rates Pandarinath al. (2018. Saxena et al. 2022
3.2
Inference of SRNNs
We use VI to learn the parameters of model. Here, we maximize the posterior:
(e.4(21:7) h1.7 (81.r)
(lt::rll1)))))))11. //1.7 (1.1)
(5)
where qo (hly) is the posterior of continuous hidden states given the observation, & the parameters of our inference network. pe(z)h, 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 (h1rrly1:r) We show the structure of our inference network in Figure B3 we build an inference network to model go (hly) as the following:
T
T
= II (heluir)) = II Nhe (hellu,, t=1
(6)
1=1
where a and a are the mean and variance Gaussian distribution. Here, we use the reparameteriza- tion trick to infer the continuous hidden states Kingma and Welling (2013):
#e cØ2
(7)
where (0, 1) is a Gaussian noise parameter. The mean μ and variance a are typically learned via training neural networks Kingma and Welling [203). However, here, we do not find apparent improvement with learning the a in our work. Therefore, as done in previous studies, we set the variance a to be constant reduce the complexity of the model and only optimize the Dong et al.] [2020] Ganguly and Earp [2021]; we use 0.0001 for a. Since our generative model is based on RNNs, we make 14 depend on h4-1 1 to match the recurrence property. Thus, we have
Pr f ht-1)
(8)
We pass V1.7 to a bidirectional recurrent neural network followed by forwarding the output to standard recurrent neural network to model Equation 8
Evidence Lower Bound and po (= hI:T, yI.T) Typically, to learn the parameters via VI, maximize the evidence lower bound (ELBO). Here, the ELBO is defined as:
ELBO E Bgg (z Aurelner) [log po (yVeT, hir, 21.)) log (Zir,
(9)
With Equation I and Bayes' rule, we get
ELBO
Ey (hier yeer)a huraa)))) [log pe (VI.7. h1rr (2
log (fi:r (y1.r) (21 (h:7 V1:r)]
(10)
2 Eg (hier wirr) [log pe (y1eT, /1.r) - log go ( (hirlvviv)
(11)
where the discrete hidden states 21.7 have been marginalized out given KK p(2: = k) = (Murphy
and Russell [2001] Dong al. 2020 )

---

Initialization Previous studies have found that a generative model trained by VI may become stuck in a single discrete state (Alemi et al. 2 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 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 each of the states. Practically, during a first phase of training, we add 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, remove this penalization and train the model with our loss as described in Equation 15
weights corresponding
Optimization We use gradient descent to optimize the parameters 0 and ś. Given the fact that use reparameterization to model 9 (1.l)) and gs (h1.r1111r)) is independent to 0, the gradient can be approximated following:
V ELBO Ve. log ())(()) V log (o1: (11:))
(12)
where (o) is the output of inference network, therefore, is dependent on ø. Since E. p(=)=1, therefore, log po(yI.T, (m)) = >. p(z) log po h1]=) = Ep(x) log po hlz), and with Bayes'
rule pe (y, h1)2 po(y.h.2)
Equation 12 is equivalent to:
p(a)
V ELBO 2 Ep(s) Vo.6 log log hITT, #ITT - V log go
(13)
In Equation 13 the first term can be rewritten by the Markov Chain.
E Ep(ELr) Vo. log po hi.r, 21.T) Σ p A) k) k)p(hal22 k))((r/hn)
T K J
ΣΣΣ p(z = k, 21 -) loggllz = - = j)p(he)lt) = =k)((ulh)) (14) 22
where use the posterior probability of the discrete hidden states, e.g., p(=1 k|hi.r, (1:)) and p(2 = k 1 jlhir, })rr), which can be obtained easily through Baum Welch algorithm. The transition, dynamics. and emission are defined in Equation Equation 3 and Equation respectively. Equaaionnrrsspectvvvll The second term of Equation can be optimized via backpropagation through time (BPTT) of the inference network. In summary, with Equation IT Equation 13
and Equation 14, our loss function is:
K
p = k) logp(zs = k)p(h1 = k)p(vi(h1)
T
K
p(z = k, log = -k/2-11 = h1 =k)p(s/)h.)
2k=jj=1
H(go(hi:lyy:rr)
(15)
where we use the fact that (Eg. log us (/1:r//1:1)) is the entropy H of (g ..rr). 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://sa wtwggtt.uu...... com/ sww.nalal-aaabaaa Silo-neur//Sii
3.3 Analysis of SRNNs
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 test set) in order to capture neural dynamics across the entire data. We use 10, 000 initial hidden states uniformly sampled in grid to compute the flow fields.
Dynamical Features:
5

---

A. Desired Simulated Lorenz Attractor and Reconstruction by SANN
B Desired Simulated Lorenz Attractor States and States Recovered SRNN
Continuous Hilden States of SRNN Trained Lorenz Attractor
Groud Truth
SRNNN
Figure 2: Lorenz Attractor: (A) Reconstruction using an SRNN on simulated data (thick red the original data, thin black is the reconstruction by SRNN. The reconstruction is almost overlapped with the ground truth). (B) Comparison between ground truth discrete states and -eNN-recovered discrete states. (C) Plot of latent dynamics in each discrete state shows recovery of dynamics.
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ı.t, with which we can infer discrete hidden states #1:t and continuous hidden states h1z-- Next, we sample zt K and he through transitions of respectively. Finally,
Neural Activity Prediction:
discrete states and dynamics of continuous states, i... using Equations 2nd [ we predict the neural activity in the future time points 3tt++K using the emissions network, i.c., Equation 4 Here, we change the bidirectional layers of our inference network to standard forward recurrent layers to match the prediction tasks (see Figure F.). We explore the prediction performance by giving the model an imput in different lengths (t to. T), here use tu 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 simulated dataset and three diverse experimental datasets ranging from monkey electrophysiology to mouse calcium imaging. On each dataset, we do --lold coos-validatiin, 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 al. 2017 ) Switching Non-linear Dynamical Systems (SNLDS) (Dong et (2020), ). and multi- regional Switching Dynamical Systems (mrSDS) "Karniol-Tambon et al. [2023] ). Additionally, we provide comparisons with Latent Factor Analysis with Dynamical Systems (LFADS) in Appendix
Figure D.6 Lastly, we show example curves of the training loss, reconstruction MSE on validation data, and discrete states recovery error on validation data for SRNNs across different epochs for all three experimental datasets in Figure D.7
4.1 Simulated data: Lorenz Attractor
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 identify the underlying dynamics in each discrete state.
4.2 Experimental data: Electrophysiology Recordings during a Monkey Reaching Task
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. 20022]. We have 5 different behaviorally- 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.
states: (1) resting; (2) delay before the "go" cue; (3)
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 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 A and Figure A.I. The neural dynamical states recovered

---

A
Curved Reach
B
Error between Behavioral States
C
MSE between Neural Activity Reconstruction (Reaching)
De
and Recovered States (Reaching)
States Recovared SRMN (Durs)
States Rocevered by SLDS
0.8
Mean Seuared Error
SALDS
States Necovered ISLDS
SINN (Quna)
0.6
States Recovered by SNLDS
States Recovered by mV505
0.4
SLUS
D
MSE between Neural Activity
Prediction (Reaching)
Time
0.2
SALDS
milts
SPAN (Dun)
SUPS
2
16 Latent State Dimensionalitt P
2
16
Mean Squared Error
10-1
Latent State Dimensionality
SAAN Statel |Restnn!
SRNN State2 (Delay Before Go Cuc)
SRAN Stane3
SRAW Stated (Reaching)
SRNI State5 dlgg)
20"2
0.1
0.2
0.3
0.4
Noural Activtyy F
Seconds Aheed
Nan Caucal
Thue SL
True DS
Tue SNI
True mr5DS
True SANN (Ours)
-0. 0.0
0.6
-12 -066
-1.2 0.0 0.6
-066
-1.2
-0. 0.0 0.6
Movement
Figure 3: Curved Reaching: (A) Examples of comparison between behavioral states and neural dynamical states recovered by SRNNs, SLDS, rSLDS, SNLDS, and mrSDS. (B) Error between rehaviorally-relevnn states and recovered neural dynamical states; each dot represents one session, with the triangle as the mean of all sessions. (C) Mean squared error (MSE) between neural activity and the reconstruction.
(D) MSE between neural activity and the prediction forward in time. (E) Flow fields of latent dynamics of SRNNs in different discrete states. (F) Reconstruction of neural activity and the corresponding ground truth for example neuron for three models, the top color bars represent the performance on 0.1 seconds ahead prediction. More examples are shown in Figure D.2
All results for direct reaching are shown in Figure D.I
by SRNNs 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 BA, 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 one class, the error of the classification measures the difference between the neural and behavioral states. Specifically, 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 still be considered as behaviorally- states, which can help us understand the cognitive function implemented in the brain. We show the error for different hidden states in Figure B. which illustrates that SRNNs outperform SLDS, rSLDS, SNLDS, and mrSDS on recovering the behavior-relevant neural dynamical states for different values for the latent state dimensionality P. We determine the number of discrete latent states via 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 /' constant. We show a comparison in D.3J. We found that 61% of SRNNs with 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 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 platcau: 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 the minimum number of discrete states as it takes for the model to perform well. We
7

---

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. [2088 Karniol-Tambon
al. (023)); ); 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 stereatyped tasks or experiments, such as reaching, there may not be 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-rrlevan states than K = 4 and = 6 (i.e., 0.098 for = 5, 0.384 for K 4, and 0.282 for 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 BD. 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 Figure D.8 Moreover, we show the reconstruction of neural activity along with one example neuron with the corresponding ground truth in Figure 3F. with more example neurons shown in Figure D.2AABB. Here, the color bars represent performance of neural prediction. Furthermore, as a non-switching comparison model, we train an LFADS network (Pandarinath et al. (2018) on the reaching datasets. We show reconstruction MSEs in Figure (D.A,, where use P 16 for SRNN, SLDS, and rSLDS, and P = 64 for LFADS keep similar number of overall parameters in the hidden layers. We found that all three switching models have lower reconstruction error than LFADS (Figure D.
We visualize the neural dynamics by plotting the flow field of the latent states h in each discrete SRNN state in Figure where we used p = 16 based on Figures 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 E:: 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 Fig and we found that the neural dynamies between eurved 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.
4.3 Experimental data: Widefield Calcium Imaging during Mouse Decision-Makinn
We show the performance of SRNNs on decision-making makinn 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 al. (2019). In this data, mice have 5 sehaviorally-relevar states, i.c., (T) from bascline 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 behavioralll-relevan states described above. Here, we use 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 where we found that SRNNs are able to identify 5 Dehavioralll-reeevaa 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 Figure D.4 we also found that SRNNS with K = 6 converge to lower number of discrete states. Additionally, SRNNs with K- = 5 perform better in the ee-omooohhig test in Figure D.4C. Indeed, we found that SRNNs with K = 4 have better reconstruction performance. However, in Figure D.4p, have much lower variability on recovered behaviorally- states than K = 4 and K- = 6. Therefore, given the visualization of recovered states, co-smoothing test, and variance of recovered states across different pseudo-sessions, we consider K = 5 as the best hyperparameter. Furthermore, we show the reconstruction performance of all models in Figure
we found that SRNNs with K
8

---

A
Rehavicral States
B Error between Behavioral States
C
MSE between Neural Activity
Reating
hodding
Daly
and Recovered States (Decisian Making)
and
Reconstruction (Decision Making)
States Recovered by
SAAN SLDS
1SLO5 SALDS
mI505
SRAN
FSLOS SMLDS
muios
1.0
SLUS
Mean Squared Error
0.8
+
'4:
0.6
14
Recovered
0.4
-5
-4
Time Ueking
D
MSE between Neural Activity and
0.2
Prediction (Decision Making)
Mean Squared Error
101
ITLDS
SLDS
SALDS
0.0
2
Latent State Dimensionalityy P
4
8
16
16 Latent State Dimensionality P
10
SRMN 1Ount
E
SRAN Statel |Restin]
SRNN State2 (Holding)
SRNN State3 (Stimulus)
SPAN State4 (Delsy Before Licking)
SRNI State5
(Licking)
10-1
20-2
10
0.1
0.3 Seconds Ahead
0.2
0.4
F
Nan Cauroll
Noural Actouity
-6 -3
-6
Time from Novement
Figure 4: WFCI Decision-Making: (A-F) Legend similar Figure legend. The top color bars in
(F) represent 0.33s ahead prediction accuracy; more examples shown in Figure
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 In addition, we show the reconstruction of neural activity for an example component with the corresponding ground truth in Figure 4 with more results shown in Figure 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.
in the Appendix. Furthermore, we visualize
4.4 Experimental data: Widefield Calcium Imaging during Mouse Self-initiated Lever Pull
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 very short time because the mice pull the lever very fast, Mitelut et al. [0222] 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 behaviorally- relevant states: (1) from resting to the starting of inhibitory dynamics (as found in Mitelut et al.)
(2022)) ); (2) from the starting of inhibitory dynamies to behavior onset: (3) lever- 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 in the datasets above. SRNNS also outperform in recovering the behaviorally-relevaa While SRNNs can recover the inhibitory dynamics of neural activity from around seconds before the behavior onset, SLDS and rSLDS can only identify 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 where we found that SRNNs with
neural dynamical states; an example is shown in Figure 5AA.
2 are not able to detect the switches before the behavior onset. Additionally, SRNNs with K = 3 have lower reconstruction error in Figure D.5B and perform better in the co-smoothing test in Figure D.3C. Moreover, we found that SRNNs with K- = 2 have relatively lower variability on recovered behaviorally- states than K 3 and K = Figure However, as demonstrated the state visualizations in Figure D.5 A, this is primarily due to the majority of time points being classified as single state.

---

A
Reheoionnl Stares
Aftee Batwetor
B Error between Behavioral States
MSE between Neural Activity C
Reading
Intiobtry
and Recovered States (Lever Pull)
and Reconstruction (Lever Pull)
States Recovored by
SLOS
SNLOS
51WV (Oust
10
States Recovered
Mean Squared Error
0.B
Error
0.6
10
Recovered by misos
0.4
-10
-8
-4
D
MSE between Ncural Activity and
0.2
10-3
Prediction (Lever Pull)
ISLDS
Log Mean Squared Error
SLPS
0.0
Latent State Dimensionalitty
4
8
P
Labent State Dimensionality
2
E
SRNN State1
SRNN State2 (Inhibitory)
SRNN State3
20-3
SRAN
(Resting)
(Atter Lever Pull)
0.33
0.67
1
1.33
F
Seconds Ahead
Nan Cauroll
Noural Actovity
True SLDS
Tree SNLDS
True mrSDS
FSLDS
SRAN (Ours)
-9
-
-6 -3 0
-9
- -3 0
6 -3
Time from Novement
Figure 5: WFCI Selt-nnit Lever Pull: (A-F) Legend similar that of Figure 4 Moreover, the reconstruction of neural activity using SRNNs performs slightly worse than other methods (Figure BC); 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 5F. with more results shown in Figure D.2p. 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 B 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 order to elucidate region-to-region communication in Figure D.9
5 Limitations and Conclusions
We developed 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- 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 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 A A100 GPUs), which is more than the amount of time taken by SL.DS and rSLDS (typically less than one hour). In conclusion, we use SRNNs to identify behaviorally-reles 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 control signals to SRNNs to explore the identification of hhaaviorally-relevan neural dynamics while providing external stimuli.
neural states, and predicting future neural activity.
neural dynamical states. We find that the neural dynamics have
behaviorally
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.
10

---

Acknowledgement
We gratefully acknowledge support from the NIH Brain Initiative 7RF1 7RF157005377 and NSF NCS 2350329. We thank Mark Churchland, Anne Churchland, Catalin Mitelut, and Timothy H Murphy for sharing experimental datasets. We also thank Nancy Padilla-Coreano for helpful discussions.
References
A. Alemi, B. Poole, I. Fischer, J. Dillon, R. A. Saurous, and K. Murphy. Fixing broken elbo. In International conference on machine learning, pages 159-168. PMLR, 2018.
C. Baldassano, J. Chen, A. Zadbood, J. W. Pillow, Hasson, and K. A. Norman. Discovering event structure in continuous narrative perception and memory. Neuron, 95(3):709-721, 2017.
S. Baldassano, D. Wulsin, H. Ung. T. Blevins, M.-G. Brown, E. Fox, and B. Litt. A novel seizure detection algorithm informed by hidden markov model event states. Journal of neural engineering, 13(3):03601 1, 2016.
C. M. Bishop and M. Nasrabadi. Pattern recognition and machine learning, volume Springer, 2006.
D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the Americun statistical Association, 1125518) 2017.
M. Breakspear. Dynamic models of large-scale brain activity. Nature neuroscience, 20(3):340-352, 2017.
M. M. Churchland, J. P. Cunningham, M. T. Kaufman, J. D. Foster, P. Nuyujukian, S. I. Ryu, and
K. V. Shenoy. Neural population dynamics during reaching. Nature, 487(7405):51-556 2012.
J. P. Cunningham and B. M. Yu. Dimensionality reduction for large-scale neural recordings. Nature neuroscience, 17(1 (((11):1500) 1509995559, 2014.
Dong, B. Seybold, Murphy, and H. Bui. Collapsed amortized variational inference for switching nonlinear dynamical systems. In International Conference on Machine Learning, pages 2638-2641 PMLR, 2020.
A. Dubreuil, A. Valente, M. Beiran, F. Mastrogiuseppe, and Ostojic. The role of population structure in computations through neural dynamics. Nature neuroscience, 25(6):783-794, 2022.
L. Duncker and M. Sahani. Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings. Current opinion in neurobiology, 70:163-170, 2021.
D. Durstewitz, G. Koppe, and M. I. Thurm. Reconstructing computational system dynamics from neural data with recurrent neural networks. Nature Reviews Neuroscience, 24(11):6933-) 1):69--710, 2023.
E. Fox, E. Sudderth, M. Jordan, and A. Willsky, Nonparametric bayesian learning of switching linear dynamical systems. Advances in neural information processing systems, 21, 2008.
A. Ganguly and S. W. Earp. An introduction to variational inference. arXiv preprint 2021.
A. Ganguly, S. Jain, and U. Watchareeructai. Amortized variational inference: systematic review. Journal of Artificial Intelligence Research, 18:167, 2023.
M. D. Golub and D. Sussillo. Fixedpointinder: A tensorflow toolbox for identifying and charac- terizing fixed points in recurrent neural networks. Journal of Open Source Software, 3(31):1003, 2018.
M. He, P. Das, G. Hotan, and P. L. Purdon. Switching state-space modeling of neural signal dynamics. PLoS Computational Biology, 191 13955 2023.
D. Hernandez, A. K. Moretti, Z. Wei, S. Saxena, J. Cunningham, and L. Paninski. Nonlinear evolution via spatially-dependen linear dynamics for electrophysiology and calcium data, arXiv preprint arXiv.1811.02459, 2018

---

M. J. Johnson, D. K. Duvenaud, A. Wiltschko, R. P. Adams, and S. R. Datta. Composing graphical models with neural networks for structured representations and fast inference. Advances in neural information processing systems, 29, 2016.
O. Karniol-Tambon D. M. Zoltowski, E. M. Diamanti, L. Pinto, C. D. Brody, D. W. Tank, and J. W. Pillow. Modeling dependennn communication between brain regions with switching nonlinear dynamical systems. In The Twelfth International Conference on Learning Representations, 2023.
M. T. Kaufman, M. M. Churchland, S. I. Ryu, and K. V. Shenoy. Cortical activity in the null space: permitting preparation without movement. Nature neuroscience, 17(3):440 448, 2014.
D. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv: I 1312.1114 2013.
S. Linderman, M. Johnson, A. Miller, R. Adams, D. Blei, and L. Paninski. Bayesian learning and inference recurrent switching linear dynamical systems. In Artificial intelligence and statistics, pages 914-922. PMLR, 2017.
V. Mante, D. Sussillo, K. V. Shenoy, and W. T. Newsome, Context-dependent computation by recurrent dynamics prefrontal cortex. nature, 503(7774):78-84) 2013.
C. C. Margossian and M. Blei. Amortized variational inference: When and why? arXiv preprint arXiv:2307.100/8, 2023.
L. Masaracchia, F. Fredes, M. W. Woolrich, and D. Vidaurre. Dissecting unsupervised learning through hidden markov modeling in electrophysiological data. Journal of Neurophysiology. 130 (2):364-379, 2023.
F. Mastrogiuseppe and S. Ostojic. Linking connectivity, dynamics, and computations in low-rank recurrent neural networks. Neuron, 9993)660-6663, 2018.
McKenna, McMullen, and Shlesinger. The brain as dynamic physical system. Neuroscience, 60(3):587-605, 1994.
A. Miri, C. L. Warriner, J. S. Seely, G. F. Elsayed, J. P. Cunningham, M. M. Churchland, and T. M. Jessell, Behaviorally selective engagement of short-latency effector pathways by motor cortex. Neuron, 95(3):683-696 2017.
C. Mitelut, Y. Zhang. Y. Sekino, J. D. Boyd, F. Bollanos, N. V. Swindale, G. Silasi, S. Saxena, and Murphy. Mesoscale wontex-wide neural dynamies predict self-initiated actions in mice several seconds prior to movement. Elife, 1 11:e76506, 2022.
L. Mosley, A balanced approach to the multi-class imbalance problem. 2013. URL https //api //app.semmaee www.seeeaansaaaaeeeeee Degamatttcccccolaa.....oo/oaaaaa 59664519
K. Murphy and Russell. Rao-blackwellised particle filtering for dynamic bayesian networks. In Sequential Monte Carlo methods in practice, pages 499-515. Springer, 2001.
S. Musall, M. T. Kaufman, A. L. Juavinett, S. Gluf, and A. K. Churchland. Single-trial neural dynamics are dominated by richly varied movements. Nature neuroscience, 22710):167--1666 2019.
C. Pandarinath, D. J. O'Shea, J. Collins, R. Jozefowicz, S. D. Stavisky, J. C. Kao, E. M. Trautmann, M. T. Kaufman, S. I. Ryu, L. R. Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10)805-881) 2018.
M. G. Perich and K. Rajan. Rethinking brain-wide interactions through multi-region 'network of networks' models. Current opinion in neurobiology, 65: 15511 2020.
M. Perich, C. Arlt, S. Soares, M. Young, C. P. Mosher, Minxha, E. Carter, U. Rutishauser, P. H. Rudebeck, C. D. Harvey, et al. Inferring brain-wide interactions using data-constrained recurrent neural network models. BioRxiv, pages 2020-12, 12, 2020.
M. Rigotti, O. Barak, M. R. Warden, X.-J Wang, N. Daw, E. K. Miller, and Fusi. The importance of mixed selectivity in complex cognitive tasks. Nature, 497(7451):585-599 2013.
12

---

T. Salimans, D. Kingma, and M. Welling, Markov chain monte carlo and variational inference: Bridging the gap. International conference on machine learning. pages 1218-1226. PMLR. 2015.
O. G. Sani, H. Abbaspourazad, Y. T. Wong, B. Pesaran, and M. M. Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroscience, 24 (1):140-149, 2021.
S. Saxena and J. P. Cunningham. Towards the neural population doctrine. Current opinion in neurobiology, 55: 103-111, 2019.
S. Saxena, I. Kinsella, S. Musall, S. H. Kim, J. Meszaros, D. N. Thibodeaux, C. Kim, J. Cunningham, E. M. Hillman, Churchland, et al. Localized semi-nonnegative matrix factorization (locanmf) of widefield calcium imaging data. PLoS computational biology, 16(4))e1007791, 2020.
S. Saxena, A. A. Russo, J. Cunningham, and M. M. Churchland. Motor cortex activity across movement speeds is predicted by network-level strategies for generating muscle activity. Elife, 11: c67620, 2022.
C. Y. Song, H.-L. Hsich, B. Pesaran, and M. M. Shanechi. Modeling and inference methods for switching regime-dependent dynamical systems with multiscale neural observations. Journal of Neural Engineering. 19(6):066019, 2022.
D. Sussillo and O. Barak. Opening the black box: low-dimensional dynamics in high-dimensional recurrent neural networks. Neural computation, 25(3):626-649, 2013.
A. Valente, J. W. Pillow, and S. Ostojic. Extracting computational mechanisms from neural data using low-rank rnns. Advances in Neural Information Processing Systems, 35:24072- -20066 2022.
Q. Wang, S.-L. Ding, Y. Li, J. Royall, D. Feng, P. Lesnar, N. Graddis, M. Naeemi, B. Facer, A. Ho, et al. The allen mouse brain common coordinate framework: 3d reference atlas. Cell, 181(4): 936-953, 2020.
L. West, J. D. Aronson, L. S. Popa, K. D. Feller, R. E. Carter, W. M. Chiesl, M. L. Gerhart, A. C. Shekhar, L. Ghanbari, B. Kodandaramaiah, al. Wide-field caleium imaging of dynamic cortical networks during locomotion. Cerebral Cortex, (2(12))2668-22) 2687, 2022.
B. M. Yu, J. P. Cunningham, G. Santhanam, S. Ryu, K. V. Shenoy, and M. Sahani. factor analysis for oow-dimensional single-trial analysis of neural population activity. Advances in neural information processing systems, 21, 2008
process
13

---

A Simulated data: Lorenz Attractor
The equations of Lorenz Systems are:
dz
oly x))
(16)
dt
x((p 2) dy
(17)
dl
dz
zy A2
(18)
dt
where y, represent 3 dimensions of the system, and p are three parameters. Here we use commonly used values for the parameters, i.e., a 10. p = 28, B 2.667.
B Experimental Datasets
Electrophysiology Recordings during 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 'go' cue after 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. (012)) 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. 20122. The data is recorded for 180 neurons and has 236 time points representing 2.36 seconds.
Widefield Calcium Imaging during Mouse Decision-Making Behavior In self-initiated decision making dataset, mice initiated trials by touching either of two handles and hold a handle for 1 followed by sensory stimuli. The sensory stimulus was presented for 600ms, and after 500ms pause with no stimulus and then the stimulus was repeated for another 600ms. After the second stimulus, 1000ms delay was imposed, then the mice are required to lick one of two spouts, the mice were rewarded with drop of water if they licked the spout twice on the same side as the stimulus. More details can be found in Musall al. 2019]. Here, we average across multiple trials recorded from the same mouse to create pseudo- sesuudssssssons... Specifically, we averaged around 80 trials of 516 trials in one mice produce N of the data can be found in Musall et al. [2019). The data is also processed via LocaNMF [Saxena al. 22020), and the 210 temporal components returned by LocaNMF are used. The data has 189 time points representing 6.3 seconds.
speudo-sessions and fit the trial- averaged data using our models. The details
Widefield Calcium Imaging during Mouse Self-initiated Lever Pull Behavior In self-initiated lever pull data, mice were trained to pull lever and hold it at an angle (for > 100ms) in order to receive a water supplement. We then apply LocaNMF Saxena et al, (220) while spatially aligning the imaged neural activity with the Allen mouse brain coordinate framework Wang et al. (2000) using affine transformations. as previously performed in Musall et al. [219]. and take 16 components as identified by LocaNMF, which form input signals, with each input dimension from brain region. In this work, we focus on the signals around the lever pull, i.e., from 10 seconds before lever pull 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 al. 2022
C Region-to-region Communication WFCI
We also use our model to uncover the relationship between different brain regions using self- initiaeed 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 FigureD.9/4 we find that the region-to-region communication is highly different between different states. Moreover, the
14

---

self-communicatoon during rest state (diagonal of SRNNI) 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 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.
D Appendix Figures
Curved Reach
Error between Behavioral States and Recovered States (Reaching)
MSE between Ncural Activity and Reconstruction (Reaching)
States Recovcred SLDS
0.8
10
mises
States Recovered ISLDS
SKMNN (Ouns)
0.6
Squared
Grates Recovered SWins
States Recovere mest3s
0.4
SLOS
ISLDG
AA-
D
MSE between Neural Activity and Prediction (Reaching)
miss
11
(DuN)
SLDS
2
16
Latent State Dimensionalitty
Dimensionality
Mean Squared Error
SALDS
SANN Statel (Resting:
SRNN State?
SRNN State3 Litter Go Cue)
SRNN Stated (Reaching)
SRNN Stote5 (Holding)
[Delay Before Go Cue)
10-1
10-2
0.1
0.2
0.3
0.4
Seconds Ahead
Nat Cousal
Neural Activtyy
True
True
Tre
True
True
(Durs)
-1.2
-0.6
0.0
0.6
-1.2 -0.6
0.0
0.5
-1.2 -0.6 0.0 Novement
0.6
-1.2 -0.6
0.0
-1.2
-0.8
0.0
0.0
Figure D.1: Direct Reaching: (A) Examples of comparison between behavioral states and neural dynamical states recovered by SRNNs, SLDS, and rSLDS. (B) Error between behaviorally-releva states and recovered neural dynamical states; each dot represents one session, with the triangle as the mean of all sessions. (C) Mean squared error (MSE) between neural activity and the reconstruction. (D) MSE between neural activity and the prediction forward in time. (E) Reconstruction of neural activity and the corresponding ground truth for one example neuron for three models, the top color bars represent the performance on 0.1 seconds ahead prediction. More examples are shown in Figure D.2 (F) Flow fields of latent dynamics of SRNNs in different discrete states.
15

---

A
Reconstruction vs True Neural Activity (Curved Reaching)
Neural Activity
1911
HEI
-1.2
-0.6
0.0
0.6
-1.2
-0.6
0.0
0.6
0.6 Time from Movement (s)
-1.2
-066
0.0
-1.2
-0
0.0
0.5
-1.2 -0.6
0.0
0.6
Time from Movement (s)
Time from Movement
Time from Movement (s)
from Movement
B
Reconstruction vs True Neural Activity (Direct Reaching)
5W
SALOS
Neural Activity
Time from Movement (s)
Time from Movement (s)
Time from Movement (5)
Time from Movement (s)
Time from Movement (s)
B
Reconstruction vs True Neural Activity (Decision Making)
M/y
Neural Activity
-2
-6
-4
-2
-6
-4
-2
Time from Licking
Time from Licking
Time from Licking
Time from Licking
Time from Licking (s)
C
Reconstruction vs True Neural Activity (Lever Pull)
"ue
Neural Activity
19-4t
Mein
Mbat
-10 -8 -6 -4
-2
-10 -B
-B
-4
-2
-10 -8 -6 -4
-2
0
-10 -8
-6
110 -8 -6 -4 -2
Time Lever (s)
Pull
from Lever (s)
from (s)
Figure D.2: Reconstruction of neural activity and the corresponding ground truth in different brain region/neuron for (A) curved reaching, (B) direct reaching. (C) self-initiated decision making, and (D) self-initiated lever pull. Red color represents SLDS, green color represents rSL.DS, blue color represents SRNN, purple color represents SNLDS. and orange color represents mrSDS.
16

---

A
B
MSE between Neural Activity C
D Variance of Behavioral States Recovered by SRNN between Different K (Reaching)
Reconstruction (Reaching)
and Reconstruction (Rraching)
1e-3
le-3
SUAR [Deamastinng)
Curved
SRNN
Resbro
Delay
Go Gun ne holdrn
Mean Squared Error
Mean Squared Error
3
States Reccvoree by Sttit (-4)
Variance
States Recovered by SANNECC-G)
-1.5
1.0
0.5
0.0
6.5
0.2
0.1
Km4 Km5
K-6
K=4
K=5
K=6
K=4K=5K=6
Figure D.3: Curved Reaching: (A) Examples of comparison between behavioral states and neural
dynamical states recovered by SRNN with K = 4, K 5, and K = (B) Mcan squared error (MSE) between neural activity and the reconstruction by SRNN with K 4, K 5, and K 6. (C) Mean squared error (MSE) between neural activity and the reconstruction by SRNN with = 4, K and K 6 using the co-smoothing test. (D) Variance of behavioral states recovered by
SRNN with K = 4, K 5, and K =6 6.
A
B MSE beiween Neural Activity and Reconstruction (Decision Making)
CMS between Neural Activity and D Variance Behavioral States Reconstruction (Decision Making)
Recovered by SRNN between Different (Decision Making)
1e-3
1e-3
Dehavioral States
Beding
levding
Debyy
tirm
Mean Squared Error
Mean Squared Error
States Recovered by Sihh (K+5)
by
Variance
States Recovered by SRNN (C
SRNN
0.2
from Licking
emooming!
K=4
K=5 K=6
K=4 K=5 K=6
K=4K=5K=6
Figure D.4: WFCI Decision-Making: (A) Examples of comparison between behavioral states and neural dynamical states recovered by SRNN with K 4, K 5, and K 6. (B) Mean squared еттог (MSE) between neural activity and the reconstruction by SRNN with K 4, K = 5. and
K 6. (C) Mean squared error (MSE) between neural activity and the reconstruction by SRNN with
K 4, K = 5, and K = 6 using the co-smoothing test. (D) Variance of behavioral states recovered
by SRNN with K 4, K 5, and K 6.
17

---

A
Activity Reconsttuction Pull)
CSE between Neural Activity and D Variance of Behavioral States
Reconstruction (Lever Pull)
Recovered SANI between Different (Lever Pull)
1e-4
SRAA (E-emostting)
Bchavioral
SRNN
Besting
Intibatrry
Mean Squared Erro
Mean Squared Error
States Becmoore GRMNIX-2)
Variance
2
States Recovered SRNNIK=))
-10
-5
-4
0.10
K=2
K=3
K=4
K=2 K=3 K=4
K=2K=3K=4
Figure D.5: WFCI Self-Initiated Lever Pull: (A) Examples of comparison between behavioral states and neural dynamical states recovered by SRNN with K = 2, K 3, and K 4. (B) Mean squared error (MSE) between neural activity and the reconstruction by SRNN with K = 2. = 3, and K 4. (C) Mean squared enror (MSE) between neural activity and the reconstruction SRNN with = 2, K = 3, and K = 4 using the co-smoothing test. (D) Variance of behavioral states
recovered by SRNN with K 2, K 3, and K 4.
Reconstruction vs True Neural Aclivity Using LFADS
(Curved Reaching)
(Direct Reaching)
PM6
PM81
le- 3
Comparison of Reconstruction
LEASS
UScS
10
SLDS
Neural Activ-
PMc11
8
rSLDS
MSE
SRNN
6
LFADS
4
Curved
Direct
0.6
Time from Movement
Time from Movement
Figure D.6: (A) Comparison of neural activity reconstruction among SLDS, rSLDS, SRNN and LFADS in reaching dataset. (B) Reconstruction of neural activity and the corresponding ground truth four example neurons for LFADS.
18

---

A
C
1e-2
0.8
0.7
10
0.6
Loss
MSE
0.5
2
0 0.4
-10-
0.3
0.2
-10
1000 2000 300D 4000 5000
1000 2000 3000 4000 5000
0
1000 2000
3000 4000
Epochs
Epochs
Epochs
D
1e
F
10a
0.5
Loss
10
MSE
0.4
-10
0.2
-102
1000 2000 4000
4000 5000 Epochs
1000 2000 3000-4000
4000 5000 Epochs
1000 2000 3000 4000
Epochs
G
0.8
10
10
0.6
Loss
10a
MSE
0.4
-10?
0.24
-109
1000 2000 3000 4000 5000
1000 2000 4000 Epochs
0.
1000 2000 3000 4000 5000 Epochs
Epochs
Figure D.7: First column: Plot of SRNN training loss across 5000 epochs for example in (A) reaching dataset, (D) self-initiated lever pull dataset, and (G) self-initiated decision making dataset. Second column: Plot of MSE between desired neural activity in validation and the corresponding reconstruction by SRNN for (B) reaching dataset, (E) self-initiated lever pull dataset, and (H) self- initiated decision making dataset. Third column: Plot of error between behavioral states in validation set and the corresponding recovered state by SRNN for (C) reaching dataset, (F) self-initiated lever pull dataset, and (I) self-initiated decision making dataset.
Behavioral Recovered
SRNN Bid/rectional Inference Layers
Forward Inference Layers
SAAN recional intrenon
(didimctionll infernne)) SRA. (Forward infeneece)
0.6
SRAN disiree infenecc)
SANM (Fonward infennece)
0.6
SAA.N saadd infereneet
0.4
0.4
0.2
0.2
Curve
Direct
Lever Pull
Decision Making
MSE between Desired Neurao Activity and Reconstruction SRNN rectional Inference
SRNN with Forword inference Layers
1e-3
E
2
1o_4
10-3
3PAM reciioaa intrrraal
(SU) recilonal Infereeccl
MSE
MSE
MSE
1.0
Curve
Direct
Lever Pull
Decision Making
Figure D.8: Error between behavioral states and recovered states by SRNN with bidirectional recurrent inference layers (blue) SRNN with forward recurrent inference layers (gray) in (A) reaching dataset, (B) self-initiated lever pull dataset, and (C) self-initiated decision making dataset. MSE between desired neural activity and reconstruction by SRNN with bidirectional recurrent inference layers (blue) vs SRNN with forward recurrent inference layers (gray) in (D) reaching dataset, (E) self-ine initiinit lever pull dataset, and (F) self-in ceff-initiated decision making dataset.
19

---

A
Region-to-region Communication Map Uncovered by SRNN (Lever Pull)
Resting
Inhibitory
After Behavior
Other
Other
Other
0.5
SSp
SSp
SSp
Visual
Visual
Visual
Motor
Motor
Motor
Motor
Visual
5Sp
Other
Motor
Visual
SSp
Other
Motor
Visual
SSp
Other
-0.5
h++1
h++1
ht+1
C
1e-4
3
SRNN
SRNN
0.6
Error
2
MSE
0.4
1
0.2
Lever Pull
Lever Pull
Figure D.9: (A) Region-to-region communication map uncovered by constrained SRNN trained on self-initiated lever pull dataset. We focus 4 regions, i.e., "Visual", "Motor", "Somatosensory (SSp)", and "Others". (B) Error between behavioral states and recovered states by constrained SRNN on self-initiated lever pull dataset. (C) MSE between desired neural activity and reconstruction by constrained SRNN on self-initie relf-initiated lever pull dataset.
20

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We claim that our abstract and introduction accurately reflect our contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discussion the limitations of our model in the Conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
21

---

Justification:
We provide equations on the main text for our theory.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We cite all the source of our datasets in the main text and include the details of our experiments in Method and Result section, and we will provide our code for reproduction upon publication.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
22

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
Polooyoollyyyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.ccpublle pwwppppguuuiieeeeep hetpi///uuidessoeeeuuubiiii SooodeSubiissi.e ssiiiioicccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We provide details 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, confi-
dence intervals, statistical significance tests, 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 with given experimental conditions).
23

---

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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide details in the Methods section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration 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 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 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 particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
25

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "fBuXPJ3x0b", "reviewer": "Reviewer_8VtM", "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.", "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).\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.", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you very much for giving us a positive rating, and for you very helpful comments. Below we respond to specific weaknesses pointed out.\n\n**Weakness 1**: *Lack of comparison with other methods, for example, ARHMMs and their extensions, as well as SNLDS and mrSDS*\n\nWe have now included the two suggested baseline models, mrSDS and SNLDS, in the PDF attached above. To summarize, as above, while recovering behaviorally-relevant states, our model SRNN outperforms SNLDS and has comparable performance with mrSDS, whereas in reconstructing neural activity, SRNN outperforms mrSDS and has comparable performance with SNLDS . **Therefore, SRNNs have overall better performance on recovering behaviorally-relevant states and reconstructing neural activity**. Additionally, we would like to point out that we used the SSM package to install 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, already included in the initial submission.\n\n**Weakness 2**: *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*\n\nWe thank the reviewer for suggesting the noisy high dimensional Lorenz Attractor and the common simulated data NASCAR. We now include results on testing SRNNs using the suggested setup for the noisy high dimensional Lorenz Attractor and NASCAR in the PDF attached above (Figure 9 and Figure 10). To simulate the Lorenz Attractor, we project the 3-dimensional Lorenz into an 8-dimensional space and add additive Gaussian noise. We found that SRNNs with P=3 can successfully recover the butterfly structure of a Lorenz Attractor using the noisy high dimensional data. Moreover, we found that our SRNNs can successfully recover the ground truth of the trajectories and states of the NASCAR dataset.\n\n**Q1)** *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*\n\nWe thank you for your question. One failure mode may be when the training data is limited. To identify this, we trained our model on a single condition in the reaching dataset instead of 26 conditions; we found that the test reconstruction accuracy is lower than for SLDS and rSLDS (see attached PDF, Fig 5). Another failure mode may be the commonly-encountered problem of overfitting: we found that SRNNs may overfit to the training data if we increase the number of parameters (number of states K and number of latents P) to be large. Lastly, there may be a small problem of identifiability as pointed out by Reviewer Kpfw, i.e., if we increase the number of latents to be very large, the neural data can be successfully modeled by a single RNN, but of course this fails to recover interpretable switches in neural dynamics. We identify these tradeoffs and provide a discussion in the rebuttal for Reviewer Kpfw.\n\n**Q2)** *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*\n\nWe agree with the reviewer that behavioral data may provide the interpretability that we desire while modeling. We tested SRNNs on poses tracked using DeepLabCut, specifically the nose and paw positions, as recorded by a behavioral camera in the head-fixed decision-making task (see Fig. 4C in Rebuttal PDF). Interestingly, the states recovered were not as accurate as when using the neural data (see Fig. 4 of original submission). In this task, a visual stimulus was presented either to the left or the right, not visible in the camera. Specifically, the model trained on behavioral pose data failed to identify a separate state for ‘stimulus presentation’, presumably since the nose and paw location was not noticeably different here than during the delay period. This exercise points towards the possibility of augmenting behavioral observations with neural activity in future work, to better inform switching states. \n\n**Q3)** *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\nOne important parameter is the number of discrete states K, and we show a comparison between different numbers of hidden states in the PDF attached and include a thorough discussion in the **global rebuttal**. Another very important hyperparameter is the number of hidden states P of SRNN, we show a comparison between different numbers of hidden states in the original submission. We included the discussion about determining K and P in the global rebuttal as well as in the official comment below. For other hyperparameters, such as the number of hidden units in our inference network and learning rate, we used cross validation and kept the parameters with best performance, but we found that these other parameters did not affect the results very much. It is possible to sweep over the hyperparameter space, however, this may require training in parallel with large resources, for example, in the reaching dataset, training one model with P=16 and K=5 takes 9 hours 32 mins on a NVIDIA A100 GPU, and training models with (P=2, K=5), (P=4, K=5), and (P=8, K=5) takes 8 hours 01 min, 8 hours 30 mins, and 8 hours 54 mins respectively.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "doTdTsAQtM", "reviewer": "Reviewer_hFy4", "summary": "The authors propose to model time series neural population activity using switching recurrent neural networks. The generative model includes discrete latent states", "questions": "- 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?", "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.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed method does appear to outperform related switching linear dynamical systems approaches in certain contexts.", "weaknesses": "High-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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you deeply for your time and attention in reading our paper, and for your valuable comments. Below is our response to specific weaknesses and questions.\n\n**Weakness 1**: *Comparison between SNLDS and mrSDS*\n\nThank you for raising this weakness. As also mentioned in the global rebuttal, we have now implemented these two models, the SNLDS and mrSDS. We show their performance on the Reaching Dataset in the attached PDF. To summarize, while recovering behaviorally-relevant states, our model SRNN outperforms SNLDS and has comparable performance with mrSDS, whereas in reconstructing neural activity, SRNN outperforms mrSDS and has comparable performance with SNLDS. **Therefore, SRNNs have overall better performance on recovering behaviorally-relevant states and reconstructing neural activity**. We would like to point out that mrSDS was published recently by ICLR 2024 in May, two weeks before our original submission, and to date, we have not found any code released with the paper. However, we implemented the code based on our reading and understanding of the paper to the best extent possible.\n\n**Weakness 2**: *Determining the appropriate number of discrete states. One approach to this might be co-smoothing*\n\nWe thank you for your comments and for the excellent suggestion of ‘co-smoothing’. We have now implemented co-smoothing and find that SRNNs are successful at latent state modeling using this metric, and that this method is helpful for determining the appropriate number of discrete states. In addition, we now provide a comprehensive set of metrics to determine this important hyperparameter, including **co-smoothing** in the **global rebuttal** as well as in the official comment below. To summarize, we use appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance and co-smoothing method, 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.\n\n**Detail Weakness 1**: *The mathematical details and notation are often unclear. For example, eq 2: Shouldn't this instead be a categorical distribution or similar? Relatedly, f is also used in eq 8*\n\nYes, you are right. It should be a softmax activation function and categorical distribution, which is what we implemented in the original paper. We apologize for the typo and we have revised both equation 2 and equation 8 in the paper.\n\n**Detail Weakness 2**: *The authors should more clearly describe the cross-validation techniques used for each dataset*\n\nThank you for this comment and we apologize for the confusion. We are training the model on N-1 conditions/sessions/subjects and testing it on the 1 held-out condition/session/subject. In this process, we train N different models, one for each one of conditions/sessions/subjects on hold, and show the results for each model as one dot in Figures 3-5. We had detailed for each experimental dataset whether we are considering conditions or sessions or subjects to be held-out data, but have now moved it to the main text in the interest of clarity. Specifically, we hold out data from one condition at a time in the reaching dataset, where there are 18 different curved reaching conditions. We hold out data from one session at a time in the decision-making dataset, where there are 8 total sessions, and we hold out data from one subject at a time in the lever pull dataset, where there are 6 total subjects. \n\n**Q1)**: *Are the predictions in Fig 2 cross validated?*\n\nWe thank you for your question and apologize for the lack of clarity. Yes, all models in this paper are cross-validated using held-out test sets. Specifically for Figure 2, we have one Lorenz attractor as the underlying dynamical system. We train SRNNs on 15 trials (representing different initial conditions), and test on 1 held-out trial. We have now clearly included this information in the paper.\n\n**Q2)**: *In Fig 3, are the authors modeling single-trials or condition averages (ie PSTHs)?* \n\nWe are modeling condition averages in the paper. In addition to the analyses in the paper, we have now also trained SRNNs on simulated single trial data to assess if our method works on this kind of data. Specifically, we implemented the thinning algorithm using the time-rescaling theorem to simulate single trial spiking data based on the condition averaged firing rate [Brown et al., 2002]. We simulated 200 single trials from 5 condition averages and computed the binned spike counts, which we used to train the SRNN models. We found that the SRNNs successfully represent this data with P=16 and K=5; we show the results of training on simulated single trials in the attached PDF (Figure 6).\n\n**Q3)**: *Why does SRNN perform worse than SLDS and rSLDS in Fig 5CD?*\n\nIn Figure 5C, the SRNNs indeed have worse reconstruction performance than SLDS and rSLDS. As we mentioned in our paper, ‘Indeed, all three methods have acceptable reconstruction.’: we didn’t find any significant difference by visualizing the reconstruction (Figure 5E of original paper). A potential reason causing the worse performance may be the small size of this dataset. As our response to Review 8VtM, a failure mode of SRNNs is when the model is trained in the low data regime, e.g., single conditions (see attached PDF). However, SRNNs have better recovery of behaviorally-relevant states. This may explain the finding that SRNNs have better prediction performance in 0.67, 1, and 1.33 seconds ahead in Figure 5D of the original paper; as the prediction window increases, the prediction performance may depend both on the recovered dynamics and the recovery of correct behavioral states. We believe that prediction of future activity provides a better test of a dynamical model, and here, SRNNs have overall better performance in prediction of neural activity.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mDNXRHhvug", "reviewer": "Reviewer_Kpfw", "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.", "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?\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?", "limitations": "The discussion adequately acknowledges limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The results shown in panels A of Figs 3, 4, and 5 are nice and convincing.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you very much for giving us a positive rating, and for your extremely relevant  comments. Below we respond to specific weaknesses pointed out.\n\n**Summary of Weakness 1**: *Like many other deep learning based approaches, the model is not particularly interpretable. 2D flow fields poorly capture the true dynamics of the system.*\n\nYes, you are absolutely correct, a 2D flow field is not able to accurately capture the true dynamics when P>2, and simply provides a visualization. We are able to also calculate slow points and other features of the dynamics in the full dimensional space, provided here as an example in Figure 8 of the attached PDF. \nWhen comparing the dynamics in different discrete states of our model, we use the same principal components to project the higher-dimensional dynamics to 2D. Therefore, the dynamics can technically be compared with each other. Additionally, where possible (e.g., in Figure 5F of the original submission), we plot the flow fields for SRNNs with P=2; here, the visualization and structure found is of the true dynamics. We have also run this model with a different random seed and recovered similar features of the dynamics in Figure 8 of the attached PDF, showing identifiability of the dynamics.\n\n**Summary of Weakness 2**: *It is unclear how to choose tunable parameters on datasets without ground truth, or at least good educated guesses.*\n\nWe thank you for your comments, and comprehensively discuss this in the **global rebuttal** as well as in the official comment below. To summarize, we use appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance and co-smoothing method, 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.\n\nAdditionally, for selecting the number of continuous hidden states P, we included comparisons in the paper, we think an appropriate P should perform well on both recovering states constantly and reconstructing behavior, meanwhile, the P should be as small as possible.\n\n**Summary of Weakness 3**: *How to choose between a SRNN and a model with higher P.*\n\nYes, you are definitely right. According to the theory of universal function approximation, a model with higher dimensions can do the same work as multiple models with low dimensions. To empirically test this, we trained a standard RNN with P=32, which has a similar number of parameters on recurrent weights but more on emission weights, and found that standard RNNs are also able to reconstruct neural activity well (see attached PDF Figure 7). However, neural dynamics are broadly thought to be low-dimensional, and we would like this dimensionality to be as low as possible to achieve interpretability in the dynamics as well as the discrete states (also discussed above). Indeed, we also found that the switches in the neural dynamics are  relevant to different behavioral and stimulus states; therefore, if we use a higher dimensionality but a non-switching model, the model may not capture this interpretability.\n\n**Q1)**: *Typo in Equation (2)*\n\nWe thank you for pointing this out, it is a typo, and we used softmax in implementation. We will revise the equation.\n\n**Q2 and Q3)**: *It wasn't obvious to me whether you are allowing the continuous hidden state to impact the transition for the discrete state. In Figure 1B, what is the meaning of the red dashed arrow? Does that carry any difference to the black arrows?*\n\nWe thank you for your question. The value of the transition probability depends on $h_t$. Our model is similar to rSLDS, where the transition probabilities depend on continuous states. We plot the red dashed arrow to show that the transition probabilities depend on previous continuous states (they are red since the states z are not directly computed using $h_t$). The black arrows between z represent the transitions between z.\n\n**Q4)**: *Are the good results shown in panel A of Figs 3, 4, and 5 due solely to differences in the initialization procedure across models?*\n\nWe thank you for your question. Initialization might be the most common and powerful method to overcome the difficulty of training in this area. We used the SSM package to implement SLDS and rSLDS where HMMs are also used as an initialization.\n\n**Q5)**: *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\nWe thank you for your question. We agree that we should see the same structure in the dynamics across different random seeds. To show a simple example, we changed the random seed and trained the SRNN on lever pull data with P=2. We show a visualization in the attached PDF in Figure 8. We found similar flow fields with a comparable fixed point structure. Specifically, the new fixed points seem to be rotated from the old ones in the same direction.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1VAJQ1p6ad", "reviewer": "Reviewer_uaif", "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.", "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?\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.", "limitations": "The authors have addressed limitations in the last section of the paper, and I do not envision any societal impact of this work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. 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.", "weaknesses": "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.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you deeply for your time and attention in reading our paper, and for your valuable comments. Below is our response to specific weaknesses and questions.\n\n**Summary of Weakness 1**: *comparison to SNLDS.*\n\nWe thank you for pointing out this weakness. In addition to the existing SLDS, rSLDS, and LFADS, we implemented two additional baseline models that reviewers mentioned, the SNLDS [Dong et al., 2020] and mrSDS [Karniol-Tambour et al., ICLR 2024]. We show their performance on the Reaching Dataset in the attached PDF. To summarize, while recovering behaviorally-relevant states, our model SRNN outperforms SNLDS and has comparable performance with mrSDS, whereas in reconstructing neural activity, SRNN outperforms mrSDS and has comparable performance with SNLDS . Therefore, SRNNs have overall better performance on recovering behaviorally-relevant states and reconstructing neural activity. \n\n**Q1)**: *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\nWe thank you for allowing us to elaborate on the differences between the two models. (1) Indeed, the transitions in our model are inspired by rSLDS, i.e., in SRNNs, the transition depends on continuous states $h_t$. However, we observed no significant performance difference between transitions based on the observations $y_t$ as compared to $h_t$, as $y_t$ is a linear transformation of $h_t$. However, $h$ typically has a lower dimensionality, requiring fewer parameters to model the transition, e.g., in the context of reaching, the dimensions of the continuous latent variables, $h$, P=16, vs. the dimensions of neural activity, $y$, R=180. (2) SNLDS is able to be tweaked in order to make predictions by using a causal inference network, which is the approach utilized in our prediction work. The original SNLDS paper used non-causal inference, which involves inferring states from the entire sequence. Consequently, the original SNLDS method is unsuitable for predicting neural activity, as it requires information from future neural activity or even the entire neural activity sequence for state inference, but again, this is able to be tweaked for the purpose of prediction. (3) Lastly, SNLDS and SRNNs have different initialization processes, which we suspect is contributing to SRNNs’ better performance in modeling neural data (see response to Weakness 1). While SNLDS uses regularization in the objective function to ensure uniform state utilization, SRNNs use HMM as an initialization, again inspired by previous models that have been successfully applied towards modeling neural data.\n\n**Summary of Weakness 2**: *It is not clear if SRNN is capable of inferring the # of underlying states.*\n\n**Summary of Question 2**: *how results vary across different # of discrete states across all models. This is for two main reasons: i. In a new dataset we might not know ground-truth behavioral segmentations. ii. The fact that rSLDS does fine in predicting data but infers discrete states inaccurately.*\n\nWe thank you for your questions. We now provide a comprehensive set of metrics to determine this important hyperparameter in the **global rebuttal** as well as in the official comment below. To summarize, we use appropriate selection metrics for the number of discrete states: (1) convergence / reuse of discrete states, (2) reconstruction performance and co-smoothing method, 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.\n\nMoreover, we also show reconstruction performance across all models with different K in the attached PDF (Figure 2).\n\nLastly, we agree that rSLDS does fine in reconstruction of data but infers discrete states inaccurately. Please see below (response to Weakness 4) for a discussion on this topic. \n\n**Summary of Weakness 3**: *The paper will benefit from some editing by the authors.* \n\nWe thank you for pointing out this weakness and we apologize for the inconvenience in locating the figures when you review our paper. We have fixed the reference formatting and the figure references. Additionally, we have added some simulation and experimental details in the main text, and removed some commonly-used VI details to the Appendix. \n\n**Summary of Weakness 4**: *Prediction + reconstruction performance across all models is similar. However, the discrete states being inferred look hugely inaccurate for SLDS and rSLDS.*\n\nWe thank you for pointing this out; this has puzzled us as well. We believe that accurate reconstruction of data does not necessitate accurate inference of discrete hidden states or dynamics, since the models are trained purely for maximizing the ELBO, i.e., here, reconstruction accuracy. However, accurate prediction of neural activity many time steps ahead requires simultaneously (a) accurate reconstruction, (b) accurate inference of discrete states, and (c) accurate inference of dynamics. Here, we see concrete differences in the prediction capability of SRNNs vs. SLDS / rSLDS (e.g., see Figure 4D of the original submission). We believe that prediction of future activity provides a better test of a dynamical model, and here, SRNNs have overall better performance in prediction of neural activity many time steps ahead.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Reviewers,

We would like to thank you for providing constructive feedback that helped us improve the paper. As a reminder, in this submission, we propose ‘Switching Recurrent Neural Networks’ (SRNNs) for discovery of switching neural dynamics that leads to behaviorally-relevant discrete states, and leads to more accurate reconstruction and prediction of neural data than baseline methods. Here, we detail responses to suggestions made by multiple reviewers, and address individual reviews below. If any questions are unanswered or our responses are unclear, we would appreciate the chance to engage further with you. Additionally, **please find a PDF with helper figures attached**. These are referenced and described in our responses.

### **1. Additional baseline models**
In addition to the **existing SLDS, rSLDS, and LFADS**, we implemented and have updated our paper with two more baseline models that reviewers mentioned, the **SNLDS** [Dong et al., ICML 2020] and **mrSDS** [Karniol-Tambour et al., ICLR 2024]. We show their performance on the Reaching Dataset in the attached PDF (Figure 1). To summarize, while recovering behaviorally-relevant states, our model SRNN outperforms SNLDS and has comparable performance with mrSDS, whereas in reconstructing neural activity, SRNN outperforms mrSDS and has comparable performance with SNLDS. **Therefore, SRNNs have overall better performance on recovering behaviorally-relevant states and reconstructing neural activity**. We would like to point out that mrSDS was published recently by ICLR 2024 in May, two weeks before our original submission, and to date, *we have not found any code released with the paper*. However, we implemented the code based on our reading and understanding of the paper to the best extent possible.

### **2. Determining relevant hyperparameters**
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 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 (Fig 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 (Fig 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 (Fig 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.29 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.

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

### **3. Additional analyses**

In response to **individual reviewers**, we also implemented the following which have made our paper much stronger.

(1) SRNN trained on single reaching conditions to explore failure modes of our model

(2) SRNN trained on simulated single trials of binned spike counts

(3) Comparison between SRNN and a standard RNN with a higher dimensional continuous latent state to assess identifiability of the models

(4) Flow fields with a different random seed to assess stability of recovered dynamics

(5) SRNN on two more simulated datasets: high dimensional noisy Lorenz Attractor and NASCAR

### **4. Typos in Equations and Editing of Paper**
We appreciate all comments and suggestions on pointing out the typos and errors of the paper. We apologize for the inconvenience during the review process.

We thank all reviewers for the suggestions and comments and provide a detailed rebuttal to each reviewer below.

Meta-review: 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.
Decision: Accept (poster)