Identifying neural dynamics using interventional state space models

The theory is clear and well written, but is a straightforward combination of existing analytic approaches.

In the first paragraph of our general response we’ve tried to address this concern.

Assumption 3.4 of causal inference is a very strong assumption, and it is not clear how to provide perturbation in systems neuroscience that adhere to it, and it is not discussed how to check this assumption in data.

We’ve included a discussion on this point in our general responses under section 4 and 5 in our general response; we hope it can partially address this comment.

Only low-d latent spaces are considered, but data could be high-d depending on task.

We’ve added new experiments with varying dimensions in Fig 3 and 4. As we discuss in the general response the optimal dimension depends on the amount of data and number of interventions. With few interventions on a few dimensions we cannot hope for recovering all the dimensions of true latents, but we can identify the system for the dimensions that we’ve observed data.

Relatedly, selection of number of dimensions is not discussed.

To address this, we’ve performed new cross-validation results in the revised manuscript. Please see our general response under section 2.

Compute time of algorithm not discussed? Useful in real-time?

The algorithm is linear in time and the number of dimensions. The algorithmic complexity follows the best practices in variational inference for SSMs which is quite tractable. Our codes are implemented in jax which is compatible with GPU and our experiments took mostly a few minutes to run.

Corr. coeff is used, and the values look very low to me.

We have included comparisons against SSM to have a baseline to compare against (section 1 of our general response). In addition, we replaced the error bars to represent standard deviation (std) as opposed to standard error of the mean (sem) to make the comparisons clearer. Regarding the visuals, we will improve the font sizes and other aesthetics for the camera-ready paper if we get accepted.

Lots of plots, the relevance or results of which are poorly explained.

We would appreciate it if you point us to specific instances where we can focus on.

There is no comparison to other methods whatsoever.

Thank you for the suggestion, we agree that we’ve been missing a critical comparison which is the comparison against SSM, which we’ve now included in the revised manuscript. Please also see section 1 of our general response.

No software availability as far as I could tell.

We have a code repository which is implemented according to best practices and standards of academic software development. We will publicly release the code package upon acceptance. Here, we include an anonymized version of the code package for your consideration.

No validation on non-interventional data (as far as i can tell) as another benchmark.

We agree, please see our earlier response to your comment on comparisons and section 1 of our general response.

No numerical assessment of identifiability in the data.

This is a great suggestion, for experiments presented in Fig. 3, 4 we’ve run the model multiple times with random initialization and calculated the variance of the (aligned) stimulation matrices $\boldsymbol{B}$ across seeds. Since our latents and inferred $\boldsymbol{B}$ matrices are permutation invariant, the alignment procedure finds the best permutation that matches the columns of $\boldsymbol{B}$ across different runs. Our results suggest that iSSM’s inferred parameters are robust across different runs while SSM does not have this property. We’ve included these results in Fig. 3G and Fig. 4J.

Lack of comparison: the authors

Thank you for pointing this out, please see section 1 of our general response.

What do those blue dots mean ...

The blue dots correspond to interventions. They represent the state changes when the intervention is applied to the system. We’ve added this to the figure caption for clarity.

What do the inferred flow fields look like in Figure 2?

This is an interesting question. There are two ways for computing the flow fields, (1) Using the generative model where the flow fields are given by the linear system that governs our latent dynamics followed by the emission model (2) Using the recognition model to infer the latents followed by the emission transform. The flow fields computed based on (2) indeed resemble the true flow fields as illustrated in Fig. 5 in the appendix. The discrepancy between (1) and (2) could be due to optimization, the mismatch between generative and inference function families, and the ELBO gap. This is a common issue for all statistical models that use variational inference [1,2]. To investigate this gap, we’ve included trajectories generated by the generative and recognition models in Fig. 5 of the appendix.

How do you define input u?

We denote the interventional inputs by $\boldsymbol{u}$. Normally the magnitude of $\boldsymbol{u}$ is determined by the experimentalist. In our case we’ve set the magnitude to be in the same order as the maximum value of neural responses during interventions. Notice that the magnitude change can be recovered by $\boldsymbol{A}$ and $\boldsymbol{B}$ matrices making the model identifiable up to re-scaling.

How do you choose your parameters?

We’ve added new cross-validation experiments addressing this. Please see section 2 of our general response.

In the first equation on page 4, does

