One Model to Train Them All: A Unified Diffusion Framework for Multi-Context Neural Population Forecasting

• Questions:

1. Thank you for the great question. Typically, models focused on forecasting primarily report behavior forecasting results, as it serves as an easier and more direct metric. In these cases, behavior forecasting often reflects how well the model forecasts neural activity, especially since those models are not jointly trained. Therefore, behavior forecasting becomes a proxy for neural activity forecasting performance. This makes it challenging to find direct comparisons, particularly because our model uses smoothed spikes, whereas models like SeqVAE, which also use smoothed spikes, do not report neural activity values. However, we acknowledge the lack of results regarding neural activity and have added relevant plots in Fig. 3.E to address this gap.

2. We would like to clarify that POYO is not capable of forecasting; instead, it performs filtering, reporting results in 500 ms windows. This means it predicts the next time step of behavior variables based on a history window of neural activity and slides this history window until the desired length of predicted behavior is achieved. To the best of our knowledge, ours is the only unified model designed for forecasting. We compared the performance of our model with the most recent state-of-the-art models: XFADS (for single-session) and SeqVae (for multi-session). However, based on your feedback, we also adapted our model to perform filtering to evaluate its performance in this context. The results show that our model is as competitive as NDT2 and POYO, despite filtering not being its primary focus. The design choices made for our model were not optimized for filtering, as diffusion models are traditionally better suited for full-sequence generation rather than next-token prediction, as is the case with transformers. The new experiments have been added to the manuscript in Fig. 3.C and Fig. 5.D. A more detailed comparison between our approach and other models can be found in Appendix A.

3. These are very relevant questions: (1) Yes, we tested this and have added the results to Fig. 3.B. Additionally, we conducted an ablation study comparing the model's performance with and without behavioral data (Fig. 6, first plot). We found that including or excluding behavioral data did not impact the model’s performance in terms of neural activity forecasting. However, we would like to clarify that the difference observed in behavioral performance stems from the fact that, when trained solely on spiking data, we need to use a linear decoder (ridge regression in this case) to map the predicted neural activity to behavior values. As a result, this performance metric reflects the quality of neural activity prediction, not the model’s ability to forecast behavior directly. (2) We thank the reviewer for the suggestion to conduct a lambda ablation study. We trained the model with different lambda values and reported the performance in terms of both neural spiking activity and behavior. The lambda term in the loss function significantly influences model performance. Our benchmarks indicated that a value of 0.5 worked well for both behavior and spiking activity. However, this value may vary depending on the specific brain area or task being studied (Fig. 6, last plot). This dependency is expected, as increasing the attention the model gives to behavior can reduce its accuracy in predicting neural activity, since not all neurons are behavior-relevant. Conversely, focusing more on neural activity can decrease behavior forecasting accuracy. (3) We do not believe that our model excessively focuses on one modality, especially in this case. For the forecasting results, we followed the preprocessing strategy used by the models we compared against, which involved aligning the data to movement onset. This means that, in the history provided to the model for forecasting what happens after movement onset, the behavior variables are simply set to 0’s. As a result, the model is not biased toward behavior but instead relies on the neural activity to make predictions after the onset.

4. Once again, we thank the reviewer for the suggestion to run ablation studies. We conducted several experiments with different time bins, fixing the smoothing filter, as well as with fixed time bins and varying smoothing filters. Our results show that the model is robust to time bins larger than 5 and across different smoothing filters. The choices of both the time bin and smoothing filter used in the paper were based on those applied by our baseline models.

We thank the reviewer for acknowledging our efforts to improve the manuscript and for the updated score. Below, we address the remaining questions:

1. We thank the reviewer for highlighting the importance of comparing both behavior and neural forecasting, a point we initially thought was less critical. We also appreciate the reminder regarding the existence of several studies that report these values, which we had overlooked. In response, we have now included the reporting values for single-session baselines for neural activity forecasting and next-step prediction (Fig. 3B and C in the new version), alongside the previously reported behavior forecasting results.

