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

Authors: Ziyu Lu, Wuwei Zhang, Trung Le, Hao Wang, Uygar Sümbül, Eric Todd SheaBrown, Lu Mi
Primary area: applications to neuroscience & cognitive science
Keywords: neuronal dynamics, dynamical connectivity, interpretability, attention mechanism, transformer

Abstract:
Neuronal dynamics are highly nonlinear and nonstationary. Traditional methods for extracting the underlying network structure from neuronal activity recordings mainly concentrate on modeling static connectivity, without accounting for key nonstationary aspects of biological neural systems, such as ongoing synaptic plasticity and neuronal modulation. To bridge this gap, we introduce the NetFormer model, an interpretable approach applicable to such systems. In NetFormer, the activity of each neuron across a series of historical time steps is defined as a token. These tokens are then linearly mapped through a query and key mechanism to generate a state- (and hence time-) dependent attention matrix that directly encodes nonstationary connectivity structures. We analyze our formulation from the perspective of nonstationary and nonlinear networked dynamical systems, and show both via an analytical expansion and targeted simulations how it can approximate the underlying  ground truth.  Next, we demonstrate NetFormer's ability to model a key feature of biological networks, spike-timing-dependent plasticity, whereby connection strengths continually change in response to local activity patterns. We further demonstrate that NetFormer can capture task-induced connectivity patterns on activity generated by task-trained recurrent neural networks. Thus informed, we apply NetFormer to a multi-modal dataset of real neural recordings, which contains neural activity, cell type, and behavioral state information.  We show that the NetFormer effectively predicts neural dynamics and identifies cell-type specific, state-dependent dynamic connectivity that matches patterns measured in separate ground-truth physiology experiments, demonstrating its ability to help decode complex neural interactions based on population activity observations alone.

Decision: accept (spotlight)
Reviewer scores: 8, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1729977417846

thread_type:
official_review

<thread_object>
creator: reviewer_ovpw

<content>
summary:
This paper introduces a model (NetFormer) for estimating the time-varying connectivity between neurons given access to recordings of their activity. At each timepoint, the approach maps the recent history of neural activity to an estimated connectivity matrix based on a linearized attention mechanism. The authors apply their method on a combination of synthetic tasks and real-world datasets, demonstrating that it can successfully capture nonstationary effects, as well as recapitulate experimental measures of connectivity in realistic settings. They show that their approach outperforms several baselines, both in predicting activity, as well as in estimating connectivity matrices.

strengths:
Drawing inspiration from the attention mechanism used in transformers, the authors propose a method for estimating connectivity weights when the underlying dynamics are nonstationary. The proposed method is simple and interpretable, and is shown to outperform baselines with static connectivity, especially in synthetic tasks where there is good control over ground truth. Further, they demonstrate the utility of their approach for auxiliary tasks, such as cell type inference from learned positional embeddings. The majority of methods for connectivity inference in the neuroscience literature are not well-suited for the nonstationary setting, so this work does indeed address an important question.

weaknesses:
The estimated connectivity at time $k$ is constrained to be of the form $\mathbf{A}_k = \tilde{\mathbf{X}}_k\mathbf{W}_Q\mathbf{W}_K^\top \tilde{\mathbf{X}}_k^\top$. This can be intepreted as a kind of generalized covariance of the recent history of neural activity, where $\mathbf{W}_Q\mathbf{W}_K^\top$ performs some form of temporal mixing/weighting. Given this, I suspect that the proposed method may be particularly susceptible to spuriously estimating connectivity between unconnected neurons that nonetheless have highly correlated activity (as discussed in Das & Fiete, 2020), even in a time-lagged sense. It is unclear to me whether any of the experiments presented test this scenario. One possibility is to test how NetFormer performs on a synthetic network with Mexican hat connectivity and white driving noise, as done in Das & Fiete, 2020. 

My biggest reservation is that the baselines used in the experiments were rather weak. In particular, no other baselines that could account for time-varying connectivity were compared to. While admittedly the vanilla static RNNs used as baselines are standard in computational neuroscience, one could easily imagine a time-varying linear model of the form $x_{k+1} = W_kx_{k}+b$, where $W_k$ is estimated via OLS over a time window $[k-H,...k]$; such models are common in the signal processing literature. One could also imagine a nonlinear version $x_{k+1} = \tanh(W_kx_{k}+b)$.