Let’s consider a simple case where $\boldsymbol{B}$ is equal to identity, and the interventions are performed on individual neurons (therefore the latent and observed dimensions are the same $D=N$). When there’s an interventional input in a certain dimension of $\boldsymbol{u}$ say $\boldsymbol{u}^{(i)}$ that makes $\boldsymbol{B}\boldsymbol{u}^{(i)}$ non-zero at that time point which makes the corresponding dimension in $\boldsymbol{A}\boldsymbol{x}_{t-1}$ to zero. Therefore the latents at time $t$ is given by $\boldsymbol{x}^{(i)}_t = \boldsymbol{u}^{(i)}_t$ which means that the interventional input replaces the activity at time $t$ and $\boldsymbol{x}^{(i)}$ is decoupled from its parents in the network. We hope this clarifies our argument, let us know if further clarification is needed.

Some equations are not labeled.

Thank you for pointing this out, we will fix this in the camera-ready paper upon acceptance.

References

[1] Yao, Yuling, et al. "Yes, but did it work?: Evaluating variational inference." International Conference on Machine Learning. PMLR, 2018.

[2] Yao, Yuling, and Justin Domke. "Discriminative calibration: check bayesian computation from simulations and flexible classifier." Advances in Neural Information Processing Systems 36 (2024).

Strengths

Thank you for the encouraging words.

Model assumption 2 is too strong

Please check section 4 and 5 of our general response for related discussion to this point.

This setup is essentially equivalent ...

The control theory literature normally models the interventional inputs as additive. The main novelty of our work is to incorporate interventions in a causal manner. The intuition of why this is helpful is that every time a new intervention is performed, a submatrix of the dynamics matrix senses the data which makes model identification easier compared to classical control theoretical results. We originally had a discussion on connections with control theory in the paper. Since our work connects to a few separate fields we removed that section to make it easier to follow.

Missing technical details in the inference

Thank you for the suggestion, the variational scheme we’ve presented is quite popular and standard. We’ve referenced a paper that discusses the details of this variational scheme for SSMs in the paper. Still if you think it’s helpful to have the details in the paper we’re happy to include a section in the supplementary material in the camera-ready paper if we get accepted.

For inference results, reconstruction correlation

This is a good suggestion, please see section 1 of our general response.

From the derivation in Session 3.3

The addition of input terms mimics a do-intervention. If there is no intervention, i.e. $\boldsymbol{B}\boldsymbol{u_t} = 0$, then each node of the latent dynamical system evolves as usual $\boldsymbol{x_{t+1,j}} = \boldsymbol{A_j} \boldsymbol{x_{t,j}} + \boldsymbol{\epsilon_{t,j}}$. when there is intervention on a node $j$, that is $\boldsymbol{B}\boldsymbol{u_t} \ne 0$, then the j-th node is decoupled from the other nodes and is set to the value $\boldsymbol{B}\boldsymbol{u_t}$, namely $\boldsymbol{x_{t+1,j}} = \boldsymbol{B}\boldsymbol{u_{t,j}} + \boldsymbol{\epsilon_t}$. In the meantime, all other nodes evolve as usual $\boldsymbol{x_{t+1, j'}} = \boldsymbol{A_j} \boldsymbol{x_t} + \boldsymbol{\epsilon_{t,j'}}$.

While we in general do not know which latent node is intervened, this decoupling of the intervened node and the un-intervened nodes under do-interventions offers signatures for identifying each latent node.

