NetFormer: An interpretable model for recovering dynamical connectivity in neuronal population dynamics

References

[1] 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.

[2] Fox, E., Sudderth, E., Jordan, M., & Willsky, A. (2008). Nonparametric Bayesian learning of switching linear dynamical systems. Advances in neural information processing systems, 21.

[3] Linderman, S., Johnson, M., Miller, A., Adams, R., Blei, D., & Paninski, L. (2017, April). Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial intelligence and statistics (pp. 914-922). PMLR.

Thank you for your thoughtful review! We have revised our manuscript (with revisions marked in purple). We hope the revision and the following responses could address some of your concerns:

• MSE as performance metric: Thanks for the question. We wanted to clarify that we used MSE to evaluate our dynamics prediction, not connectivity inference. During training, we have used the MSE between the true and predicted next-step activity to guide model parameters optimization. In Table 1, the MSE metric is for evaluating next-step activity prediction on the test set, instead of evaluating connectivity inference. Please refer to Table 1 caption for more details.

• Spurious correlation: Thanks for the insightful question. We compared NetFormer with cross-correlation (Table 1, Figure 5) for connectivity strength inference and with Granger causality (Appendix A7.2, Figure 12) for binary connectivity inference. In both cases, NetFormer demonstrated better performance, leading us to conclude that our model can capture effective connection beyond mere spurious correlation. Meanwhile, we would also like to highlight that we compared the connectivity estimated from our method with the cell-type level connectivity measured from the patch-clamp experiment. It involves stimulating one cell and recording from another, thereby establishing causal connectivity in neural circuits. Additionally, the partial observation experiment in Figure 5 also shows the robustness of our approach to hidden confounding variables. We appreciate the reviewer's suggestion to evaluate NetFormer on alternative datasets. We believe this is a valuable direction for future work.

• Baselines accounting for time-varying connectivity: Thanks for the suggestions. We have added two baselines that account for trial-(state-) varying connectivity: LtrRNN ([1] Pellegrino et al., 2023, as suggested by reviewer 9MUd) and AR-HMM ([2], [3]) in section 4, Figure 6, and Appendix A8. Specifically, LtrRNN models trial-varying connectivity using a low rank tensor $\mathbf{W} \in \mathbb{R}^{N \times N \times K}$, where $K$ denotes the number of trials. AR-HMM assumes that there are discrete latent states switching underlying the observed activity, and each state admits unique dynamics through a different connectivity matrix. We observed that compared to weights inferred by LtrRNN, PCA on NetFormer-inferred weights yields a cleaner separation between the running and stationary states. The states inferred by AR-HMM for each timestep shows some similarity with true states but has higher noise. Moreover, compared to LtrRNN and AR-HMM, inferred connectivities from NetFormer’s attention are in better agreement with the cell-type level patch-clamp measurements. We would also like to point out that NetFormer is more efficient than LtrRNN: NetFormer uses shared parameters across time/trials, while LtrRNN trains trial-specific parameters, leading to an increase in parameters as the number of trials grows. Also, discrete state change modeled by AR-HMM is not ideal for capturing continuous connectivity changes, and heavily relies on the user-specified number of states.

• A single cell type-averaged ground truth connectivity: We included additional visualizations of neural activity traces, mouse behavioral states, and cell-type level attention weights for each state within a block of time in Figure 13 of Appendix A8. We observed that the inferred intra-state connectivity is more similar compared to connectivity across different states. We would like to clarify that the patch-clamp cell-type connectivity data and the multi-modal neural activity data were obtained from separate experiments, meaning that there does not exist an exact ground-truth connectivity structure underlying each behavioral state or timestep. Nevertheless, even for state-dependent connectivity inference (Table 6 in Appendix A8), we believe that the patch-clamp experimental results still provide a reasonable reference, with the following reasoning: As the neural data was collected during spontaneous activity, even though there are some periods with high locomotion (as in the running state), there is no learning or perturbations that could alter the wiring of the neural population drastically. Meanwhile, our results are also consistent with the findings of cell-type specific state-dependent connectivity reported in Fu et al. 2014: VIP cells are more activated during locomotion, while SST cells are more inhibited, suggesting “a circuit in which VIP cells increase activity of neighboring excitatory cells by inhibiting their inhibitory input from SST cells” (Fu et al. 2014). In Fig 6d top row, NetFormer-inferred attention shows that the presynaptic VIP column (rightmost column) is darker (more inhibition) in the running state compared to the stationary states and the presynaptic SST column (second column from the right) is lighter (less inhibition).

