Parsing neural dynamics with infinite recurrent switching linear dynamical systems

Thank you for your review! We hope what follows answers your remaining questions.

Weaknesses: We start by first acknowledging the weakness raised. Indeed, the partitioning into different linear DS you mentioned is an important intuition. It is present in our model through the recurrent ($x_n \to z_{n+1}$) connections, but with the effective partition being time-varying. For rSLDS it is indeed a fixed global nonlinear DS, which induces a weaker prior over non-stationary dynamics. Nonetheless, the resulting system can still support non-stationary dynamics, which we believe makes comparisons on non-stationary systems still justified. A simple extension to the rSLDS that would allow for better modeling of nonstationarity would be to include an additional dependence on time $n$ through external inputs. This way, this partition over the $\boldsymbol x$ state space would be allowed to also vary with time. If deemed relevant, we would be happy to include this extension analysis as a supplement in the final version!

Questions: Addressing your questions:

• Extended NASCAR: We assume you are referring to the sample continuous states in Figure 2E. The posterior mean trajectories in Figure 2C are smooth, but in Figure 2E we plot the sampled continuous state trajectories, that will by definition be noisy.
• Thank you for pointing this out. We have now updated Table 1 in the main text to incorporate the dynamical MSE, as well as Table 2 in the Appendix.

Finally, could you expand on this sentence?

It would be more convincing to see the discrete state inference from continuous recording for real nonstationary system.

Thank you again for your feedback. Let us know if we can better answer any remaining questions you may have!

We are glad to read our reply addressed your concerns!

Thank you for clarifying. Indeed, there are considerable practical considerations to the full continuous recordings since we'd be dealing with an order of magnitude longer sequences. There are also considerations on the modeling side, as modeling the continuity in x and z between trials has to be handled carefully. For instance, in the decision-making data considered here, there are inter-trial breaks (elapsed times) that might make the latent states discontinuous. To tackle this, one approach is to define different dynamical transitions for the latents at those time steps from one trial to another (see e.g. (Roy et al, Neuron, 2021)). Thank you for raising this point, we will include a comment on this aspect in the limitations section of the paper.

Thank you for your careful review! We appreciate the time and dedication towards highlighting the strengths and valuable feedback to make the paper better. We hope below to have summarised faithfully the weaknesses and questions you raised.

First, addressing the weaknesses point by point:

• Polya-gamma: Thank you for the suggestion on the Polya-gamma augmentation, we will add some of the suggested details on the time-dependent sizes in the main text for the final version.

• Extended NASCAR: The rSLDS states equivalent to Figure 2E-F can now be found in Appendix Figure 5 -- we apologize for their omission in the original submission. Indeed as you predicted, the sampled continuous states do not look as good for the rSLDS, whereas the samples are much more in line with the posterior mean for the irSLDS. Furthermore, we also provide the equivalent to Figure 2F for the rSLDS by plotting the most likely $\hat z_{1:T}$ trajectory, and we see that the rSLDS model mistakenly maintains only $K=4$ states after the switch to the alternate dynamics (where $K=6$). Together these results reinforce the performance of the irSLDS over the rSLDS in this non-stationary case. We can include these details in the main text for the final version.

• Empirical considerations: Finally, thank you for bringing up concerns on the practitioner, empirical, side. First, we hope the provided running times in the updated manuscript (see Appendix D.2.) help give a better sense of the computational cost. They support, for instance, how a grid search over $K$ for the rSLDS would be quite computationally expensive compared to our irSLDS, which has more flexibility with respect to the inferred $K$. Furthermore, recall that we provided empirical evidence (in Section 4.2 on the IBL data) towards the irSLDS model's ability to uncover fluctuations in the number of discrete states $K$ on a trial-by-trial basis. The rSLDS provided a slightly worse fit, and did not uncover those fluctuations as much. This indicates that even if the grid-search over $K$ with the rSLDS was computationally efficient and preferred, it would still yield a constant $K$ that might miss those within-dataset fluctuations.

As for the additional questions/comments:

• Computational complexity: Thank you for pointing this out, this is highlights a good need for clarification! For the parallel scan, we focus on its use for the computation of the full transition matrices $W_{1:T}$. Since these are matrix dynamics, and not vector dynamics, they do require $\mathcal{O}(K^3 T)$ ops, from $T$ sequential matrix-to-matrix multiplications. We agree that the text was misleading, making it seem that we use it for the computation of the vector sufficient statistics $\boldsymbol w_n$. This latter computation only really arises if we want to generate samples from the model, in which case it is indeed sequential -- but so are all other models. We provided some clarification in the updated manuscript, and will pay attention to those details for the final version!
• Thank you for the relevant reference, we've included it in the introduction.

Thank you again for your feedback. We look forward to answering any further questions you have!

Thank you for your positive review and feedback! We hope what follows answers your remaining questions.

To address the weakness raised, handling non-stationarity is indeed often challenging. Fortunately, the discrete-state process is defined by only a few key parameters, which helps in terms of interpretability and can help with hyperparameter optimization during training. Furthermore, our implementation of the scan operation does significantly reduce the computational complexity (see empirical running times added in Appendix D.2) where it can be deployed.

As for your questions:

• Indeed we set $K=8$ as an upper bound. For this specific visualization of inferred discrete states, namely the most likely sequence, indeed it does not use state #2. This state is still however well defined for the model, and other analyses (such as for instance smoothing marginal posteriors $p(z_t \mid \hat x_{1:T}, y_{1:T})$) might leverage that state.
• Thank you for pointing that out! We provide the same plot as Figure 2E-F for the rSLDS in Figure 5 in the Appendix. As an overview, we observe that the sampled continuous states do not look as good for the rSLDS, whereas the samples are much more in line with the posterior mean for the irSLDS. As for the equivalent to Figure 2F for the rSLDS, with the most likely $\hat z_{1:T}$ trajectory, and we see that the rSLDS model mistakenly maintains only $K=4$ states after the switch to the alternate dynamics (where $K=6$). Together these results reinforce the performance of the irSLDS over the rSLDS in this non-stationary case. We can include these details in the main text for the final version.
• There would be no limitations to considering Poisson instead of Gaussian observation process for the irSLDS framework, with this specific implementation. Indeed, we do not leverage the Gaussian observation process directly, which one could do for instance to do efficient (even exact if given the $z_{1:T}$ sequence) marginalization of the continuous states in SLDS models. Here, all of our implementation is based on the SSM package, so all computation is numerical and readily generalizable to other emission processes.

Finally, thank you for pointing typos/clarifications, we will integrate them in the final version. We look forward to answering any further questions you have!

We thank you for the review! Specifically, we appreciate bringing to light the influence of hyper-parameters. Addressing the weaknesses point by point:

• Smoothness: Indeed influence from ``nearby'' states is mediated by the diffusion parameter $\gamma$, or in the FDM implementation, $\beta \propto \gamma$. For the results reported in the submission we fixed $\beta = 1.0$ (for all tasks) determined from a quick hyper-parameter search. We agree an exploration of smoothness, per task, would be insightful -- we propose to include the log-likelihood search as well as example generated trajectories in the final version.

• Time duration and dependence: Similarly to the smoothness described above, discrete-state duration can be mediated by the $\kappa$ hyperparameter. We used $\kappa = 0.4$, determined from a quick hyper-parameter search. As with the smoothness, we could include in the final version details on this hyper-parameter search. Regarding the question raised on time dependence: this is where our accent on the inclusion of input-driven transitions, either from external inputs $\boldsymbol u_n$ or continuous internal states $\boldsymbol x_n$ (recurrence), can shine through. Indeed one could include the time $n$ as an external input, for instance by defining the sufficient statistic dynamics as

$

\boldsymbol w_{n+1} = U \boldsymbol w_{n} + \kappa\mathbf{1}_{i} + R \boldsymbol x_n + r + g(n,i)

$

for $g(n,i)$ some parametrized function (e.g. affine $Q^{(i)} n + q$). Concretely for the Lorentz system mentioned, this modeling could leverage the continuous latent $\boldsymbol x_n$ to dictate one discrete state per side of the attractor, or use the time dependence to do so.

• Number of states: For the synthetic experiments, $K$ was set to the true known value ($K=4$ NASCAR, $K=6$ extended NASCAR). For the IBL task, we used $K=8$ and provided in Appendix section D results on the log-likelihood search over different values of $K$. To address your final question, we will provide in the final version an analysis for scaling $K$, showing for instance how the different models fare in ($K$) over-parametrized regimes.

Thank you again for your feedback. We look forward to answering any further questions you have!

Thank you for your thorough review, we appreciate the time and dedication towards making the paper better. We hope to have summarised faithfully the weaknesses and questions you raised, and we address them point-by-point below:

• Comparison with other models: We agree that our analysis is limited to the switching linear dynamical system and its direct extensions. Our primary goal was to convey how a carefully designed prior (in our case the heat equation that captures the dist-CRP and induces a state geometry) with select hyperparameters can maintain similar performance/efficiency, while increasing interpretability. In line with this goal, we focus on only one model family, with and without this prior. We appreciate the feedback on reinforcing the analysis and will provide scalability results in the final version with a more thorough analysis beyond the 4-8 states!

• Heat equation prior: We agree that there are alternative prior formulations, and we do not make claims that this is necessarily the best prior to uncover switching. These possible alternate priors are not discussed thoroughly in the current manuscript, and we will make sure to include a discussion element on that regard. Nonetheless, our claim is that this prior encapsulates important properties, and as such offers a good prior for model fitting and data-driven analysis. With regards to your question "Why does the heat equation prior directly enable discovery of greater fluctuations/switching in the data?": as mentioned in the text, we attribute this capability to the fact that the heat equation prior encapsulates the dist CRP (point #2, end of section 3.2), yielding a model with greater capacity to calibrate the number of states used from the data.

• Parallel scan: Thank you for raising questions on this topic, we realize that there is a need for clarification! First, just so there is no doubt, we use the parallel scan operation for the computation of the transition matrices $W_n$, not for the computation of the dynamics of $\vec x_n$. Hence our dynamics matrix of study is the FDM matrix $U$, not the time-varying linear matrices for the continuous state dynamics. Now, indeed, the previous work mentioned using the scan operation has focused on dynamics with diagonal matrices, whereas our dynamics matrix is not (here, tri-diagonal). Still, the parallel scan operation can be deployed for any recurrent system with affine dynamics. The diagonalization mentioned is crucial for efficient computation as $K$ scales, but this is not our focus. Our main focus is in making the computation more efficient as $T$ scales, and the parallel scan allows that, by bringing the sequential $T$ to $(T/L + \log L)$ ops. In our updated manuscript, we make explicit the dynamics of $W_n$ in Appendix C.3 for reference, and provide empirical running times in Appendix D.2 to showcase the increased efficiency from the scan.

Thank you again for your feedback. We look forward to answering any further questions you have!