Generalized and Invariant Single-Neuron In-Vivo Activity Representation Learning

Dear Reviewer,

We sincerely thank you for your valuable time and insightful feedback on our manuscript. We particularly appreciate the acknowledgment of the paper's significance and originality in addressing the critical problem of generalizability in computational neuroscience.

Below, we address your specific questions.

1. On the Encoding of Batch Labels

We treat the task of identifying the batch as a multi-class classification problem. The batch discriminator, $D_{\phi}$, is trained to predict the batch label (e.g., Animal ID or Stimulus ID) from the neuron's embedding $z_i$. The loss function for this task is the standard cross-entropy loss. We clarify this mechanism further in the main text of the revised manuscript to ensure this key detail is more accessible.

2. On Hyperparameter Selection via Cross-Validation

You requested more details on the cross-validation process for choosing hyperparameters.

Our goal was to select the optimal adversarial weight, $\lambda$[cite: 204, 628], to maximize the model's ability to generalize across different experimental batches. This entire process was performed using only the training dataset, ensuring the final test set remained completely unseen.

To accomplish this, we employed a standard k-fold cross-validation procedure within the training data. The process is as follows:

• Partitioning the Training Set: If our training set consisted of data from k distinct batches (for instance, in the cross-animal experiment, we train on 3 animals, so k=3), we perform a k-fold cross-validation.
• Iterating Through Folds: For each potential value of $\lambda$ in our grid search, we iterate through k folds. In each fold, we train the model on data from k-1 batches and validate its performance on the single, held-out batch. The performance metric is the accuracy on the primary downstream task (e.g., cell type classification).
• Averaging Performance: After completing all k folds, we calculate the average validation accuracy for that specific value of $\lambda$.
• Selecting the Best Hyperparameter: We repeat this for every value of $\lambda$. The value that yields the highest average cross-validation accuracy is chosen as the optimal hyperparameter.

We update the methods section of our manuscript to reflect this clear and standard procedure.

3. On the Convergence of Adversarial Training

You asked if we observed differences in training stability or convergence difficulty among the different representation encoders.

This is an excellent and practical question, as training stability can indeed differ across model architectures, and our approach was adapted accordingly.

For most of the models we tested, including the end-to-end supervised LOLCAT and the VAE-based framework, convergence with the adversarial loss was generally stable. We were able to train these models end-to-end with both the primary and adversarial losses from the start.

However, the NeuPRINT model, which learns an implicit representation, required a more tailored approach. We observed that jointly optimizing its delicate primary objective alongside the adversarial objective from scratch was challenging for convergence. To address this, we adopted a two-stage training strategy specifically for NeuPRINT:

1. Stage 1: Pre-training. We first trained the NeuPRINT model solely on its self-supervised objective until it achieved a stable, converged state.
2. Stage 2: Adversarial Fine-tuning. After the model had converged, we introduced the adversarial loss and fine-tuned the entire network to learn batch-invariant representations.

We will add these crucial details about our training methodology to the revised manuscript.

4. On the Discriminator Architecture

We thank the reviewer for this excellent suggestion. Prompted by your feedback, we conducted a new set of ablation studies to systematically investigate the impact of the discriminator's architecture on model performance.

We first increased the depth of the discriminator to three and four hidden layers. We found that this additional complexity had a negligible impact on the final generalization accuracy across our models.

We then replaced the MLP with a much simpler linear model (a logistic regression classifier). In this configuration, we observed a consistent and notable decrease in performance. This suggests that a non-linear classifier is necessary to effectively capture the complex, subtle artifacts that constitute batch effects.

Adding further complexity yields no additional benefit, while a simpler linear model is insufficient. This analysis validates our architecture as a robust and efficient choice for this task. We add a summary of this new ablation study to the appendix of our revised manuscript.

Once again, we thank you for your constructive comments. We are confident that by incorporating these clarifications, the final version of our paper will be stronger and more complete.

Dear Reviewer,

Thank you for your thorough review and the constructive feedback on our manuscript. We appreciate the time and effort you dedicated to our work. We agree that clarity is paramount and have undertaken a significant revision to enhance the manuscript's readability and logical flow, making our core contributions even more accessible.

We found your specific suggestions invaluable. In response, we have made the following improvements:

• Clarified Key Terminology: We have revised the Abstract and Introduction to provide explicit, context-specific definitions for crucial concepts such as "batch effects," "invariant representations," and "out-of-domain generalization."

• Restructured for Readability: Following your excellent advice, we have merged the "Dataset" and "Experiment Setup" sections into a unified "Benchmark" section. Each experiment is now presented as a self-contained unit, creating a clearer, more logical narrative for the reader.