• Positional encoding: Thanks for the question and we are more than happy to clarify. The positional encodings are learnable parameters of NetFormer to reflect individual neuron’s identity, and are not based on the physical locations of recorded neurons. The original transformer paper [6] has argued the importance of positional embedding in transformer models: The transformer architecture treats input tokens in parallel, without any built-in mechanism to understand their positions in a sequence, and the positional embeddings enable the model to distinguish between tokens at different positions. In our case, recall that each token is the activity of one neuron. So we assign each neuron a unique learnable positional embedding vector, which can help represent the identity of each neuron within the population and help the model better understand their roles. Without positional embeddings, the model would treat each neuron identically. As for using position embeddings for cell-type decoding, it has been shown in [7] that learned positional embeddings can represent neuron identities and be used for cell-type classification.

• Dataset specificity and transferability: We agree with the reviewer that compared to the current dataset, electrophysiological recordings can make a better dataset with higher temporal resolution. Nevertheless, given our focus on recovering cell-type level connectivity, currently we are not aware of any publicly available electrophysiological neural activity recordings with simultaneous cell-type labels and connectivity ground truth as an alternative to the current imaging dataset. We look forward to testing our model on such datasets once they are available.

• Code availability: We have attached our code in the supplementary material. We will also make this code publicly available.

References

[1] 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.

[2] Fox, E., Sudderth, E., Jordan, M., & Willsky, A. (2008). Nonparametric Bayesian learning of switching linear dynamical systems. Advances in neural information processing systems, 21.

[3] Linderman, S., Johnson, M., Miller, A., Adams, R., Blei, D., & Paninski, L. (2017, April). Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial intelligence and statistics (pp. 914-922). PMLR.

[4] Grigsby, J., Wang, Z., Nguyen, N., & Qi, Y. (2021). Long-range transformers for dynamic spatiotemporal forecasting. arXiv preprint arXiv:2109.12218.

[5] Kim, T. D., Luo, T. Z., Pillow, J. W., & Brody, C. D. (2021, July). Inferring latent dynamics underlying neural population activity via neural differential equations. In International Conference on Machine Learning (pp. 5551-5561). PMLR.

[6] Vaswani, A. (2017). Attention is all you need. Advances in Neural Information Processing Systems.

[7] Mi, L., Le, T., He, T., Shlizerman, E., & Sümbül, U. (2023). Learning time-invariant representations for individual neurons from population dynamics. Advances in Neural Information Processing Systems, 36, 46007-46026.

Thank you for your thoughtful review! We have revised our manuscript (with revisions marked in purple). We hope the revision and the following responses could address some of your concerns:

• Visualization of the attention weights and neural trajectories: We have added visualizations of neural activity traces, mouse behavioral states, and cell-type level attention weights for each state within a block of time in Figure 13 of Appendix A8. We observed that the inferred intra-state connectivity is more similar compared to connectivity across different states.

• More recent and relevant models: We appreciate the reviewer’s input in brainstorming relevant baselines. We have added two new baselines capable of capturing nonstationary connectivity: LtrRNN ([1] Pellegrino et al., 2023, as suggested by reviewer 9MUd) and AR-HMM ([2], [3]) in section 4, Figure 6, and Appendix A8. Specifically, LtrRNN models trial-varying connectivity using a low rank tensor $\mathbf{W} \in \mathbb{R}^{N \times N \times K}$, where $K$ denotes the number of trials. AR-HMM assumes that there are discrete latent states switching underlying the observed activity, and each state admits unique dynamics through a different connectivity matrix. We found that NetFormer is able to capture connectivity changes between behavioral states more clearly. Moreover, compared to LtrRNN and AR-HMM, NetFormer-inferred connectivities from neural recordings are in better agreement with the cell-type level patch-clamp measurements (Table 6 in Appendix A.8). We also want to comment that NetFormer is more efficient than LtrRNN: for a given neuronal population, NetFormer uses shared parameters across time/trials, while LtrRNN trains trial-specific parameters, leading to an increase in parameters as the number of trials grows. For AR-HMM, its discrete state switching mechanism makes it less suitable for capturing continuous connectivity changes, and its state discovery heavily relies on the user-specified number of states.

On the other hand, the following methods seem less applicable to our focus:

