1 Introduction↩︎

Systems neuroscience experiments are growing in complexity with a cascade of technical advances. In addition to recording neuronal activity from hundreds to thousands of neurons in parallel from multiple brain areas [1]–[4], experimenters are also simultaneously acquiring behavioral data including eye positions and body movements [5], [6]. Neuroscientists seek models for these rich datasets to obtain insights into how neural circuitry is involved with behavior. However, complex datasets and experimental designs present challenges for model development, particularly in incorporating the various modalities, and subsequently creating models for a host of different applications.

In recent years, deep neural networks (DNNs) have demonstrated their potential in modeling neural activity and circuitry [7]–[9]. Brain-inspired DNNs have exhibited responses similar to visually evoked neural responses in visual cortical areas [7], [10]–[12]. These models were trained in a goal-driven approach, such as the classification of images. Then intermediate representations in the model were compared to observed neuronal activities in biological brains. In other work, autoencoder-based latent variable models have been applied to the analysis of neuronal activities, with the assumption that the single-trial spiking activity of a behavior task depends on underlying dynamics [13]–[16]. These approaches can provide insight into principles of neural circuitry, but often entail inductive biases that are not always ideal. General purpose, multimodal models can provide a complementary approach.

To develop models for correlating neuron activities, sensory input, and behavior, we sought a powerful multimodal architecture and training objective. The transformer architecture has demonstrated its flexibility in modeling data from various domains, finding widespread adoption in applications including vision, language, and sound [17]–[22]. In particular, transformer-based language-vision models exhibit high performance for learning multimodal representations for text and images, enabling text-to-image generation [20], [23]. Key to this is that a) the transformer is scalable; i.e it does not suffer from the vanishing gradient problem [24] b) it makes few assumptions about the input data (minimal inductive bias) [25], [26] and c) it can contextualize information within long-duration multimodal contexts. Here we present a transformer-based model for multimodal datasets in systems neuroscience, including neuronal activity, stimuli, and behavior. We developed a generative model which firstly aligns and subsequently fuses neuronal activities, input signals (e.g., visual stimuli) and behavior.

Our key contributions are as follows: a) Introduce an approach that reframes the population response inference problem as a causally-masked, spatiotemporal autoregressive generation problem. b) Show how to employ a generative pretraining paradigm for neural data. c) Investigate the potential of multimodal, multitask, generative transformer models to analyze and parse neuroscience data, and obtain insights.

2 Related Work↩︎

2.0.0.1 Representational and Functional Similarities between DNNs and Mammalian Brains

There are similarities between the hierarchical representations of convolutional DNNs and the visual and inferior-temporal cortices [7], [8], [10]–[12]. More recently, parallel-path networks trained using a contrastive-predictive objective were shown to separate into pathways that mirrored the functional specializations of the dorsal and ventral streams in mouse visual cortex [27]. Similarly, there are parallels between human speech and language models [28], [29]. Lastly, reports have shown the plausibility of transformers replicating specific neural functions and circuits [30], [31]. In our work, establishing a similarity between a DNN and the brain is not our focus. Instead, we are focused on developing an analytical and functional tool. This tool, Neuroformer, can exhibit similarities in representational and functional characteristics with biological neural circuitry, and we explore that here.

2.0.0.2 Learning Latent Dynamics of Population Activity

Variational auto-encoders have been used to infer latent dynamics embedded in complex, high-dimensional neural data [13], [32]. Recently, a transformer architecture has been used in place of RNNs for the same application [14], and another project has trained a latent model using a contrastive objective [16]. These empirical findings show that this approach avoids over-fitting and outperforms previous methods. In our work, we integrate these paradigms, by employing a contrastive objective at the front end of a generative model.

2.0.0.3 Large Multimodal and Multitask Models

The transformer has proven effective at processing a variety of modalities including vision, languange, and sound [17]–[19]. Subsequently, the field introduced multimodal models. For example, models for interactions between images and text [33]–[35]. Some models attempt to align the modalities with separate feature encoders using contrastive learning approaches [20], [21], and alignment and fusion [22]. We experimentally validate these methods on neural data, by using our own generative and contrastive learning objective. Furthermore, architecture innovations have resulted in models that are agnostic to the input modalities and can be applied to a large variety of tasks [25], [26]. Lastly, it has been recently observed, that large models pretrained on massive datasets of text exhibit emergent properties. This has been observed in Large Language Models (LLMs) as in [36]–[40], but also in scientific domains [27], [41]. We take inspiration from these lines of work to build a framework for modelling neural data that can efficiently scale to large inputs and an arbitrary number of modalities.

