Identifying Neural Dynamics Using Interventional State Space Models

We sincerely appreciate you and the other reviewers for your time and thoughtful evaluation of our work. We found the feedback to be highly constructive, as well as both fair and encouraging. Below we address your specific questions:

Nonlinearity limitation: We completely agree with the reviewer that linear dynamics introduce a limitation into the model. However, developing theoretical results for more complex models is quite challenging. To our knowledge this is the first paper that develops such results for a non-LDS noisy SSM. We hope this brings the causal inference and neuroscience communities closer together and inspires the development of more theoretical results for more complex models.

Complexity of implementation: Thanks for pointing this out, we have included a new section in the supplementary material detailing the inference and its computational complexity. To summarize the section we find that our algorithm scales as $\mathcal{O}(TN^2 + TD^2)$ where $T$ is the length of the trajectories, $D$ is the latent dimension, and $N$ is the number of neurons. Therefore our inference is linear in time and quadratic in the number of neurons.

Limited score of interventions: The existing evidence for interventions forcing the neural state outside of the observational trajectories is discussed in both prior theoretical and experimental papers [1, 2, 3]. That said, we believe that collectively as a field our understanding of how interventions affect the neural states is quite limited. The model we present here is among the first models that incorporate interventions into the modeling framework and therefore we believe it’s critical to develop interventional models to better understand how interventions influence neural dynamics.

[1] Jazayeri, M. and Afraz, A. Navigating the neural space in search of the neural code. Neuron, 93(5):1003–1014, 2017.

[2] O’Shea, D. J., Duncker, L., Goo, W., Sun, X., Vyas, S., Trautmann, E. M., Diester, I., Ramakrishnan, C., Deisseroth, K., Sahani, M., et al. Direct neural perturbations reveal a dynamical mechanism for robust computation. bioRxiv, pp. 2022–12, 2022.

[3] Daie, K., Svoboda, K., and Druckmann, S. Targeted photostimulation uncovers circuit motifs supporting short-term memory. Nature neuroscience, 24(2):259–265, 2021.

Expressivity of the emissions: While we agree that linear dynamics is a limiting factor of the model, we would like to point out that adding nonlinear emissions to the linear dynamics changes its expressivity quite drastically. We reference the Koopman theory in the paper, which posits that a (possibly infinite dimensional) linear dynamical system followed by a nonlinear emission is sufficient to express "any" autonomous dynamical system. The catch here is of course "infinite dimensional" but intuitively as the latent dimension grows larger the expressivity of the model also increases.

Increasing latent dimension: You are absolutely correct that if we assume low-dimensionality of the data then increasing the latent dimension is counter-intuitive. We would like to clarify that while we think the optimal latent dimension crucially depends on the task and dataset, we think of interventions as kicking the state of the system outside of its low-dimensional manifold increasing the dimensions spanned by the trajectories. Therefore the conjecture we've included in the discussion section is only valid for large datasets where diverse and numerous interventions are performed. But we have only presented this argument as a conjecture and we agree that experimental validation is needed for confirming it.

Discussion of the implementation: Thank you for the suggestion, we added a new section in the supplementary material on the details of our variational inference algorithm. We have also included our code package in the submission which follows standard code development practices.

LSTMs, hyperparameters: Using LSTMs in the variational posterior has become a standard practice in the field (e.g. [1]). We did experiment with vanilla RNNs but they drastically underperformed LSTMs. We will include these results if the paper gets accepted. Regarding the hyperparameter search, our two hyperparameters are the sparsity of $\boldsymbol{B}$ denoted by $s$ and the latent dimension $D$. The cross-validation and the range of parameters are included in Fig. 3C, 4D. We used a logarithmic axis for s to cover a wide range of possibilities and for $D$ we increased it up to a point where we see a drop of validation performance. Regarding how the hyperparameters change w.r.t. the task at hand, this question is highly problem specific. We think these value highly affected by the details of the experimental setup such as the task, type of recording, animal model, etc.

