Meta-Dynamical State Space Models for Integrative Neural Data Analysis

Figure 1: It would be helpful if panel B (or some other panel) could include $\Omega$ and the inference networks proposed by the model. For generation, this seems fine, but that's only capturing half the model.

We thank the reviewer for the suggestion. We are working on an illustrative figure for the inference framework and will include it in the updated manuscript.

Is there a ground truth for Figure 2? It's hard to judge what one should be looking for here.

We apologize if the figure is confusing. Since we train the model jointly on multiple datasets (M=2 or M=20), and consequently, multiple underlying dynamical systems, it is difficult to generate a single ground truth figure. We could show separate figures for different datasets, however, it is challenging to discern variations in velocity in vector field plots.

References

[1]. Krishnan, R. G., Shalit, U., & Sontag, D. (2015). Deep kalman filters. arXiv preprint arXiv:1511.05121.

[2]. Pandarinath, C., O’Shea, D. J., Collins, J., Jozefowicz, R., Stavisky, S. D., Kao, J. C., ... & Sussillo, D. (2018). Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10), 805-815.

[3]. Herrero-Vidal, P., Rinberg, D., & Savin, C. (2021). Across-animal odor decoding by probabilistic manifold alignment. Advances in Neural Information Processing Systems, 34, 20360-20372.

[4]. Azabou, M., Arora, V., Ganesh, V., Mao, X., Nachimuthu, S., Mendelson, M., ... & Dyer, E. (2024). A unified, scalable framework for neural population decoding. Advances in Neural Information Processing Systems, 36.

[5]. Galanti, T., & Wolf, L. (2020). On the modularity of hypernetworks. Advances in Neural Information Processing Systems, 33, 10409-10419.

We thank the reviewer for taking the time to provide thoughtful and constructive feedback on our submission. We are appreciative of the reviewers’ overall positive reception to the work in terms of presentation and experiments, as well as the recognition of the potential impact for the field. We provide clarifications on the points raised by the reviewer below.

The model makes a number of design choices that are unclearly or incompletely motivated. For instance: does the dynamical VAE matter? Should one use the Bowman et al. inference scheme or the SMC-based autoencoder models of [1, 2, 3]? What about the bidirectional RNNs for $\mu_{\alpha}$ and $\sigma_{\alpha}^2$?

We agree that it is important to expand on the relevance of the chosen inference scheme. In the paper, we followed the standard encoder parameterization introduced in [1] for simplicity, as this is a standard baseline across the dynamical VAE literature. We note that our approach is independent of the parameterization of the encoder and we do not expect the relative performance values reported on the paper to depend on this choice. We will update the manuscript to make this clear.

The experimental comparisons are done for alternative approaches to learning the embeddings, but there aren't really comparisons with other models.

We agree with the reviewer that in general, comparing across different model classes is important. However, we consider the same model class in our baseline comparisons so as to not conflate the effect of the specific inference scheme with the proposed generative model, which is the novel contribution of our work. It is likely that an improvement in the inference scheme would lead to better performing single-session models. But the results suggest that the corresponding multi-session embedding conditioned generative model would still outperform it in few-shot regimes.

Lastly, although a minor point, we do consider shared dynamics (Session Stitching in LFADS [2], and similar to [3]) and Single session model baselines as well, which do not include embeddings.

At least one of the metrics tabulated is k-step ahead $R^2$, which could surely be compared with LFADS on single-session data or POYO [4] on multi-session data.

• We could indeed compute this metric by using LFADS on single-session data but we consider the same base model due to the reason mentioned in our previous response. Since LFADS is a special case of a dynamical variational autoencoder, it should be possible to incorporate our generative model into that framework; for instance, by including an additional embedding encoder, and adapting the generative network conditioned on the inferred embedding.
• Since POYO [4] is trained by minimizing a supervised velocity decoding loss, we would not be able to compute the forecasting $R^2$. For the motor cortex experiments, we used the trained generative model to forecast neural observations starting from movement onset time, and used these forecasted values to decode behavior.

I recognize that this paper has goals in addition to pure prediction, but since the authors do consider that metric, it would be good to see how the model proposed here compares to other predictive models.

Again, we agree with the reviewer that this would be a good addition. However, we want to highlight that this is the first approach for learning a dynamical systems model from neural recordings that considers non-trivial variations across recordings, to the best of our knowledge.

Likewise, it would be good to know how to assess (other than qualitatively) what constitutes a "good" learned latent space in the absence of ground truth.

