CellPainTR: Contrastive Batch Corrected Transformer for Large Scale Cell Painting

Thank you for acknowledging the additions to our paper and for raising your score. We appreciate your continued engagement and the constructive feedback.

Training Time and Computational Efficiency

We understand the concern regarding computational inefficiency. To clarify, the reported training time reflects a setup that was not optimized for speed (e.g., using a single Quadro 4000 GPU, without distributed training or mixed precision). These timings are provided solely indicatively for users running the current publicly available code in similar conditions. While computational efficiency is important, our focus in this paper is on the method’s design and its efficacy in addressing the dual objectives of batch correction and biological signal preservation in the Cell Painting context (which does not impact batch correction and biological signal preservation). A clearer discussion of potential optimizations (e.g., improved hardware, data processing, distributed training) has now been added to the revised manuscript.

Intermediate Stages and Practical Guidance

We have expanded the discussion to provide clearer guidance on when intermediate stages—such as CellPainTR(1) or CellPainTR(2)—might suffice, depending on the user's goals:

• CellPainTR(1): Suitable for applications requiring foundational representation learning but retains significant batch effects, particularly across data sources. (Appendix - Visualization - Figure 7)
• CellPainTR(2): An excellent choice for analyzing single-source datasets, offering strong intra-source batch correction and biological signal retention.
• CellPainTR: The final model provides the most robust inter-source batch correction, making it ideal for large-scale, multi-source analyses, albeit with minor biological signal trade-offs compared to CellPainTR(2). (Appendix - Visualization)

We have emphasized these trade-offs and their implications in the revised discussion, along with visual and quantitative evidence in the appendix to demonstrate why the additional training for CellPainTR is justified for specific use cases. This nuanced explanation aims to offer better practical guidance and transparency. Thank you once again for the valuable feedback, which has helped us refine the paper further.

Dear Reviewer z5XD,

To further address your concern, we conducted additional experiments using the Bray et al. dataset. We would like to kindly remind you that the Bray et al. dataset was entirely unseen during training. In order to properly train the model, we would typically need to separately train the source token specific to this dataset. However, due to the time constraints of this revision, we opted to identify the most suitable pre-trained token to serve as a proxy for the fully trained one. Specifically, we found that the source token for "Source 10" performs effectively as a proxy.

Additionally, due to differences between the older version of CellProfiler (which generates 1,383 features) and the version used in this study (which generates 4,765 features), only 275 features common to both versions were used for this experiment. While there are minor differences in naming conventions, this subset represents the intersection of compatible features. For this experiment, we did not apply any curation to align compound diversity across datasets.

To ensure consistency, we applied the same preprocessing and evaluation pipeline uniformly across all methods, including ComBat, Harmony, and CellPainTR. The results are presented as follows.

Results

Interpretation

We would like to highlight that CellPainTR demonstrated strong generalization to this unseen dataset, even under a suboptimal setup, which included:

• The use of a source token not specifically optimized for this dataset.
• A limited subset of features (275 out of 3,251) from those it was originally trained on.

Despite these constraints, CellPainTR achieved the highest overall performance across both batch correction and biological metrics compared to other methods. This underscores its robustness and effectiveness under challenging conditions. Notably, the model's design—particularly its use of context tokens—appears to facilitate superior generalization compared to traditional methods like ComBat and Harmony.

These findings provide strong evidence of the model's adaptability to new datasets, even when faced with significant limitations and time constraints. We hope this additional experiment has fully addressed your concerns.

We sincerely thank reviewer ciaD for their thoughtful follow-up question, which allows us to clarify important aspects of our work.

Regarding Cross-Dataset Evaluation

While we understand the value of replicating the cross-dataset evaluation from [1], our current approach faces some methodological constraints that make direct comparison challenging:

1. Our model uses fixed source context tokens that are specifically trained to learn and align source-specific distributions through supervised contrastive learning. New sources require fine-tuning of new context tokens before inference - a limitation we now discuss more thoroughly in the appendix.

2. Replicating [1]'s experimental setup with Bray et al. would effectively make it a 14th source in our framework, providing similar insights to our existing 13-source evaluation rather than a truly independent cross-dataset validation.

Architectural and Feature-Space Distinctions

We appreciate the observation about CellProfiler's image-based nature. However, there are fundamental differences between our approach and [1] that warrant consideration:

