SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Christos Sourmpis, Carl C.H. Petersen, Wulfram Gerstner, Guillaume Bellec Brain Mind Institute,
School of Computer and Communication Sciences and School of Life Sciences,
École Polytechnique Fédérale de Lausanne (EPFL),
CH-1015 Lausanne, Switzerland,
{firstname_lastname}@epfl.ch
Simultaneous behavioral and electrophysiological recordings call for new methods to reveal the interactions between neural activity and behavior. A milestone would be an interpretable model of the co-variability of spiking activity and behavior across trials. Here, we model a mouse cortical sensory-motor pathway in a tactile detection task reported by licking with a large recurrent spiking neural network (RSNN), fitted to the recordings via gradient-based optimization. We focus specifically on the difficulty to match the trial-to-trial variability in the data. Our solution relies on optimal transport to define a distance between the distributions of generated and recorded trials. The technique is applied to artificial data and neural recordings covering six cortical areas. We find that the resulting RSNN can generate realistic cortical activity and predict jaw movements across the main modes of trial-to-trial variability. Our analysis also identifies an unexpected mode of variability in the data corresponding to task-irrelevant movements of the mouse.
Over the past decades, there has been a remarkable advancement in neural recording technology. Today, we can simultaneously record hundreds, even thousands, of neurons with millisecond time precision. Coupled with behavior measurements, modern experiments enable us to better understand how brain activity and behavior are intertwined [1]. In these experiments, it is often observed that even well-trained animals respond to the same stimuli with considerable variability. For example, mice trained on a simple tactile detection task occasionally miss the water reward [2], possibly because of satiation, lack of attention or neural noise. It is also clear that there is additional uncontrolled variability in the recorded neural activity [3, 4, 5] induced for instance by a wide range of task-irrelevant movements.
Our goal is to reconstruct a simulation of the sensory-motor circuitry driving the variability of neural activity and behavior. To understand the generated activity at a circuit level, we develop a generative model which is biologically interpretable: all the spikes are generated by a recurrent spiking neural network (RSNN) with hard-biological constraints (i.e. the voltage and spiking dynamics are simulated with millisecond precision, neurons are either inhibitory or excitatory, spike transmission delay takes  $2 - 4\mathrm{ms}$ ).
First contribution, we make a significant advance in the simulation methods for data-constrained RSNNs. While most prior works [6,7,8] were limited to single recording sessions, our model is
1Our code is available at github.com/EPFL-LCN/pub-sourmpis2023-neurips (code DOI hosted by zenodo).
constrained to spike recordings from 28 sessions covering six cortical areas. The resulting spike-based model enables a data-constrained simulation of a cortical sensory-motor pathway (from somatosensory to motor cortices responsible for the whisker, jaw and tongue movements). As far as we know, our model is the first RSNN model constrained to multi-session recordings with automatic differentiation methods for spiking neural networks [8, 9, 10].
Second contribution, using this model we aim to pinpoint the circuitry that induces variability in behavior (asking for instance what circuit triggers a loss of attention). Towards this goal, we identify an unsolved problem: "how do we enforce the generation of a realistic distribution of neural activity and behavior?" To do this, the model is fitted jointly to the recordings of spiking activity and movements to generate a realistic trial-to-trial co-variability between them. Our technical innovation is to define a supervised learning loss function to match the recorded and generated variability. Concretely the trial matching loss function is the distance between modeled and recorded distributions of neural activity and movements. It relies on recent advances in the field of optimal transport [11], [12], [13] providing notions of distances between distributions. In our data-constrained RSNN, trial matching enables the recovery of the main modes of trial-to-trial variability which includes the neural activity related to instructed behavior (e.g. miss versus hit trials) and uninstructed behavior like spontaneous movements.
Related work While there is a long tradition of data fitting using the leaky integrate and fire (LIF) model, spike response models [14] or generalized linear models (GLM) [6], most of these models were used to simulate single neuron dynamics [15, 16] or small networks with dozens of neurons recorded in the retina and other brain areas [6,7,8,17]. A major drawback of those fitting algorithms was the limitation to a single recording session. Beyond this, researchers have shown that FORCE methods [18] could be used to fit up to 13 sessions with a large RSNN [17, 19, 20]. But in contrast with back-propagation through time (BPTT) in RSNNs [9, 10, 21], FORCE is tied to the theory of recursive least squares making it harder to combine with deep learning technology or arbitrary loss functions. We know only one other study where BPTT is used to constrain RSNN to spike-train recordings [8] but this study was limited to a single recording session.
Regarding generative models capturing trial-to-trial variability in neural data, many methods rely on trial-specific latent variables [22, 23, 24, 25, 26]. This is often formalized by abstracting away the physical interpretation of these latent variables using deep neural networks (e.g. see LFADS [22] or spike-GAN [27]) but our goal is here to model the interpretable mechanisms that can generate the recorded data. There are hypothetical implementations of latent variables in RSNNs, most commonly, latent variables can be represented as the activity of mesoscopic populations of neurons [25], or linear combinations of the neuronal activity [26, 28, 29]. These two models assume respectively an implicit grouping of the neurons [25] or a low-rank connectivity matrix [26, 28, 29]. Here, we want to avoid making any structural hypothesis of this type a priori. We assume instead that the variability is sourced by unstructured noise (Gaussian current or Poisson inputs) and optimize the network parameters to transform it into a structured trial-to-trial variability (e.g. multi-modal distribution of hit versus miss trials). The optimization therefore decides what is the network mechanism that best explains the trial-to-trial variability observed in the data.
This hypothesis-free approach is made possible by the trial matching method presented here. This method is complementary to previous optimization methods for generative models in neuroscience. Many studies targeted solely trial-averaged statistics and ignored single-trial activity, for instance methods using the FORCE algorithm [17, 19, 20, 30, 31], RSNN methods using back-propagation through time [8] and multiple techniques using (non-interpretable) deep generative models [32]. There exist other objective functions which can constrain the trial-to-trial variability in the data, namely: the maximum likelihood principle [6, 15] or spike-GANs [27, 33]. We illustrate however in the discussion section why these two alternatives are not a straightforward replacement for the trial matching loss function with our interpretable RSNN generator.
2 Large data-constrained Recurrent Spiking Neural Network (RSNN)
This paper aims to model the large-scale electrophysiology recordings from [2], where they recorded 6,182 units from 12 areas across 18 mice. All animals in this dataset were trained to perform the whisker tactile detection task described in Figure 1; in  $50\%$  of the trials (the GO trials), a whisker is
2of which we use 3,810 units from 6 cortical areas in this study
Figure 1: Modeling trial-variability in electrophysiological recordings. A. During a delayed whisker detection task, the mouse should report the sensation of a whisker stimulation by licking to obtain a water reward. Neural activity and behavior of the mouse are recorded simultaneously. B. A recurrent spiking neural network (RSNN) of the sensorimotor pathway receives synaptic input modeling the sensory stimulation and produces the jaw movement as a behavioral output. C. The stimuli and the licking action of the mouse organize the trials into four types (hit, miss, false alarm, and correct rejection). Our goal is to build a model with realistic neural and behavioral variability. Panels A and C are adapted from [34].
deflected and after a 1 s delay period an auditory cue indicates water availability if the mouse licks, whereas in the other  $50\%$  of trials (the No-Go trials), there is no whisker deflection and licking after the auditory cue is not rewarded. Throughout the paper we attempt to create a data-constrained model of the six areas that we considered to play a major role in this behavioral task: the primary and secondary whisker somatosensory cortices (wS1, wS2), whisker motor cortices (wM1, wM2), the primary tongue-jaw motor cortex (tjM1) and the anterior lateral motor cortex (ALM), also known as tjM2 (see Figure 1A and 3A). While we focus on this dataset, the method described below aims to be broadly applicable to most contemporary large-scale electrophysiological recordings.
We built a spiking data-constrained model that simulates explicitly a cortical neural network at multiple scales. At the single-cell level, each neuron is either excitatory or inhibitory (the output weights have only positive or negative signs respectively), follows leaky-integrate and fire (LIF) dynamics, and transmits information in the form of spikes with synaptic delays ranging from 2 to  $4\mathrm{ms}$ . At a cortical level, we model six brain areas of the sensory-motor pathway where each area consists of 250 recurrently connected neurons (200 excitatory and 50 inhibitory) as shown in Figure 3A, such that only excitatory neurons project to other areas. Since the jaw movement defines the behavioral output in this task, we also model how the tongue-jaw motor cortices (tjM1, ALM) drive the jaw movements.
Mathematically, we model the spikes  $z_{j,k}^{t}$  of the neuron  $j$  at time  $t$  in the trial  $k$  as a binary number. The spiking dynamics are then driven by the integration of the somatic currents  $I_{j,k}^{t}$  into the membrane voltage  $v_{j,k}^{t}$ , by integrating LIF dynamics with a discrete time step  $\delta_t = 2$  ms. The jaw movement  $y_{k}^{t}$  is simulated with a leaky integrator driven by the activity of tjM1 and ALM neurons, followed by an exponential non-linearity. This can be summarized with the following equations, the trial index  $k$  is omitted for simplicity:
$$
v _ {j} ^ {t} = \alpha_ {j} v _ {j} ^ {t - 1} + (1 - \alpha_ {j}) I _ {j} ^ {t} - v _ {\text {t h r}, j} z _ {j} ^ {t - 1} + \xi_ {j} ^ {t} \tag {1}
$$
$$
I _ {j} ^ {t} = \sum_ {d, i} W _ {i j} ^ {d} z _ {i} ^ {t - d} + \sum_ {d, i} W _ {i j} ^ {\text {i n}, d} x _ {i} ^ {t - d} \tag {2}
$$
$$
\tilde {y} ^ {t} = \alpha_ {j a w} \tilde {y} ^ {t - 1} + (1 - \alpha_ {j a w}) \sum_ {i} W _ {i} ^ {\mathrm {j a w}} z _ {i} ^ {t} \tag {3}
$$
$$
y ^ {t} = \exp (\tilde {y} ^ {t}) + b \tag {4}
$$
where  $W_{ij}^{d}, W_{ij}^{\mathrm{in},d}, W_{i}^{\mathrm{jaw}}, v_{\mathrm{thr},j}$ , and  $b$  are model parameters. The membrane time constants  $\tau_{m} = 30$  ms for excitatory and  $\tau_{m} = 10$  ms for inhibitory neurons define  $\alpha_{j} = \exp \left(-\frac{\delta t}{\tau_{m,j}}\right)$  and  $\tau_{jaw} = 50$  ms define similarly  $\alpha_{jaw}$  which controls the velocity of integration of the membrane voltage and the jaw movement. To implement a soft threshold crossing condition, the spikes inside the recurrent network are sampled with a Bernoulli distribution  $z_{j}^{t} \sim \mathcal{B}(\sigma (\frac{v_{j}^{t} - v_{\mathrm{thr},j}}{v_{0}}))$ , where  $v_{0}$  is the temperature of the sigmoid  $(\sigma)$ . The spike trains  $x_{i}^{t}$  model the thalamic inputs as simple Poisson neurons producing spikes randomly with a firing probability of  $5 \mathrm{~Hz}$  and increasing their firing rate when a whisker or auditory stimulation is present (see Appendix A). The last noise source  $\xi_{j}^{t}$  is an instantaneous Gaussian noise  $\xi_{j}^{t}$  of standard deviation  $\beta v_{\mathrm{thr}}\sqrt{\delta t}$  modeling random inputs from other areas (  $\beta$  is a model parameter that is kept constant over time).
Session stitching An important aspect of our fitting method is to leverage a dataset of electrophysiological recordings with many sessions. To constrain the neurons in the model to the data, we uniquely assign each neuron in the model to a single neuron from the recordings as illustrated in Figure 2A and 3A. There are 3,810 recorded neurons in the dataset but we take a subset of 250 neurons per area to keep the same number of excitatory and inhibitory neurons in all areas (we were limited by the 290 regular spiking neurons recorded in wS1, and computation resources, mainly, the RAM of the GPU), we ignore the remaining data when fitting the model. This bijective mapping between neurons in the data and the model is fixed throughout the analysis and defines the area and cell types of the neurons in the model. The area is inferred from the location of the corresponding neuron in the dataset and the cell type is inferred from the action potential waveform of the cell (for simplicity, fast-spiking neurons are considered to be inhibitory and regular-spiking neurons as excitatory). Given this assignment, we denote  $z_{j}^{\mathcal{D}}$  as the spike train of neuron  $j$  in the dataset and  $z_{j}$  as the spike train of the corresponding neuron in the model; in general, an upper script  $\mathcal{D}$  always refers to the recorded data. With this notation the tensor of recorded data  $z_{j}^{\mathcal{D}}$  only exists if neuron  $j$  is recorded during session  $S$  (denoted  $j \in S$ ), and the activity of neurons not belonging in session  $S$  ( $j \notin S$ ) is considered to be unknown. Consequently, two neurons  $i$  and  $j$  might be synaptically connected in the model although they correspond to neurons recorded in separate sessions. This choice is intended to model network sizes beyond what can be recorded during a single session. Our network is therefore a "collage" of multiple sessions stitched together as illustrated in Figure 2A and 3A. This network is then constrained to the recorded data by optimizing the parameters to minimize the loss functions defined in the following section. Altogether, when modeling the dataset from Esmaeili and colleagues [2], the network consists of 1,500 neurons where each neuron is assigned to one neuron recorded in one of the 28 different recording sessions. Since multiple sessions are typically coming from different animals, we model a "template mouse brain" which is not meant to reflect subject-to-subject differences. In our view, this is coherent with the classical assumption that is made when recording this type of dataset: the features of the circuit being studied are shared among subjects that participate in the experiment. To study differences between sub-groups of subjects (e.g. age group) one could separate the data into two session groups and compare two fitted models. This is beyond the scope of this paper and below we fit a single model to the entire dataset.
3 Fitting single-trial variability with the trial matching loss function
We fit the network to the recordings with gradient descent and we rely on surrogate gradients to extend back-propagation to RSNNs [9, 10]. At each iteration until convergence, we simulate a batch of  $K = 150$  statistically independent trials. We measure some trial-average and single-trial statistics of the simulated and recorded activity, calculate a loss function, and minimize it with respect to all the trainable parameters of the model via gradient descent and automatic differentiation. This protocol is sometimes referred to as a sample-and-measure method [8] as opposed to the likelihood optimization in GLMs where the network trajectory is clamped to the recorded data during optimization [6]. For the sound evaluation of the model, we separate the recorded trials of every session into a training set (75% of the total trials) and a testing set (25%), that is never used in the fitting procedure. The
split is done in a stratified manner so both sets have the same distribution of trial types. The full optimization lasts for approximately one to three days on a GPU A100-SXM4-40GB.
Trial-average loss We consider the trial-averaged activity over time of each neuron  $j$  from every session  $\mathcal{T}_{\mathrm{neuron},j}$ , sometimes referred also as a neuron's peristimulus time histogram (PSTH). This is defined by  $\mathcal{T}_{\mathrm{neuron},j}(z) = \frac{1}{K}\sum_{k}z_{j,k}*f$  where  $f$  is a rolling average filter with a window of  $12\mathrm{ms}$ , and  $K$  is the number of trials in a batch of spike trains  $z$ . The statistics  $\mathcal{T}_{\mathrm{neuron},j}(z^{\mathcal{D}})$  are computed similarly on the  $K^{\mathcal{D}}$  trials recorded during the session corresponding to neuron  $j$ . We denote the statistics  $\mathcal{T}_{\mathrm{neuron},j}'$  after normalizing each neuron's trial-averaged activity, and we define the trial-averaged loss function as follows:
$$
\mathcal {L} _ {\text {n e u r o n}} = \sum_ {j} \| \mathcal {T} _ {\text {n e u r o n}, j} ^ {\prime} (\boldsymbol {z}) - \mathcal {T} _ {\text {n e u r o n}, j} ^ {\prime} \left(\boldsymbol {z} ^ {\mathcal {D}}\right) \| ^ {2}. \tag {5}
$$
It is expected from [8] that minimizing this loss function alone generates realistic trial-averaged statistics like the average neuron firing rate.
Trial matching loss: fitting trial-to-trial variability Going beyond trial-averaged statistics, we now describe the trial matching loss function to capture the main modes of trial-specific activity. From the previous neuroscience study [2], it appears that population activity in well-chosen areas is characteristic of the trial-specific variability. For instance, intense jaw movements are preceded by increased activity in the tongue-jaw motor cortices, and hit trials are characterized by a secondary transient appearing in the sensory cortices a hundred milliseconds after a whisker stimulation. To define statistics for a single trial  $k$  which can capture these features we denote the population-averaged firing rate of an area  $A$  as  $\mathcal{T}_{A,k}^{S}(z) = \frac{1}{|A\cap S|}\sum_{j\in A\cap S}(z_{j,k}*f)$  where  $|A\cap S|$  is the number of neurons in area  $A$  recorded in session  $S$ , the smoothing filter  $f$  has a window size of  $48~\mathrm{ms}$  and the resulting signal is downsampled to avoid unnecessary redundancy. We write  $\mathcal{T}_{A,k}^{\prime S}$  when each time bin is normalized to mean 0 and standard deviation 1 using the training data from session  $S$ . We, therefore, use  $\mathcal{T}_{A,k}^{\prime S}(z)$  as a feature vector to characterize the recorded area  $A$  during each trial  $k$  of session  $S$ , and we build a richer feature vector including all recorded areas and the behavior with the concatenation  $\mathcal{T}_{\mathrm{trial},k}^{\prime S} = (\mathcal{T}_{A1,k}^{\prime S},\mathcal{T}_{A2,k}^{\prime S},\boldsymbol {y}_k*\boldsymbol {f})$ , where  $A1$  and  $A2$  are the two areas recorded in session  $S$ , and  $\boldsymbol {y}_k$  is the generated jaw trace from trial  $k$ . The same feature vector can be computed with the simulated jaw and simulated neural activity by sub-selecting the neurons from areas  $A1$  and  $A2$  from session  $S$  to compare generated and recorded data.
The challenging part is now to define the distance between the recorded statistics  $\mathcal{T}_{\mathrm{trial},k}^{\prime S}(z^{\mathcal{D}})$  and the generated ones  $\mathcal{T}_{\mathrm{trial},k}^{\prime S}(z)$ . Common choices of distances like the mean square error are not appropriate to compare distributions. This is because the order of trials in a batch of generated/recorded trials has no meaning a priori: there is no reason for the random noise of the first generated trial to correspond to the first recorded trial – rather we want to compare unordered sets of trials and penalize if any generated trial is very far from any recorded trial.
Formalizing this mathematically we consider a distance between distributions inspired by the optimal transport literature. Since the plain mean-squared error cannot be used, we use the mean-squared error of the optimal assignment between pairs of recorded and generated trials: we select randomly  $K' = \min(K, K^{\mathcal{D}})$  generated and recorded trials ( $K$  and  $K^{\mathcal{D}}$  are respectively the number of generated and recorded trials in one session), and this optimal assignment is formalized by the integer permutation  $\pi : \{1, \ldots, K'\} \to \{1, \ldots, K'\}$ . Then using the feature vector  $\mathcal{T}_{\mathrm{trial},k}'^S$  for any trial  $k$ , we define the hard trial matching loss function as follows by the sum  $\mathcal{L}_{\mathrm{trial}} = \sum_{S} \mathcal{L}_{\mathrm{trial}}^S$  with:
$$
\mathcal {L} _ {\text {t r i a l}} ^ {S} = \min  _ {\pi} \sum_ {k} \left\| \mathcal {T} _ {\text {t r i a l}, k} ^ {\prime S} (\boldsymbol {z}) - \mathcal {T} _ {\text {t r i a l}, \pi (k)} ^ {\prime S} \left(\boldsymbol {z} ^ {\mathcal {D}}\right) \right\| ^ {2}. \tag {6}
$$
We compute this loss function identically on all the recorded sessions and take the summed gradients to update the parameters. Each evaluation of this loss function involves the computation of the optimal trial assignment  $\pi$  which can be computed with the Hungarian algorithm [35] (see linear_sum_assignment for an implementation in scipy). This is not the only way to define a distance between distributions of statistics  $\mathcal{T}_{\mathrm{trial},k}^{\prime S}$ . In fact, this choice poses a potential problem because the optimization over  $\pi$  is a discrete optimization problem, so we have to assume that  $\pi$  is a constant with respect to the parameters when computing the loss gradients. We also tested alternative
Figure 2: Artificial Dataset. A. Session stitching: every neuron from the recordings is uniquely mapped to a neuron from our model. For example, an excitatory neuron from our model that belongs in the putative A1 is mapped to an excitatory neuron "recorded" in A1. In our network, we constrain the connectivity so that only excitatory neurons can project across different brain regions. B. The first area (A1) responds equally in a hit-like and a miss-like trial, while the second area (A2) responds only in hit-like trials. A model that does not use trial matching cannot capture the bimodal distribution of A2. C. Distribution of max firing rate of the population average of A2 from each trial. Only the trial matching algorithms retrieve the bimodal behavior of A2.
choices relying on a relaxation of the hard assignment into a smooth and differentiable bi-stochastic matrix. This results in the soft trial matching loss function, which replaces the optimization over  $\pi$  by the Sinkhorn divergence [12, 13] (see the geomloss package for implementation in pytorch [13]). In practice, to minimize both  $\mathcal{L}_{\mathrm{trial}}$  (either the soft or hard version) and  $\mathcal{L}_{\mathrm{neuron}}$  simultaneously we optimize them in an additive fashion with a parameter-free multi-task method from deep learning which re-weights the two loss functions to ensure that their gradients have comparable scales (see [36] for a similar implementation).