We would also like to offer additional context on two other points you raised: our choice of evaluation metrics and our focus on single-neuron functional activity. We hope this clarification will better highlight the motivation and significance of our work.

On the Contribution and Focus on Single-Neuron Data: In computational neuroscience, many models have recently emerged to learn neuron "identity representations" from in-vivo, single-neuron activity data. These representations are highly valuable for downstream tasks such as predicting a neuron's cell type or its anatomical location. However, while this direction is promising, a critical challenge often overlooked is that these models exhibit poor generalization across different experimental batches (e.g., different subjects or stimulus sets). Our work is the first to systematically identify, benchmark, and provide an effective solution for this crucial generalization gap. This specific focus is the core of our contribution.

On the Framework and Evaluation Metrics: To address this challenge, we proposed a simple, model-agnostic adversarial training strategy. We believe a key strength of our approach is its flexibility; it is an intuitive framework that can be easily incorporated into a wide variety of existing and future model architectures.

Given this practical goal, we chose downstream classification performance as our primary evaluation metric. This is the established standard for gauging the utility of representations in our field because it directly measures what neuroscientists need: the ability to predict meaningful biological labels. Our extensive experiments, conducted across two different datasets and four distinct model architectures, convincingly demonstrate that our method achieves significant gains in out-of-distribution generalization. This robust validation directly supports our central claim that our framework produces more useful and reliable representations for scientific discovery.

In summary, our manuscript:

• Identifies and provides a solution for a critical generalization problem within the emerging field of single-neuron representation learning.

• Introduces a simple, effective, and model-agnostic framework with broad applicability.

• Demonstrates robust improvements in out-of-distribution performance across multiple datasets and model architectures, confirming the practical utility of our approach.

We have substantially revised the manuscript to improve its clarity based on your feedback. We hope that with this additional context, the value of our contributions is clear. We believe our work offers the field a practical and powerful tool for building more generalizable models, and we would be grateful if you would reconsider it.

We sincerely thank the reviewer for their positive and constructive feedback. We are particularly grateful that you found our approach to be a "simple and flexible way to achieve greater out-of-distribution generalization" and appreciated the intuitive motivation and clear presentation of our work.

We agree with the weaknesses you have identified and have performed additional analysis and revisions to the manuscript to address them.

---

On Technical Novelty

We agree entirely with your assessment. The core contribution of our work is not the development of a new adversarial algorithm, but rather the novel application of this established technique to address a critical and previously un-benchmarked problem in computational neuroscience: the poor generalization of single-neuron representation models across different subjects and stimulus conditions. We believe that formally identifying this gap and demonstrating that a straightforward, model-agnostic framework can effectively solve it is a significant and practical contribution for neuroscientists who wish to apply these models in real-world experiments.

Efficacy with Fewer Training Batch-Level Attributes

Following the your suggestion, we trained the LOLCAT model with an adversarial loss using only two stimulus conditions and evaluated its performance on a held-out third condition. The results consistently demonstrated lower performance compared to the model trained on all three stimulus conditions, which indicates that employing a less diverse training set leads to performance degradation. For future work, we suggest exploring the impact of the number of batch-level attributes in the training set to further enhance the model's robustness. These discussions have been incorporated into the Discussion section.

Efficacy in boost in-distribution performance

We also tested the scenario you alluded to, where there is no out-of-domain challenge (i.e., all stimulus types and animals are present in the training data, with only a standard random test split). We did not observe significant and consistent improvement in performance, whether adversarial training was applied with respect to stimulus types or animal subjects.

Performance for each neuronal subtype

To provide a more in-depth evaluation beyond aggregate accuracies, we analyzed the performance for each neuronal subtype. The results from NeuPRINT are in the tables below.

Cross-Stimulus Generalization by Cell Type

Cross-Animal Generalization by Cell Type

We have added this detailed analysis to the manuscript. Thank you for helping us strengthen our work.

Breakdown of results for each of the four mice (SB025, SB026, SB028, SB030) as the test set

Data presentation in tables 1-3

We have recalculated the percentages in the third column "cross-attribute + adversarial training" to show the decrease relative to the first column's "In-domain" value. Thank you for helping us make our work more clear.

We thank the reviewer once again for their constructive feedback, which has helped us significantly improve the clarity and impact of our manuscript. We are confident that the revised paper now more robustly supports our contributions and provides a solid foundation for future work on generalizable single-neuron representation learning. In light of these revisions, which we believe address the initial concerns, we would be grateful if you would reconsider your assessment of our work.