[1] Krishnan, Rahul G., Uri Shalit, and David Sontag. "Deep Kalman Filters." arXiv preprint arXiv:1511.05121 (2015).

Thanks again for your time; we look forward to your final assessment.

We sincerely appreciate you and the other reviewers for your time and thoughtful evaluation of our work. We found the feedback to be highly constructive, as well as both fair and encouraging. Below we address your specific questions:

Correct vs. incorrect trials: Please notice the latents closer to the end of the trial where the correct and incorrect latents start to diverge. We noticed that that’s the time when the animal is performing its action selection. This is relevant for us because the area where the recordings are performed ALM is thought to involve action selection. Therefore our results suggest that the causal latents discovered by the model distinguish between the action selection mechanism for the correct vs. incorrect trials.

Generalization to unseen interventions: As described in the last paragraph of section 4.3, we divide the data into train and test trials with different interventions performed in each set. Our results show that (Fig. 4 C,D,E) that the test reconstruction accuracy is on par with training reconstruction accuracy, supporting the argument that the model is able to generalize to test interventions.

Compare against other baselines: Thank you for pointing at the weakness of the paper. We want to clarify that our contribution is not developing a new latent variable model, rather it’s adding the interventional components to the existing SSMs. To our best knowledge, this is the first causal dynamical systems model that incorporates latents and observational noise. Therefore we are not aware of other alternatives to use for our comparisons, we are happy to include new comparisons if the reviewer has specific methods in mind.

Thanks again for your time; we look forward to your final assessment.

Further responses to reviewer xGnX

Explanation of dynamic attractor: We borrowed the term "dynamic attractor" from the previous literature on this toy example [1]. In our understanding, a ring attractor in the neuroscience literature implies a specific latent mechanism where the spatial information is encoded in the connectivity pattern of the network. Here the dynamic attractor has non-rotational latents that are nonlinearly transformed to generate rotational dynamics in the observational space. Indeed the dynamical structure generated in the observational space is a ring attractor but we remained consistent with the terms in the literature.

Indeed you are correct that in the non-noisy setting $\boldsymbol{x}=(0,1)$ or equivalently $\boldsymbol{y}=(0,0)$ is an unstable fixed point of the model. So the model will not move away from it. However, except for that pathological case, all other settings of $\boldsymbol{x}$; $\frac{d\boldsymbol{x}}{dt}$ will be nonzero and the model does not require perturbation to move along the circle.

[1] Galgali, A. R., Sahani, M., and Mante, V. Residual dynamics resolves recurrent contributions to neural computation. Nature Neuroscience, 26(2):326–338, 2023.

Positive observations in Fig. 2C: Thank you for your observation. This is a plotting issue, and the y-axis in the plots is not meaningful. We will fix that in the final revision (if accepted). We min-subtract the data in each dimension to stack different axes of the latents on top of each other. What’s shown in the state plots (Fig. 2 A, B) shows the exact numbers corresponding to the trajectories.

Sparsity penalty on $\boldsymbol{B}$: We mention in line 132-134 (right column) that the sparsity penalty is implemented by placing a Laplace prior on the elements of $\boldsymbol{B}$ with a scale parameter $s$. In our real data experiments we do cross-validation over $s$ (the columns of Fig. 3C and Fig. 4D). For synthetic data we don’t use a sparsity prior.

Injective $f$: While we don't explicitly enforce that the function f is an injective function, this function maps from a low-dimensional latent space for x to a high-dimensional observational space. So if this function fits the data well, then this function is often injective without explicitly enforcing it, as long as none of the latent dimensions collapses and are not used --- namely many different observations are mapped to the same latent variable value. We can often empirically check that this issue does not occur.

Reconstruction accuracy: The reconstruction accuracy is defined as the correlation coefficient between true and reconstructed signals (latents or observations). Notice that except for the case of stimulation count is 0, the observational reconstruction accuracy is always close to 1 for both models. However, the latent reconstruction accuracy which directly measures the identifiability of the latents is only close to one for iSSM as the stimulation count increases. This is the main signature that the added interventional component to the model indeed leads to identifiability. The latent dimension for all the results in this figure is set to the ground truth 2.

