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
Venue: ICLR 2025 Spotlight (2025)
Keywords: neuronal dynamics
dynamical connectivity
interpretability
attention mechanism
transformer
Primary area: applications to neuroscience & cognitive science
TL;DR: We proposed NetFormer, an interpretable transformer-inspired model to capture neuronal population dynamics and uncover dynamical inter-neuron connectivity.

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.

Published as a conference paper at ICLR 2025
NETFORMER: AN INTERPRETABLE MODEL FOR RECOVERING DYNAMICAL CONNECTIVITY IN NEURONAL POPULATION DYNAMICS
Ziyu Lu Wuwei Zhang Trung Le",
Hao Wang", Uygar Sümbül', Eric Shea-Brown't Lu Mi
University of Washington
Princeton University 'Rutgers University
Allen Institute for Brain Science "Georgia Institute of Technology
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. NetFormer, the activity of each neuron across 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, demonstrate NetFormer's ability model a key feature of biological networks, 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 tusk-trained recurrent neural networks. Thus informed, we apply NetFormer to a multi-moda dataset of real neural recordings, which contains neural activity, cell type, and behavioral state information. We show that 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.
plasticity, whereby
INTRODUCTION
Inferring the underlying connectivity of network from observations of the activity of its units long-standing challenge. In the brain, this challenge is exacerbated by (i) different nonlinear dynamics present in individual neurons, (ii) the difficulty of experimentally sampling the full neuronal population simultancously, and (ii) dynamic reconfiguration of effective connectivity, mediated by both synaptic plasticity and neuromodulation. This last issue carries significant practical importance in studying behavioral dynamics, learning and memory (Bargmann, 2012; Tyulmankov al., 2021; Marder, 2012; Liu et al., 2021; Aitken & Mihalas, 2023). As such, it poses (harder) generalization of the classical problem where the connectivity should no longer be considered as a static unknown, rather as a dynamical variable that needs to be inferred and tracked over time.
surrogate, but not sufficient, measure of success in unsupervised inference of connectivity is the inferred network's success in fitting the observed dynamics. While traditional linear dynamical
Equal contribution (co-first authors).
Equul contribution (co-senior authors). Correspondence 1mi 78gatech. edu.
1uzi wuzziyueu..ee edu, w214 1 @pr inceton edu, etsbθuw. edu,

---