2. We apologize for the unclear reporting of the ablation study results and thank the reviewer for pointing this out. We have now added a description for each panel in Fig. 6. Regarding the spikes vs. spikes+behavior results, we appreciate the reviewer’s feedback on the confusion. The discrepancy arose because we reported results for forecasting over a 500ms window, where the better performance of the spikes-only model was due to ridge regression compensating for forecasting errors in spike predictions, leading to improved behavior accuracy. To address this, we have updated the results in the new version to reflect the filtering mode of the model. Panel A of Fig. 6 (in the updated version) compares our joint model with the two-step model, where behavior is predicted using ridge regression based on next-step spiking activity. The results show that the joint model outperforms the two-step model in behavior prediction while preserving spiking activity accuracy. This demonstrates that the joint model effectively integrates spiking and behavioral variables, optimizing both simultaneously. Unlike the two-step model, which treats these tasks separately, the joint model better captures their interactions, resulting in improved performance without sacrificing spike prediction accuracy.

Panels B-E show the effects of time bin size, smoothing filter, history window, and forecasting window on $R^2$ for behavior forecasting. Specifically, for each variable (e.g., time bin size), we evaluated the model's behavior performance while keeping all other variables constant.

• Panel D: We examined the behavior accuracy using different history forecasting windows: 700 ms, 500 ms, 200 ms, and 100 ms. As expected, we observed a significant performance drop at 100 ms, which aligns with our finding that the most relevant neural activity for behavior occurs 200-250 ms before movement onset.

• Panel E: We explored the impact of different forecasting windows on model performance, assessing how the forecast horizon influences behavior prediction.

• Panel F: We trained the model with three different lambda values and reported the resulting behavior and spiking activity forecasting performance.

We appreciate the reviewer’s feedback and would like to address the concerns raised.

1. Fig. 6 and Filtering vs. Forecasting: It is not accurate to state that Fig. 6 shows filtering results in all panels. While panel A focuses on filtering, the remaining panels explicitly focus on forecasting, as noted in the manuscript. This approach aligns with the discussion in lines 512–520 of the manuscript. Filtering, equivalent to one-step-ahead forecasting, minimizes error propagation and isolates the model’s ability to predict neural activity accurately. We chose to highlight one-step forecasting results (20 ms) rather than longer windows (e.g., 500 ms) because, over extended horizons, the ridge regression in the two-stage approach compensates for neural forecasting errors. Ridge regression decodes behavior from already forecasted neural activity, rather than directly forecasting behavior from prior neural activity. This creates an unfair comparison, as ridge regression smooths over the limitations of the two-stage model. The degradation in forecasting accuracy with longer windows is evident in panel E and is expected. We could add a figure to further illustrate this phenomenon, showing how the two-stage method's performance remains relatively stable due to ridge regression, while the joint model performance decreases in line with neural forecasting degradation, with the increase of the forecasting window. However, we emphasize that Fig. 6 remains aligned with the paper's primary focus on forecasting and the evaluation of joint training’s benefits.

We recognize that the M1 dataset used in Fig. 6 may not optimally highlight performance improvements with joint training. The goal of joint training is not solely to outperform the two-stage approach but to address its assumptions, learning low-dimensional latents that are predictive of behavior while preserving neural fit. While performance gains may appear minimal in datasets with near-linear dependencies (e.g., M1), joint training excels in interpretability and is particularly effective in modeling nonlinear dependencies, as demonstrated in the PMd experiments.

We hope that our clarifications address the reviewer’s observation that "joint spike+behavior training yields worse prediction than training on spikes alone." Our results suggest joint training preserves neural fit and improves behavior forecasting in some scenarios, with weaker performance only observed in certain datasets (e.g., M1) when ridge regression compensates for long forecasting windows. However, we observed benefits in shorter windows. We would greatly appreciate it if you could consider revisiting your evaluation in light of these clarifications.

We would like to direct the reviewer’s attention to the overall comment section, where we provided a detailed summary of the improvements made to both the main manuscript and the appendix. We have made significant revisions to address the reviewer’s concerns, including a more comprehensive discussion of the results, a clearer comparison with existing models (such as POYO, NDT, and EIT), and an expanded explanation of our network architecture. Additionally, we have refined the mathematical descriptions for clarity.