3 Model↩︎

3.0.0.1 Workflow

As input, our model takes action potential data, i.e., spike trains, from multiple neurons recorded or imaged simultaneously, and breaks them up into the Current State \(\;I_{c}\), and past state \(\;I_{p}\)according to specified temporal context windows (\(w_c, w_p\)) (Fig. 1, blue and purple, respectively). We do this to discretize them and precisely align the neural representations with features from other modalities. Model inputs, including neural, visual, and other features, are optionally processed by feature encoders. The modalities are first aligned via a multimodal contrastive module [20], [42], [43] which aims to learn matching representations for the neural features and the features from other modalities. Subsequently, a cross-modal transformer fuses the current neural response \(\;I_{c}\)with all other features. We then decode our fused representation using a causally-masked transformer decoder. All training is self-supervised and requires no data labelling.

3.0.0.2 Feature Backbones

We assign an ID (number) to each neuron in our dataset and model it as a token. At each step, a learnable look-up table projects each neuron spike contained in our Current and Past States onto an embedding space \(E\), resulting in vectors (\(T_{c}, E\)) and (\(T_{p}, E\)), where \(T_{c}, T{p}\) are the corresponding state’s sequence length (number of spikes + padding). For video frames, we optionally use a single layer of 3D Convolutions [44] before splitting them into patches of size (\(\;T_{f}, H, W, E\)), as is standard for visual transformers (Fig. 1, Feature Encoding, green) [45]–[47]. We then feed these patches into the cross-modal layers.

3.0.0.3 Multimodal Contrastive Representation Learning

Drawing from recent developments in vision and language representation learning [20], [22], [48], our approach leverages a contrastive learning objective before fusing our feature sets (Fig. 1, Contrastive Matching). Our architecture allows alignment of an arbitrary number of modalities; however, here we focus on the three-modal case. We use Visual features \(F_{p, c} \sim (B, T, H, W, E)\), Current State neural embeddings \(I_{c} \sim (B, T_{c}, E)\), and Behavior features \(A_{c} \sim (B, T_{a}, E)\) as our modality representatives, where \(B =\) Batch Size and \(T =\) sequence length. These are projected to a latent dimension \(d\) using three linear layers \(\boldsymbol{g_{f}}, \boldsymbol{g_{i}}, \boldsymbol{g_{a}}\). We calculate pairwise similarities between the modalities, yielding three matrices of dimensions \((B_{f}, B_{i}), (B_{i}, B_{a}), (B_{a}, B_{f})\) where \(B_f = B_i = B_a = B\). The diagonal entries correspond to ground-truth pair similarities. These scores are then softmax-normalized and used to compute the contrastive loss via cross-entropy H, where the probability \(\boldsymbol{p}\) should be 0 and 1 for the negative and positive pairs \(\boldsymbol{y}\), respectively. We generalize the computations with pair notation (m,n) where m and n are any of the modalities (either F, I, or A):

\[\begin{align} s(m,n) &= g_m(m_{p,c})^Tg_n(n_{c}) & p_{k}^{mn} &= \frac{\exp(s(m, n_{k})/\tau)}{\sum_{k=1}^{K} \exp(s(m, n_{k})/\tau)} \end{align}\]

We then apply these equations to each pair (F,I), (I,A), (A,F), and compute the contrastive loss:

\[\begin{align} L_{vna} = \frac{1}{3} \displaystyle \mathop{\mathrm{\mathbb{E}}}_{(F, I, A) \in d} [H(\boldsymbol{y}^{fi}(F),\boldsymbol{p}^{fi}(F)) + H(\boldsymbol{y}^{ia}(I),\boldsymbol{p}^{ia}(I)) + H(\boldsymbol{y}^{af}(A),\boldsymbol{p}^{af}(A))] \end{align}\]

By explicitly imposing alignment, we encourage the model to uncover the representational commonalities between biologically relevant features [16]. This is particularly helpful in low-data settings, mitigating the over-fitting that can occur with variable neural signals. We observe consistent performance improvement upon adding the contrastive objective (see Fig. 6, 10.)

3.0.0.4 Feature Fusion