Published as a conference paper at ICLR 2025
models struggle to capture the essential nonlinear mechanisms of leaky integration and firing (Gerstner & Kistler. 2002) in biological neurons, more sophisticated nonlinear models typically suffer from a lack of interpretability, making it difficult to identify the underlying connectivity (Pandarinath et al., 2018; Le & Shlizerman, 2022; Ye et al., 2023). Moreover, traditional approaches often adopt a static perspective connectivity (Tank et al., 2021; Löwe al., 2022), failing account for the nonstationary interactions, such as those produced by plasticity and modulation at synapses.
Here we propose an interpretable nonlinear and nonstationary dynamical model to represent inter- actions between neurons (Figure 1), based on the fast weight programming nature of the attention mechanism (Schlag et al., 2021). Prior research has suggested that the attention mechanism can reveal information about the underlying structure of a system (Singh & Buckley, 2023; Lu et al., 2023). We further removed the softmax activation function in the attention mechanism, as the constraint of attention weights summing up to one is not biologically meaningful because neither the in-degrees nor the out-degrees of neuronal connectivity (nor their counterparts incorporating synaptic strength) are invariant across neurons (Santuy al., 2020). We first demonstrated with both mathematical analysis and simulation study that even without the softmax activation, the core part of the attention mechanism the dot-product between queries and keys is capable of capturing nonstationary and nonlinear structural information. Next, we applied this novel approach to a wide range of simulated networks including nonstationary and/or complex nonlinear connectivity patterns, and showed that it can recover ground truth connectivity information. We then applied it to a large-scale, publicly available dataset of neuronal activity recordings. Importantly, this dataset includes the genetic cell type of individual neurons, enabling us to compare our predictions for cell-type level connectivity patterns against known ground truth values from independent experiments. Taken together, this shows the potential of our method for recovering interpretable connectivity information, even in the presence of complex nonlinear and nonstationary network dynamics.
Our main contributions are as follows: (i) We formulated transformer the NetFormer, for which the core of the attention mechanism - the dot-product between queries and keys directly encodes nonstationary and nonlinear structure of networks; (ii) On a simulated network with the spike- berke-timing-deee suming-dependett plasticity mechanism, we demonstrated that the inferred time-varying weights from attention aligned with the underlying changes in connectivity; (iii) On activity generated by task-trained recurrent neural networks, we demonstrated that attention can capture task-induced connectivity patterns; (iv) We applied the NetFormer model to population activity recorded from mouse visual cortex, and showed that attention can recover experimentally measured synaptic connectivity, while benchmarking it with standard recurrent models and other common statistical metrics. Additionally, we demonstrated that attention can naturally reflect state- dependent modulations in the inferred cell-type level connectivity, even more effectively compared to two other specialized bascline methods.
network model,
1.1 RELATED WORK
Dynamical models of neuronal activity. Dynamical models have been a powerful tool for high- dimensional neural data analysis (Vyas et al., 2020). Generalized linear models (GLMs), known for both interpretability and desirable convexity properties (Paninski et al., 2007), have been widely used to model ncuronal population activity well inter-ncuronal interactions (Pillow et al., 2008; Das & Fiete, 2020). Nevertheless, unless stacked with an explicit state switching mechanism (Escola et al., 201 in GLMs the weight matrix describing interactions among neurons is typically stationary across time (Li et al., 2024). Recurrent neural networks (RNNs) have been a popular alternative (Barak, 2017: Perich al., 2020); while the connectivity in (trained) RNNs is typically given by static connectivity matrix "W", variants including long short-term memory networks (LSTMs) (Hochreiter & Schmidhuber, 1996) and gated recurrent neural networks (GRUs) (Cho al., 2014), do include nonstationarities at the level of individual neural units. While this can enhance the model's 's expressivity and performance predictive tasks (Salinas et al., 2020; Lai et al., 2018), it also introduces challenges for interpretability (Tank et al., 2021). Recently, transformer models ( (Vaswani et al., 2017) have been observed to outperform RNNs in various time series forecasting tasks (Zhou et al., 2021; Wu et al., 2021), but their deep layered structures and nonlinear attention mechanisms also raise challenges in interpretation with respect to underlying connectivity structures in the original data (Jain& & Wallace, 2019; Abnar & Zuidema, 2020), as discussed more below. A closely related approach to the present work is the switching linear dynamical systems (Fox et al., 2008; Linderman et al., 2017a). These models have nonstationary connectivity matrices which switch among number of discrete values
2

---

Published as a conference paper at ICLR 2025
Attention Map
excitatioo
X/
&2
4x
Ax = Qx-Ke'
inhibition
Spr1
Population Dynamics
N x
N x
[ + M)
N
Kg'
D x
Wir
Figure 1: Overview of NetFormer. NetFormer learns to predict neural dynamics and infer dynamical connectivity through a lincarized attention mechanism. The model takes in activities X₃ of N neurons across H history timesteps, and predicts their activity 24+1 timestep &o 1. Queries Qs and keys K₁, are lincarly mapped from X2 (X, concatenated with positional embedding E), using We and WK. The time-dependent linearized attention matrix A& is computed as Qu Ku which leams the neuron-level dynamical connectivity, Red and blue colors in the attention map indicate excitatory and inhibitory interactions, respectively.
according to a Markov process. Nevertheless, in vivo experiments have revealed that cortical activity is more likely to go through a continuum of states instead of discrete switching (Harris & Thiele, 2011). This motivates us to propose a model capable of capturing continuous changes in connectivity.
Interpretability of the attention mechanism. Attention weights and positional embeddings provide opportunities to understand the inner working of the transformer models. However, the interpretability of these components is still subject of debate. Findings supporting certain level of interpretability. such as correlation to linguistic features, are common in the literature, with specialized metrics developed to quantify their interpretability (Clark et al., 2019; Abnar & Zuidema, 2020). However, caution should be taken when equating attention with explanation (Jain & Wallace, 2019), considering the lack of identifiability (Brunner al., 2019) and the wide variety of underlying architectures and implementations (Wang& & Chen, 2020). In this work, we seek to avoid these confounding aspects by focusing on the linearized attention mechanism (Schlag et al., 2021).
Predicting activity from connectivity. The reverse direction, predicting activity from connectivity, an allied approach for studying the complexities relating functional and structural information. A prominent line of study has focused on the worm C. elegans as its synaptic connectome was the first available among all species (White et al., 1986). Using this, generative models of activity have been proposed (Mi et al., 2021). Nevertheless, decades of electrophysiological analyses have emphasized the strong additional role neuromodulators in shaping activity (Randi et al., 2023; Marder, 2012). As result, the synaptic connectome alone predicts only partial information about recorded population dynamics (Bargmann, 2012; Randi & Leifer, 2020).
AN INTERPRETABLE MODEL FOR RECOVERING DYNAMIC CONNECTIVITY
We consider N-dimensional dynamical system
dt (t) = (((((t))))
(I)
where a=( E : R.N and ((I) is N x N matrix whose entries may vary across time. (t) prescribes how the i-th variable (i) is driven by the j-th variable x00) at time
Let 2 be observations of the system at discrete timesteps 44. For each k, we train the NetFormer model (Figure 1) predict Bi++1 based on X₁ [#k-H+1 step history of the system up to timestep k. To encode neuronal identities, a learnable positional embedding matrix E' E R. X is concatenated to X giving X = =X E] E M)) The queries Qx and keys K₁ are obtained through linear transformations of X₁, and their product gives the linearized attention matrix Ag:
σ&] R-xHH the recent H-
Q: X Wq NXD. K; = W NXD. A = Q.K NXN.
(2)
follows that entry (i.j) of Ar is computed from the history of (9) and (f), and thus describes the relationship between the i-th and j-th variables. To predict # 1, we take mg be the values v%
3

---

Published as a conference paper at ICLR 2025
and employ the residual connection (He et al., 2016), obtaining prediction
お*+1 = 22 + A 2/ + A
(3)
which similar to the update rule we would get Equation were simulated using the classical forward Euler method with step size 8
2k+1 24 + WW((.))))
(4)
Since neuronal connections can be either excitatory or inhibitory, but neither effect can arbitrarily large, choose f to be sigmoidal function with
f(0) 0, f(E) f(0) + (O( + O(5") for & within some interval (-e, around 0!
Equation can thus be written as
(5) Comparing equations 3 and 5, we deduce that the lincarized attention matrix As learned by the NetFormer may capture the true interactions between different variables W, by approximating (0) Wk. especially when 1 and the first order term plays the most significant role. It is not hard to see that this hypothesis also extends to systems in the form of
2k+1 z + (f( 24 6f =f +
+0(x))
d x = (
(6)
which includes the decaying effect that is commonly present in neural dynamics (Gerstner & Kistler, 2002) (Appendix I) We provide empirical evidence for this hypothesis on both forms of systems in the subsequent sections. Our code is https: //ggttgttttt wwg.tuu.....g...g com Cwom:/NeurrAAHHb/NeeeFFoO..
EXPERIMENTS ON SYNTHETIC DATA
3.1 NONLINEAR AND NONSTATIONARY SYSTEMS SIMULATION
We first considered four simplified simulated systems, with variations in the inclusion of nonlinearity and nonstationarity;
dæ
Wa, (b) dæ di
Manh(Wan (c) dæ di
dl W W (d) dæ tanh( W (())x).
(a)
dl
where (ar) Wo + T. Simulation details are in Appendix A.2.1. All trained NetFormer models are able to make accurate one-step-ahead predictions (R² = 1.000, Figure 2a-d left). Vi- sually, the average linearized attention matrix across timesteps, A k > As, provides good characterization of the average ground-truth dynamical association matrix across timesteps,
W = Σk= W(x))) (Figure 2a-d right). As baseline, we consider AoLS from the linear ordi- nary least squares regression Ø = AoLsæE. We used the Spearman's rank correlation coefficient (p) between the off-diagonal entries of A AOLS and W to quantify how faithfully the learned connectivities reflect the ground-truth. A achieved comparable performance as AoLS in systems (a) (b), but significantly outperformed ÁoLS in systems (c) (d), both visually (Figure 2a-d right) and quantitively (Appendix A.2.2). Moreover, in the nonstationary systems (c) (d), the linearized attention matrix is able to track the majority of changes in W(x) across timesteps (Figure 2e-f and Appendix A.2.3). This ability to capture nonstationarity also explains why NetFormer can outperform the linear regression model which only accounts for static connectivity.
3.2 TIIKE-TIMIG-- DEPENNEE PLASTICITY (STDP) SIMULATION
Next, we tested NetFormer in a more neurobiological realistic setting by considering a leaky integrate- and-fire (LIF) neuron (Gerstner Kistler, 2002) with ppike-timing-deeeee
plasticity (STDP). STDP is fundamental and widely studied synaptic modification scheme in neuroscience (Bi & Poo, 1998; Abbott & Nelson, 2000; Gerstner et al., 1996; Song et al., 2000) where the synaptic connection strength between two neurons depends the relative timing of the spikes they fire. In
Appendix A.2.4 for discussion the radius of convergence of this series representation
4

---

Published as a conference paper at ICLR 2025
a. Linear stationary
b. Nonlinear stationary
dx
dx lamb( (Xx) dr
dr
helParmer 1,000)
NarPPrmmr (.004)
10-10000
-0.25
-0.0
5
200
4U5
600
timestops
Emeestes
C. Linear nonstationary
3
d. Nonlinear nonstationary
Wxxxx
dx (an)(WWs))),,
di
(X)
(W6 + xx00)
W(x) (Wr + xwo
Ned
0.5
- 0:999
-0 3.0171
0.650)
-0.2
limestepe
Connectivity across timesteps
f.
Linea
Nethor mer
(2. 12.13
(2, 1)
[2.4]
f1 yt ccw1.00
12, 33
12, 31
12,
12, 5
7.5
5.0
2.5
com3
2.5
-5
2000
2000
2000
2000
2000
timosteps
timesteps
Figure 2: NetFormer provides accurate dynamics predictions, and recovers ground truth connectivity matrices. a-d Left:: Test set predictions of NetFormer. Predicted trajectories one-step-ahead predictions. Predicted (dashed) trajectories overlap with the true (solid) trajectories. a-d Right: True and inferred connectivity from NetFormer's linearized attention matrix or linear regression weight matrix. For systems c, d, connectivity is averaged across test set timesteps for visualization. For NetFormer, the inferred connectivity shown is the one whose Spearman correlation is closest to the average across 10 random seeds. Colorbars: scale of off-diagonal entries. Diagonal entrics are masked in grey. Inferred connectivity matrices were rescaled by the reciprocal of simulation stepsize for visualization. e: True and inferred temporal evolution of four example connections in the nonstationary system c. Timesteps used as test set are shaded in grey. f: True and inferred temporal evolution in system d.
obtained by concatenating all
typical models, when the presynaptic neuron fires before the postsynaptic neuron, their connection strengthens, increasing the postsynaptic response to future spikes. Conversely, presynaptic firing after postsynaptic weakens the connection. The amount of change in synaptic strength depends on the time interval between pre- and postsynaptic spikes. An example relationship is illustrated in Figure 3a.
We simulated postsynaptic LIF neuron receiving excitatory inputs from 100 presynaptic neurons, following Neuromatch Academy (2023) (see Appendix A.3 for simulation details). Since we only need the spike times of presynaptic neurons, instead of simulating their dynamics, we directly modeled their spike trains with independent Poisson processes (Figure 3b). When spike arrives at certain synapse, the membrane potential of the LIF neuron is increased by an amount proportional to the synaptic strength, and this potential is reset once the firing threshold is reached (see Appendix A.3 for precise equations). The strength, or weight, of each synapse is modulated following STDP (Figure 3a) within some boundaries. The weight evolution of ten example synapses across the simulation timespan is shown in Figure 3c.
We trained NetFormer to predict the next-step membrane potential of the LIF neuron based on its present and past potential and the spikes it received. This relationship is captured by the 1 x 101 linearized attention matrix of NetFormer. After training, NetFormer is able to capture the dynamics well (test set MSE=0.06 + 0.008, =0.911 + 0.013, mean + std across 5 random seeds), outperforming the linear regression model (test set MSE=0.138, =2=0.777). Visualization of their predictions (Figure 3d) shows that NetFormer effectively captures the nonlinear reset mechanism, whereas the linear regression model fails.
We further extracted the learned mee-oo-possyynnn matrix NetFormer and compared it against the ground-truth synaptic weights. Unlike the toy systems in Section 3.1 (Figure 2e, f), recovering individual synaptic weights each timestep from the attention scores does not seem viable. This is not surprising. though, given that the weight differences between synapses are much smaller compared to the membrane potential, that synapse only involved in dynamics prediction when it has a spike, and the strong nonlinearity in the membrane spiking dynamics. Nevertheless, found that for most synapses, the long-term trends in
neuron relationship from the linearized attention
5

---

Published as a conference paper at ICLR 2025
a
STOP kernel
0.02
0.825
Synapee kd
0.01
bus
0:22
0.00
1ere Faoss (m))
Time (ms)
Time (ms)
100000
85200
Timn (MM)
83250
02
0.4
06
0.
H
e
true weights
inferred we gght from attention
CC 0.957
0.000
0.800
1.000
8.000
0.000
0.032
20000 40000 60000 80000
20000 40000 60030 80000
Timaalags
Figure 3: NefFormer captures the effect of -pip---iiingedeeeedeentt plasticity (STDP), a. STDP temporal kernel used in the simulation. Change in the synaptic weight (AW) is function of the delay between pre- and post-synaptic neuron spikes (fge - ppos) b. Raster plot of spikes received by the LIF neuron the first 10 synapses during the first 200 timesteps. e. Evolution the first 10 synapses' weights across the simulated timespan. Each color represents a synapse. d. True and predicted membrane potential across 100 timesteps in the test set. Membrane potential was z-scored before all model fitting. Predicted traces were constructed by concatenating one-step-ahead predictions. e. True and inferred synaptic evolution span, after smoothing. Both true and inferred evolution trajectories were smoothed with a sliding window. Each row represents a synapse, and the rows were sorted by the correlation between the smoothed true and inferred trajectories. The correlation coefficients corresponding to each row are shown on the left. For visualization, each row rescaled to [0, through min normalization. f. Weight evolution trajectories of four example synapses, which correspond to the 1st, 31st, 61st, and 91 rows from were smoothed but not min-max normalized. Insets: Raw, unsmoothed trajectories in sliding window.
simulation time
indicated by arrows. The trajectories
the corresponding attention scores across timesteps are consistent with the true trends in synaptic weights driven by STDP (Figure 3e, f). Across the 100 synapses, the median of the correlation coefficients between smoothed synaptic weight evolution trajectories inferred from attention and the true weight trajectories is 0.608 + 0.028 (mean + std across 5 seeds). This demonstrates that the simple structure of NetFormer is capable of capturing nonstationary connectivity in spiking neuronal networks with dynamic synapses.
3.3 TASK-DRIVEN POPULATION ACTIVITY SIMULATION
We further examined whether NetFormer can identify task-driven connectivity patterns in neural populations. As single neuron level connectivity in dask-performing laboratory animals hard to measure, we resorted to task-trained recurrent neural network (RNN) models. The RNN models are trained perform tasks which mirror the ones laboratory animals are trained to perform, and the activity of their recurrent hidden units has been widely adopted in studies of task-driven neural representation and computation (Mante et al., 2013; Sussillo al., 2015; Yang et al., 2019; Duncker et al., 2020). Here we considered three representative tasks from the NeuroGym toolkit (Molano- Mazon et al., 2022): a. Perceptual Decision Making (Britten et al., 1992). Go-Nogo (Zhang et al., 2019), and c. Delay Comparision (Barak et al., 2010). For each task, we trained RNN model with hidden dynamics
(7) Here h& E RN denotes the activity of hidden units, while E R represents the stimulus input at timestep Wres E N specifies how current activity shaped by past activity, and W Wsti E NNS captures the effect of the present stimulus input. bree E N batim E N correspond to baseline activity and background input. We used RNN models with 4, 8, and 12 hidden units for tasks a, b, c, respectively. Training details are provided in Appendix 1. For each task, applied the trained RNN model to perform 1000 trials, and recorded its hidden units activity across timesteps on every trial. Then trained NetFormer to predict the next-step hidden units activity h441 based on the present and past hidden activity (ha- ........ hx) and stimulus inputs ($ 1a1------ 8x+1) on 800 of those trials, and held out the remaining 200 trials for evaluation. Compared to the linear
= tanhhhW>> Stim + bstim + W Wrech brec).

---

Published as a conference paper at ICLR 2025
a PerreptualDeeddiiontaaannn-ng
GoNogo-vO
DelayComparison-v0
10
Linarr
1.3
(ow 0.743)
10 0 -0.566)
1o 0.620)
0 0.5181
10 0.8111
50 D.625)
pred
1.
W
055
0.5
-10
17.5 25
Lim5
Limssops
Figure 4: NetFormer captures task-induced connectivity patterns. a. Top: True connectivity matrix (Wr.c) from RNN trained to perform the Perceptual Decision Making task, followed by inferred connectivity from NetFormer's 's linearized attention matrix or linear regression weight matrix. For NetFormer, connectivity shown is the one whose Spearman correlation p is closest to the average p across 5 random seeds. Colorbars: scale of off-diagen off-oiagonal entries. Diagonal entries are masked in grey. Inset: Scatter plot of off off-diagonal entries in the visualized inferred versus true connectivity matrix. Bottom: True (solid line, open circles) and NetFormer- preeeeeereeeeedddeeeeeeee (dashed line, filled circles) activity for example hidden units held-out trial. Predicted traces were obtained by concatenating one-step-ahead predictions. Hidden units are distinguished by colors. b-c. Same as a, for RNNs trained to perform the Go-Nogo task and Delay Comparision task, respectively.
inferred
regression model, NetFormer attains higher accuracy in both dynamics prediction and connectivity recovery (Figure and Appendix A.4.2).
CONNECTIVITY-CONSTRAAI SIMULATION AND NEURAL DATA
Neurons form synapses based, in part, on factors controlled by their genetic and morphological cell types. The transmission of information through these synapses is influenced by activity history and behavioral states. To test its ability to handle these and allied complexities, we applied the NetFormer to recent multi-modal dataset from (Bugeon et al., 2022), containing both simultaneously recorded neuronal activity and cell type information in the mouse primary visual cortex. After training the NetFormer predict the neural activity, used the time-averaged attention matrix an inferred connectivity strength between neurons. We compared this inferred connectivity against an independent experimental measurement that serves cell-type level ground truth. This is the cell-type averaged postsynaptic potential (PSP) measured directly using paired patch-clamp experiments (Campagnola et al., 2022), in which the postsynaptic voltage responses of individual "downstream" neurons are recorded in response to spikes elicited in specific "paired" neurons. As an additional test, we observed that the inferred dynamical connectivity more similar within the same behavioral state and more distinct between different states. Furthermore, following the measured simulation to produce neural
cell-type level connectivity, we developed connectivity.constaall dynamics with fully specified ground truth connectivity. On this simulated dataset, we assessed the NetFormer's 's ability to infer both individual-neure level and cell-type level connectivities. More details are in Appendix A.5.
4.1 DATASETS
Multi-modal in-vivo neural recording. The publicly available dataset (Bugeon et al., 2022) includes spontaneous population activity recorded from the mouse primary visual cortex (V1) via two-photon calcium imaging. We trained NetFormer models on data from one subject (SB025), which includes recordings of 2481 neurons. The dataset also provides single-cell spatial transcriptomies data, enabling identification of excitatory (EC) and three inhibitory neuron subclasses (Pvalb, Sst, Vip).
Connectivity-constriiie simulation. We generated activity of synthetic neuron population with 200 neurons whose cell-type level connectivity is constrained by the aforementioned patch clamp experiments (Campagnola et al., 2022). Specifically, we simulated the leaky-integratee system
d x() = xtt) + tanh (tt(t) + b)
(8)
dt
7