This is indeed an important question for latent variable models of neural data in general. One way of demonstrating that we have learned a good generative model (which includes a good latent space) is being able to forecast activity on novel recordings given a few samples. While good performance on forecasting indirectly implies a “good” learned latent space, this is still very much an open question.

There are some experiments to demonstrate that the parameterized low-rank embeddings outperform linear adapters or using embeddings as inputs, and these are convincing, but the intuition supplied for this is pretty terse. Why should one expect interference in these cases? Why does the rank-1 update protect against this?

We did not expect these results before running the experiments so we found these observations quite interesting as well. Theoretical results in Galanti & Wolf, 2020 [5] suggest that learning conditional functions using hypernetworks affords improved expressivity compared to a fixed input-based parametrization, and is one possible explanation for our observation. However, this requires further investigation and is beyond the scope of the current work. We will include this in the Discussion section of the updated manuscript.

We thank the reviewer for taking the time to provide thoughtful and invaluable feedback on the manuscript. We are incredibly excited about applying our method to various neuroscientific problems and therefore find the reviewers’ comments highly encouraging. We also really appreciate the reviewers’ positive impressions on the writing and overall presentation of our work. We address the points raised by the reviewer and provide further clarifications below.

But how, then, are the different conditions handled in the notation y_{1:T}^{1:M}?

We thank the reviewer for bringing our attention to missing details about how trials are handled.

• We omitted trial labels in the equations in section 3.2 for ease of notation but we are modeling individual trials. We do not include any condition related information in the model (although it can be incorporated as an input), or use the information to get the dynamical embedding. In sec 5.2, the results correspond to performance on left-out trials for each session.
• Since we focus on differences in dynamics across sessions, we learn a single dynamical embedding per session. We indeed aggregate trials (or batches) that belong to the same session when sampling from $q_{\alpha}$ to get this value.

We apologize for the confusion and are updating section 3.2 to include these important details.

Fig 5B is difficult to interpret; some more annotations or a clearer caption would be appreciated.

We will update the figure and caption, as well as provide more details on how it was generated in the Appendix in the updated manuscript.

It would be interesting to see some of the metrics computed as a function of the number of sessions used for training the model (and then, for example, testing few-shot performance on held-out sessions/subjects).

We agree that it would be important to do this analysis. We did include the results of pretraining with 4 sessions, relative to the 44 sessions used in the main text, which supports the relevance of large-scale pre-training on generalization. Specifically, the model pre-trained with 4 sessions generalized quite well to test trials from those sessions (Figure 17B) but it required more trials from novel sessions to achieve the same performance as the “larger-scale” model (Figure 17C vs Figure 6B). We plan to do a more fine-grained comparison of generalization performance as a function of the training data, and data diversity.

L390, "Fig" missing before 17? would be helpful to show the true dynamics in fig 12

We thank the reviewer for pointing out the typo as well as the suggestion. We will incorporate these in the updated version.

We thank the reviewer for taking the time to provide extensive and thoughtful feedback on our submission. We also appreciate the overall positive reception to our work, especially in terms of presentation, and experimental evaluation. We address the reviewers’ concerns below.

For instance the shared dynamics motifs (Driscoll et al, 2024), how would they best fit within this framework? Alternatively, the embedding analysis of (Cotler et al., 2023), how would it fare on SOTA unsupervised models of neural activity compared to your analysis directly on the data?

This is a great point and we agree that it is important to make the connection between our approach and similar ideas that have emerged in RNNs.

• There are subtle differences between our work and the dynamical motifs in Driscoll et al., 2024. In the latter, each motif corresponded to a distinct fixed point structure or topology and the context cue was a pre-specified input that could push the dynamics to regions of state space corresponding to task-relevant dynamics. In our case, we want to capture both topological and geometrical differences, and the “context” or embedding is learned from data. In our case, the embedding can induce changes in dynamics even in the same region of state space, i.e., given the same initial condition, dynamics can evolve differently, conditioned on the embedding. However, the broad idea of dynamical structure re-use is similar in both works because after pretraining, rapid learning is facilitated by inferring the embedding of new data, and re-using the corresponding dynamical system model. It would be interesting to reverse engineer the dynamics model learned via our approach to identify potential sub-networks that might more directly relate to the motifs identified in Driscoll et al., 2024 but it is beyond the scope of this paper and we leave such investigations to future work.
• The embedding analysis in Cotler et al., 2023 is quite similar to our main idea since they also observe that embeddings that are “close” have similar dynamical properties. We believe that the application of both approaches would lead to similar conclusions given sufficient data with high signal-to-noise ratio. However, when there are trial-limited or noisy datasets, learning individual dynamics models for each dataset, as done in Cotler et al., 2023, might be challenging, and leveraging the collective statistical power of all datasets to learn a common dynamical systems model (or meta-model in Cotler et al., 2023) as done in our approach would be preferable.