Following alignment, the masked Current State, \(I_{c}\) is fused with all other features, in our case the visual \(F\) and neural Past State features, \(I_{p}\) by cascading cross-attention modules. This stage of the model is similar to the [25], [26], where cross-attention between a smaller latent array (\(T_{c}, E\)), in our case the Current State and larger byte arrays (\(T_{f},s E\)), in our case the Past State neural features, is iteratively carried out. This architecture decision has two main advantages: a) It introduces a bottleneck imposed by the smaller \(I_{c}\) array, of length size \(T_{c}\) that upon operating with a larger feature array of length \(T_{f}\), reduces the complexity of attention from \(\mathcal{O}(T_{f}^{2})\) to \(\mathcal{O}(T_{c}T_{f})\). When \(T_{c}\) is much less than \(T_{f}\) (expressed mathematically as \(\lim_{{T_{c} \to 0, T_{f} \to \infty}} T_{c}/T_{f} = 0\)), the complexity scales linearly with \(T_{f}\). This enables the use of finer-grained features, like large-resolution visual features or in our case raw video frames, which are essential for modeling lower-level visual areas. b) The iterative nature of the algorithm enables us to fuse an arbitrary number of modalities with our Current State.

3.0.0.5 Causal Spike Modelling

After it has been fused with all other modalities, the resulting Current State is decoded using a causally masked transformer [18]. At each prediction step, the model predicts both which neuron will fire, and within which sub-interval (\(\boldsymbol{dt(n)} = n \in w_{c}\)), using two linear projections from the last self-attention layer (Fig. 1, AR Decoding) mapping to the corresponding (\(\boldsymbol{p_{i}}, \boldsymbol{p_{dt}}\)) label distributions. Notice that two Spikes can occur at the same time, and therefore by predicting bins within which neurons fire, we can sequentially predict spikes that occur simultaneously. We optimize for this objective, by minimizing the cross entropy between the predicted and ground truth distributions:

3.0.0.6 Loss Function

The model is jointly optimized end-to-end using the contrastive objective \(L_{vnc}\), and cross-entropy loss for both \(I_{c}\) and \(dt_{c}\), (\(L_{ce(I)}, L_{ce(dt)}\)). The resulting loss is a weighted average of the three losses: \[\begin{align} L = (\gamma)L_{vnc} + (\mu)L_{ce(I)} &+ (1 - \gamma - \mu)L_{ce(dt)} \end{align}\]

3.0.0.7 Inference - Generating Simulations

At inference time, the model autoregressively predicts the distributions of the \(n\)th step, (\(I_{c}(n), dt_{c}(n)\)), outputting an end-of-state [EOS] token to indicate it has completed the Current Target State \(W_{c}\), in a similar fashion to work in image-text generation [23], (Fig. 1, AR Decoding) . To sample from the respective distributions (\(\boldsymbol{p_{I}}, \boldsymbol{p_{dt}}\)) we use nucleus sampling [49]. The predicted Current Target State neuron IDs are then incorporated into the Past State according to the Past State Context Window \(w_{p}\) together with the predicted sub-intervals \(dt_{c}\), which act as additive temporal embeddings (Fig. 1, dotted arrow from AR Decoding).

3.0.0.8 Finetuning - Transferring to New Tasks

A primary motivation for redefining the spike inference problem in a manner compatible with contemporary pretraining paradigms is the ability to utilize a broad spectrum of techniques developed within this framework. To transfer to a new task, whether it be behavioral prediction or any other downstream task, we adopt a straightforward heuristic. First, if necessary, the data distribution is discretized into \(n\) classes. Then, a mean-pooled representation from our final latent layer is projected onto the appropriate feature space \(p_{v} \in n\). The loss is calculated using the standard multi-class cross-entropy for classification [50]–[52] of mean-squared error for regression.

3.0.0.9 Training

We trained all Neuroformer models with 8 layers and 8 heads for the decoder and each of the cross-attention modules, with an embedding dimension of \(256\) and an \(80/20\) train/test split. V1+AL and Visnav models equated to a total of \(40,000,000\) and \(100,000,000\) parameters respectively. We emprically observed an increase in performance up until the specified parameter sizes. We used AdamW optimizer [53] in PyTorch. Learning rate was 2e-4, with linear warmup and cosyne decay schedules. Batch size was kept at 224 for pretraining and 32 for fine-tuning.