---

Published as a conference paper at ICLR 2025
using the forward Euler scheme with step size ò 1, obtaining in discrete timesteps
1 Lanh W 2t + b) + 6.
Here e stands for Gaussian observation noise, represents the baseline activity, and W denotes the neuronal connectivity. In our simulation, 76% of neurons are excitatory, with the remaining 24% being inhibitory. The inhibitory neurons are further evenly subdivided into three cell types (Pvalb, Sst, and Vip). We used cell-type specific means and variances of PSPs measured in the patch clamp experiments define Gaussian distributions of connection strengths between each pair of cell types, and sampled the connection strength between individual cells accordingly. We also provide an additional simulation in Appendix A.7, where the tanh replaced with a sigmoid nonlinearity.
4.2 BASELINES AND EVALUATION METRICS
We benchmarked NetFormer against multiple baselines: (i) a linear recurrent model (referred as "linear regression") ) where Fk++1 = W zΙ. + b; two variants of nonlinear recurrent neural networks (referred to as "RNNs"): : (ii) k tanh a + b), which matches the coamectivity-constraine simulation and serves as an "oracle" type model giving a corresponding upper bound on performance, (iii) x exp (W 2y b). of the form of commonly used generalized linear models (GLMs) in neuroscience (Pillow et al.. 2008), with nonlinearity mismatched to that of the simulation itself. We also considered standard statistical metrics, including crose-correl,tion, covariance, mutual information, and transfer entropy (with details in Appendix A.6). We evaluated activity prediction using mean squared error (MSE), coefficient of determination (R²), and the Pearson correlation coefficient. We assessed the correlation between inferred and true connectivity using Pearson and Spearman correlation coefficients, both at the N x N neuron level (N: number of recorded neurons) and the K K cell type-level (K = 4 includes one excitatory and three inhibitory types: Pvalb, Sst, Vip). More details are provided in Appendix 1..2..
4.3 TIME-AVERAGED CONNECTIVITY INFERENCE
NetFormer outperforms baselines in inferring time-averaged connectivity. As shown in Table 1 and Figure 5a, on both simulated and real neuronal activity data, NetFormer outperforms other baseline models in predicting activity and inferring connectivity in the majority of the evaluations. We hypothesize that NetFormer's 's advantage relative to other baseline methods mainly stem from two key aspects. The first that the underlying connectivity in the NetFormer model is nonstationary rather than static, as assumed by the other baseline models and methods. The second is the absence of a need to specify an activation nonlinearity. Notable, the monnectivity-consrramme NetFormer achieves comparable performance to the "oracle" model (RNN with tanh nonlinearity). By contrast, the RNN with an exponential nonlinearity, misspecified with respect to the simulation, performed significantly worse. This contrast illustrates strength of the NetFormer model: NetFormer does not involve specification of an activation function, avoiding the fragility that this can entail. We further tested NetFormer's robustness to spurious correlations following Das & Fiete (2020) (see details in Appendix A.9). NetFormer shows some resilience to spurious correlations, but would also struggle when the recurrent weights are too strong too weak, similar other methods studied in Das & Fiete (2020), including the GLM (Pillow et al., 2008), logistic regression (Lee et al., 2006), and the Ising model (Roudi al., 2009).
simulation,
Excitation cell type can be decoded from learned positional embeddings. Following Mi et al. (2024), we also observe that the learned positional embeddings E in the NetFormer model can be used to decode an aspect of cell class information. Specifically, using learned embeddings of neurons in the training set, we trained a binary classifier via logistic regression classify neurons as excitatory or inhibitory, a coarser grouping which subsumes the senetic/morpological categories considered above. When applied to classify held-out neurons on the test set, the trained classifier attains 1 100% top-1 accuracy in simulation data, 66.67% top-1 accuracy and AUROC score of 0. 700 in real data (see confusion matrices in Figure 5b). This shows that the learned positional embeddings are linearly separable according to this aspect of neurons' cell type identity.
NetFormer demonstrates robustness against partial observation. For most biological neuronal only access the activity of subset of all neurons. To evaluate its
circuits, each recording robustness against such partial observation, we fitted NetFormer to a randomly selected subset of neurons in the simulated dataset. In Figure 5c, show that the performance in recovering the
8

---

Published as a conference paper at ICLR 2025
NetFormer
linear
RNN
RNN
cross
covariance
mutual information
transfer
regression
w/ tanh
w/ exp
correlation
entropy
connectivity
Pearson 0.869=0.002 0.81 0817:00002 0.905 =0.000 0.581=0011
0.823
-0.029
0.339
0.600
simulation
X
Spearman 0.002-00.00 0. 5.50044 5.0-4.0001
0,546 10.000 0.191-0.009
0.519
00.05
0.262
0.339
connectivvty
Pearson 0.879=0.001 0 0.885500.001 =0.00000000
0.87=00008
0,88
-0,438
0.371
0,419
K x
Spearman 0.860
0.852.10.005
0,866
10.002
0.022.0.025
0.732
0.334
0.018
0.353
connectivity
Pearson 0.777=0:047
-0.395=0:020
00.395:000036
00.407:000006
0.017
.0162
.0176
in-vivo recording
K X K
Spearman 0.847=0.063 -0,409=0.051
-0.343=0,1055
+0,191=0.300
-0.080
-0.190
-0.061
0,233
MSF 0.404=0.004
0.443=+00001
0.560=0.001
0.476=0.003
activity prodiction
Pearson 0.740=0.003 0.720::0.001 0.639=00000 0,699=0.002
R
0.548=0.004 0.515=+0,001 0.386#0.001 0,478=0.004
Table 1: NetFormer outperforms classical baselines methods in dynamics prediction and connectivity inference. Results from both connectivity consttriie simulation and neural recording. An asterisk (*) indicates that the RNN with tanh activation serves as the oracle model (upper bound for performance data). Results from mutual information and transfer entropy compared against the absolute values of ground truth connectivity. Simulation data has ground truth for both neuron-level (N x N) and cell type-level (K x K) connectivity. Patch-clamp results serve as the ground truth for real neural data cell-type level connectivity. Connectivity inference is assessed with Spearman's 's and Pearson's 's correlations. Next-step activity prediction is evaluated with MSE, Pearson's 's coefficient, and H the test set.
the simulation
EC/IN Cull Type Classification
BC Paoto Sot Via
Simutartin Data
Neural Data
NetFormar
Unnar Regression
RNN Tann
RNN Cop
Cross Corretation
1:60
0.00
3:50
Simulattin Data
0.00
1.00
2.31
3.60
Partomappe
Partial Oisservation on Simullation
Connectivity Inference on Noural Date
1.0
1.9
Real House Data
66-
mariic Spramen
022
Premy
Monder of Ubserved noursne
Segge Madia
Figure 5: a. Visualization of true and inferred connectivitics at both neuron level and cell-type level in simulation data and neural data. NetFormer is benchmarked with linear regression, RNNs, and standard statistical metrics. A positive lincar transformation has been applied standardize all matrices the same range for better visualization. b. Confusion matrices of excitatory (EC) and inhibitory (IN) cell type classification on both intneccciittycconstaaneed simulation and neural data. e. Experiment on different levels of partial observability in the crmmectiviit-conttaaree simulation. Connectivity inference is evaluated at both neuron level and cell-type level. d. Connectivity inference in real neural data with NetFormer models trained on a single session and multiple sessions from the same subject. Errorbars: Spearman correlations across random seeds. neuron-level connectivity does not decrease significantly even with only half of the neurons observed. Cell-type level connectivity inference is more robust against such partial observations, highlighting the potential of NetFormer to effectively derive cell-type level connectivity from real neural data.
NetFormer can fit neural recordings across sessions. RNN-type models, designed to recover connectivity at the level of individual neurons, cannot easily incorporate data of varying population sizes across experimental sessions. In contrast, NetFormer promotes scalability by allowing parameter sharing (Wq. Wx) across sessions. These parameters are defined in the temporal dimension and thus do not increase with the number of neurons. Figure 5d shows that NetFormer achieves comparable performance for connectivity inference when both fitting a single session and multiple sessions.
4.4 NONSTATIONARY CONNECTIVITY INFERENCE
Building on NetFormer's ability to capture nonstationary connectivity changes over time in Section we extended this evaluation to real neural data. In this dataset (Bugeon et al., 2022), activity at each timestep labeled by one of the three behavioral states: running, stationary desynchronized, and stationary synchronized. Figure 6a shows that neurons of different cell types can exhibit different activity patterns across these states, suggesting that neural activity is informative of behavioral states. Thus informed, we explored whether the time-varying attention could capture connectivity changes among these states (see details and visualizations in Appendix A.8). We compared NetFormer with two baseline methods capable of capturing nonstationary connectivity: low-tensor-rank RNN (LurRNN) (Pellegrino et al., 2023) and an autoregressive Hidden Markov Model (AR-HMM) (Fox al., 2008; Linderman al., 2017b). LuRRNN model the connectivity each "trial" combination of rank-1 I matrices. Since our neural data is measured during spontancous activity without an explicit trial structure, to apply LiRNNN, we constructed "trials" using sliding windows on the data. Unlike
9

---