We hope this offers more insights into the connections between these approaches. We will include a succinct summary of these points in the Discussion section of the updated paper.

Empirical comparisons, implementing the other models mentioned in Section 4 using the way they would typically handle the variability across datasets (such as inputs or indicator variables).

• Could the reviewer specify which additional methods they would like us to incorporate? The Shared Dynamics baseline incorporates the design choices from Pandarinath et al., 2018 and Herrero-Vidal et al., 2021. The transformer-based approaches only indirectly incorporate temporal information.

• We agree that it is important to compare against different model classes, however, we did not want to confound the effect of the specific inference scheme with our proposed generative model, which is the novel contribution of our work. Moreover, we chose a general sequential variational autoencoder inference framework, but our method can be adapted to work with specific cases as well, for instance, LFADS (Pandarinath et al., 2018).

Figures 3 and 4 lack detail on how they were generated, specifically as to which parameters (/embedding values $e^i$) are being used.

We will include these details in the updated manuscript.

Section 3.1. takes unnecessary turns in presenting the hierarchical model in my opinion.

We thank the reviewer for the suggestion. We were trying to tie our model to more classical hierarchical state-space modeling approaches (such as Linderman et al., 2019). We are working on incorporating your comments to make the presentation more concise and streamlined, and will post the updated manuscript soon.

Minor (/typos) L155: I am unsure about the use of the term "inductive bias". \theta and \epsilon^i capture shared and dataset-specific structure by construction -- I see the "bias" here as referring to how dynamical structures are similar to a mean structure with \theta.

We will rephrase and clarify the points mentioned by the reviewer in the updated version.

How does the read-in network $\Omega$ handle different data dimensionalities?

This network is learned per-dataset (eq. 16, line 223) similar to the likelihood function. As suggested by Reviewer 6MTE, we are working on an illustrative figure for the inference network and will include it in the updated manuscript.

We thank the reviewer for their feedback and for acknowledging that the proposed meta-learning approach provides a promising direction for neural data analysis. We address the concerns brought up by the reviewer below.

Benchmarks/Related Works

We apologize if our presentation of these works in the Related Works section was in any way misleading. We believe that this is more of a misunderstanding that we can clarify with better terminology. Broadly, we are interested in building a generative model for neural time series data and there are fundamental differences between our approach and the works mentioned by the reviewer.

Schneider et al. specifically learn the underlying dynamics rather than a priori prescribing them, as other algorithms do (e.g., SIDS, LFADS).

The reviewer is correct in noting that the authors in CEBRA do not assume a data-generating process since the goal of their work is extracting latent representations for efficient downstream task performance, for instance, decoding behavior. We are instead interested in learning a dynamical systems model [2] of diverse recordings from data. While the self-supervised version of CEBRA conditions the encoding model on time during training, it does not learn a dynamical systems model. We will make this distinction more clear in the revised manuscript.

CS-VAE assumes the same shared underlying dynamics and uses switching linear dynamical systems (SLDS) to discover them.

In CS-VAE [3], the authors are interested in learning latent variable models from high-dimensional multi-subject behavioral data. We agree that the authors assume shared latent structure and exploit this for training a common RNN behavioral decoder, which broadly aligns with our motivation. However, our focus is on building a generative model for data by learning the underlying dynamical system, which is entirely different. The authors use SLDS post-hoc to find discrete behavioral motifs within trials by treating the learned behavioral latents as observations. Crucially, this analysis is done separately for each session/animal rather than in a unified space. We will update the Related Works section to better highlight these differences.

Lastly, although this is a minor point, both CEBRA and CS-VAE are applied on recordings from different subjects performing the same task. In this case, assuming shared latent dynamical systems also works. We are interested in the case when there are non-trivial changes in the underlying dynamical system, as noted in Section 2 of our manuscript.

and the authors should update their citation to: https://elifesciences.org/reviewed-preprints/88602.

We thank the reviewer for pointing out the correct citation – we will update it in the revised version.

Furthermore, this paper does not provide any benchmarking against models that they themselves acknowledge have the same goal: to create robust encoders across multiple datasets for downstream tasks (e.g., decoding). Therefore, they should report their results in comparison to Ye et al., 2023; Zhang et al., 2023; Caro et al., 2024; Azabou et al., 2024.