• Neural Data Transformers (NDT): Using multi-layer, multi-headed transformer models, NDT can produce accurate dynamics forecasts. However, for models with more layers and parameters, their attention weights can be hard to identify or interpret. Moreover, with the softmax activation applied in their attention computation, NDT’s attentions cannot model that the effect of one neuron on another can be either excitatory or inhibitory, whereas NetFormer employs a linearized attention mechanism without softmax to account for this fact. Resolving the ambiguity in interpretability of NDT’s attention is beyond the scope of the current study.

• Temporal GNNs: While some GNN models are able to describe both spatial and temporal relationships, they have two major limitations, as also pointed out in [4]: 1. They often require a pre-defined graph, while in the neuroscience setting, connectivity information in the neuronal population recording is rarely available. 2. The graph they learned is often static, which will not be able to account for the nonstationary connectivity that we are interested in modeling. Therefore, we believe that GNNs models are not suitable for our goal of discovering nonstationary connectivity from fitting dynamics.

• Neural ODEs: Neural ODEs can be a powerful model for fitting neural dynamics as shown in [5]. Nevertheless, their focus is not to recover dynamical neuronal connectivity, and thus we do not think it fits into the current study.

• Alternative metrics for connectivity evaluation, such as causality metrics or time-lagged correlations: Thanks for the suggestions. In fact, one of our existing baselines, cross correlation, is exactly based on time-lagged correlations, and has been applied on both connectivity-constrained simulation and neural data, as shown in Table 1, Figure 5, Table 5, and Figure 11. We also added Granger causality (GC) on connectivity-constrained simulation as an additional baseline. As shown in Figure 12 and Appendix 7.2, NetFormer achieves better AUROC than GC in inferring binary connectivity. Meanwhile, we would like to emphasize that GC is commonly used for measuring binary connectivity (connected or not), but in both our simulations and neural data, the ground-truth connectivity has signs (excitatory or inhibitory) and magnitude (strength), which cannot be captured by Granger causality.

Thank you for your thoughtful review! We have revised our manuscript (with revisions marked in purple). We hope the revision and the following responses could address some of your concerns:

• Model performance across different nonlinear activations: We performed a second set of connectivity-constrained simulation, where a sigmoid activation is used in place of the tanh activation. Then we used NetFormer, along with other baselines, to reconstruct connectivity at both individual-neuron level and cell-type level. The newly added results in Table 5 and Figure 11 in Appendix A7 show that NetFormer’s performance is consistent across both tanh and sigmoid activations. We would like to reiterate that activations with a sigmoidal shape, such as tanh and sigmoid, are of particular interest in the current neuroscience application. As we have also mentioned in Section 2, sigmoidal activations comply with the biological fact that the effect of one neuron on another can be either excitatory or inhibitory, but neither effect can be arbitrarily large.

• Adapting the model to bigger datasets where partial observability is a problem: While we look forward to testing the scalability of our model on larger real-world datasets, the current dataset we used (Bugeon et al. 2022) is the largest public dataset that provides both activity and cell type information of simultaneously recorded neurons. Meanwhile, in the connectivity-constrained simulation (section 4, Fig 5c), we have assessed the robustness of NetFormer against partial observations at both individual-neuron level and cell-type level. We found that individual-neuron level connectivity recovery remains robust even with only half the neurons observed, and more robustness against partial observations could be achieved if we focus on cell-type level connectivity recovery.

• Connecting attention weights to recognized patterns of connection or brain phenomena: We would like to first highlight that the attention weights that we recovered from the mouse cortex recording is in good agreement with the patch-clamp measured cell-type level connectivity strength (Campagnola et al., 2022). Moreover, we also observed changes in cell-type level connectivity estimated from attention weights that can depend on behavioral states, which is consistent with experimental reports in Fu et al. 2014. It is reported that VIP cells are more activated during locomotion, while SST cells are more inhibited, suggesting “a circuit in which VIP cells increase activity of neighboring excitatory cells by inhibiting their inhibitory input from SST cells” (Fu et al. 2014). In Fig 6d top row, NetFormer-inferred attention shows that the presynaptic VIP column (rightmost column) is darker (more inhibition) in the running state compared to the stationary states and the presynaptic SST column (second column from the right) is lighter (less inhibition). Moreover, we also evaluated the effectiveness of our approach on synthetic datasets which incorporate biological features, such as STDP (section 3.2), neuroscience task training (section 3.3), or cell-type level connectivity constraints (section 4).