4 Results↩︎

To evaluate performance of Neuroformer, we used both artificial data sets, with known ground truth, and real neural data sets, with authentic variability and real-world relevance. We analyzed the performance in terms of inference and features.

4.1 Data↩︎

4.1.0.1 Simulated Dataset

We simulated a spiking neural network characterized by directional connectivity with three hub neurons using Brian2 simulator [54] (Hub-net simulation). The network has \(N=300\) leaky integrate-and-fire neurons with stochastic current. The membrane potential \((V_{i})\) of each neuron was simulated as a stochastic differential equation: \[\begin{align} \frac{dV_i}{dt} = \frac{I - V_i}{\tau} + \sigma \sqrt{\frac{2}{\tau}} \cdot x_i \end{align}\] The membrane time constant was set to \(\tau=10\)ms as a Gaussian random variable with zero mean and 1 standard deviation. We scaled the random variable by \(\sigma=0.5\). The total dataset equated to 1 million tokens.

4.1.0.2 Two-photon calcium imaging datasets

We assessed the performance of Neuroformer using real neural datasets. These consisted of neuronal activity recorded from awake mice, which were either passively viewing visual stimuli (Passive-Stim data), or actively engaged in a visually-guided navigation task (Visnav data). The activity was recorded using two-photon calcium imaging, [1], [55], [56] and neuronal spikes were deconvolved from calcium traces via the suite2p calcium imaging processing pipeline, [54] complemented by a Bayesian spike inference method [57]. The Passive-Stim data comprised recordings from 386 neurons in the primary visual cortex and a higher visual area (V1 + AL). These neurons responded reliably to drifting grating stimuli and a naturalistic movie. Two Visnav datasets were derived from recordings in the lateral visual cortex (2022 neurons) or L2/3 medial (1905 neurons), spanning V1 and multiple higher visual areas. During these recordings, the mice were engaged in a visually-guided navigation task within a virtual environment, with their movements tracked by a ball system [58] that controlled the first-person view within the virtual environment. The Visnav neurons conveyed information about both the visual input and the speed of the animal’s movements. The resulting dataset sizes equated to \(80\times10^{3}\) , \(800\times10^{3}\) and \(1\times10^{6}\) tokens.

4.2 Attention Reveals the Functional Connectivity of Simulated Neuronal Circuits↩︎

To test our hypothesis that attention can provide insights into neural circuitry, we first trained the Neuroformer on a simulated dataset with known ground-truth connectivity. This dataset is unimodal, so we simply omitted the visual cross-attention module, essentially reducing to a vanilla GPT architecture. After training, we ran masked inference on the data, aggregating the attention values that are assigned to each neuron, and counting how many times neurons attend to each other during this process. This results to two square matrices \(N_{a}\), \(N{z}\) of size (\(N, N\)), where \(N\) is equal to our neuronal population’s size \((300)\). By dividing (\(N_{a}\)) by (\(N{z}\)) we compute the average attention \(a_{av.}(i, i)\) that neurons assign to one another. Our results recovered the hub connectivity structure of the ground truth neuronal circuit (\(\;I \in 50, 150, 200\)) using only the spiking activity (Fig. 2). This attention analysis identified the neurons that were driving circuit activity. Compared to a simple correlation matrix, the attention analysis matrix provides directional, potentially causal, information.

4.3 Realistic Autoregressive Simulation of Multimodal Neuronal Dataset↩︎

To validate Neuroformer as a viable analysis tool for neuronal population datasets we tested its ability to model real neuronal data. We trained Neuroformer to predict visually evoked neuronal activity in a population of 386 neurons, recorded using two-photon calcium imaging in a mouse. The neurons were recorded simultaneously, and spanned two visual cortical areas: primary visual cortex (V1) and the anterolateral (AL) higher visual area. After training, Neuroformer is able to autoregressively generate whole-trial simulations (96 seconds) of the neuronal population that closely resemble held-out trials (Fig. 3).

4.3.1 Neuroformer Outperforms GLM↩︎