questions:
1. The claimed performance in estimating experimentally measured cell type-averaged postsynaptic potential (Table 1) is quite remarkable, yet I was left confused on some of the details here. For example, the fitted model appears to predict different cell type-averaged activity for different behavioral "states" (Fig. 5),  yet it appears there is a single cell type-averaged ground truth connectivity derived from patch clamp experiments that all models in Table 1 are being compared to. Is it simply the case that "connectivity" measured under patch-clamp conditions happens to match an average over these "states"?  It would be useful to help further evaluate the claimed performance on this dataset if the authors provided a supplementary figure showing additional details on these fits, as done in other figures (e.g. fitted connectivity vs the ground truth "connectivity" for NetFormer and the various baselines, fitted activity traces, etc.)

2. Minor: the MSE between the average connectivity and the ground truth is used throughout as a performance metric for NetFormer. This is sensible when average connectivity is used as the final readout of estimated connectivity. Does this metric significantly differ from the average MSE between the time-varying connectivity and the ground truth? This is arguably a more accurate measure of how much the model deviates from the ground truth, as deviations at all times should be penalized.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711225650

thread_type:
official_review

<thread_object>
creator: reviewer_ipsn

<content>
summary:
The paper introduces NetFormer, a Transformer-based model designed to capture dynamic and nonstationary connectivity structures in neural systems. Through linearized attention mechanisms that bypass the softmax constraint, NetFormer encodes neuronal interactions that evolve over time. The model's performance is tested on synthetic data and real neural recordings from the mouse visual cortex, showing promising results in predicting neural activity and recovering cell-type connectivity.

strengths:
1.  The conceptualization of linear Transformers for dynamically evolving connectivity is well-conceived, bridging gaps left by traditional static models. It provides a valuable alternative by proposing a method to infer state-dependent connectivity in neural circuits, which addresses the challenge of nonstationarity.
2.  The paper demonstrates the model's robustness by evaluating it on both simulated data and real biological recordings, including calcium imaging. The comparison with an oracle baseline further highlights the method's effectiveness and provides a strong benchmark for its performance across different neural contexts.

weaknesses:
1. Model Interpretability in Practice: Although the model is interpretable in its mathematical formulation, it remains challenging to track how specific interactions are captured in real neural data, which may affect its transparency in biological applications. Visualization of the attention weights and their correspondence to neural trajectories could enhance interpretability.

2. Limited Comparison with Advanced Baselines: The model's comparisons primarily involve classical methods or simpler RNN architectures. Recent models like Neural ODEs and dynamic GNNs, which have been shown to capture temporal and spatial interactions effectively, are not included. This absence weakens the claims of state-of-the-art performance 

3. Positional Encoding: The rationale behind applying positional encoding across the spatial dimension of time series data from multiple neurons is unclear and warrants further clarification. While positional encoding is typically well-suited for distinguishing between different time steps in a temporal sequence, the authors should explain the motivation and justification for using it in the spatial context of calcium imaging data. If the encoding is based on the physical locations of neurons, it should be specified, though this does not appear to be the case in this paper. Additionally, the reasoning for using positional embeddings to categorize neurons as excitatory or inhibitory requires a more detailed explanation.  

4. Dataset Specificity and Transferability: The real neural data focuses on calcium imaging with lower temporal resolution than electrophysiological recordings. This choice may limit the model's applicability in higher-temporal-resolution neural data, which would be better suited for capturing fast synaptic interactions. Additionally, the use of positional embeddings in calcium imaging data needs further justification, as it might not correlate well with physical neuron placement. 

5. Reproducibility: The authors have not provided any code repository or made their scripts publicly available, which raises concerns about the reproducibility of their results. Although pseudocode is included in the supplementary materials, it does not fully address the need for accessible implementation details to enable independent verification of their findings.

questions:
1. Could you incorporate visualizations that link the learned attention weights to specific neural trajectories or connectivity structures, perhaps using case studies or examples from the experimental dataset?