In the paper, we do acknowledge the similarity with foundation models, specifically in lines 271-272; our approach shares the same broad goal of pretraining a single model for rapid adaptation on downstream recordings. However, as noted above, we are not interested in creating robust encoders across multiple datasets for downstream tasks (e.g. decoding). Instead, we want to rapidly learn a dynamical systems model for downstream recordings and consequently, in lines 272-274, we follow up with these methods leverage transformer-based architectures that encode temporal information indirectly via positional embeddings. Since transformers do not have recurrence in their architectures, they do not learn a dynamical systems model [2]. We will update the manuscript to better explain this point.

The effectiveness of the proposed framework depends on the availability of large-scale pretraining data. The sources show that models trained on smaller datasets tend to learn specialized solutions that may not generalize well to new recordings.

While we investigate generalization performance when pretraining with 4 recording sessions, relative to the 44 used in the main experiment, we observe that our method is still able to effectively capture diverse dynamical systems and generalizes well to sessions used for training (Figure 17 A, B), but needs relatively more trials from novel recordings for generalization (Figure 17 C). The improvement in generalization with increase in pretraining data or data diversity is unsurprising and has been observed in other works [4][5][6]. We do not necessarily see this as a weakness of the approach – as data collection techniques are improving, there needs to be concomitant development of methods that can extract meaningful structure from large-scale recordings under varied conditions.

Dear authors, thanks for the response. I agree several misunderstandings can be clarified in the text, but I still think there is a large value in benchmarking other methods and clarifying/quantifying what precisely you mean to say you discover the dynamical systems model -- i.e., my understanding is that you are attempting to uncover pθ(), in pθ(zt|zt−1) = N(zt|fθ(zt−1),Q), but your metric remains a common downstream task, predicting/forecasting. Let's take Figure 9; the GT and the prediction look wildly different, but your R2 metric is still close to 1; shall I believe you truly learned the dynamics, or rather just a good SSM/embedding? That is why I am asking you to compare to other methods that share a common subgoal -- discovering pθ() i.e. the latents.

"rapidly learn a dynamical systems model for downstream recordings and consequently"

• sorry, what do you mean by rapid here? I can't find any information on time to run your method.

We certainly agree with the reviewer that it would be a good addition to apply our model on recordings from other brain areas. However, we would like to emphasize that ours is the first approach that can learn diverse dynamical systems in a unified model directly from data. Through extensive experiments on large-scale motor cortex recordings, we have tried to assess the limitations and applicability of our approach. We are excited to apply our approach to recordings from other regions as well, or studying differences in dynamics across areas in the future.

I just have to disagree here; the motor cortex data is simple and from highly similar tasks across all subjects - at the end of the day, it's motor cortex recordings while animals make simple arm movements.

learning the dynamical system and discovering the latents involve separate parts of the framework

Yes, exactly. Thanks for being clear here, and this is why I (and other reviewers too) are pushing you to benchmark your solution for finding the latents with other methods/architectures; you could still learn dynamics in this second step. Namely, I’m not convinced your method is the only way for latent variable modeling that should be tested in your final estimation of $p_{\theta}( )$.

To note, if your solution worked “on top of” other methods, that would be even more exciting for the field…

References

[1] Pandarinath, C., O’Shea, D. J., Collins, J., Jozefowicz, R., Stavisky, S. D., Kao, J. C., ... & Sussillo, D. (2018). Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10), 805-815.

[2] Schneider, S., Lee, J. H., & Mathis, M. W. (2023). Learnable latent embeddings for joint behavioural and neural analysis. Nature, 617(7960), 360-368.

[3] Naesseth, C., Linderman, S., Ranganath, R., & Blei, D. (2018, March). Variational sequential monte carlo. In International conference on artificial intelligence and statistics (pp. 968-977). PMLR.

[4] Karl, M., Soelch, M., Bayer, J., & Van der Smagt, P. (2016). Deep variational bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint arXiv:1605.06432.

[5] Driscoll, L. N., Shenoy, K., & Sussillo, D. (2024). Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. Nature Neuroscience, 27(7), 1349-1363.

[6] Cotler, J., Tai, K. S., Hernández, F., Elias, B., & Sussillo, D. (2023). Analyzing populations of neural networks via dynamical model embedding. arXiv preprint arXiv:2302.14078.

[7] Pellegrino, A., Cayco Gajic, N. A., & Chadwick, A. (2023). Low tensor rank learning of neural dynamics. Advances in Neural Information Processing Systems, 36, 11674-11702.