The Generalized Linear Model (GLM) [59], [60] is a versatile and widely used statistical modeling framework that generalizes linear regression to accommodate non-normal response variables and non-linear relationships. It has been widely used by neuroscientists as a performant tool for stimulus-response population decoding. Although useful, it has some significant disadvantages, in that it requires heavy inductive bias (linear kernel, non-linearity, Poisson spiking). This limits both its generality and flexibility in simultaneously integrating many modalities. We compared the predicted, trial-averaged responses from a GLM to Neuroformer on the V1+AL dataset and found that our model’s predictions where more closely correlated with the ground-truth (Fig. 3c). Our results revealed that our model could better capture the variability across the diverse set of stimuli in the dataset. We empirically observed that a single GLM could perform well on either the diffraction gratings or naturalistic movies, but not both at the same time.

4.3.2 Biologically Relevant Features Emerge from Stimulus-Response Cross-attention↩︎

Again, we turned to the attention parameters in the Neuroformer, as we had for the hub neuron analysis. We plotted the spatial locations on visual stimulus frames where cross-attention was allocated, i.e., a sort of attention map. These maps exhibited localized features and structure. These maps are related to the concept of receptive fields, but are specific to the full context of ongoing visual stimuli and neuronal activity. Thus, they can reveal novel relationships and insights into potentially causal relationships among neuronal activity and stimuli (Fig. 3).

4.3.3 Highly Accurate and Few-shot Prediction of mouse Behavior↩︎

In this section, we introduce the notion of finetuning [50] for neuronal data. By reframing the neuronal modelling problem to be compatible with modern NLP approaches, we can leverage similar techniques to solve tasks that can benefit from neuronal pretraining. One such task is predicting the behavior of a mouse from its neuronal responses. This is a common thing to do in neuroscience analysis, where researchers want to investigate the relationship between neuronal activity and behavior. Furthermore, such paradigms could be useful in brain-computer interface applications. For this analysis, we used neural and behavior data from a mouse running in a virtual reality setup. Neuroformer can support both classification with a mutli-class cross-entropy objective or regression using mean-squared error, which can be chosen according to the nature of the task. Here we report results using regression (Fig. 4, Table 1****) and have included classification and regression inference results for all models in the appendix (Figs 11, 12). Similarly to how multitask models in the field of vision and language work, finetuning a model for a new task involves projecting the features of the last layer onto the new feature space. [50]

r0.5

Despite the inherent stochasticity in the speed of the mouse, the predictions of our model when trained on the full dataset are highly correlated with the held-out ground-truth. Additionally, Neuroformer significantly outperformed conventional approaches, achieving pearson \(r=0.97\), with the held-out ground truth for the respective lateral and medial datasets (Table 1). To determine the effectiveness of the masked neuronal prediction pretraining we compared the performance of pretrained vs. randomly initialized models on 10% and 1% of the finetuning data. The pretrained few-shot models quickly approached the performance of the model trained on the full dataset (Fig. 5b,c). Crucially, we observed that the pretrained model that was finetuned on 1% of the data outperformed the randomly-initialized model that was finetuned on 10% of the data. This suggests that pretraining enables the model to begin learning representations that go beyond the optimization objective as has been observed in generative models within natural language generation. [61], [62]

5 Ablations↩︎

To explore the impact of each component of Neuroformer, we implemented a range of model versions. Starting from just decoding from the Current State, we progressively incorporated Past State, Video, Behavioral modalities, and Alignment using contrastive learning. With each addition, we noted an enhancement in the model’s ability to generate realistic neuronal responses (Fig. 6, 10).

Of particular interest, was the Lateral dataset where the behavioral modality we incorporated for the ablation study was eye position (see Fig. 13 for predictions). We hypothesized that by learning to extrapolate from the eye position to the visual field of the mouse, the model could improve its performance, which is what we observed. While more investigation needs to go into understanding exactly how Neuroformer incorporated eye position within this context, this result highlights the potential of our approach.

By observing the fluctuation in performance as components within the model are selectively incorporated, scientists can examine hypotheses about the relevance of various input modalities or phenomenological behaviors to different brain regions. Ultimately, the development of modular and performant foundation models of the brain will create new avenues for the scientific exploration of the principles that underpin neuronal function and intelligence.

6 Conclusion↩︎

We have developed a new method to enable generative pretraining of multimodal neural data, and have demonstrated its utility accross four novel neural datasets with diverse structure. Our initial observations showed that pretraining not only adopted representations that mirrored the causal patterns found in neuronal circuitry, but was also able to effectively enhance performance in downstream applications, enabling rapid few-shot adaptation to behavior-related tasks.