Regarding the insights, we would like to emphasize that forecasting neural activity serves as a way to confirm that the model has learned a good approximation of the true underlying neural dynamics. While a model might excel in next-step decoding (filtering), this does not necessarily mean it has learned the correct dynamics. The use of session embeddings enables the alignment of multiple contexts by learning a shared, common representation while capturing session-specific "styles" through these embeddings. This allows our model to better generalize across different sessions.

Finally, further motivation for the benchmark dataset can be found in the section titled "Tools for Foundation Models: Benchmark Datasets," with additional details provided in Appendix C to E. In brief, there is currently no benchmark designed specifically for multi-session models. The most similar dataset is NLB, which was designed for single-session models. However, for a fair comparison of large pre-trained models, they should be trained on the same dataset. While we acknowledge that NLB can always be used for performance comparison, using different datasets for model pretraining would not lead to a fair evaluation. Therefore, we argue that a large-scale and accessible dataset is essential for advancing the field and enabling proper comparisons across models.

We hope the reviewer found our rebuttal helpful. If there is anything we can further clarify or do to improve the reviewer's overall assessment of our work, please let us know.

• Weaknesses:

We thank the reviewer for the suggestion to explore what the model is learning. In response, we conducted an occlusion sensitivity study, which allowed us to identify the behavior-relevant neurons and determine the specific time windows during which neuronal activity encodes movement-related information.

Additionally, we carried out the recommended ablation study to compare the model's performance with and without behavioral data (Fig. 6, first plot). The results show that the joint model outperforms the two-step model in behavior prediction while preserving spiking activity accuracy. This demonstrates that the joint model effectively integrates spiking and behavior variables, optimizing both simultaneously. Unlike the two-step model, which separates these tasks, the joint model better captures their interactions, leading to improved performance without sacrificing spike prediction accuracy.

We have added visualizations of the session embeddings for the multi-animal case (Figure 4.B). These embeddings cluster according to behavioral strategy, even though the model was not explicitly trained to do so. Two distinct strategies were identified: (1) fast reaction times with higher rates of failed trials, and (2) slower reaction times with higher task success. Notably, the shared session ID is positioned between the two clusters, indicating that the model successfully learned a shared representation.

• Questions:

That's an interesting question! Currently, we don't have data on naturalistic tasks, with the least structured task being the RTT, where the model performed well (Fig. 5.D). In the context of naturalistic tasks, we speculate that the current session embedding approach might not be directly applicable. However, we believe that we could adapt this approach by conditioning it to distinguish between different animal movements (e.g., left/right).

Our hypothesis is that an initial step would involve training an embedding specifically suited for freely moving animals, which could help the model adapt to such scenarios. For example, using CLIP-style embeddings, where we provide descriptions of the animals' movements, and applying a contrastive learning approach to correlate these descriptions with neural activity and behavior. This could potentially align with techniques like CEBRA to capture the relationships in naturalistic settings.

We hope the reviewer found our rebuttal helpful. If there is anything we can further clarify or do to improve the reviewer's overall assessment of our work, please let us know.

We would like to highlight that we conducted additional experiments to emphasize the differences between our proposed approach and previous unified models.

We ran new experiments to address the previous concerns:

Diffusion Training: To test whether the improved performance is due to our architecture or the training strategy, we trained the model without it (details in Appendix B.2). Our results show that the architecture itself is inherently superior, and that training with diffusion further enhances its performance.

The poor performance of NDT2 with a causal mask was expected, since recent work on transformer architectures has shown that learning global dependencies in temporal tokens leads to poor prediction results on multivariate time series datasets [1]. Thus, this model is not competitive in terms of forecasting and, in fact, that was not its original purpose, nor does POYO's architecture allow for forecasting, so our model is the only unified model proven to be able to so competitively with state-of-the-art neural data forecasting models.

In Fig.1 of the Appendix B.6 we show that training without diffusion required substantially more training iterations and consistently yielded lower prediction accuracy. These findings align with previous studies emphasizing the importance of intermediate noise steps in diffusion training for enhancing model prediction quality [2].