2. Have you considered including comparisons with more recent and relevant models, such as Neural Data Transformers, Neural ODEs, or temporal GNNs? Additionally, would exploring ablations with models specifically designed for connectivity tasks, like GLMs with state-switching mechanisms, help further validate NetFormer’s performance claims?

3. Would it be beneficial to use alternative metrics for connectivity evaluation, such as causality metrics or time-lagged correlations, to provide more biologically relevant insights?

4. Could you clarify the rationale behind applying positional encoding across the spatial dimension of time series recorded from multiple neurons? While positional encoding is typically justified for temporal sequences, how does it apply to the spatial dimension of calcium imaging data? Additionally, what is the reasoning for using these positional embeddings to classify neurons as excitatory or inhibitory?

Overall, if the authors are able to address some of these concerns I would be happy to increase my initial score accordingly.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731009125703

thread_type:
official_review

<thread_object>
creator: reviewer_hfca

<content>
summary:
In order to overcome the shortcomings of static and non-interpretative models, the research presents NetFormer, a novel model for assessing dynamic connectivity in neural networks. In order to infer nonstationary and nonlinear connection, NetFormer employs a linearized attention mechanism. Without softmax restrictions, the model uses a time-dependent attention mechanism and treats brain activity across time as tokens to achieve interpretability, which makes it more biologically plausible. Spike-timing-dependent plasticity (STDP) simulations, synthetic data, and actual brain recordings are used to evaluate the model, which shows that it can accurately predict neural activity and capture dynamic neuronal connections.

strengths:
Originality: A genuine gap in the present neuroscience tools is filled by the creation of an interpretable model for nonstationary neural networks based on linearized attention.

Quality: The robustness of the model is supported by numerous theoretical and experimental validations, including actual brain data.

Clarity: Although more illustrations would help the paper's complex theoretical content, the model's structure is explained in detail.

Significance: The work is very important since it offers a tool that may improve our comprehension of how neurons communicate, which could have ramifications for different applications in neuroscience.

weaknesses:
Complexity: Some parts, especially the ones about the mathematical foundations, can be too difficult for a wider range of neuroscience readers.

Limitations of Interpretability: Although NetFormer can be interpreted in comparison to more conventional models, some results might be strengthened by further biological contextualization.

Evaluation Scope: The model performs well under partial observability, but it might require additional testing in a variety of scenarios found in real-world neuroscience datasets.

questions:
Nonlinear Dynamics: Could the authors elaborate on whether using various non-linear activation functions in the simulation tests significantly alters the model's performance?

Scalability: How does the model adapt to bigger datasets, especially in real-world scenarios when partial observability is a problem?

Interpretability: Could the writers elaborate on how particular attention weights might be connected to recognized patterns of connection or brain phenomena?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731103329385

thread_type:
official_review

<thread_object>
creator: reviewer_9mud

<content>
summary:
This work aims to recover the time-dependent connectivity structures in neural recordings using a transformer based model with a simple attention mechanism. The authors validate their approach on simulated RNNs and later apply to a published neural dataset.

strengths:
- The simulations are mostly convincing (see below) of the approach's validity, providing compelling evidence for its use in biological experiments.

- To my knowledge, apart from [1], this is one of the first actionable use cases of transformer models in systems neuroscience.

weaknesses:
I believe this work may benefit from at least one cycle of revisions, with new experiments added to strengthen the validity claims. Specifically, the weaknesses below are the main reason for my score, though a substantial revision can change my mind.

- Missing baseline: The authors claim that existing work mainly focus on stationary weight matrices and compare with respect to those. [2] is a relevant work that does not seem to be cited, which should also be one of the baselines.

- Missing validity experiments: I believe the validity of the approach should be further strengthened by considering task trained RNNs with short-term synaptic plasticity (claimed in Figs. 2 and 3).  See [3] for an example of how to validate with task-trained, not randomly connected, RNNs. [4] shows how to incorporate STSP into firing rate models and [5] does the same thing for the spiking network models. 

- The experiment in Fig. 4 should use [7], which I believe is a more appropriate baseline.