Published as a conference paper at ICLR 2025
a.
Tarning Battereryyogtchhonee
Pwil
Nethormms
Running
Desynchronized
Syncheentred
Connectietty
C
NO 19BI Intrred strne
Phato
Fione Mey
Tua stiks
Teaysppe
Figure 6: Connectivity inference across behavioral states. Ncural activities across cell types in three behavioral states. b. Projection of NetFormer attention maps and LuRNNN-ifeered connectivity weights onto top two principal components. Each dot represents different timestep, colored by behavioral states. PCA on NetFormer attention maps shows greater similarity between the two stationary states compared to the running state. c. Comparison of RR-HMM inferred states with true behavioral states. d. Comparison of inferred sute-specific connectivities across NetFormer, LirRNN, and AR-HMM. NetFormer-infen connectivities are in better agreement with the patch-clamp result.
NetFormer which shares parameters across time/trials, the number of parameters LiRRNN has to learn scales linearly with the number of trials. As a second baseline, we considered the AR-HMM, which assumes that there discrete latent states switching underlying the observed activity, and each state admits unique dynamics through a different connectivity matrix. Due to its discrete state switching mechanism, it is not suitable for capturing continuous connectivity changes, and its state discovery relies heavily on the user-specified number of states.
While NetForreer-iffereed cell-type level connectivity is good agreement with the patch-clamp experimental ground truth, connectivities inferred by LtrRNN and AR-HMM bear little correlation this ground truth (Figure 6d). A quantitative comparison is provided in Table 6, Appendix A.8. Notably, consistent with prior experimental observations (Fu et al., 2014). attention also witnesses an increase inhibitory activity from presynaptic Vip neurons and a decrease in inhibition from presynaptic Sst neurons during the running state, as seen by the darker Vip column in the running state compared the stationary states and the lighter Sst column (Figure 6d, top row). Figure 6b further demonstrates that state information is implicitly captured by NetFormer's attention maps, showing greater similarity within states and clear distinctions between states. Notably, stationary desynchronized and stationary synchronized states show more similarity to each other than to the running state. Moreover, compared to weights inferred by LtrRNN, PCA on inFFommr--nfferee weights yields cleaner separation between running and two stationary states. When tasked to find three states from the neural data, those inferred by the AR-HMM largely align with the three behavioral states, albeit with higher noise (Figure 6c).
CONCLUSION AND DISCUSSION
Experience, activity, and adaptation change the effective connectivity of biological neuronal networks via mechanisms including synaptic plasticity and neuromodulation, all playing out across various dynamical variable that should be tracked, rather
timescales. This perspective poses connectivity than inferred once. Here, we propose the NetFormer as a light-weight model for dynamical connec- tivity inference. We began with mathematical analysis that relates nonlinear and nonstationary dynamics to its linearized attention mechanism. We further demonstrated, representative simulated and in- vivo neural datasets, the strength of our model through comparison against various baselines to predict nonlinear neural dynamics and to capture the underlying dynamical connectivity.
This said, our method has several limitations: (i) Partial observability neuronal population dynamics has been a major confounding factor for connectivity inference, and our method is no exception. (ii) As our model learns the forward dynamics through hiseere-de linearization the system in a local temporal neighborhood, its ability to capture nonlinear or nonstationary systems is limited compared to fully nonlinear or layered transformer presents step forward to addressing the long-standing challenge of extracting nonstationary neuronal network structure from complex functional data, and brings new insights into the interpretability of the transformer model and its applicability in modeling nonstationary dynamical systems.
models. Despite these limitations, our work
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the NSF grant NCS-FO #2024364 and the NIH grant IRF IRFIDA055669 (ZL and ESB) and the Shanahan Foundation Fellowship (WZ and LM). TL acknowledges the support in part by A3D3 NSF grant OAC-211 11997 and the Department of Electrical and Computer Engineering the University of Washington. HW is partially supported by Amazon Faculty Research Award, Microsoft AI & Society Fellowship, NSF CAREER Award IIS- IIS-2340125, NIH grant ROICA297832,
and NSF grant IIS- 212998..
REFERENCES
Larry Abbott and Sacha B Nelson. Synaptic plasticity: taming the beast. Nature neuroscience, (11)1117-1118,, 2000.
Samira Abnar and Willem Zuidema. Quantifying attention flow in transformers. arXiv preprint arXiv:2005.00928, 2020.
Kyle Aitken and Stefan Mihalas. Neural population dynamics of computing with synaptic modulations. Elife, 12:e83035, 2023.
Omri Barak. Recurrent neural networks versatile tools of neuroscience research. Current opinion in neurobiology. 46:1-6, 2017.
Omri Barak, Misha Tsodyks, and Ranulfo Romo. Neuronal population coding of parametric working memory. Journal of Neuroscience, 30(28):9424 9430, 2010.
Cornelia Bargmann. Beyond the connectome: how neuromodulators shape neural circuits. Bioessays, 34(6):458 465445586666666 2012.
Guo-qiang Bi and Mu-ming Poo. Synaptic modifications cultured hippocampal neurons: depen- dence spike timing, synaptic strength, and postsynaptic cell type. Journal of neuroscience, 18 1444:004441 1998.
Kenneth H Britten, Michael N Shadlen, William T Newsome, and J Anthony Movshon. The analysis of visual motion: comparison of neuronal and psychophysical performance. Journal of Neuroscience, 12(12):4745 1992.
Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wattenhofer. On identifiability in transformers. arXiv preprint arXiv: 190::::0000011111 2019.
Stéphane Bugeon, Joshua Duffield, Mario Dipoppa, Anne Ritoux, Isabelle Prankerd, Dimitris Nicoloutsopoulos, David Orme, Maxwell Shinn, Han Peng, Hamish Forrest, Aiste Viduolyte, Charu Bai Reddy, Yoh Isogai, Matteo Carandini, and Kenneth D. Harris. A transcriptomic axis predicts state modulation of cortical interneurons. Nature, 607, 2022. ISSN 1476-4687.
Luke Campagnola, Stephanie C Seeman, Thomas Chartrand, Lisa Kim, Alex Hoggarth, Clare Gamlin, Shinya Ito, Jessica Trinh, Pasha Davoudian, Cristina Radaclli, et al. Local connectivity and synaptic dynamics in mouse and human neocortex. Science, 375(6583))eab)5881. 2022.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406. 1078, 2014.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does bert look at? 040,0
an analysis of bert's attention. arXiv preprint
Abhranil Das and Ila R Fiete. Systematic errors in connectivity inferred from activity in strongly recurrent networks. Nature Neuroscience, 23( (3110) 2020.
Lea Duncker, Laura Driscoll, Krishna V Shenoy, Maneesh Sahani, and David Sussillo. Organizing recurrent network dynamics by task-computation to enable continual learning. Advances in neural information processing systems, 33:1 33:14387-14397, 2020.

---

Published as a conference paper at ICLR 2025
Sean Escola, Alfredo Fontanini, Don Katz, and Liam Paninski. Hidden markov models for the stimulus-respon response relationships of multistate neural systems. Neural computation, 23(5):1071- 1132, 2011.
William Falcon and The PyTorch Lightning team. PyTorch Lightning, March 2019.
Emily Fox, Erik Sudderth, Michael Jordan, and Alan Willsky. Nonparametric bayesian learning of switching linear dynamical systems. Advances neural information processing systems, 21, 2008.
Yu Fu, Jason M. Tucciarone, J. Sebastian Espinosa, Nengyin Sheng, Daniel P. Darey, Roger A. Nicoll, Z. Josh Huang, and Michael P. Stryker. A cortical circuit for gain control by behavioral state. Cell, 156(6):1139- 1199--1522 1122, Mar 2014. ISSN 0092-8674, doi:
Wulfram Gerstner and Werner M Kistler. Spiking neuron models: Single neurons, populations, plasticity. Cambridge university press, 2002.
Wulfram Gerstner, Richard Kempter, J Leo Van Hemmen, and Hermann Wagner. A neuronal learning rule for sub-millisecond temporal coding. Nature, 383(6595)776-78, 1996.
Kenneth D Harris and Alexander Thiele. Cortical state and attention. Nature reviews neuroscience, 12(9))509- 000999552222200 2011.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Lstm can solve hard long time lag problems. Advances in neural information processing systems, 9, 1996.
Sarthak Jain and Byron C Wallace. Attention is not explanation. arXiv preprint 1902: 2019.
Guokun Lai, Wei-Cheng Chang, Yiming Yang, and Hanxiao Liu. Modeling long-and short-term temporal patterns with deep neural networks. In The st international ACM SIGIR conference on research & development in information retrieval, pp. 95-104, 2018
Trung Le and Eli Shlizerman. Stndt: Modeling neural population activity with spatiotemporal trans- formers. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 17926-17939. Curran Associates, Inc., 2022.
Su-In Lee, Honglak Lee, Pieter Abbeel, and Andrew Y Ng. Efficient |" regularized logistic regression. In Aaai, volume pp. 401-408, 2006.
Chengrui Li, Soon Ho Kim, Chris Rodgers, Hannah Choi, and Anqi Wu. One-hot generalized linear model for switching brain state discovery. In The Twelfth International Conference on Learning Representations, 2024.
Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial intelligence and statistics, pp. 914-922. PMLR, 2017a.
Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian Learning and Inference in Recurrent Switching Linear Dynamical Systems. In Aarti Singh and Jerry Zhu (eds.), Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learning Research, pp. 914-922. PMLR, 20-22 Apr 2017b.
Yuhan Helena Liu, Stephen Smith, Stefan Mihalas, Eric Shea-Brown, and Uygar Sümbül. Cell- type-specific neuromodulation guides synaptic credit assignment in a spiking neural network. Proceedings of the National Academy of Sciences, 118(51):e211 182111111111111111 18, 2021.
Sindy Löwe, David Madras, Richard Zemel, and Max Welling. Amortized causal discovery: Learning to infer causal graphs from time-series data. In Conference on Causal Learning and Reasoning, 2022. pp. 509-525. PMLR, 2022.
12

---

Published as a conference paper at ICLR 2025
Ziyu Lu, Anika Tabassum, Shruti Kulkarni, Lu Mi, Nathan Kutz, Eric Shea-Brown, and Seung- Hwan Lim. Attention for causal relationship discovery from biological neural dynamics. arXiv preprint arXiv:2311.06922 2023.
Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome, Context-dependent computation by recurrent dynamics in prefrontal cortex. nature, 503(7474):78 84, 2013.
Eve Marder. Neuromodulation of neuronal circuits: Back to the future. Neuron, 76(1):1-11, 2012.
ISSN 0896-6273, doi: https://doi.or org/ 1111110.11100 nww.:/i..cuoon21112220
Lu Mi, Richard Xu, Sridhama Prakhya, Albert Lin, Nir Shavit, Aravinthan Samuel, and Srini- vas C Turaga. counectome-constraaned latent variable model of whole-brain neural activity, In International Conference on Learning Representations, 2021.
Lu Mi, Trung Le, Tianxing He, Eli Shlizerman, and Uygar Sümbül. Learning time-invariant representations for individual neurons from population dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Manuel Molano Mazzon Joao Barbosa, Jordi Cittor-Caurra,, Marta Fradera, Ru- Yuan Zhang, Jeremy Forest, Jorge del Pozo Lerida, Li Ji-An, Christopher J Cueva, Jaime de Rocha, et al. Neurogym: An open resource for developing and sharing neuroscience tasks. 2022.
Neuromatch Academy. Biological neuron models tutorial, 2023. Accessed: 2024-09-28.
Chethan Pandarinath, Daniel J O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et al. Inferring triale-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10):805 -15, 2018.
Liam Paninski, Jonathan Pillow, and Jeremy Lewi. Statistical models for neural encoding, decoding, and optimal stimulus design. Progress in brain research, 165:493-507, 2007.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Arthur Pellegrino, N Alex Cayco Gajic, and Angus Chadwick. Low tensor rank learning of neural dynamics. Advances in Neural Information Processing Systems, 36:11674-117022 2023.
Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton P Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter H Rudebeck, Christopher D Harvey, al. Inferring brain- wide interactions using cata-constr recurrent neural network models. BioRxiv, pp. 2020-12, 2020.
Jonathan W Pillow, Jonathon Shlens, Liam Paninski, Alexander Sher, Alan M Litke, EJ Chichilnisky, and Eero P Simoncelli. Spatio-temppral correlations and visual signalling in a complete neuronal population. Nature, 454(7207):995 9999 2008
Francesco Randi and Andrew M Leifer. Measuring and modeling whole- neural dynamics in caenorhabditis elegans. Current opinion neurobiology, 65:1677175, 2020.
Francesco Randi, Anuj K Sharma, Sophie Dvali, and Andrew M Leifer. Neural signal propagation atlas of caenorhabditis elegans. Nature, 623(7986):4064444 414, 2023.
Yasser Roudi, Joanna Tyrcha, and John Hertz. Ising model for neural data: model quality and approximate methods for extracting functional connectivity. Physical Review E EStaatistiaal, Nonlinear, and Soft Matter Physics, 79(5):051915. 2009.
David Salinas, Valentin Flunkert, Jan Gasthaus. and Tim Januschowski. Deepar: Probabilistic forecasting with autoregressive recurrent networks. International journal of forecasting, 36(3): 1181-1191, 2020.
13