Thank you, authors, for your time and effort running these follow-up experiments and answering the questions I had. I would encourage the authors to add these new details to their manuscript and/or supplemental material.

Regardless of the fact that the adversarial training method is not a novel contribution, I remain excited about the practical utility of this approach in this domain, where cross-subject and cross-task generalization is challenging.

We sincerely thank the reviewer for their valuable feedback and for appreciating the straightforward nature of our work—identifying an important generalization problem and applying a direct, effective solution. We agree with the points raised regarding clarity and presentation and have revised the manuscript to address them.

Clarifying the Paper's Core Contribution and Goal

We agree entirely that the paper's primary goal—learning robust representations for neuronal identity, not predicting neural activity—was not stated clearly enough in the original manuscript. This lack of clarity is the most critical point to fix, and we thank you for highlighting it.

To address this, we have significantly rewritten parts of the abstract and introduction. The revised text now explicitly states that our focus is on "representation learning for neuronal identity from in-vivo activity." This clarification sharpens the paper's core contribution: identifying a critical generalization gap in this emerging area and demonstrating that an established method offers a robust, model-agnostic solution. This also makes it clear why the literature on predictive neural activity is not the focus of our work.

Improving Presentation and Coherence

Thank you for rightly pointing out the weakness and lack of coherence in the "Theoretical Motivation" section. We agree completely that the original section and its corresponding appendix were uninformative and detracted from the paper's quality.

We have substantially revised this part of the paper to improve clarity and coherence:

• We have removed the uninformative appendix on theoretical motivation.
• We have retitled the original Section 4.5 to "Intuitive Explanation" to more accurately reflect its purpose.
• We have completely rewritten this section to provide a clear, intuitive walkthrough of how the adversarial method works, grounding the explanation in the visual illustration provided in Figure 2, rather than making unfulfilled promises of a deep theoretical dive.

We have also merged the "Dataset" and "Experiment Setup" sections into a single, unified "Benchmark Experiments Setup" section. This improves the paper's flow, as each experiment is now a self-contained subsection that introduces its dataset and protocol together, preventing the reader from having to switch back and forth.

---

We believe these revisions make the manuscript clearer, more coherent, and more accurately focused on its primary contribution. We thank you again for your constructive feedback, which has been invaluable in helping us strengthen the paper. We would be grateful if you would reconsider your assessment of our work.

Dear Reviewer:

We sincerely thank you again for your constructive and positive feedback.

Best regards,

Authors

We extend our sincere gratitude to the reviewer for their thorough and insightful feedback. We are encouraged that the reviewer recognized our work as "technically solid, clearly written," and a "worthwhile addition to the literature" that "addresses a practical painpoint for neuroscientists." We were particularly pleased that the reviewers appreciated the novelty of our contributions, noting that this is the "first paper... to marry domain-adversarial training with single-neuron time-series models" and "to quantify OOD gaps systematically across stimulus and animal factors."

We have carefully considered all comments and have revised the manuscript accordingly. Below, we address each of the points raised.

---

Weaknesses

Response: Limited Datasets and Cross-Lab Claims

Only two mouse datasets are used; no check on primate or synthetic benchmarks, so the claim of "cross-lab variation" is still speculative.

We thank the reviewer for this important point. We agree that our study, being based on two mouse datasets, represents the first crucial step in addressing generalization, and that demonstrating robustness across labs and species is a critical future direction. We have explicitly acknowledged this in the Limitations section of our manuscript. We agree that the claim of "cross-lab variation" is speculative at this stage. Accordingly, we have revised the manuscript to confine the discussion of cross-lab variation to the Future Work section, ensuring our claims accurately reflect the scope of our current results. Our immediate future work will focus on extending this framework to more diverse, multi-lab, and multi-species datasets, such as the Allen Institute’s Visual Coding Neuropixels dataset and the Steinmetz dataset. This will allow for a rigorous evaluation of the cross-lab and cross-species generalization capabilities of our proposed method.

Response: Alternative Methods and Hyperparameter Sensitivity

No results on alternative batch-invariance methods (e.g., mix-style, batchnorm-free) or sensitivity to $\lambda$ (adversarial weight).

We thank the reviewer for raising these two important points.

