Paper: Parsing neural dynamics with infinite recurrent switching linear dynamical systems (ICLR-2024)

Reviewer comment (Experiments):
rSLDS has the intuition of partitioning the latent space into regions that are governed by different linear DS individually, and the resultant nonlinear DS globally. It's not very suitable to do comparison on a nonstationary system.

Author response:
**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!

Reviewer comment (Experiments):
In the extended NASCAR example, why is the trajectory not smooth? Is the bias too large?

Author response:
**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.

Reviewer comment (Writing):
The paper would benefit from providing a bit more context/intuition behind the Polya-gamma approach in Section 3.1 and its weaknesses. Just reading this section currently does not make it as clear why the approach in section 3.2 is necessary. Section C.1 in the appendix is helpful. Perhaps a bit more details about the Gaussians with time-dependent sizes could be added to the main paper to make this point more clear?

Author response:
**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.

Reviewer comment (Evaluation):
The experimental results are the biggest weakness of the paper. It is not currently obvious that a practitioner should clearly reach for irSLDS as the preferred method.

Author response:
**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.

Reviewer comment (Experiments):
Can the extended NASCAR example be made more extreme (e.g. more states) to make the differences between irSLDS and rSDLS more apparent? It is not clear that the differences in test log-likelihood are that drastic (and there do not appear to be underlying states of rSDLS displayed to compare to).

Author response:
**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.

Reviewer comment (Experiments):
It is claimed in Section 3.3 that irSLDS has a similar time complexity to rSLDS. It would be nice to see some kind of empirical runtime comparison to support this. How does this change when choosing different upper bounds of \(K\)?

Author response:
**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.

Reviewer comment (Experiments):
Related to the previous points, it seems two potentially simple baselines would be an rSLDS with a large number of states or alternatively performing a grid search with rSLDS over different state sizes. Given the flexibility of irSLDS, it should be able to achieve better performance than these two baselines in one-shot. Instead, it seems the same amount of grid search was done for each method over the number of states. Alternatively, if either of these two baselines still takes less time/cost than running irSLDS (this is where the empirical runtime comparison comes in), then perhaps rSLDS would be preferable for a practitioner in many cases. Empirical evidence to resolve these doubts — in terms of empirical runtime under different scenarios and the time/cost per method to achieve strong results — would strengthen the paper.

Author response:
**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.

Reviewer comment (Theory):
Unless I am missing something, I believe the last paragraph of Section 3.3 has an error regarding the computational complexity of the sequential vs parallel versions of the algorithm. It is stated that the sequential scan requires \(\mathcal{O}(K^3T)\) ops, but shouldn't this be  \(\mathcal{O}(K^2T)\) ops? This is because it requires \(T\) sequential matrix-vector multiplications (compared to the cubic cost of matrix-matrix multiplications of the parallel version of the scan). Thus the choice of parallel vs sequential scan would be dependent on the state size and sequence length, since the cubic cost could grow quite large for larger state sizes.

Author response:
**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!

Reviewer comment (Novelty):
This paper https://arxiv.org/abs/2111.01256 seems relevant to infinite switching states and could be cited as related work.

Author response:
Thank you for the relevant reference, we've included it in the introduction.

Reviewer comment (Evaluation):
In Figure 3 F lower, irSLDS identifies the state 2, although it does predict (as I understand from the absence of orange dots) that the discrete state of the system was never at that state through the trials. Isn’t that strange? Shouldn’t the framework have estimated 7 states in total for this dataset?

Author response:
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.

Reviewer comment (Theory):
What are the limitations for considering Poisson instead of Gaussian observation process for the irSLDS framework?

Author response:
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.

Reviewer comment (Writing):
Conclusion has various typos, i.e. first sentence of conclusion: “a”-> “an”, missing ’s’-es.

Author response:
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!

Reviewer comment (Experiments):
How is the smoothness prior determined? In the paper, the authors state “Each time a discrete state j is chosen, ‘nearby’ states also become more probable at the next time step”, it may be better to show the comparison with no smoothness or other smoothness formulations.

Author response:
**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.

Reviewer comment (Theory):
The prior based on dist-CRP is sensitive to time, what if there is a state with a long duration in the real dynamics? Or what if the dynamics are influenced by time a lot, such as in a Lorenz attractor?

Author response:
**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
\begin{equation}
    \boldsymbol w_{n+1} = U \boldsymbol w_{n} + \kappa\mathbf{1}_{i} + R \boldsymbol x_n + r + g(n,i)
\end{equation}
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.

Reviewer comment (Experiments):
How was the upper bound of K in irSLDS determined? Since the method is called infinite rSLDS, can the K be very large?

Author response:
**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.

Reviewer comment (Experiments):
Most of the comparisons are done between rSLDS and irSLDS, yet the results are not convincing and substantially novel in the experiments conducted. If the authors are adamant about just comparing these two models, then it would be nice to see comparisons at greater scales instead of just testing 4-8 states for NASCAR and the IBL experiments. The results at just these scales do not appear novel.

Author response:
**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!

Reviewer comment (Experiments):
Both experiments show that the irSLDS models uncovered switching at task-relevant states, yet it isn't clear why adding the heat equation prior itself is the best choice in discovering this switching. The authors discuss their reasons for using the heat equations but don't explain further into infinite priors that possibly could have achieved the same if not more.

Author response:
**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.

Reviewer comment (Theory):
The efficiency of the parallel scan is not clear. Other works that utilize scans to efficiently compute the transition probabilities such as S4D (Gu et al., 2022), DSS (Gupta et al., 2022), S5 (Smith et al., 2023) are done on diagonalized transition matrices. The connection to the parallelization of the linear dynamics in Equation 8 is unclear since the transition matrix in this work is not diagonalized, so this work does not reap the same computational benefits that they claim, and their algorithm cannot scale effectively since it incurs cubic cost.

Author response:
**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.

Reviewer comment (Theory):
Smith et al. is cited in reference to the use of a scan operation to compute transition probabilities. However, S5 assumes a diagonal state matrix to efficiently compute the linear recurrence which takes a different form than the time-varying transition matrices in this work. How can this scan method scale efficiently?

Author response:
**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.

Reviewer comment (Experiments):
Why does the heat equation prior directly enable discovery of greater fluctuations/switching in the data? Can you attribute the performance gains to this specific infinite prior?

Author response:
**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.

Reviewer comment (Experiments):
Why not compare to other linear state space models and generalize outside of rSLDS and irSLDS that can allow for input-dependent transitions and recurrent relations without the explicit use of the influence functions as a geometric prior?

Author response:
**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!