When $\boldsymbol{B}\boldsymbol{u_t} = 0$, then all latent nodes evolve following a usual linear dynamical system $\boldsymbol{x_{t+1}} = \boldsymbol{A}\boldsymbol{x_t} + \boldsymbol{\epsilon_t}$. Without interventions, one can find sets of latent variables that can perfectly describe the observed data. In particular, suppose the data is generated by $\boldsymbol{y_t} \sim P(\boldsymbol{y_t} | f_{\theta}(\boldsymbol{x_t}); \boldsymbol{x_{t+1}} = \boldsymbol{A} \boldsymbol{x_t} + \boldsymbol{\epsilon_t}$ with the true latent nodes $\boldsymbol{x_t}$, then it can also be equally well described by another set of latents $\boldsymbol{\hat{x_t}}$ that is linear transformation of the true latent, i.e. $\boldsymbol{\hat{x_t}} = \boldsymbol{M} \boldsymbol{x_t}$ for all t: $\boldsymbol{y_t} \sim P(\boldsymbol{y_t} | \hat{f_{\theta}}(\boldsymbol{\hat{x_t}})); \boldsymbol{\hat{x_{t+1}}} = \boldsymbol{\hat{A}}\boldsymbol{\hat{x}_t} + \boldsymbol{\hat{\epsilon}_t}$, where $\hat{f} = f \circ \boldsymbol{M}^{-1}, \boldsymbol{\hat{A}} = \boldsymbol{A} \cdot \boldsymbol{M}^{-1}$. Hence the latent $\boldsymbol{x_t}$ is not identifiable under only observational data, i.e. when $\boldsymbol{B}\boldsymbol{u_t}=0$.

Assumption 3.2 assumes the observation

While Assumption 3.2 assumes piecewise linear, it can be replaced by other assumptions on the mixing function $f$. The key utility of Assumption 3.2 is to achieve linear identifiability of the latent nodes. It can be achieved via piecewise linear mixing function (as in Assumption 3.2) or polynomial function (as in Assumption 4.2 in [1]) or injective differentiable functions (as in Assumption 1 of [2]).

Assumption 3.4 seems to be too strong

Please see section 4.2 in our general response.

References

[1] Ahuja, Kartik, et al. "Interventional causal representation learning." International conference on machine learning. PMLR, 2023.

[2] Buchholz, Simon, et al. "Learning linear causal representations from interventions under general nonlinear mixing." Advances in Neural Information Processing Systems 36 (2024).

Thank you for the prompt response and apologies if our description wasn’t clear.

First, we make the assumption that each latent causally links to a sparse set of neurons. As we discussed in the general response, this assumption (referred to as sparse coding in neuroscience literature) has neuroscientific roots and therefore is well-motivated. The sparsity assumption translates into the stimulation matrix $\boldsymbol{B}$ being sparse since each column corresponding to a latent has a sparse set of active rows.

Second, when we intervene on individual neurons, only their corresponding latents (which again is a sparse set) are activated. Therefore, interventions on individual neurons translate into do-interventions on subsets of latents. Our theoretical result—identifiability of the latents, dynamics, and emissions using do-interventions—is equivalent to block-identifiability if the interventions activate sparse subsets of neurons. Therefore, if we designate multiple latents per group of functionally coherent neurons, the model identifies those groups of latents. This is still a valuable result since compressing neural dynamics into a smaller set of disentangled latent groups offers interpretability, reproducibility, and the potential of scientific discovery.

Third, given the sparse coding assumption, enforcing the sparsity on matrix $\boldsymbol{B}$ allows for its recovery as well as identification of the latents and model parameters. This is supported by our new results on the consistency of inferred $\boldsymbol{B}$ matrices across random initializations.

That said, as we mentioned in our general response, the model still allows for directly uncovering the causal relationships between neurons by setting the $\boldsymbol{B}$ matrix to identity. In this case having access to interventional datasets where individual neurons are intervened on is equivalent to performing do-interventions on individual latents. We added a new result to the appendix of the paper (Fig. 7) to showcase this application on synthetic models of working memory.

Figures could be improved

Thank you for these suggestions, we’ve fixed these in the revised manuscript and we’ll further improve them in the camera-ready paper upon acceptance.

No way to capture time varying interventional dynamics

Short-term temporal effects: If the interventions have temporal effects consistent with the generative model, then theoretically the model should capture it. When the interventional input is omitted, its effect still persists through the value of latents that are affected by it, and will continue propagating in the network through the dynamics and emission models (or the approximate recognition model). Therefore the statement that “the model does not capture time-varying interventional dynamics” is not entirely correct.

Long-term temporal effects: The reviewer is indeed correct that if the effect of intervention is long-lasting with complex long-term dynamics (e.g. if interventions cause plasticity effects that change the long-term behavior of the system) the model is not designed to capture it. However, we argue that in that case fitting multiple iSSM models to different chunks of data splitted by the expected time course of learning can be one way to bypass this issue.

That said, we agree that capturing complex interventional temporal effects is still an open question and we hope to study this in more depth in future work.

How does the amount of interventional data ...

In the experiment in Fig. 2 the total amount of data (i.e. number of trials and the length of each trial) is kept constant across all runs. We change stim_count on the x-axis while keeping the amount of data fixed.

Observations reconstruction seems to fall

The effect is not statistically significant (notice that the error bars represent standard error of the mean). Our conclusion of Fig. 2 G-P is that both observational and interventional data allow for a good reconstruction of data, but only interventional data allows for capturing the true latents.

In the simulations, adding interventions improves recovery

This is a great suggestion, we’ve done a new experiment to address this. In the experiment we fit iSSM to the interventional data from the two systems, and input the emission models with a grid of latents. If the models faithfully capture the flow fields, the generated trajectories for the given input should be different. This is indeed the case as we present in Fig. 6 in the appendix of the revised manuscript.

Strengths

Thank you very much for the encouraging words, our thinking is very much aligned with these statements.

A key, unaddressed limitation of the work is the assumption that ...

This is a great suggestion, please see section 4 and 5 of our general response which might partially address this comment.

An assumption of the SSM framework is that high dimensional ...

Indeed modeling interventions as affecting latents vs. neurons has been a critical question for us during the development of the model. Please see 3 and 5 in our general response for a more detailed discussion.

The switching mechanism proposed by the authors

Thanks for pointing this out. As we discuss in section 5 of our general response, in principle the model can be used to capture individual neural interventions through considering one latent per neuron and setting $\boldsymbol{B}$ to identity. But that requires datasets where individual neurons are perturbed in near perfect isolation from other neurons. At least we are not aware of such a dataset yet. We would like readers to think of latents as “groups of neurons” or “micro-circuits” that are functionally similar. In a sense a causal latent variable model such as iSSM can discover such grouping in an identifiable way. Our theoretical results suggest that in the cases where $\boldsymbol{B}$ is unknown but sparse, enforcing sparsity leads to its identification. Inspired by this, we add a sparsity penalty to $\boldsymbol{B}$. We explain this further in section 4 of our general response.

There is no way to bake in known ...

Section 4.1 of our general response tries to address this comment.

Can this framework be used to analyze ...

In section 5.2 of our general response we discuss this. We also add a new experiment in the appendix (Fig. 7) to address this comment.

Assumptions of the proof

Indeed our current proof requires that the interventions are required to be separate for different latents; i.e. the interventions happen on each latent dimension in isolation. If each intervention involves do-interventions on multiple latent nodes, the proof can be extended, under assumptions on the diversity of nodes being intervened in each intervention (e.g. assumption 3 in [1]). We are actively working on these extensions and hoping to present it in future work. Please see also the “entanglement of the latents” section in our generic response.

Lack of comparison to non-interventional PfLDS model

Thank you for this suggestion, we’ve added comparisons against PfLDS (referred to as SSM) in the revised paper. Please see section 1 of our general response.

Showing that generative dynamics are consistent with inferred latents

While this is a really interesting suggestion, we think this is an unsolved problem in SSMs that use variational inference schemes. When we condition the model on the observations $\boldsymbol{y_{1:T}}$, $\boldsymbol{u_{1:T}}$, the recognition network provides us with latents at all time points $\boldsymbol{x}_{1:T}$ which are used to generate sequences. When generating sequences from the generative model, we only have access to $\boldsymbol{A},\boldsymbol{B}$ which might not be sufficient to capture all the nonlinearity. We’ve done the experiment you’ve asked for, the generated sequences from the inference network and generative model have similarities, but the gap is still considerable (as is for other SSM models [2,3]). We’ve included this result in Fig. 5 in the appendix.

References

[1] Ahuja, Kartik, Amin Mansouri, and Yixin Wang. "Multi-domain causal representation learning via weak distributional invariances." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

[2] Yao, Yuling, et al. "Yes, but did it work?: Evaluating variational inference." International Conference on Machine Learning. PMLR, 2018.

[3] Yao, Yuling, and Justin Domke. "Discriminative calibration: check bayesian computation from simulations and flexible classifier." Advances in Neural Information Processing Systems 36 (2024).

References

$1$ Jazayeri, Mehrdad, and Arash Afraz. "Navigating the neural space in search of the neural code." Neuron 93.5 (2017): 1003-1014.

$2$ Moran, Gemma E., et al. "Identifiable deep generative models via sparse decoding." arXiv preprint arXiv:2110.10804 (2021).

$3$ Lachapelle, Sébastien, et al. "Synergies between disentanglement and sparsity: Generalization and identifiability in multi-task learning." International Conference on Machine Learning. PMLR, 2023.

$4$ Barth, Alison L., and James FA Poulet. "Experimental evidence for sparse firing in the neocortex." Trends in neurosciences 35.6 (2012): 345-355.

$5$ Qian, William, et al. "Partial observation can induce mechanistic mismatches in data-constrained models of neural dynamics." bioRxiv (2024): 2024-05.

3. Low dimensionality of the dynamics

• Indeed as we discuss in the introduction and schematize in Fig. 1, we think of interventions as forcing the neural state out of its low-dimensional attractor manifold. Given this, it is a valid question of how we are still assuming low dimensionality. Our argument is the following: (1) While the observational data is low-dimensional, we think that interventional data is higher dimensional but still lives in a much lower dimension than the number of neurons. This is simply because certain configurations of neural states are not biologically possible $1$ . (2) In our model, we place a centered Laplace prior on $\boldsymbol{B}$ with the scale hyperparameter $s$ to encourage its sparsity which is critical for identifiability as discussed in the literature $2,3$ . The optimal latent dimension for our experiments depends on the sparsity of the $\boldsymbol{B}$ matrix as well as the amount of data that we have. With larger datasets and more interventions, our conjecture is that increasing the number of latents always provides better test accuracy.

4. Entanglement of the latents

• 4.1 Constant vs. inferred $\boldsymbol{B}$: Our code package allows for both inferring the stimulation matrix $\boldsymbol{B}$ and setting it to a previously known constant matrix. If we have precise information about the relative location of perturbation sites with respect to neurons, we can encode that information into the matrix $\boldsymbol{B}$ by setting $\boldsymbol{B}_{ij} = D(\text{stim}_i, \text{neuron}_j)$. With this approach, every time a neuron or a group of neurons are perturbed, their corresponding latent is activated which in turn silences or activates downstream neurons in isolation of their other parent neurons.

• 4.2 The sparsity of $\boldsymbol{B}$: If we don’t have prior information about $\boldsymbol{B}$, our approach is to add a sparsity penalty by considering a Laplace prior over the elements of $\boldsymbol{B}$. In this case, as we increase the effect of the prior, the latents become more and more disentangled while still fitting the data. We intentionally used a weak prior for the $\boldsymbol{B}$ matrix in our experimental results since we used a small number of latents. With a sparser $\boldsymbol{B}$ in order to fit the data well we need more latents. We added a new experiment where we increased the latent dimensions and used a stronger prior for $\boldsymbol{B}$. The new results are presented in Fig. 3H and Fig. 4H. Our main argument is that if there are multiple iSSM models consistent with interventional data, assuming sparsity of $\boldsymbol{B}$ will result in identifiability. From a neuroscientific perspective, this means that the latents are activated by sparse (and non-overlapping subsets of neurons) which is a sensible assumption for some use cases $4$ .

5. Intervening on latents vs. neurons

• 5.1 Neuroscientific motivation: Please note that many causal experiments in neuroscience (both optical and electrophysiological) still do not perform interventions at the single neuron level. In some cases, the light-gated proteins (e.g. Channelrhodopsin) are expressed in all the neurons of a subtype, and broad-field illumination is used to activate all neurons of that certain subtype. In other experiments, light-gated proteins are expressed in a sparse subset of neurons, and two-photon lasers are used to activate only those neurons. In these cases, due to the spillover effect of lasers, it’s impossible to activate a dense set of neurons with enough spatial precision. Similarly, for electrophysiological techniques such as micro-stimulation, each electrode targets several neurons located in its vicinity as opposed to a single neuron. Due to these (mostly) experimental limitations, we think modeling interventions as affecting latents is a good first step. Therefore, we think given an intermediate latent dimension the model provides a natural causal grouping of neurons into subsets that are behaviorally relevant.

• 5.2 Modeling causal relations between neurons: That said, our model still allows for capturing the effect of intervening on individual neurons. In the case where we’re interested in investigating causal connections between neurons, one approach is to consider one latent per neuron. If the experimental setup allows for stimulating individual neurons, then we can set the $\boldsymbol{B}$ matrix to identity. To showcase this specific use case of the model we added a new experiment identifying the connectivity matrix to test between two models of working memory inspired by $5$ . In this new experiment, we’ve set the matrix $\boldsymbol{B}$ to identity and used it to discover the causal relationships between neurons (Fig. 7).