QuantFormer: Learning to Quantize for Neural Activity Forecasting in Mouse Visual Cortex

TL;DR

QuantFormer addresses the challenge of forecasting neural activity from two-photon calcium imaging by reframing the regression problem as classification through latent space vector quantization. It combines self-supervised pre-training via masked auto-encoding with neuron- and stimulus-specific tokens to scale to arbitrary neuronal populations, and adapts the encoder for two downstream tasks: activation prediction and response forecasting. The approach achieves state-of-the-art performance on the Allen Brain Observatory dataset, demonstrating strong generalization across stimuli and subjects and offering interpretable codes and neuron embeddings that reflect activation dynamics. This work lays a foundation for scalable, real-time neural signal prediction applicable to online interventions in the mouse visual cortex.

Abstract

Understanding complex animal behaviors hinges on deciphering the neural activity patterns within brain circuits, making the ability to forecast neural activity crucial for developing predictive models of brain dynamics. This capability holds immense value for neuroscience, particularly in applications such as real-time optogenetic interventions. While traditional encoding and decoding methods have been used to map external variables to neural activity and vice versa, they focus on interpreting past data. In contrast, neural forecasting aims to predict future neural activity, presenting a unique and challenging task due to the spatiotemporal sparsity and complex dependencies of neural signals. Existing transformer-based forecasting methods, while effective in many domains, struggle to capture the distinctiveness of neural signals characterized by spatiotemporal sparsity and intricate dependencies. To address this challenge, we here introduce QuantFormer, a transformer-based model specifically designed for forecasting neural activity from two-photon calcium imaging data. Unlike conventional regression-based approaches, QuantFormerreframes the forecasting task as a classification problem via dynamic signal quantization, enabling more effective learning of sparse neural activation patterns. Additionally, QuantFormer tackles the challenge of analyzing multivariate signals from an arbitrary number of neurons by incorporating neuron-specific tokens, allowing scalability across diverse neuronal populations. Trained with unsupervised quantization on the Allen dataset, QuantFormer sets a new benchmark in forecasting mouse visual cortex activity. It demonstrates robust performance and generalization across various stimuli and individuals, paving the way for a foundational model in neural signal prediction.

Figures (13)

• Figure 1: Comparison of encoding, decoding, and forecasting tasks. Encoding methods take a stimulus and behavioral variables at time $t$ to predict neural spikes at the same time point. In contrast, decoding methods work do the opposite, using spike responses at time $t$ to predict behavioral variables for that time step. Neural forecasting differs from both, as it uses the stimulus at time $t$ and raw fluorescence traces at time $t-1$ to predict neural responses at time $t$.
• Figure 2: QuantFormer architecture. During pre-training we employ a self-supervision quantization strategy that learns to reconstruct the randomly-masked patches along a quantization scheme. For response forecasting, [NEURON] and [STIM] tokens are prepended to the input, and neuronal response patches are masked; the model predicts for the masked patches quantized codes that are converted, through the quantization decoder learned during self-training, to a continuous signal. For activation classification, an additional [CLS] token is included in the sequence, and its output embedding is fed to the activation classifier.
• Figure 3: Qualitative analysis of stimuli response forecasting performance by QuantFormer and its competitors: forecasting examples for each type of stimuli: drifting gratings (top-left), static gratings (top-right), natural scenes (bottom-left) and locally sparse noise (bottom-right). More examples can be found in Section D of the Appendix.
• Figure 4: Attention maps for all stimulus types. Each row corresponds to a predicted code, while columns represent the $[STIM]$ token, $[NEURON]$ token, pre-stimulus patches, and the predicted codes. Color intensity indicates attention strength, with yellow denoting the most attended tokens and darker shades indicating less attended ones. This visualization highlights the model's selective focus across different components of the input and predicted outputs.
• Figure 5: Interpretability of codes. (a) t-SNE of a codebook, with patterns representation in the same scale. We can appreciate along the first axis the amplitude variation. (b) Same as before, but with normalization to appreciate differences in patterns. (c) Effect of sequence.