questions:
I believe I understood most of the work adequately, but I do have one question on Fig. 5: I am not exactly sure how fig. 5 is consistent with the findings from Fu et al. 2014. Could you explain a bit more, or maybe add a comparison panel to the figure highlighting this claim? Also, as a neuroscientist, to me, a large-scale dataset is one that has 10,000 to millions of neurons [6]. The dataset considered seems to have at most hundreds of neurons. 

To address my concerns, you can perform the following revision:

1) Can you add [2] as a baseline to your experiments?

2) Can you please replace Fig. 2 (which can become supplementary) with a new figure that is based on task-trained RNNs? It is now well known that weights drawn randomly from a normal distribution may have some desirable properties, masking the true abilities of models like Netformer to capture the underlying connectivity [3].

3) Can you please replace Fig 3 with a task-trained RNN as well? If you wish, and this would be a spotlight level contribution, you can show the superiority of netformer in finding stsp patterns in both firing rate based [4] and spiking based [5] models.

4) For Fig. 4, I believe [7] would be a relevant baseline. 

I believe, for a borderline score, the suggestions 1, 2 and 4 should be addressed. 

Edit: The authors have performed substantial revisions to address my concerns 1; 2, and 4. Moreover, the addition of the mexican hat experiment makes this paper a solid contribution as a poster in my opinion.

[1] Kozachkov, Leo, Ksenia V. Kastanenka, and Dmitry Krotov. "Building transformers from neurons and astrocytes." Proceedings of the National Academy of Sciences 120.34 (2023): e2219150120.

[2] Pellegrino, Arthur, N. Alex Cayco Gajic, and Angus Chadwick. "Low tensor rank learning of neural dynamics." Advances in Neural Information Processing Systems 36 (2023): 11674-11702. 

[3] Qian, William, et al. “Partial Observation Can Induce Mechanistic Mismatches in Data-Constrained Models of Neural Dynamics.” NeurIPS 2024, OpenReview, 25 Sept. 2024

[4] Masse, N. Y., Yang, G. R., Song, H. F., Wang, X. J., & Freedman, D. J. (2019). Circuit mechanisms for the maintenance and manipulation of information in working memory. Nature neuroscience, 22(7), 1159-1167.

[5] Mongillo, Gianluigi, Omri Barak, and Misha Tsodyks. "Synaptic theory of working memory." Science 319.5869 (2008): 1543-1546.

[6] Manley, Jason, et al. "Simultaneous, cortex-wide dynamics of up to 1 million neurons reveal unbounded scaling of dimensionality with neuron number." Neuron 112.10 (2024): 1694-1709

[7] Sani, Omid G., Bijan Pesaran, and Maryam M. Shanechi. "Dissociative and prioritized modeling of behaviorally relevant neural dynamics using recurrent neural networks." Nature Neuroscience (2024): 1-13.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732254632733

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Thanks for your constructive feedbacks

comment:
We thank all reviewers for their thoughtful suggestions. We are much encouraged to see that our work has been deemed “well-conceived, bridging gaps left by traditional static models” (ipsn) and our novel model “is one of the first actionable use cases of transformer models in systems neuroscience”(9MUd), “does indeed address an important question” (ovPw), and “a genuine gap in the present neuroscience tools is filled by the creation of an interpretable model for nonstationary neural networks based on linearized attention” (HFca). 

Based on all reviewers’ comments, we have carefully revised our manuscript (with revisions marked in purple). We have uploaded our code in Supplementary Material, and we will make it publicly available. We kindly invite all reviewers to review our revised manuscript, and we also provide an overview of the major updates below:

* **Task-driven population activity simulation (section 3.3, Fig 4).** As suggested by reviewer 9MUd, we further tested whether NetFormer can identify task-driven connectivity patterns in neural populations. Specifically, on three representative neuroscience tasks from the NeuroGym toolkit [1], we trained recurrent neural networks (RNNs) to perform the tasks, and applied NetFormer to learn the connectivity of trained RNNs from the activity of their recurrent units. NetFormer achieves higher accuracy in both dynamics prediction and connectivity recovery compared to linear models (Table 4 in Appendix A.4.2, also attached below).

     |   | NetFormer (MSE)  | NetFormer ($R^2$) | Linear (MSE)   | Linear ($R^2$) | NetFormer (Spearman) | NetFormer (Pearson) | Linear (Spearman) | Linear (Pearson) |
     |-------------|------------------|-------------------|----------------|----------------|-----------------------|---------------------|-------------------|------------------|
     | **a**       | **0.000** ± 0.000 | **0.998** ± 0.001 | 0.014          | 0.920          | **0.694** ± 0.163     | **0.760** ± 0.120   | 0.566            | 0.548            
     |
     | **b**       | **0.010** ± 0.000 | **0.972** ± 0.001 | 0.013          | 0.960          | **0.622** ± 0.006     | **0.722** ± 0.004   | 0.518            | 0.553            
     |
     | **c**       | **0.001** ± 0.000 | **0.997** ± 0.001 | 0.034          | 0.897          | **0.811** ± 0.011     | 0.633 ± 0.105       | 0.626            | **0.650**        
     |

* **More advanced and well-recognized baselines designed for modeling nonstationary connectivity (section 4, Fig 6).** As suggested by reviewer 9MUd, ipsn, and ovPw, we further introduced two baselines that can account for trial- and (state-) varying connectivity: low-tensor-rank RNN (LtrRNN, [2]) and autoregressive Hidden Markov Model (AR-HMM, [3, 4]). 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, connectivity inferred from NetFormer's attention is in better agreement with the cell-type level patch-clamp measurements. (Table 6 in Appendix A.8, also attached below).

    | Context               | Condition                  | Metric     | NetFormer                | LtrRNN                   | AR-HMM                   |
    |-----------------------|----------------------------|------------|--------------------------|--------------------------|--------------------------|
    | **Running**           | K x K                     | Pearson    | **0.591** ± 0.204        | 0.007 ± 0.182            | 0.199 ± 0.027            |
    |                       |                            | Spearman   | **0.601** ± 0.209        | -0.036 ± 0.189           | 0.274 ± 0.048            |
    | **Stationary Desync** | K x K                     | Pearson    | **0.662** ± 0.176        | -0.003 ± 0.181           | -0.320 ± 0.012           |
    |                       |                            | Spearman   | **0.723** ± 0.173        | -0.057 ± 0.174           | -0.206 ± 0.017           |
    | **Stationary Sync**   | K x K                     | Pearson    | **0.713** ± 0.148        | 0.000 ± 0.186            | 0.145 ± 0.041            |
    |                       |                            | Spearman   | **0.767** ± 0.151        | 0.069 ± 0.175            | 0.151 ± 0.009            |

* **Visualization of NetFormer’s dynamics prediction on neural data (Fig 13, Appendix A.8).** As suggested by reviewer ipsn and ovPw, we have added visualizations of true and predicted neural activity as well as connectivity in different behavioral states in Fig 13, Appendix A.8.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732255091590

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
**References**

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

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

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

[4] 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732688979704

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Thank you for your constructive feedback

comment:
We appreciate all reviewers’ feedback and comments on our work; our manuscript has benefited substantially from the many valuable suggestions we have received during this rebuttal period. In the latest revision, we have included new experiments in Appendix 9 on applying NetFormer to networks showing spurious correlations. In summary, we found that NetFormer shows some resilience to spurious correlations, while can also struggle when the recurrent connections are too strong or too weak, similar to other methods studied in Das & Fiete (2020). Nevertheless, since this is a well-recognized challenge for models aiming to infer connectivity from fitting activity, we see it as an opportunity for future improvement. Our latest revision is marked in red, and our code for these experiments are also attached in supplementary materials. We kindly invite all reviewers to review our revised manuscript.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors present a novel approach for reconstructing time-varying neural connectivity from data based on a transformer-like self-attention model, and show good results on model and real data. The authors were unanimous that it is a good paper worth publishing, and I agree. As this is a topic which has been widely studied by the ML community for years, it might be a good idea to include some more baselines, such as classic work on state space models (in the sense of HMMs and Kalman Filters with GLM observation models, not the more modern sense of Mamba, S5 etc) e.g. Pfau, Pnevmatikakis and Paninski (NeurIPS 2013), but it is not necessary for the paper to be accepted.