1. Feature Space Choice:

• [1] learns features directly from images
• Our method operates on CellProfiler's engineered features

2. Design Philosophy:

• Deep learning features likely capture richer context but at the cost of interpretability
• CellProfiler features (intensity-based, texture-based, etc.) maintain direct biological interpretability
• Our approach deliberately preserves this interpretability through the CWMM loss while addressing batch effects

Complementary Rather Than Competing Approaches

We view our work as complementary to image-based approaches like [1] rather than competing with them:

• Our focus is on establishing foundations for interpretable batch correction within well-understood feature spaces
• This lays groundwork for future work bridging deep learning's representational power with biological interpretability requirements
• Direct benchmarking against image-based methods would misrepresent the distinct goals and constraints of each approach

We hope this clarifies our methodological choices and positions our contribution in the broader context of the field. We have expanded our discussion of these points in Appendix D of the revised manuscript.

We sincerely thank the reviewer for their feedback.

W1. Concerns about the novelty of incorporating context tokens.

A. While we appreciate the reviewer's perspective, we would like to again emphasize the following novelties in the several key aspects:

1. Training Strategy & Purpose:

   • ContextViT: Single-phase training for visual domain variations
   • CellPainTR: Novel three-phase training strategy:
     • Phase 1: Unconstrained pretraining for general pattern learning
     • Phase 2: Intra-source fine-tuning using supervised contrastive learning for source-specific batch effects
     • Phase 3: Inter-source training for refined batch effect understanding

2. Context Token Implementation:

   • ContextViT: Complex dynamic inference through pooling and transformation
   • CellPainTR: Efficient codebook lookup system where:
     • Each source has a dedicated context token
     • Tokens are directly retrieved via source ID
     • Simple concatenation with input

3. Problem-Specific Innovation:

   • Different modality: Cell profiles vs. images
   • Different challenge: Batch correction vs. visual domain adaptation
   • Specific requirements:
     • Known, fixed sources
     • Preservation of biological signal while removing technical variation
     • Need for interpretable source-specific corrections

While we acknowledge that both approaches use the term "context token," the fundamental differences in implementation, purpose, modality, and training strategy make CellPainTR's approach a distinct and novel contribution to computational biology and batch correction. The confusion may arise from the similar terminology, but the approaches serve different purposes and solve different problems in their respective fields.

W2. Request for comparison with alternative methods to Hyena operators.

A. The reason to decide to use Hyena operator was the long token size in Cell Painting problems. For example, in our case the number of feature was 3251, which is a bit longer than the direct use of naive BERT architecture. The use of VQ-VAE could reduce the token size to generate smaller latent space, but we found that the resulting compression makes the training unstable, implying that the raw data (rather than compressed latent) is important for batch correction. That being, we will discuss this in more detail in the final version of the paper.

W3. Request for significance measures given small improvements.

A. Regarding Statistical Testing for the Overall Score We computed the Overall Score as suggested by the reviewer. This score is the mean of two aggregate scores—Batch Correction Aggregate and Biological Metrics Aggregate—which themselves are means of their respective categories, each capturing a distinct aspect of performance.

W4. Request for ablation study on correlated features.

A. Thanks for the interesting suggestion. The reviewer is kindly reminded that our goal is not to remove the correlation. The suggestion of removing highly correlated features reveals a fundamental misunderstanding of our goals:

1. Correlation in features is highly context-dependent:

   • Varies by dataset based on:
     • Batch effects
     • Compound treatments used
     • Cell line responses
     • Phenotypic perturbation patterns
     • Distribution of perturbations (extreme vs. normal)

2. While feature correlation removal is useful for single experiments to address dimensionality challenges, our goal is fundamentally different:

   • We aim to create a foundational model ready for any present or future experiment
   • The model must handle varying correlation structures across different experimental contexts

3. Making our feature set dependent on correlation analysis would:

   • Limit the model's generalizability
   • Make it dataset-dependent
   • Reduce its utility as a universal tool for batch correction Our approach instead allows the model to learn and adapt to these varying correlation structures while maintaining independence from any specific dataset's characteristics.