• On alternative batch-invariance methods: We agree that a comparative analysis with other methods is an important area for future research. Our primary contribution is to be the first to identify, systematically benchmark, and provide a viable solution to the critical generalization problem in the emerging field of single-neuron identity models. While other batch-invariance methods exist, we believe that demonstrating the effectiveness of a well-established technique like adversarial training provides a strong and essential baseline for this new domain. We have now expanded our Discussion section to include these alternative approaches as a promising avenue for future investigation.
• On sensitivity to the adversarial weight $\lambda$: We apologize for not making this clearer in the main text. The adversarial weight $\lambda$ was not arbitrarily chosen but was selected through rigorous cross-validation, as detailed in our appendix. Specifically, we tested 8 values on a logarithmic scale from $10^{-5}$ to $100$. We observed that downstream task accuracy consistently increased with the adversarial weight before stabilizing for values of 0.01 and higher, with no subsequent drop in performance. This demonstrates the robustness of our approach to this hyperparameter. We have clarified this selection process in the revised manuscript.

Response: Preservation of Trial-Level Tuning Curves

While embeddings preserve coarse cell-type groups, the paper does not test whether trial-level neural tuning curves remain interpretable after adversarial training.

This is an excellent point that touches on a more fine-grained level of analysis. Our current study focuses on mitigating batch effects at the level of entire stimulus conditions (e.g., natural scenes vs. gratings) and animals, which we argue is the foundational step for building generalizable models. Analyzing trial-to-trial adaptation within a stimulus block—such as the response decrements observed with repeated stimuli in works like Kehl et al. (2024)—is a fascinating, orthogonal problem. This would require a more complex experimental design where individual trials or trial blocks are treated as distinct batches. While this is outside the scope of our current work, we recognize its importance. We have added a discussion of this topic to our Future Work section to inspire further research in this direction.

Kehl, M.S., Mackay, S., Ohla, K. et al. Single-neuron representations of odours in the human brain. Nature 634, 626–634 (2024). https://doi.org/10.1038/s41586-024-08016-5

Response : Accessibility of Implementation Details

Key implementation details (optimizer, schedule, discriminator capacity) sit in Appendix D; surfacing them earlier would aid replication.

We appreciate the reviewer’s suggestion to improve the paper's accessibility and reproducibility. To address this, we have surfaced the key implementation details by creating a new Appendix A.5: Implementation Details. Furthermore, we have added a prominent reference to this new section in the first paragraph of the main Section 5: Experiment Setup. This change ensures that researchers can more easily find the necessary information to replicate our work.

---

Questions:

Response : Question1

Thank you for pointing out the missing reference to Figure 1; we have now added it. The figure is a schematic to illustrate our out-of-distribution testing principle, not a literal depiction of our test set. Our actual experiment did not use one single "unlike" test set. Instead, we used a leave-one-out protocol with the four diverse stimulus types in the dataset (natural scenes, drifting gratings, spontaneous activity). We trained on three conditions and tested on the held-out one, repeating for all four. Your intuition about an "extreme distribution shift" is correct. We found that models struggled most when generalizing from complex stimuli (like natural scenes) to simple ones (like spontaneous activity). The likely biological reason is that complex stimuli are required to reveal the unique computational roles of different neurons, making the functional signatures learned during training much richer than the simple patterns observed in spontaneous activity.

Response : Question2

That is a critical question. We believe the success of our method in primates will depend on the evolutionary conservation of the features being studied.

What we expect to transfer: The method should generalize well for broad, conserved features. This includes classifying major brain regions (e.g., isocortex vs. cerebellum) or high-level cell types (e.g., excitatory vs. inhibitory), as their fundamental roles are shared across species.

What we expect to fail: Generalization will likely be challenging for highly specific neuron subtypes or anatomical sub-regions that have diverged significantly between mice and primates. Our adversarial method cannot force a common representation if the underlying biological link between a neuron's identity and its activity patterns has fundamentally changed.

Response : Question3

Our work provides a framework for understanding the specific function of an individual computational unit (a biological neuron) by isolating its intrinsic properties from confounding variables. This has a direct analogy in mainstream AI.

One could apply our adversarial framework to an artificial neural network (ANN) to learn a robust representation for each individual node or neuron. By training the model to perform its primary task while simultaneously making it impossible for a discriminator to tell which training batch or input type an individual node's activation came from, we could uncover the node's core, invariant function.

This would help us understand the "division of labor" within large AI models, improving interpretability by revealing the specific role of their most basic components.

We thank the reviewer once again for their constructive feedback, which has helped us significantly improve the clarity and impact of our manuscript. We are confident that the revised paper now more robustly supports our contributions and provides a solid foundation for future work on generalizable single-neuron representation learning. In light of these revisions, which we believe address the initial concerns, we would be grateful if you would reconsider your assessment of our work.