QuantFormer: Learning to quantize for neural activity forecasting in mouse visual cortex

Thank you for your detailed and thoughtful feedback. We appreciate the opportunity to address your concerns and provide clarifications. Thank you for bringing the cited works to our attention. We will include Zhang et al. (2024), Zhu et al. (2022), and Ye & Pandarinath (2021) in the literature review to provide a more comprehensive context for our contributions. Regarding Zhu et al. (2022), which focuses on calcium imaging data, we will conduct a direct comparison to evaluate our model against this baseline. For the ablation study on the importance of stimulus tokens, we are currently running the comparison. In the meantime, we encourage you to refer to our interpretability analysis in the appendix, which provides insights into how the model utilizes stimulus tokens during forecasting.

Thank you also for pointing out inaccuracies. We will rephrase the relevant sentences to correct and clarify these points. First, we acknowledge that some cited works use deconvolved calcium traces, which involve additional preprocessing compared to raw fluorescence traces. To the best of our knowledge, raw fluorescence data without deconvolution are not commonly used in these works. We are actively working to extend our approach to include deconvolved traces for a more direct comparison. However, to the best of our understanding, Antoniades et al. (2023) (Neuroformer) focuses exclusively on spike data, as outlined in the workflow section of their paper. While some approaches (e.g., Turishcheva et al. 2024a;b, Li et al. 2023, Xu et al. 2023a) do not rely on trial-averaged responses during training, they do use repeats during evaluation or focus solely on visual stimuli for prediction. The total number of unique neurons used in our work is 3,989, considering the partial overlap between sessions. The mentioned 230,000 traces refers to the total number of trials, not neurons. We will ensure the work accurately reflects existing contributions and highlights our approach's unique focus on trial-specific neural forecasting with minimal preprocessing and we will correct this terminology to avoid any misunderstanding.

Regarding your questions: For subject generalization, we pretrained the model on data from 10 mice and fine-tuned it on the 11th mouse. During fine-tuning, neuron-specific information was incorporated to adapt the model to the test mouse. Importantly, the neurons used during the fine-tuning phase were not included in the pretraining phase. For stimulus generalization, we agree that static and moving gratings share similarities, particularly for neurons that are less sensitive to motion. However, the conditions still differ because parameters like orientation and spatial frequency vary, creating a meaningful distinction for evaluating generalization. All results were averaged over 44 runs, representing fine-tuning experiments across different combinations of mice and stimulus types. The STIM token is derived from a representation of each stimulus, and what is learned is a linear projection that maps stimulus features to the embedding size. Specifically, we employ the following stimulus features:

• Natural Scenes: We use the features from the last layer of a ResNet-50 as input.
• Gratings: We encode information about spatial frequency, temporal frequency, orientation, and phase into four real-valued numbers.
• Locally Sparse Noise: We encode the $x, y$ positions of white and black points into a vector. For gratings and locally sparse noise, we prefer our encoding method because ResNet-50 is pretrained on natural images, and thus, its representation of these stimuli most probably will not be significant.

Thanks a lot for clarifications. I have one more question

For subject generalization, we pretrained the model on data from 10 mice and fine-tuned it on the 11th mouse

Did you freeze some parts of the model during fine-tuning?

Thank you for your detailed and insightful feedback.

As a general consideration, the choice of 32 codes results from careful experimentation, balancing effectiveness within our methodology. Increasing codes poses challenges akin to regression-based approaches, with most codes representing baseline activity, complicating classification.

If "image computable" refers to image-based features, our tokenization process accommodates features from CNNs or low-level image features, allowing us to handle diverse stimuli. We appreciate your feedback regarding the finite set of stimuli and the potential overlap between training and testing data. While pretraining includes the same neurons and stimuli, the model captures trial-specific dynamics beyond average stimulus-driven activity. Preliminary evidence is shown in the supplementary material (Fig. R1), comparing our average response to the training average.

Regarding neuron identities, not using them during pretraining encourages the model to focus on common features across neurons rather than neuron-specific patterns. This enables the model to learn generalizable patterns during pretraining. When neuron identities are later introduced during fine-tuning, the model combines general and neuron-specific information, achieving its full predictive capability.

More specifically regarding your questions,