In conclusion, while we understand the reviewer's interest in an ablation study of correlated features, such an approach would fundamentally compromise our goal of creating a generalizable foundation model. Furthermore, removing highly correlated features would actually work against our objectives by making the model dependent on specific dataset characteristics and potentially biased by dataset-specific batch effects. Our design choice to maintain all features and let the model learn appropriate correlations is essential for ensuring broad applicability across different experimental contexts and future datasets.

We sincerely thank the reviewer for their thoughtful follow-up and for seeking clarification regarding the interpretability of features in our work. This is an essential point, and we appreciate the opportunity to elaborate on it.

Clarifying "Interpretability"

You are correct that the term "interpretability" would benefit from a more precise definition in the context of this discussion. As the reviewer noted, features extracted by tools like CellProfiler—such as Zernike moments or texture measurements—may not inherently correspond to specific biological phenomena. However, these features are measurable, understandable metrics (e.g., intensity, shape, texture) that, individually or in combination, can provide insights or hypotheses about underlying biological mechanisms. For instance, in laboratory settings, CellProfiler features are routinely used in Cell Painting assays to link cell morphology to specific mechanisms of action.

In contrast, deep learning-derived features, while powerful, are often difficult to interpret, even within an image analysis framework. This is one reason why engineered features remain a robust choice in scenarios where biological interpretability is crucial.

Supporting Evidence

The importance of engineered features for interpretability is exemplified in works such as Cross-Zamirski et al., 2022 ([1]). In this study, the authors aim to predict Cell Painting images from brightfield inputs. Crucially, the evaluation of their model relies on comparing the cell profiles of generated images to those of the ground-truth images. This approach underscores the value of engineered features as a shared, interpretable metric for assessing biological relevance across imaging modalities. By contrast, representations learned directly from images (e.g., deep learning embeddings) do not inherently guarantee alignment with biological phenomena. Engineered features, by design, are tied to measurable aspects of cell morphology, which in turn relate to cellular conditions and states—the foundational hypothesis behind Cell Painting.

Scope of Our Paper

While we acknowledge the importance of linking features to specific biological phenomena, we emphasize that this is not the focus of our current work. Our paper builds on the established premise of Cell Painting: that morphological measurements provide meaningful insights into cellular conditions. Thus, demonstrating the interpretability of specific features in the context of biological conditions falls outside the scope of this study. Addressing this question would require a broader, more foundational investigation, which we believe aligns more closely with basic science or philosophical inquiries into the principles of Cell Painting.

Additional Resources

For further context, we point the reviewer to several seminal papers that address the interpretability and utility of Cell Painting and traditional feature extraction:

• [2]: Introduces Cell Painting and highlights how extracted features relate to morphology and biological phenomena.
• [3]: Explores trade-offs between deep learning and traditional features, emphasizing interpretability.
• [4]: Examines the utility of image-based profiling strategies, contrasting interpretability in traditional vs. computational approaches.
• [5]: Discusses challenges of linking machine-learning-derived features to biological meaning compared to traditional methods.
• [6]: Explores large-scale image-based profiling and highlights interpretability as a key strength of traditional features.

We hope these references provide additional clarity and demonstrate how "interpretability" is understood and leveraged within the broader Cell Painting community.

Concluding Remarks

To conclude, while we recognize the fundamental importance of the question raised, we believe it is beyond the immediate scope of our paper. Instead, our work assumes the established premise of Cell Painting and engineered features as interpretable metrics. We hope our response has addressed your concerns and provided clarity on this matter.

[1] Cross-Zamirski, J. O., Mouchet, E., Williams, G., Schönlieb, C. B., Turkki, R., & Wang, Y. (2022). Label-free prediction of cell painting from brightfield images. Scientific reports.

[2] Bray, Mark-Anthony, et al. "Cell Painting, a high-content image-based assay for morphological profiling using multiplexed fluorescent dyes." Nature protocols

[3] Christiansen, Eric M., et al. "In silico labeling: predicting fluorescent labels in unlabeled images." Cell

[4] Caicedo, Juan C., et al. "Data-analysis strategies for image-based cell profiling." Nature methods

[5] Grys, Ben T., et al. "Machine learning and computer vision approaches for phenotypic profiling." Journal of Cell Biology

[6] Bougen‐Zhukov, Nicola, et al. "Large‐scale image‐based screening and profiling of cellular phenotypes." Cytometry Part A