We sincerely appreciate you and the other reviewers for your time and thoughtful evaluation of our work. We found the feedback to be highly constructive, as well as both fair and encouraging. Below we address your specific questions:

Generalizing to other SSMs: Thank you for bringing this up. Indeed we present our theoretical results for linear dynamics and nonlinear observations. The algorithmic part of the paper is readily applicable to more complex generative models. Since we’re using variational inference which is agnostic to the underlying model, changing the components of the generative model–i.e. dynamics, emissions, and noise model–will not impact the inference. For identifiability results, the key idea of the proof lies in independence testing: The intervened latents at time t+1 are independent of the others at previous time step t, and this independence leads to the identifiability of the latents. The linearity assumption in the latent space does make this argument easy to make, since we only need to consider the covariance. That said, this argument is readily generalizable to other settings where $x_{t+1}$ and $x_{t}$ are nonlinearly related (e.g. polynomial function or ReLU function up to sparsity constraints), as long as the statistical independence of the independent latents is still sufficient to identify the latents. That said, our goal for this paper is laying the foundation of applying interventional models to neural datasets and we leave these developments for future work.

Hallmark of short-term memory: We thank the reviewer for noticing this. While we agree that multiple mechanisms are proposed for working memory, here we’re using "persistent activity" loosely without any implications of the underlying mechanism. Notice that even in Goldman 2009 the proposed mechanism still explains the observed persistent activity. We will clarify in the revised manuscript that the observed persistent activity should not be confused with the underlying mechanism generating it.

Reproducibility: Thank you for pointing this out. While we have not detailed the specifics of our inference scheme, we’re using standard variational inference tools. That said, we added a new section to the supplementary about the details of inference and computational complexity for completeness. Additionally, we realize that we’ve omitted some important reproducibility details, therefore we added a new section to the supplementary on this. Table 1, 2 summarize these results.

Recovering true latents: Thank you for this great observation. The arrows presented are all from the ground truth model, we have not shown the inferred dynamics by the model. We will include this in the revised manuscript. That said, notice that it’s not quite straightforward to visualize the inferred dynamics. The dynamics inferred by the model are encoded in the variational LSTM and the fitted emission. There are two ways for visualizing the inferred dynamics. (1) Using the inferred A: This is not ideal because of the variational gap. (2) Using the emission & variational posterior. This approach requires us to generate a grid in the observation space. However, due to the use of LSTMs the variational posterior is context dependent. We’re open to the reviewer’s suggestions for visualizing the dynamics, but so far we haven’t found an appropriate visualization.

Definition of iSSM: The SSM and iSSM in all experiments share the exact same characteristics–i.e. same latent dimension, emission & inference architecture, step size, etc.–except for the interventional components where SSM models interventions as an additive term. We did this to focus on the utility of interventional models while accounting for all other confounding factors and bringing the two models to the same footing. In the codes, there is an interventional: bool flag which if set to False will model the interventions in an additive way. For the Poisson experiment (Fig. 4) the model is equivalent to Gao et. al. 2016, except that they did not include any inputs to the model. For the Normal experiment (Fig. 2, 3) the model departs from Gao et. al. 2016 in terms of the observational noise model.

Low-d synthetic experiments: The synthetic experiment in Fig. 2 is mainly presented to provide intuition. It’s an example where latents are clearly defined from a behavioral standpoint and the nonlinear emission is required to produce rotational trajectories. We think it’s an intuitive example for the computational neuroscience audience because it’s a toy version of two prevalent hypotheses about rotational dynamics in the motor cortex. The synthetic example in Supp. Fig. 7 uses higher dimensional latents.

We used the remaining space in the responses to reviewer 3PQB and kBw8 to address your remaining concerns. Thanks again for your time; we look forward to your final assessment.