1. The clear benefit of our architecture lies in the use of quantization, which enables us to reformulate the forecasting problem as a classification task.
2. The Allen dataset was chosen because it provides fluorescence data making it a benchmark for studying neural dynamics in response to visual stimuli. Temporal leakage is carefully avoided by splitting session data along the temporal dimension: training and testing data from the same session are separated by a gap of at least 10 min, ensuring that the model cannot directly learn from overlapping temporal patterns within a session.
3. We cannot conduct experiments with completely separate neurons because our model relies on learning neuron identities, which are session-specific. Regarding different stimuli, we cannot separate stimuli into training and testing sets within the Allen dataset without shuffling data along the temporal dimension, which would introduce information leakage. For example, test responses might inadvertently be included in the training set as "baselines." This limitation arises from the construction of the Allen dataset, not from our approach. However, our method does not rely solely on simple stimulus IDs, but rather on stimulus features. For example, for natural scenes, we use features extracted by a pretrained ResNet, which are then projected into the model's d-dimensional space via a trainable linear layer. We will take your suggestion to evaluate the model on the SENSORIUM 2022 dataset, in future updates.
4. (and 7) Derivative-based normalization is particularly effective for sparse signals characterized by distinct peaks because it mathematically emphasizes the sharp changes associated with these peaks while preserving their essential features. For a sparse signal $x(t)$ consisting of distinct peaks, the signal can be approximated as: $x(t) = \sum_i A_i \delta(t - t_i),$ where $A_i$ represents the peak amplitudes and $t_i$ are the peak locations. The derivative of this signal is given by: $x'(t) = \sum_i A_i \delta'(t - t_i),$ which highlights sharp changes at the peak locations due to the presence of the derivative of the Dirac delta function, $\delta'(t)$. To normalize the signal using its derivative, we define: $x_{\text{norm}}(t) = \frac{x(t)}{\lvert x'(t) \rvert}.$ This normalization amplifies the regions with significant changes (i.e., the peaks) while diminishing other less significant parts of the signal.
5. For each trial, we calculate the baseline activity using a 2-second window before the stimulus onset. We then compute the ΔF/F (change in fluorescence) relative to this baseline, which allows us to normalize the responses. We consider a neuron to be active if its ΔF/F exceeds 10%.
6. We use sparsity in a broad sense, meaning that neural activity is rare compared to base activity, which aligns with the structure of fluorescence data. We chose not to use deconvolved traces or spiking data because these require additional preprocessing steps, and our goal is to avoid such steps to remain consistent with a real-time scenario. We agree that deconvolved traces are indeed more sparse and could potentially provide additional benefits in certain contexts, but our focus is on minimizing preprocessing to streamline applicability. That said, we will test our model on the SENSORIUM 2022 dataset to evaluate its performance under different conditions. Regarding spike data, their fundamentally different nature makes quantization less meaningful, as spiking activity is already discrete and does not benefit from this approach.

Thank you for the detailed and insightful feedback. We greatly appreciate the opportunity to address your concerns and clarify the motivations, methodology, and results presented in our work. Below, we provide detailed responses to your comments.

1. Vector quantization was chosen to address the critical challenge of data sparsity in neuroscience. VQ compresses neural data into discrete representations, enabling robust pattern discovery (as shown in Fig. A-9) while preserving essential temporal structure. This transformation allows us to reformulate the response forecasting problem as a classification task, which is more effective than traditional regression methods that struggle with sparsity. We acknowledge that the motivation for this choice may not have been sufficiently emphasized early in the paper, and we will clarify it in the introduction to avoid confusion.
2. We are currently working to evaluate comparisons with NDT1 and MtM. However, this comparison is not straightforward because these methods were designed for calcium imaging spiking activity, while our approach operates on raw fluorescence traces. Although this may seem like a minor modification, it is challenging because fluorescence traces are continuous and noisier compared to the discrete nature of spiking data. To address this, we are experimenting with using inferred spike probabilities as inputs to these methods and converting their predictions back into time series data. We will include these results in future updates.
3. We indeed included this analysis in the paper; please refer to lines 474–476. The baseline in Tab. 3 is our encoder backbone, a BERT-like transformer encoder, followed by a linear layer. This model was trained using cross-entropy loss for activity classification and mean squared error (MSE) for forecasting, offering a direct comparison to our approach.
4. The observed superior performance of our method is due to its ability to use quantized indices for reconstructing the original shape of the signal, which enables it to capture more detailed patterns. In contrast, other methods rely on regression-based losses that focus on modeling average activity patterns. Although we cannot compute PSTHs strictly (as we are not using spikes), we can approximate them by averaging all responses for each stimulus and using this average as a baseline for comparison at test time. We included in the supplementary materials (Fig. R1) qualitative examples showcasing both cases where our model succeeds and where it fails, providing a balanced view of its performance relative to the baselines.