Our model outperformed conventional approaches, while in comparison incorporating minimal inductive biases and retaining an identical architecture across datasets and experiments. Overall, the model was applied to 4 modalities (neural data, video, speed, eye position), and 3 decoding tasks (neuron ID, neuron dt, speed). The architecture enables easy expansion of the model to more modalities and tasks, making it compatible with almost any type of systems neuroscience dataset and decoding objective.

Concluding, the Neuroformer enables systems neuroscience to leverage the full suite of tools in generative, multitask and multi-modal transformers. It paves the way for scaling training of neuroscience models and potentially connecting them with Large Language Models. This advancement could foster models with new emergent capabilities, benefiting not just neural data analysis but also various other applications. Our work sets the stage for promising research within this domain. ²

Supplementary Materials↩︎

7 Limitations↩︎

In preliminary experiments, the Neuroformer was trained on simulated data, revealing potential in identifying hub network connectivity. However, its reliability varied across datasets. For instance, in a multi-channel simulation, neuron attention sometimes accurately identified the correct channel, but at other times was noisy. In another test with neurons responding to a naturalistic video via 3D gabor filters, the method’s reliability in uncovering receptive fields was inconsistent. We think this technique could benefit from more advanced methods for calculating the attention like Grad-CAM.

Our model, the first of its kind using an auto-regressive decoder-based generative pre-training approach, contrasts with encoder-based methods like LFADS or CEBRA. While these methods explicitly exploit the intrinsic dimensionality of neuronal populations, the Neuroformer, inspired by the GPT objective, takes a different approach. Out attempts to use our data with representation learning approaches yielded poor results (see Fig. 10). These methods are more typically applied to motor control tasks for BCI applications. Direct comparisons with methods like LFADS are needed, as they might currently offer better representation learning due to explicit low-dimensional constraints. In contrast, our method is more akin to an in-silico model of the brain.

While transformers, like the Neuroformer, offer advantages like being modality-agnostic and contextualizing information over extended sequences, traditional methods like the GLM and CNNs have their merits. They are grounded in neuroscientific assumptions, offering interpretability. Although the Neuroformer showed promise in decoding tasks, the reasons for its performance advantages remain unclear. Factors like overparameterization, ability to model non-linearities, and effective representation learning might contribute.

Our aim with the Neuroformer was two-fold: 1) to introduce a tool for neuroscience to flexibly incorporate data from increasingly diverse and complex experimental settings and 2) to explore generative pre-training for such data. The true potential of this method might be realized with increased scale, given the vast data available in systems neuroscience. The Neuroformer marks a step towards this, but challenges persist. We’re keen on scaling these methods and finding synergies between large language and neuroscience models.

8 Ethical Statement↩︎

This project involves data from mouse experiments. Our AAALAC (Association for Assessment and Accreditation of Laboratory Animal Care)-accredited institution is committed to the humane care and use of animals in research and education. Our Institutional Animal Care and Use Committee (IACUC) oversees animal care and use, facilities, and procedures. The attending veterinarian and IACUC ensure that research activities involving mice meet the ethical and legal requirements as set by relevant policies and guidelines.

9 Experimental Setup↩︎

10 Hyperparameters↩︎

10.1 Neuroformer↩︎

11 Comparison Model Details↩︎

11.1 Population Response↩︎

11.1.1 GLM↩︎

A generalized linear model (GLM) for neuron data was adapted from [63]. The instant firing rate (\(r_{i}(t)\)) of individual neurons was modeled by a linear-nonlinear cascade, including linear stimulus components (\(k*x\)) and spike history components (\(h*y\)) followed by an exponential nonlinearity. \(\mu\) represents the baseline firing rate of the neuron. The stimulus kernel (\(k\)) and history kernel (\(h\)) are modeled by cosine bases. The parameters are optimized by minimizing the log-likelihood for spike prediction.

\[\begin{align} r_{i}(t) = exp(k*x + h*y+\mu) \end{align}\]

11.2 Behavior Decoding↩︎

11.2.1 MLP↩︎

The MLP used for the behavior comparisons was comprised of 5 stacked MLP layers, with a hidden state of (\(h=1024\)), GeLU non-linearity, dropout \(p=0.01\), layer-normalization between each of the MLPs. No improvement was observed upon adding more layers.

11.2.2 GRU↩︎