• Adding Sani et al. 2024 as a baseline: After a careful reading of Sani et al. 2024, we found their focus to be very different from ours, and thus might not be a suitable baseline for comparison. While their work emphasizes capturing behaviorally relevant neural activity patterns, our work focuses on modeling nonstationary connectivity between neurons. We realized that the confusion may arise from our use of behavioral features in Fig 6. To clarify, during training, our model works solely with neural activity, without access to any behavioral information. Specifically, our model is trained to make next-step neural activity predictions as accurately as possible. After training, we compared the time-varying connectivity inferred by our model to the behavioral information, and found that the unsupervised patterns emerging from the time-varying connectivity are actually behaviorally meaningful. In contrast, Sani et al. 2024 uses both neural and behavioral activity as model inputs and prediction targets for training. Their focus is not on recovering connectivity between neurons, but on modeling the neural dynamics that give rise to the observed behavior. As a result, they discovered relationships between neural activity and behavior in a supervised way, and their model is less interpretable in terms of neuron-to-neuron connectivity.

Reference

[1] Molano-Mazon, M., Barbosa, J., Pastor-Ciurana, J., Fradera, M., Zhang, R. Y., Forest, J., ... & Yang, G. R. (2022). NeuroGym: An open resource for developing and sharing neuroscience tasks.

Thank you for your thoughtful review! We have revised our manuscript (with revisions marked in purple). We hope the revision and the following responses could address some of your concerns:

• Relating Fig 5 (now Fig 6) to Fu et al. 2014: Thanks for the question. It is reported in Fu et al. 2014 that VIP cells are more activated during locomotion, while SST cells are more inhibited, suggesting “a circuit in which VIP cells increase activity of neighboring excitatory cells by inhibiting their inhibitory input from SST cells” (Fu et al. 2014). In Fig 6d top row, NetFormer-inferred attention shows that the presynaptic VIP column (rightmost column) is darker (more inhibition) in the running state compared to the stationary states, and the presynaptic SST column (second column from the right) is lighter (less inhibition).

• Dataset being limited compared to more recent ones (Manley et al. 2024): Thank you for introducing this valuable dataset. In this study, one of our main focuses is to investigate cell-type specific connectivity. As such, the current dataset (Bugeon et al. 2022) is central to our analysis, as it includes both simultaneously recorded neural activity and their cell type information. As pointed out in Bugeon et al. 2022, using calcium imaging in transgenic mice (as done in Manley et al. 2024) cannot identify neurons of different cell types simultaneously, which is crucial for validating our model. Given that fact, Bugeon et al. 2022 provides the largest public dataset that contains both activity and cell type information of simultaneously recorded neurons. As described in Appendix A.5.3, this dataset yields approximately 500 neurons per session, and for the subject we considered (SB025), a total of 2482 neurons were recorded across 6 sessions. These sessions have been employed to study NetFormer’s cross-session parameter sharing capacity in Fig 5d. We have also toned down the description of “large-scale” for the dataset size.

• Adding LtrRNN (Pellegrino et al., 2023) as a baseline: Thank you for the suggestion. We have added LtrRNN (section 4 and Fig 6). We found that the inferred connectivity from LtrRNN is not in good agreement with the cell-type level patch-clamp measured ground truth (Fig 6 and Appendix A8). We further observed that compared to weights inferred by LtrRNN, PCA on NetFormer-inferred weights yields a better separation between running and two stationary states. We also want to comment that NetFormer is more efficient than LtrRNN: for a given neuronal population, NetFormer uses shared parameters across time/trials, while LtrRNN trains trial-specific parameters, leading to an increase in parameters as the number of trials grows.

• New figure on task-trained RNNs: Thanks for the inspiring question. We added new experiments on applying NetFormer to infer the connectivity between hidden units of task-trained RNNs (see section 3.3 and Fig 4). We considered three representative tasks from the NeuroGym toolkit [1], which involve integration of evidence over time, and/or maintaining working memory to make decisions after a delay period. Compared to a linear model, we observed that NetFormer achieves higher accuracy in both dynamics prediction and connectivity recovery (see Fig 4 and Table 4 in Appendix 4.2). While we agree that extending our experiments to task-trained RNNs with STSP is a valuable future direction, we believe that the two core concepts there have been verified: The new task-trained RNN results show that NetFormer can capture task-driven connectivity patterns, and our previous result on STDP shows that NetFormer can capture synaptic modification.