---

Published as a conference paper at ICLR 2025
Andrea Santuy, Laura fomás-Roca José-Rodrigo Rodríguez, Juncal González-Sorian Fei Zhu, Zhen Qiu, Seth GN Grant, Javier DeFelipe, and Angel Merchan Merehan-Perez. Estimation of the number of synapses in the hippocampus and brain-wide by volume electron microscopy and genetic labeling. Scientific Reports, 10(1); 14014, 2020.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning. pp. 9355-9366. PMLR, 2021.
Ryan Singh and Christopher Buckley. Attention as implicit structural inference. Advances in Neural Information Processing Systems, 36:24929-24946, 2023.
Sen Song. Kenneth D Miller, and Larry F Abbott. Competitive hebbian learning through spike-timing- dependent synaptic plasticity. Nature neuroscience, 3(9):919 926699922222099 2000.
David Sussillo, Mark M Churchland, Matthew T Kaufman, and Krishna V Shenoy, A neural network that finds a naturalistic solution for the production of muscle activity, Nature neuroscience, 18(7): 1025-1033, 2015.
Alex Tank, Ian Covert, Nicholas Foti, Ali Shojaie, and Emily B Fox. Neural granger causality. IEEE Transactions on Puttern Analysis and Machine Intelligence, 44(8):4267-4277 2021.
Danil Tyulmankov, Ching Fang, Annapuma Vadaparty, and Guangyu Robert Yang. Biological learning in key-value memory networks. Advances in Neural Information Processing Systems, 34: 22247-22258, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Annual review of neuroscience, 43:249-275, 2020.
Yu-An Wang and Yun-Nung Chen. What do position embeddings learn? an empirical study of pre-trained language model positional encoding. arXiv preprint arXiv:2010.04903, 2020.
John G White. Eileen Southgate, Nichol Thomson, Sydney Brenner, al. The structure of the nervous system of the nematode caenorhabditis elegans. Philos Trans R Soc Lond B Biol Sci, 314 (1165):1 1340, 1986.
Haixu Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. Autoformer: Decomposition transformers with cuto-coreelation for long-ter terettermme series forecasting. Advances in neural information processing systems, 14:22419-22430, 2021.
Guangyu Robert Yang, Madhura R Joglekar, Francis Song, William Newsome, and Xiao-Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature neuroscience, 22(2):297-306, 2019.
Joel Ye, Jennifer L Collinger, Leila Wehbe, and Robert Gaunt. Neural data transformer 2: Multi- context pretraining for neural spiking activity. In Thirth-seventh Conference on Neurul Information Processing Systems, 2023.
Xiaoxing Zhang, Wenjun Yan, Wenliang Wang, Hongmei Fan, Ruiqing Hou, Yulei Chen, Zhaoqin Chen, Chaofan Ge, Shumin Duan, Albert Compte, et al. Active information maintenance in working memory by sensory cortex. Elife, 8:e43191,
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang Informer: Beyond efficient transformer for long sequence time-series forecasting. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 11106- 1110-11115,, 2021.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 JUSTIFICATION FOR LINBARIZED ATTENTION APPLIED TO "LEAKY" SYSTEMS (SEC 2,
EQN 6)
Using the forward Euler method and step size 6, Equation 6 can be simulated as
Lk 3t + 6 Th + So (Waar))) 23. dæv + 6f(We...)
(10)
Following the same sigmoidal assumption on f. Equation 10 can be written as
E++1 3t () I)ar + 50(z)),
(11)
where is the N x N identity matrix. Therefore, the linearized attention matrix As learned from Equation may reflect the true interactions WE by approximating (( - I), and can capture the interactions between different variables (off-diagonal entries of Wx) up to scaling factor (0).
A.2 ADDITIONAL DETAILS FOR NONLINEAR AND NONSTATIONARY SYSTEMS SIMULATION
(SEC 3.1)
A.2.1 SIMULATION DETAILS
In Figure 2 a, b. ground-truth W were generated randomly, with real-part of each eigenvalue clipped ensure stability of the system. W in a, were also used as Wo in c,d, respectively. in d were picked randomly while maintaining stability of the system. In the system trajectory was simulated using the closed-form solution x(t) eWi,, where n7 is the initial state. In b-d. trajectories were simulated using the forward Euler method: 0x++1 = 2 + 6f(W, where W: = W for b, and W₂- Wn + wyww for c, All simulations consist of 3000 timesteps with stepsize 5 0.01, with the first 80% used as training set, and last 20% as test set. For c, d. ground-truth connectivity matrices were computed as the time- averged connectivity across test set 3000
timesteps W = -000 k=200 WR. Simulated trajectories are visualized in Figure 7. In all settings, the the trainining set for 1 100 epochs using the Adam optimizer in Pytorch, with H = 1, M = 5, D = 5, batch size = 80, initial learning rate 0.01. In b, d, learning rate was decayed by factor of 0.9 every 100 epochs. In learning rate was decayed by factor of 0.8 every 100 epochs.
NetFormer model was trained to minimize the mean squared error (MSE)
dx Wx
0
tanh(WX)
0
X0
-2
x12
W(x)x, W(x) (Wo x
x
x(1
0.0
x(5)
-2.5
dx 28
tanh(Wixix) w(x = (Wo + x
11
0
500
1000
1500
2000
2500
3000
timesteps
Figure 7: Simulated trajectories of toy models in Figure 2. Shaded regions represent timesteps used as test set.
15

---

Published as a conference paper at ICLR 2025
A.2.2 QUANTITIVE COMPARSION WITH LINEAR REGRESSOIN MODEL
For each toy system in section 3.1, we trained 10 NetFormer models with different random seeds (initializationss, and computed the Spearman's rank correlation coefficient (p) and the Pearson correlation coefficient (r) between the off-diagonal entries of Á and W for each trained model. In terms of p. A achieved comparable performance as AOLS in systems a, b (p > 0.3, two-sided one sample test), but significantly outperformed AOLs in systems c, (p < 10 10-8) (Figure 8 left). Similar observations can be made with (Figure 8 right). For each system, the linearized attention matrix visualized in Figure 2 the one whose p is the closet to the average p across 10 random seeds.
10
X Asd
Spearman
Pearson
0.9
0.8
0.8
nonlinear sts.
linear nonlinear
linear nonlinear nonlinear nonsta. nonsta.
sta.
nonsta.
Figure 8: Comparison between AOLS (red cross) and Á from NetFormer models with 10 different random initializations (boxplots).
A.2.3 NONSTATIONARITY CONNECTIVITY TRACKING
On toy systems with nonstationary connectivity (Figure 2c, d), evaluated how well linearized attention matrices across timesteps can track changes in the connectivity. For each pair (i,j), i,j j. we collected Ai and Wij across all test timesteps, resulting in two time-varying
1...., 5, series (t) and Wiy (l), and computed the Pearson correlation coefficient between them. Results for 10 trained NetFormer models with different random seeds are shown in Figure 9, Distributions of the temporal correlation coefficients for all off-diagonal pairs (i. (i.j) are shown as violin plots, where each violin corresponds to model trained with one random seed. The median of each distribution is marked with a black line. All medians are greater than 0.999.
linear nonsts.
nonlinear nonsta.
Tamooral comcaation
1
2
6
1
2
6
Figure Distribution of test set temporal correlation between the linearized attention matrix and the true nonstationary connectivity. Fach column shows result from NetFormer model with different random seed. Median of each distribution is marked black. All medians are greater than 0.999.
A.2.4 FURTHER DISCUSSION ON NONLINEAR DYNAMICAL SYSTEMS
In Section 2, we showed that when f is sigmoidal, A can reflect W through Taylor series approx- imation of ([W.).. Take f tanh w xb < x Vi I,...,N, N,
example. Let w, denote the i-th row of W. When
(anh( W xx) = tanh(O) + oz +(x)).
As (O) 0, the forward Euler method is
1 26 5 ttthh(Waa (wx)) ø + 5 tanh (o)nh -00x))
This analysis also applies to other sigmoidal functions f. such as arctan, with
f(0) 0, (E) (0( + O(g3) for & within some interval around 0.
16

---

Published as a conference paper at ICLR 2025
Therefore, we hypothesize that A can capture W through learning (0) for sigmoidal f. is also clear that learning W becomes more challenging when w; a does not always stay within the radius of convergence of the Maclaurin series. Nonetheless, we note that if some w outside the convergence region, [w] x)) will be constantly positive or negative, and the system will either blow up or decay to zero. Therefore, for the systems of interest here, which are those with interesting persistent dynamics, from time to time w. æ must fall within the convergence region where the Maclaurin series representation is valid. That being said, while A may still be able to capture some aspect of W, it could become less accurate, and may require more observations of the system to gather sufficient timesteps within the convergence region.
constantly
In the example nonlinear dynamical system shown in Figure 2b, w, 2 stays within the radius of convergence of tanh for all i and k, which makes the Maclaurin series approximation valid for all timesteps. In Figure 10, we provide another toy model example showing that the attention from NetFormer still bears considerable similarity to the ground-truth W even when w 2y falls out of the convergence region for some and k.
tanh
zoomed
120
200
-30
500
1000
1500 timesteps
2000
2500
3000
Eimesttpge
Ground-iruth W
0.842)
Agus 0.842)
20
20
Figure 10: Demonstration of NetFormer on nonlinear system de tanh Wx) where wi x z does not always stay within the convergence region of the Maclaurin series of tanh. a. Top row: Trajectories of the simulated system. Simulation was done using the forward Euler method with stepsize 0.1. Shaded regions represent timesteps used as test set. Bottom row: Visualization of w. x across simulated timesteps. Boundaries of the convergence region, + 2. were marked with black horizontal lines. The right column provides a zoomed-in in view of the first 500 simulation timesteps. b. Left to right Ground-truth W, average linearized attention matrix across test timesteps from NetFormer (A), AoLS fitted through least-squares regression. Colorbars indicate the scale of the off off-diagonal entries, and the diagonal entries masked in grey. A, AoLS were rescaled for visualization. Spearman's 's rank correlation coefficients (p) were computed between the off-diagonal entries of A or Aots and W. We trained 10 NetFormer models with different random initializations (p = 0.841 0.02, mean + std), and A shown is the one whose is the closet to the average / across 10 random initializations. NetFormer models achieved similar performance as AOLS (p 0.9, two- sided one sample test). All NetFormer models were trained to minimize the mean squared error on the trainining set for 600 epochs using the Adam optimizer in Pytorch, with / 1, M 5, D batch size = 80. Learning rate was initialized to 0.01, and was decayed by factor of 0.9 every 100 epochs.
17

---

Published as a conference paper at ICLR 2025
A.3 ADDITIONAL DETAILS FOR STDP SIMULATION (SEC 3.2)
Our STDP simulation is largely based on Neuromatch Academy (2023). The dynamics of the LIF
neuron follows
dV -VV - FL) (V - FFE)
Ti dt
(12)
V(t) 2 VvA = V(t)= = Vreset
(13)
where V is the membrane potential, Tm is the membrane time constant, Ez is the resting potential, Eg is the synapse reversal potential, Ven is the spiking threshold, and Vreset is the reset potential. Once (tt) crosses the spiking threshold, we say the LIF neuron emits spike, and (t) will be reset to and held at Vreaet for a refractory period byef. (t) is the total excitatory synaptic conductance, which is the total conductance of all active pre-synapses (i.e. synapses coming into the LIF neuron)
at that time:
N
gE(1) = Σ (i(()dt(t - tapk)
(14)
==1
where 5; is the delta function: d,(1 - lsph) if there is spike at pre-synapse i at time I, and 0 otherwise. The conductance of each pre-synapse gs evolves following
dg:
9: + g: - ttpk)
(15)
dt
where TE the EPSP time constant, and Di the peak synaptic conductance of pre-synapse i. All g:
are bounded between 0 and Imax.
When there is a spike arriving at the i-th pre-synapse,
Ji 9s + (L)gmaz
(16)
where (t) helps tracking the time since the last postsynaptic spike emitted by the LIF neuron.
When the postsynaptic LIF neuron spikes, all pre-synapses are updated:
9: = 9s + P.(t))mar, Vi
(17)
where Pt( helps tracking the time since the last spike at the i-th pre-synapse. STDP is enforced through (1) and P.(l). Specifically, (1) follows
dM
T- -MM
(18)
dt
and whenever the postsynaptic LIF neuron spikes,
M(t) M(1) A_
(19)
(1(() follows
dP
= - P T+ dt
(20)
and whenever the i-th presynaptic neuron spikes.
P(I) - P(I) A+.
(21)
T+ and T- specify the range of separation between pre- and postsynaptic spikes where STDP takes effect. A4, A_ are both positive, and define the maximum amount of synaptic strengthening and weakening, respectively. follows that (() S 0, P.(t) 2 Vt. (t) and (t()) effectively capture the STDP rule
AW A4e (tre (tre /oultT if Epost > Epre
(22)
(1pre AW - - (pantt/ 1post < 1pre
(23)
as shown in Figure 3a. The constant parameters and initial conditions in our simulation are set in the same way as Neuromatch Academy (2023), and are summarized in tables 3. Pre-synaptic spike trains are modeled as independent Poisson processes with rate 50Hz.
To generate data, we run the simulation for 100,000 timesteps where each timestep corresponds to mms. We used the first 80% of data for training, and the last 20% of data as the test set. Membrane potential of the LIF neuron was then z-scored using its mean and std on the training set. We fitted 5 NetFormer models using different random seeds. All NetFormer models have /t 5, M 101, D 101, and were trained to minimized the mean squared error on the training set for 20 epochs using the Adam optimizer with learning rate 0.005, batch size 64 in Pytorch. Result from one seed is visualized in Figure 3d-f. We used a sliding window of length 10,000 timesteps to smooth both true and inferred synaptic weight trajectories before computing the correlation coefficients.
18

---

Published as a conference paper at ICLR 2025
Table Constant parameters in STDP simulation
Parameter
Value
Tm
10 [ms]
Er.
-75 [mV]
Eg
0 [mV]
Vih
-55 [mV]
Vreset
-75 [mV]
tref
[ms]]
TE
[ms]
Jmar
0.024
20 [ms]
T-
20 [ms]
A+ 0.008
A_ 0.0088
Table 3: Initial conditions in STDP simulation
Variable
Initial value
-65 [mV]
9i
0.014
M
0
P
0
A.4 ADDITIONAL DETAILS FOR TASK-DRIVEN POPULATION ACTIVITY SIMULATION (SEC 3.3)
A.4.1 SIMULATION DETAILS
The hidden dynamics of RNN models follows equation 7, and at each timestep kt, the network activity read out through a linear mapping
yk
W oothh.
(24)
All RNN models are trained to minimize the cross entropy loss using the Adam optimizer in Pytorch.
In the Perceptual Decision Making task, we trained RNN with 4 hidden units for 5000 epochs using learning 0.005, batch size 32. In the Go-Nogo task, we trained RNN with 8 hidden units for 2000 epochs using learning 0.01, batch size 32. In the Delay Comparision task, trained RNN with 12 hidden units for 5000 epochs using learning 0.01, batch size 32. All trained RNNs achieve over 90% accuracy in the 1000 test trials. All trials generated using the default parameters in the NeuroGym toolkit (Molano-Mazon et al., 2022).
In each task, we recorded the hidden units activity of the trained RNN during the 1000 test trials. We then trained NetFormer to predict the next-step hidden units activity based on the present and used for training, and the remaining 200 trials were using for evaluation. In each task, we trained NetFormer models from different initializations minimize the mean squared error (MSE) on the training set using the Adam optimizer in Pytorch. In the Perceptual Decision Making task, NetFormer 7, and was trained using learning rate 0.0025, batch size 64 for 100
past H-step hidden activity and stimulus inputs. Hidden units activity during 800 trials
has H 5, N M epochs. In the Go-Nogo task, NetFormer has H/ 1, N = M = D 11, and was trained using learning rate 0.01, batch size 64 for 50 epochs. In the Delay Comparision task, NetFormer has N M D 14, and was trained using learning rate 0.005, batch size 64 for 50 epochs.
A.4.2 QUANTITIVE COMPARSION WITH LINEAR REGRESSOIN MODEL
Table 4 provides quantatitive comparison between NetFormer and the linear regression model.
19

---

Published as a conference paper at ICLR 2025
Dynamics prediction
Connectivity recovery
NetFormer
Linear
NetFormer
Linear
MSE
R²
MSE
R²
Spearman
Pearson
Spearman
Pearson
a
0.000 000
0. 0.008
0.014
0,920
0.694=0.163
0.760:E0.120
0.566
0.548
b
0.000
0.97210.001
0,013
0.960
0.622=0.0066
0.722:10.004
0,518
0,553
0.001:10.00
0.997110.001
0.034
0.897
0.811=:00011
0.633::0.105
0.626
0.650
Table 4: a, b, corresponds to the Perceptual Decision Making task, Go-Nego task, and Delay Com- parision task, respectively. MSE and R were evaluated on concatenated held-out trials. Spearman and Pearson correlation coefficients were computed between the off-diagonal entries of the inferred and true connectivity matrices. NetFormer results are the std across 5 random seeds.
A.5 CONNECTIIIY--ONSTRANNNE SIMULATION AND NEURAL DATA: DATASETS AND PREPROCESSING (SEC 4.1)
A.5.1 CONEECTVVTY-CONSTRAAAII SIMULATION
To construct the ground-truth connectivity matrix, for each cell-type pair and for every pair of neurons, we first drew a random sample from the uniform distribution between 0 and 1. Then, we used the connectivity probability from patch-clamp experiments Campagnola to determine if two neurons are connected. For connected neurons, we sampled their connection strength from a normal distribuion M(H, 0.1), where y is the measured post-synaptic potential from patch-clamp experiments. We simulated 30, 000 steps for 200 neurons, using the first 80% timesteps for training and the last 20% for testing.
(2022) as a cutoff threshold
A.5.2 PATCH-CLAMP DATASET
The dataset released in Campagnola et al. (2022) contains experimental results of connectivity probability and connectivity strength (Postsynaptic Potential (PSP)) at the cell-type level measured using patch-clamp. In each experiment, up to eight neurons were simultaneously subjected to whole- cell patch-clamp recording, mainly under current-clamp conditions, with some stimuli also tested under voltage-clamp conditions. Stimuli were applied to each patched neuron while recording the other neurons for postsynaptic responses. We mainly focus on the experimental results for layers 2/3 in mouse primary visual cortext (V1), to match the neurons recorded in the multimodal mouse datatset (Bugeon et al., 2022).
A.5.3 MULTIMODAL NEURAL RECORDING
For functional activity recordings from neuronal populations, we used recent, public multimodal dataset provided by Bugeon et al. (2022). This dataset includes spontancous population activity recordings from the mouse primary visual cortex (VI) across layers 2/3 via 2-photon calcium imaging at temporal sampling frequency of 4.3Hz across six 20-minute sessions, recording approximately 500 neurons per session. Spatial coordinates of the recorded neurons are also provided. We trained our models on data from one experimental subject (SB025), which includes recordings of 2481 neurons, with some neurons repeating across six sessions. The dataset also includes single-cell spatial transcriptomics, profiling mRNA expression for 72 selected genes to identify excitatory and inhibitory class labels of neurons. 51% of neurons in the inhibotiry class can further be identified to be one of Lamp5, Pvalb, Vip. Sncg. and Sst.
A.5.4 DATA PREPROCESSING
In the comnectivity-constraine simulation, when using RNN with an exponential activation, we added the minimum activity value all simulated neuronal activity ensure nonnegativity.
For the experimental neural recording, which is nonnegative, we normalized using the mean and standard deviation calculated across all sessions and neurons involved in training. When using the RNN model with exponential activation, we normalized the data by dividing by the standard deviation only, without first subtracting the mean.
20

---

Published as a conference paper at ICLR 2025
A.6 BASELINES
A.6.1 STATIONARY CONNECTIVITY BASELINES (SEC 4.2)
Linear Regression: We denote neural activity data as Xg. = [ xu| E RNH
recorded from N neurons and H time steps. Let 2k+1 E R denote the neuronal activity at the (k + t) time step. Given previous time step, Th, linear regression predicts current time step activity as
kb+1 = Wxx +
Recurrent Neural Network (RNN) with tanh activation: Given neuronal activity across previous , a RNN with predefined Tanh activation function predicts current
p time steps, Zj, time step activity
žk+1= =0 (11102 IR + (( Ix + ... W xx-p- + +6) o- tanh
where W0, 10, represent how the previous (-th step affects the current step activity and each element W represents how the j-th neuron the previous 1-th time step influences the ith neuron in current step. The RNN is trained by minimizing mean squared errors (MSE) of the current time step activity prediction given previous time steps. p = 1 is commonly used for RNN. For modeling both simulation data and real mouse data, chose p because the simulation data has exactly one timestep dependency. Using larger did not improve performance in real data either.
Recurrent Neural Network (RNN) with exponential activation: Similar to above, with
1 (W(0) Ie + W(1) x4-1 + ................................ W 2k xx-pti + b) exp
I&] E RNH denotes activity of N neurons across H time steps. Let 2255) E R denote the i-th and j-th neurons' activity across H time steps. For simpliticity, let a = x(()) [r :],b = x00) -r]. Cross correlation with time delay r is computed as
Cross correlation: Recall that Nx = H++
a
|a ū b1
We used 7 1 for inferring connectivity in both simulation and read data.
Covariance: Similar to above, let x(,)
R denote the i-th and j-th neurons' activity across H time steps. Covariance between neuron and j defined as
x(0
H 1
H
-
Mutual information: Computed using the Python package Pylnform.mutualing
Transfer entropy: Computed using the Python package SyInform.transfereentro
A.6.2 NONSTATIONARY CONNECTIVITY BASELINES FOR NEURAL DATA (SEC 4.4)
Low-tensor-rank RNN (LtrRNN): LtrRNN is designed to model trial-varying neural dynamics by capturing low-dimensional changes in connectivity over time. We split continuous neural activity data into overlapping sliding windows, treating each window as a distinct trial, which allows LtRRNN learn temporal variations in connectivity patterns. The dimension of the network is defined by the number of neurons. To fit LiRNNN, we followed the code provided by its authors https: // g thub. com/ a wtthurrpe//eee TEtRRN, After fitting, we extracted the trial-specific weight matrices from LirRNN and compared them with the attention maps learned by NetFormer.
Autoregressive Hidden Markov Model (AR-HMM): AR-HMM extends the traditional HMM by incorporating state-specific autoregressive dynamics, where each state has its own unique autoregres- sive model to describe the temporal dynamics of neural activity. AR- RR-HMM is able to infer the latent states of the system and estimates the connectivity matrices associated with each state. However, a limitation of AR HMM is the need to predefine the total number of states. To fit AR-HIR-HN AR-HMM, we followed the code https: / hubb. com / inderman ab/ ssm.
21

---

Published as a conference paper at ICLR 2025
excitation
EC
Pvalb Set Vip
Inhlaition
NetFormer
Regression
Exp
Correlation
Postsynaptic
c4
Presynoplic
Figure 11: Conmectivty-coosttained simulation with sigmoid activation. Visualization of ground
truth and inferred connectivity matrices at both individual- level and cell-type level. A positive linear transformation has been applied to standardize all matrices to the same range for better visualization.
A.7 MORE ON CONNECTIVITY--OOSTRANND SIMULATION (SEC 4.2, 4.3)
A.7.1
CONNECTIVITYY SIMULATION WITH SIGMOID NONLINEARITY
To test the robustness of NetFormer towards different nonlinear activations in the connectivity- constrained simulation (Section 4.2), we replaced the tanh activation with sigmoid activation and used NetFormer, along with other baselines, to reconstruct connectivity at both neuron-level and cell-type level. Table presents quantitative comparisons of the inferred connectivitives with ground truth, averaged over five random seeds. Qualitative comparison is in Figure 11. NetFormer shows competitive performance, and outperforms other baselines at neuron-level (N N) connectivity inference.
NetFormer
linear
RNN
RNN
covariance information
transfer
regression
w/ tanb
w/ exp
correlation
entropy
connectiviy
Peatson 00.00=00000 0.765 0.002 0.772=0.001 0.679
0.829
-0.022
0.447
0.329
simulation
X
Spearman 0.508110.008 0.482 1 0.0002 0.485 0.001 0.45410.000
0.512
-0.006
0.201
0.223
connectivity
Pearson 0.880=0.003 9.904 = 002 o. 0.000 0.803 -00002
0.921
-0.438
0.350
0.304
Spearman 0.860 0.003
10.011
0.856
-0.299
0.104
0.340
Table 5: Connectivity
simulation with sigmoid activation. Simulation data has ground truth for both neuron-level (N x N) and cell type-level (K K) connectivities. Connectivity inference is assessed using Spearman's 's and Pearson's 's correlation coefficients. Results of connectivity inference using mutual information, transfer entropy, and Granger causality are assessed by comparing against the absolute values of ground truth connectivity.
A.7.2 GRANGER CAUSALITY TEST
We performed Granger causality tests on connectivity simulation data with both tanh and sigmoid activations. To create a binary ground-truth neuron-level connectivity matrix, we assigned a value to all nonzero entries and 0 to zero entries. We then conducted Granger causality tests on pairwise neural activities to generate a matrix of test statistics, which is then min-max normalized to between and 1. To compare with Granger causality on inferring binary connectivity, we also performed min-max normalization the time-averaged neuven-level attention matrix. We evaluated the performance using AUROC. NetFormer outperforms Granger causality in both simulations (Figure 12).
22

---

Published as a conference paper at ICLR 2025
NedFormer AUROC 0.16
NetFormer AUROC 0.894 (hanh simulation activation)
PG Pusth Sar Vp
A [Aigmoid simulation activation)
Accurscy 0,872
(rigmoid wimusation activation)
Jann simulationn activution)
10
Postonippt
6.
Pregnaphis
Figure 12: NetFormer for binary connectivity recovery. Visualization of NetFormer and Granger causality results for inferring binary connectivity (presence or absence of connections) in connectivity- constrained simulation data.
a.
Behavional States (unseen during training)
b. Gell AltutionntAlttttonn Altution at Runninn (Zrundom time stopp)
Neuron1
Neuron 2
Neuron
200
15.00
25
Neurond
Sype-laall Amentiona Stationary Sync mandom time-stapd)
Nearon 5
Neuron7
Figure 13: a. Visualization of neural data and behavioral states for an example session. Behavioral states are plotted the top row. Predicted activity traces from NetFormer are shown in orange, while measured activity traces are shown in blue. NetFormer predictions effectively capture overall patterns of neural activity. Three blocks of time from a. are selected, each representing different state. NetFormer attention weights at two timesteps randomly selected within each block are visualized.
A.8 NETFORMER FOR NONSTATIONARY CONNECTIVITY INFERENCE ON NEURAL DATA (SEC 4.4)
We visualized neural activity traces, mouse behavioral states, and cell-type level attention weights for each state in Figure 13. While behavioral states are not provided as model inputs, they can be inferred using unsupervised methods such as PCA on NetFormer's attention weights. We also benchmakred NetFormer with LiRRNN and AR- HMM on state-dependent connectivity inference, as shown in Table The inferred connectivities are compared against the patch-clamp experimental result.
A.9 NETFORMER APPLIED TO NETWORKS WITH SPURIOUS CORRELATIONS (SEC 4.3)
Das Fiete (2020) shows that connectivity inference can obscured by spurious correlations arising from strong recurrent connections. To assess NetFormer's resilience to such spurious correlations, we performed two experiments on simulated neuron populations with strong recurrent connections. In the first experiment, generated continuous activity from the simulated network, where NetFormer directly applicable. In the second experiment, we generated spiking activity from the simulated network, which follows Das & Ficte (2020) more closely but requires modifying the objective function of NetFormer. In both experiments, we simulated neuron populations size 100. The population connectivity matrices in both experiments constructed in the way as Das & Fiete
23

---

Published as a conference paper at ICLR 2025
NetFormer
LirRNN
AR-HMM
Running
Pearson
0.591
0.007100182
0.199=0.027
K x K
Spearman
0.601=0.209
-0.036=0.189
0.274=0.048
Stationary Desync
Pearson
0.662-0.176
-0.003=0.181
-0.320+0.012
in-vivo recording
K K
Spearman
0.723=0.173
-0.057=0.174
-0.206+1.0017
Stationary Syne
Pearson
0.713 0.00010.186
0.145
K
Spearman 0.767=00151
0.069
0.151=0.009
Table 6: State-dependent connectivity inference. NetFormer attention maps and LtrRNN weight matrices grouped by state labels and averaged for each state. The inferred connectivity for each state is compared against Postsynaptic Potential (PSP) resting state amplitude obtained from patch-clamp experiments.
(2020) with a local Mexican hat profile. Specifically, the connection weight between neurons and computed as
Waj = e d2y/2oo E ae 42,/2035
(25)
where de is the distance between neurons i and j (in units of neurons), oi 6.98, 02 7, and a 1.0005.
A.9.1 NETWORK WITH CONTINUOUS ACTIVITY
The network model used to generate the data is similar equation with modifications to the feed-forward inputs b and incorporation of the recurrent weight strength r. Specifically, Gaussian noise is injected into the feed-forward inputs following Das & Fiete (2020):
Fk+1 = tanh(r Ø + b(1 + cx)).
(26)
We conducted similar experiments in Das & Fiete (2020) by varying the recurrent weight strength r. The corresponding connectivity inference results are shown Figure When recurrent weight strength is large (c.g. r 35), spurious correlations between neurons appear, seen in the synchronized activity and noise correlation. On the other hand, when the recurrent weight strength is small (e.g. 5), activity is mainly driven by noise rather than recurrent connections, making hard to infer connectivity by fitting activity. In both extremes, NetFormer will struggle to capture the underlying ground-truth connectivity, reflected by lower correlations with the ground truth. That being said, NetFormer still recovers some patterns of the ground truth, better than what noise correlation alone can capture.
A.9.2 NETWORK WITH SPIKING ACTIVITY
To adapt NetFormer for modeling spiking activity, we take the exponential of NetFormer's output as the mean rate of the Poisson distribution for spike generation. The spike generation process is constructed in the same way Das & Fiete (2020), which we describe briefly below, and refer the readers to the original paper for more details.
Let 8k, Jk. σ denote the synaptic activation, neural input, and emitted spikes of the neuron population at timestep k. respectively. gk includes both inputs from other neurons and externally injected inputs by: that
rW 8h + bu-
(27)
At timestep k. the i-th neuron will emit spike, that is, o 1, if the input it receives exceeds certain threshold, that is, > O.
To fit the NetFormer, every timestep k. have
9x NetFormer(s . 8k-H+1):
(28)
Å = cxp():),
(29)
σ Poisson(Ae).
(30)
This setup is chosen to match the setup for GLM in Das & Fiete (2020) (see Equation 9 there). One minor difference is that, as model inputs, they used past spike trains convolved with decaying
24

---

Published as a conference paper at ICLR 2025
True Connectivity
10
15 1-
50 F5
Atteention Weight Noine Correlation
Activity
Figure 14: Spurious correlation: network with continuous activity. Correlation between attention weight as connectivity inference and ground-truth recurrent weights (top row), noise correlations (second row), attention weights (third row), and time-by- by-nuooon activity traces (last row), at different recurrent weight strengths r.
exponential kernel, while we used the synaptic activation. Nonetheless, since their filter is truncated at value much larger than the exponential decay time constant, it can be easily verified that these are actually very similar.
We followed the procedure described in Das & Fiete (2020) to generate population activity different recurrent strength r. We followed the data generation code provided by the authors https://github com/ coom/Fietee wrieteLab/rae pewuura--orruutt-iieeeeeeee riken_dempe m. We adopted and spiking thresholds listed in https:://itht/i com/ htt///llbb/maaaaaa httpt/hhehhhds___pp
mat. After the simulation starts, we collected activity after the initial 10,000 simulation timesteps in order for the network activity to stabilize. total of 100,000 timesteps were collected, where the first 80,000 steps were used as the training set, and the rest 20,000 steps were used as the test set. Since this is a prototype study, the data volume we used here is much smaller than what is used in the orignal paper: their experiments are typically done 10 spikes, while there are average 5.05 x 10 spikes in our training set across all r. Nevertheless, synchronized activity patterns and high noise correlation large can also be observed in our training (Figure 15).
Fietelab/neut Aneural-ciret cerrrcitttt touuit-iifeeetie
For all r. we used NetFormer with H = 1, N = 100, D = 100. For r = 0.0225, M = 40, and M = 100 for the rest of r. All NetFormer models were trained to minimize the negative Poisson log likelihood loss on the training set, using batch size 32 for 100 epochs. We used learning rate 0.001 for 0.0025, 0.005, 0.002 for 0.0075, 0.01, 0.0125, 0.015, 0.0175, 0.02, 0.025, and 0.0005 for r = 0.0225. We used the average NetFormer attention across test set timesteps as the NetFormer-inferred connectivity. Each row in NefFormer-inferred connectivity is then rescaled to minimize the l2 distance from the corresponding row in the ground-trudd truth connectivity matrix, as also done in Das & Fiete (2020). Visually, Figure 15 shows that NetFormer-inferred connectivity can be affected by the presence of strong spurious correlation in neural activity, while being more resilient than noise correlation. Quantitively, we computed the Spearman and Pearson correlation coefficients between the off-diagonal entries of NetFFormer-inferee connectivity and true connectivity. We also computed the inference error (the lo distance between the off- off-diagonal entries of true and inferred connectivity weights) as in Das & Fiete (2020). Consistent with their findings in Figures 2 and 3, we observe U-shape curve in inference error as a function of r, and an inverted U-shape curve for correlation.
25

---

Published as a conference paper at ICLR 2025
True Connectivity
-
00.025
r-0 0075
m-0.015
/ 025
5.0003
10.0000
3.0000
0.500
-0 0030
-0.0003
0.0025
6.5000
NetFormer attention
-00011
0.0020
-0,0100
00005
0..00
-0.0000
6.10
33
0.20
Noise correlation
-2
Synaptic activations
neurans
Spearman Pearson
0.8
0.80
0.7
0.6
Inference error
Correlation
0.5
0.70
0.4
0.65
0.3
0.60
0.2
0.55
0.1
0.0025
0.005
0.0075
0.01
0.0125
0.015
0.0175
0.02
0.0225
0.025
0.
Recurrent strength
Figure 15: Spurious correlation: network with spiking activity Correlation and inference error between true and NetFormer-inferee connectivities. Errorbars: mean across three random initializations for NetFormer. In addition, rescaled NetFormer-inferred connectivity, noise correlation, and population synaptic activations four example visualized. For correlation and connectiviy matrices, colorbars indicate the scale of the off-diagonal entries, and the diagonal entries are masked in grey.
A. IMPLEMENTAION DETAILS FOR FITTING CONNECTIVITY-CONSTRRO SIMULATION AND NEURAL DATA (SEC 4)
A.10.1 MODEL FRAMEWORK
IK] N To encode neuron identities, learnable positional embedding matrix E E NxM is concatenated to X, giving X₁ [X El e IR (IH++M) The queries Q+ and keys K₁ are obtained through linear transformations of X Qx = N&Wa E N and K2 = X&WK E RNDD NetFormer model is trained to predict the next time-ste step activity 2x+1, defined as
Following Section train the NetFormer predict #x+1 based on X₁ -+k-H+1
Q&KT VD
(XaWo) X )a Tk.
Jr +
where A& is the self-attention that want use for inferring connectivity, ó is the attention acti- vation. In the standard Transformer model Vaswani et al. (2017). softmax used as the attention activation function, but here set é equal identity for better interpretability. In fact, experi- mented with different activation functions and empirically found that the identity activation yields the best results recordings from the mouse cortex.
Although not used in the current experiments, is possible to apply an additional linear transformation (Wout) on X to accommodate neuronal dynamics that depend multiple previous timesteps,
26

---

Published as a conference paper at ICLR 2025
that is,
ix+1 = A& (x + X Wouts Wont E R R
A.10.2 NETFORMER TRAINING AND EVALUATION
We first assign each unique neuron across all sessions ID, which is later used to track positional embedding for each unique neuron, because same neuron can be recorded more than one session. Then, within each session, we construct samples with window size 200 in simulation and 60 in real data, and we make sure samples in one batch should come from the same session so that the dimensions can match
We use the first 80% timesteps in all sessions for training and the last 20% timesteps for validation. The model is trained using MSE as the loss function, comparing the predicted activity for the next time step with the ground-truie truth activity. We employ early stopping criteria, ceasing training if there are 20 epochs without improvement, with hard limit of 100 epochs maximum.
After training is complete, we calculate the attention for each sample in the dataset. For each session, we aggregate the attentions from all samples to compute a single time-averaged attention. Averaged attentions from all sessions are then transformed into a final cell-type level attention. We achieve this by aggregating attention values according to their corresponding presynaptic and postsynaptic cell types and dividing by the total count of such pairs.
We also extract positional embeddings from the trained model and utilize each neuron's unique ID to determine the neuronal embedding for every unique neuron. These embeddings are then used as features for logistic regression to classify neurons as either excitatory or inhibitory (Figure5b).
For evaluation, assess the inferred cell-type level connectivity against the Postsynaptic Potential (PSP) resting state amplitude obtained from patch-clamp experiments, which serves as the experimen- tal ground-truth. Additionally, evaluate the accuracy of the binary cell-type classification using experimental data from single-cell spatial transcriptomics, which provides a classification of neurons into excitatory and inhibitory types across all sessions.
1.1 EVALUATION METIRCS
We use python libraries and built-in functions for computing evaluation metrics.
For connectivity inference, we flatten the inferred 2-dimensional N N or K K connectivity matrix and grond-truth matrix.
Pearson correlation: siipy.statspearssee
Spearman rank correlation:
For activity prediction, given the input matrix E IR B for NetFormer and input matrix E R Bx for RNN, where B is the batch size, NetFormer outputs BxXN and RNN outputs R BxN We flatten the predicted activity and the grond truth.
MSE: torch.nn. www..nn...ntiioaa....mmmmm las__oss(( Pearson correlation: ccipystatt.pearsonnf R² skleam.mettres.f2..comm
For binary classification, classifier predicts the probability for all neurons.
Top-1 accuracy: www.:meeetttterrrreeee score()
Area Under the Receiver Operating Characteristie (AUROC): sklearn.metries _ac__aac_______________ suc_score()
10.4 HYPERPARAMETERS
In connectivit--onntraimne simulation data, for training NerFormer, we use history window size 100, embedding size 200, hidden dimension of query and key matrices is 300, learning rate 10 and batch size 32. For training RNN, we use p = 1, batch size 32, and learning rate 10
In neural recording, for training NetFormer, hidden dimension of query and key matrices 90, learning rate 10-3, and batch size 32. For training RNN, use p 1, batch size 32, and learning rate 10-1
use history window size 60, embedding size 30,
27

---

Published as a conference paper at ICLR 2025
We PyTorch Paszke et al. (2017) and PyTorch Lightning Falcon & The PyToorch Lightning team (2019) for model development and training, and Adam as the optimizer.
A. A..5.5 PSEUDOCODE
We train NetFormer and extract attentions and positional embeddings for connectivity inference and binary cell-type classification. The pseudo code for model training connectivity inference and cell-type classification is provided follows:
NetFormer (x, heuron_ids)-
if constraint True:
parameters (nun_cell_typp num_cell_type)
cell_type__evel_val parameters (num_cell__tal num_cell cam_cell_type)
embeddings embedding_table (neuron_ids, input = layer_norm x, embedding ngs input[:, : :T), input [:, : T:]
poooooommccncttttt (concat (x, embeddings )
dim dim_ x. shape -11, embeddings. shape [-1]
scale (dim_x + dim_e) -0.5
logits input 8 __ input.
logits = logits * scale
if activation =- softmax:
attention softmax (logits)
elif activation == sigmoid:
attention simooid
elif activation == tanh:
attention tanh (logits)
elif activation none:
attention = logits
output = layer_norm (attention @ x + x)
if out_layer == True:
# linear_out is lienar transformation from dimension T to output linear_eeee out (output) return output, attention
else:
# Use the last column as prediction
return output[:, -1], attention
[eatDormee__Tainnnnggll__ee (al1_sapp1l
all_inputs, all_neuron_ids ids, all_GT_targett all aal_samples model = NetFormer () optimizer Adam (model, learning_rate)
all all_predictions, attentios model (all (all_aaall___aaaaa inputs, all all_neuron ids)
prediction_loss MSE (all [eall_lledictionn all_GT_targets)
loss = prediction_loss
optimizer. loss backwardd optimizer.stepl
28

---

Published as a conference paper at ICLR 2025
Connectivity_IIfrreeeee Inference (all_samples, trained
trained_NetPornee NetFormer, GT CI_connectivity)=
all all_inputs, ail_neuron_ide _ds, all_cT_targets ttragetgess = all _amp_ssssessssse
all all_predictions, all_attentions trained NetFormer (all inputs, all _e__nerro__ooe ids)
avg_attention mean (a11 [ean(all_attenttoons axis=0)
pearson_corr pearsonr (GT (crconccrtl avg_attention)
spearman_corr = spearmanr (Crconnctity,.... avg_attention)
(rrtiiiirrrrineeeeeeeFrroeeee neuron_ids, cell_types,
embedd ngs Brained_NNttoormer saabeeddin_tablee (neuron_ids)
X_train embeddings TRAIN idx]
y_train celi_types [ [RAIN_idde [ix]]
X_test [MMeddings[TEEST....
Y_test [ell_types[TSS10
# Train classifier
classifier Logistickegressionn.ff fit (X (X_train, y_train) # Test on test set
Y_pred wlwssifier. ct (X_test)
A.11 COMPUTE RESOURCES
Model training simulated systems in Section 3 was done a MacBook Pro with Apple MI chip. Using the NVIDIA A100 GPU, NetFormer model trained on connectivity-constee data took about 10min. On neural recording, NetFormer model trained the subject SB025 took about 20min, which requires at least 30 GB of RAM and 16 GB of GPU memory.
simulation
29


Reviews:
[{"review_id": "y1ABBRIY0E", "reviewer": "Reviewer_9MUd", "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.", "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. \n\nTo address my concerns, you can perform the following revision:\n\n1) Can you add [2] as a baseline to your experiments?\n\n2) 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].\n\n3) 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.\n\n4) For Fig. 4, I believe [7] would be a relevant baseline. \n\nI believe, for a borderline score, the suggestions 1, 2 and 4 should be addressed. \n\nEdit: 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.\n\n[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.\n\n[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. \n\n[3] Qian, William, et al. “Partial Observation Can Induce Mechanistic Mismatches in Data-Constrained Models of Neural Dynamics.” NeurIPS 2024, OpenReview, 25 Sept. 2024\n\n[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.\n\n[5] Mongillo, Gianluigi, Omri Barak, and Misha Tsodyks. \"Synaptic theory of working memory.\" Science 319.5869 (2008): 1543-1546.\n\n[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\n\n[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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The simulations are mostly convincing (see below) of the approach's validity, providing compelling evidence for its use in biological experiments.\n\n- 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.\n\n- 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.\n\n- 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. \n\n- The experiment in Fig. 4 should use [7], which I believe is a more appropriate baseline.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LbYB2B19ch", "reviewer": "Reviewer_HFca", "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.", "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?\n\nScalability: How does the model adapt to bigger datasets, especially in real-world scenarios when partial observability is a problem?\n\nInterpretability: Could the writers elaborate on how particular attention weights might be connected to recognized patterns of connection or brain phenomena?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "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.\n\nQuality: The robustness of the model is supported by numerous theoretical and experimental validations, including actual brain data.\n\nClarity: Although more illustrations would help the paper's complex theoretical content, the model's structure is explained in detail.\n\nSignificance: 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.\n\nLimitations of Interpretability: Although NetFormer can be interpreted in comparison to more conventional models, some results might be strengthened by further biological contextualization.\n\nEvaluation 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f0HrcT0zYq", "reviewer": "Reviewer_ipsn", "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.", "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?\n\n2. 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?\n\n3. 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?\n\n4. 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?\n\nOverall, if the authors are able to address some of these concerns I would be happy to increase my initial score accordingly.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n2.  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.\n\n2. 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 \n\n3. 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.  \n\n4. 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. \n\n5. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AUV4NU2EiL", "reviewer": "Reviewer_ovPw", "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.", "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.)\n\n2. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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. \n\nMy 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)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)