The GRU model used for the behavior comparisons had the following architecture: 5 Bi-directional GRU layers with dropout \(p=0.01\), layer-normalization, and tapering t=0.5 for each consecutive hidden state. Tapering allowed us to stack more GRU layers without undergoing collapse. More layers resulted in a decrease in performance and eventual collapse. The first hidden state of the model was initialized as the neural past state as defined in the paper, allowing the model to integrate past information. We employed Neuroformer-style tokenization which shortened the sequence of spikes and resulted in better performance.

12 Block-Sparse Causal Masking↩︎

To precisely align the corresponding neural features with all other modalities, we broke our state features into a Current State \(w_{c}\) and a Past State \(w_{p}\) according to specified temporal windows (see Table 2).

Because of this, we cannot simply model the neural features as a continuous string of spikes (or the equivalent language tokens) as in the case of a standard GPT model [61]. The unique challenge of arranging our data in this way is that each of these intervals contains a different number of spikes. Our situation is similar to [23], where during inference time, the block is autoregressively populated by the predicted tokens. The difference in our case is that each resulting "block" is of a different size. We therefore employ right-masking and padding to enable efficient training and causal inference (see Fig. 8).

Note that by just masking our Current State and projecting from the \(T_{n}\) row of the block during inference time, we are able to fuse with any other non-masked features, while preventing any backflow of information, and retaining the causal aspect of our model.

13 Calculating Connectivity With other Methods↩︎

14 Capturing Trial-to-Trial Variability↩︎

When investigating the performance of the model as we progressively incorporated each module, we examined F1 scores, and the ability of the model to capture the trial-to-trial variation across the data sets.

To do this, we computed the correlation across trial-averaged spikes for each neuron across groups of four trials. We compared a holdout group to the corresponding set generated by the Neuroformer, and to another corresponding group of trials from the data set. We observed a steady improvement in the distribution of correlations when compared to the ground-truth (Fig. 10), with the full model’s predictions closely matching those of the ground-truth sets.

15 Precise Behavior Predictions↩︎

An advantage of the Neuroformer architecture, is that modalities can be flexibly incorporated as an input or label. In the case of labels, variables can be modeled as both classes as in multi-class logistic regression using a cross-entropy objective (Fig. 11), or raw values as a linear regression using mean-squared error (Fig. 12). Our model vastly outperformed other baselines, and performed reliably across datasets and objectives, as opposed to other methods, whose performance fluctuated.

16 multitask Decoding↩︎

In addition to speed, Neuroformer is able to jointly predict many other variables. Here we show that the model could effectively predict the eye position of the mouse - even in the absence of input stimulus. We used a single model checkpoint to decode all three variables. Speed and longitudinal + latitudinal angle for eye position.

17 Multimodal Latent Variables↩︎

The contrastive part of Neuroformer, enables us to align any of the corresponding auxilliary modalities with our neural data, including visual features, behavioral features or reward. Here we show the raw features extracted from the contrastive stage that was trained to align speed with neural features.

No dimensionality reduction method (like T-SNE) was used for this visualization. These are the raw features, extracted from our contrastive alignment module, by setting \(E_{clip} = 3\).

18 Analysis of Intermediate Features and Probing Speed↩︎

18.1 Intermediate Features↩︎

The features represented in Figure 14 were derived from the contrastive alignment module by setting the \(E_{clip}\) to 3, effectively projecting the features to 3 dimensions for alignment during training, as detailed in Hyperparameters 2. Our initial hypothesis posited the need for two dimensions for direction—considering movement occurs on a 2D plane—and one dimension for magnitude. Although we did not expect intermediate representations to fully segregate according to speed without explicit supervision, the process proved beneficial for analytical purposes. As indicated by the toy model experiment in Figure 2, a convergence towards speed-based separation was observed. Additional intermediate representations were extracted from transformer blocks of a model not trained on speed, and a linear regression was performed to map these to speed. These were compared against random arrays of identical shape as a control. supplementary material. As more modalities are fused, the resulting representation becomes more linearly separable in terms of speed.

18.2 Linearly Probing Speed↩︎

19 CEBRA Behavior Results↩︎

We also attempted to compare Neuroformer behavior predictions with the SOTA in representation learning methods, CEBRA [16]. Despite attempting an extensive hyperparameter search accross batch size, temperature, learning rate, offset, and decoding method (KNN and Lasso Reg.) we were unable to get the model to generalize, and typically observed heavy overfitting as seen in Fig. 15.

References↩︎