Contribution: We would like to emphasize that our primary goal is forecasting, and previous unified models either cannot perform forecasting or their performance is not competitive with single-session neural data forecasting. To address this limitation, we have developed a distinct architecture and training strategy. Our model is specifically designed to handle both long-term and short-term dependencies for accurate forecasting. Additionally, we improved the performance and training efficiency through the use of diffusion.

Given these updates, we hope the revisions address the concerns raised and provide a clearer demonstration of the contributions and significance of our work. We would greatly appreciate it if you could consider revisiting your evaluation in light of these improvements.

[1] Liu, Y., Hu, T., Zhang, H., Wu, H., Wang, S., Ma, L., & Long, M. (2023). itransformer: Inverted transformers are effective for time series forecasting. arXiv preprint arXiv:2310.06625.

[2] Nichol, A. Q., & Dhariwal, P. (2021, July). Improved denoising diffusion probabilistic models. In International conference on machine learning (pp. 8162-8171). PMLR.

Weaknesses:

1. We thank the reviewer for the suggestion to run more baseline experiments. However, the comparisons with NDT2 and POYO are not directly feasible because our proposed model is a unified framework capable of forecasting, whereas NDT2 and POYO were specifically designed for filtering tasks—i.e., predicting only the next time step based on a history of neural activity. This distinction was the main reason for not including these comparisons, and we apologize if the original manuscript was unclear on this point. To address this, we have now provided a clearer explanation in the methods section (C.5) and in Figure 1 of the revised manuscript. Based on the reviewer’s feedback, we conducted a new experiment where we adapted our model to a filtering setting, predicting only the next time step. Unfortunately, NDT2 does not provide results for NLB datasets. In this new experiment, we compared our model against several other methods, including the Wiener filter, GRU, MLP, AutoLFADS, NDT, NDT-Sup, EIT, XFADS, and POYO on the NLB-maze dataset (Figure 1.C in the updated version). Additionally, to compare our model with other modern pretraining strategies, we ran experiments on the RTT dataset, comparing the performance of training from scratch versus training with the full dataset for NDT2 and POYO (Figure 5.D in the new version). While POYO remains the best-performing model in this comparison, our model is highly competitive, outperforming NDT2. However, it is important to note that filtering is not the primary focus of our model. A detailed, though not exhaustive, summary of the differences between our approach and others has been added to Appendix A. Finally, we would like to emphasize that POYO is not capable of forecasting, which represents an unexplored area in the NDT2 framework.

2. We agree with the reviewer that the network architecture is an important detail missing from our paper. To address this concern, we have included a full description of the model in Section C.2 and additional details in Appendix B. We hope that the updated manuscript fully clarifies any remaining questions related to the model. Nevertheless, we will address the missing details of the method here:

• Each neural population in the model has a different number of neurons, which can pose a challenge when working with multiple sessions. To address this, we standardize the input size by aligning all populations to the size of the largest population, padding the smaller populations as needed. It is important to note that, similar to transformers, the padded portions of the input are ignored during training by the loss function, ensuring that the padding does not influence the model's learning process (Appendix B.4).
• Each sample (trial) is associated with a session ID, which links it to a specific session. The session IDs are represented as a learnable lookup table, where each ID (scalar) corresponds to a high-dimensional vector. This vector is then used to guide the generation process through cross-attention, in a manner similar to how CLIP embeddings guide the text-to-image generation in diffusion models.
• We agree with the reviewer that the distinction between conditional and unconditional models was not well defined in the original manuscript. The model is considered unconditional when using the shared session ID. The equation is only applied during sampling in the conditional case, where we generate a sequence conditioned on a specific session ID. The conditioning is achieved by assigning 90% of the model’s focus to the chosen session ID and 10% to the shared one. This guidance scale, which controls how much attention the model gives to the session ID, is a hyperparameter during the sampling phase.
• The model autoregressively predicts an arbitrarily long sequence by using the previously generated sequence as input for predicting the next steps.

• Weaknesses:

1. Yes, our goal is forecasting, not filtering, as is the case with NDT2 and POYO. We acknowledge that the original manuscript did not clearly distinguish our approach from these two models. To address this, we have added a detailed comparison in Appendix A to highlight the differences between our model and NDT2/POYO.

2. We thank the reviewer for the suggestion to systematic explore the limitations of forecasting capabilities of the model. We ran several ablation studies exploring different history forecasting windows and forecasting windows, with the results presented in Figure 6 of the updated manuscript. For panel D, we saw what is the behavior accuracy using a history forecating window of 700 ms, 500 , 200 and 100 ms. The model showed a huge performance degradation for the 100 ms what was expected due to the fact that we found that the most revelant neural activity for behavior was presnet in 200- 250 ms before the movement onset. For E, we explored how good is the forecating performance for different forecating windows.

3. We conducted the recommended ablation study to evaluate the model's performance with and without behavioral data (Fig. 6, first plot). The results show that the joint model outperforms the two-step model in behavior prediction while preserving spiking activity accuracy. This demonstrates that the joint model effectively integrates spiking and behavior variables, optimizing both simultaneously. Unlike the two-step model, which separates these tasks, the joint model better captures their interactions, leading to improved performance without sacrificing spike prediction accuracy.

• Questions:

1. Thank you for the insightful question, which helps clarify our results. In the new version, we added our explanation/hypothesis: Our transfer strategies aim to learn shared features while capturing session-specific styles via embeddings. Training on one small session, however, is insufficient to build a noise-robust model, and few-shot examples from another animal fail to separate shared and session-specific features, pushing the model further from the data manifold and amplifying noise in the learned representation. With sufficient data, the model learns robust shared features, and few-shot examples capture session-specific styles, improving performance. This emphasizes the importance of adequate data for robust representations. While Azabou et al. (2023) showed that large-scale pretraining enables session identification to adapt effectively, our results highlight that without robust pretraining, this does not hold true.

2. Yes, good catch. We have added a clear distinction between the time of neural activity and the diffusion processing. The time steps we refer to are specifically the diffusion time steps.

3. We added a detailed comparison in Appendix A

We hope the reviewer found our rebuttal helpful. If there is anything we can further clarify, we are happy to do so!

Contribution: We would like to emphasize that our primary goal is forecasting, and previous unified models either cannot perform forecasting or their performance is not competitive with single-session neural data forecasting. To address this limitation, we have developed a distinct architecture and training strategy. Our model is specifically designed to handle both long-term and short-term dependencies for accurate forecasting. Additionally, we improved the performance and training efficiency through the use of diffusion.

Comparison with NDT2 with causal mask: NDT2 is not trained with a causal mask by default; rather, it is trained with a non-causal mask. For the sake of comparison and to address feedback from the reviewers, we conducted experiments to evaluate NDT2 with a causal mask. The results of these experiments are now included in Figure 3.B of the revised manuscript.

The poor performance of NDT2 with a causal mask was expected, since recent work on transformer architectures has shown that learning global dependencies in temporal tokens leads to poor prediction results on multivariate time series datasets [1]. Thus, this model is not competitive in terms of forecasting and, in fact, that was not its original purpose, nor does POYO's architecture allow for forecasting, so our model is the only unified model proven to be able to so competitively with state-of-the-art neural data forecasting models.

Diffusion Training: To test whether the improved performance is due to our architecture or the training strategy, we trained the model without it (details in Appendix B.2). We found that the architecture itself is better and that training with diffusion improves it.

In Fig.1 of the Appendix B.6 we show that training without diffusion required substantially more training iterations and consistently yielded lower prediction accuracy. These findings align with previous studies emphasizing the importance of intermediate noise steps in diffusion training for enhancing model prediction quality [2].

[1] Liu, Y., Hu, T., Zhang, H., Wu, H., Wang, S., Ma, L., & Long, M. (2023). itransformer: Inverted transformers are effective for time series forecasting. arXiv preprint arXiv:2310.06625.

[2] Nichol, A. Q., & Dhariwal, P. (2021, July). Improved denoising diffusion probabilistic models. In International conference on machine learning (pp. 8162-8171). PMLR.