Minor:

1. If stimulus information were excluded, the model would rely solely on the temporal structure and patterns in the neural activity itself for forecasting, likely reducing its predictive accuracy, particularly for responses tightly linked to specific stimuli. However, we conducted interpretability analysis to examine the attention distribution during predictions (see Fig. A-8). While the stimulus token is important, it plays a relatively minor role compared to neuron identity and has a comparable influence to the neuron's state before the stimulus onset.
2. We asked ourselves this question during preliminary experiments. The rationale is that a latent representation is meaningful if it enables recovery of the original input. When the reconstruction loss is applied only to masked tokens, the representation of unmasked tokens becomes poor, as their primary role is limited to aiding in reconstructing the missing parts. Moreover, while we did not conduct an exhaustive ablation study due to time constraints, preliminary experiments indicated that using the combined loss led to improved performance.
3. The [NEURON] embedding represents the identity of each neuron, ensuring that the model captures neuron-specific dynamics. Unlike POYO, where each spike is tokenized as an event, our approach tokenizes trace patches rather than individual spikes. Regarding the dimensions, the neuron, stimulus, and patches are embedded into vectors of dimension d, which matches the model's embedding space.
4. We visualized training and test data distributions (find attached an example in the supplementary material, Fig. R2) and observed significant overlap. As fine-tuning is always performed within the same session, we ensure the model adapts to within-session variability.
5. We conducted ablation studies on the number of indices and found that the best performance was achieved with 32 codes. However, we are not suggesting that 32 is the optimal encoding size for neural information in general—this is specific to our methodology. The relatively small number reflects how, as the number of codes increases, many of them end up representing neutral conditions (flat activity), leaving only a few for peak patterns. This imbalance creates challenges in classification, as too many similar codes make the task more like a regression problem, reducing effectiveness.

Thank you for your thoughtful and detailed feedback. Below, we address each of your concerns and clarify our approach and its motivations.

1. We agree with this observation. However, traditional encoding methods are designed to predict the average response, which does not capture variability or trial-specific dynamics. In contrast, our approach focuses on forecasting individual trial responses, which is crucial for applications like online experiments with optogenetic manipulation. In such scenarios, real-time predictions of neural responses can guide dynamic adjustments, even when the visual stimulus is controlled.
2. To demonstrate how our approach captures trial-specific dynamics beyond the stimulus-driven average response, we visually compare our forecasting model to a PSTH-based baseline. Specifically, we averaged the responses across all trials in the training set to create a baseline (train average). During evaluation, we compared this baseline to both the average response across test trials (test average) and our model's average prediction (Refer to Fig. R1 in supplementary material). This analysis highlights how our method accounts for variability beyond the stimulus-driven average response.
3. We address this point using the attention maps provided in the appendix. These maps reveal that while neuron identity is the most significant factor in determining the forecasted indices, the final predictions are also influenced by the context surrounding the onset. This demonstrates that the model does not rely solely on stimulus-driven patterns but also incorporates pre-stimulus and trial-specific variability. Additionally, although the correlation value of 0.33 in Table 2 may seem modest, it reflects the inherent challenge of capturing trial-specific variance, which extends beyond replicating the average response.
4. (5) We have not directly compared our model's performance to the baselines on the original datasets reported by their authors, as our approach is specifically tailored to work with fluorescence traces, which differ significantly from data used in these studies. For the Allen dataset, we trained the baseline models ourselves under identical conditions to ensure a fair comparison. While we cannot claim to outperform these baselines on their originally tested datasets, our results highlight the strengths of our method in modeling trial-specific dynamics in fluorescence data.