We sincerely appreciate you and the other reviewers for your time and thoughtful evaluation of our work. We found the feedback to be highly constructive, as well as both fair and encouraging. Below we address your specific questions:

Definition of reconstruction accuracy: We noticed that the definition of the reconstruction accuracy is missing in the paper. The reconstruction accuracy is defined as the correlation coefficient between the true and reconstructed latents, which can only be computed for synthetic experiments with ground truth latents. It’s a common metric used in the causal representation learning literature for assessing the identifiability of the latents [1,2,3].

[1] Khemakhem, I., Kingma, D., Monti, R., & Hyvarinen, A. (2020, June). Variational autoencoders and nonlinear ica: A unifying framework. In International conference on artificial intelligence and statistics (pp. 2207-2217). PMLR.

[2] Khemakhem, I., Monti, R., Kingma, D., & Hyvarinen, A. (2020). Ice-beem: Identifiable conditional energy-based deep models based on nonlinear ica. Advances in Neural Information Processing Systems, 33, 12768-12778.

[3] Song, X., Li, Z., Chen, G., Zheng, Y., Fan, Y., Dong, X., & Zhang, K. (2024). Causal temporal representation learning with nonstationary sparse transition. Advances in Neural Information Processing Systems, 37, 77098-77131.

Other variational distributions: This is a great suggestion, we thank the reviewer for pointing this out. LSTMs have been used in previous SSM literature [1]. We tried vanilla RNNs and LSTMs and LSTMs were significantly better than RNNs. The performance enabled by LSTMs is already sufficient for identifiability experiments we show in the paper. Unfortunately due to the limited time we will not be able to provide new results using new variational families but if we benefit from new variational families we will incorporate the new results in the final manuscript upon acceptance.

[1] Krishnan, Rahul G., Uri Shalit, and David Sontag. "Deep Kalman Filters." arXiv preprint arXiv:1511.05121 (2015).

Compare against other SSMs: In all experiments where we compare SSM with iSSM, we use the exact same model, i.e. linear dynamics and nonlinear emissions. The only difference between SSM and iSSM is the interventional formulation. In addition, the focus of our experiments are identifiability as opposed to mere reconstruction of the data. That said, one can develop new iSSMs with different types of dynamics and emissions using the same framework, such as nonlinear dynamics and nonlinear emissions. Notice that the our inference method will still be effective in other scenarios, and only the generative model should account for the new components. This work attempts to lay the foundation of applying interventional models to neural datasets. While our theoretical results currently do not support nonlinear dynamics we’ll consider developing theoretical results for more complex models in the future work.

Thanks again for your time; we look forward to your final assessment.

Further responses to reviewer xGnX

Question 1: Our main intuition is that the perturbations kick the state of the system outside of the "observational manifold" and force the system to explore regions in the state space that are not visited in the observational regime (we depict this in the schematic in Fig. 1). In addition, our theoretical results suggest that if the true generative model of the data follows our modeling assumptions, using interventional models allows for identifying all model components which theoretically leads to out of distribution generalization. However, in reality the true generative model of the data is far more complex and our model components only provide an approximation to their true counterparts.

Question 2: iSSM is equivalent to SSM in the absence of interventional data. If $\boldsymbol{u}_t = 0$ then the term that involves $\boldsymbol{u}_t$ is ignored and we recover the observational SSM model. It’s only in the presence of interventional inputs where the latent nodes dissociate from their parents in the graphical model that enable all the benefits that we exploited (enabling the model to visit new states, changing the graphical model of the data providing more information about the model, etc.).

Question 3: Yes. Using LSTMs as the variational family has become standard in the literature of SSMs. Some benefits of LSTMs include (1) it’s causal, therefore it respects the arrow of time (2) it’s context dependent, therefore it integrates information from far away time points to inform the model predictions (3) it’s data efficient and scales well with dimensionality and time. Due to these benefits and specifically since we observed sufficiently good results we did not experiment with alternatives. The only alternative we tried was vanilla RNNs which drastically underperformed LSTMs.