CellCLIP – Learning Perturbation Effects in Cell Painting via Text-Guided Contrastive Learning

We introduce CellCLIP, a framework for contrastive learning (CL) on Cell Painting data. Unlike natural image CL methods, our approach accounts for the unique structure of multi-channel cellular images and the many-to-one relationship between perturbations and image sets. CellCLIP integrates three key components:

CrossChannelFormer

To encode images, we adapt pretrained natural image models (e.g., DINOv2) by treating each Cell Painting channel as a grayscale input. We then aggregate images from the same perturbation using gated attention pooling, producing a per-perturbation profile. Finally, our CrossChannelFormer introduces lightweight cross-channel reasoning by combining pooled profiles with channel-specific embeddings and a global CLS token. This design enables efficient modeling of stain-specific structures while requiring only C+1 tokens per perturbation.

Perturbation Encoding

Instead of building modality-specific encoders, we represent perturbations as natural language prompts. These descriptions capture both cell type and perturbation details (e.g., compounds or CRISPR targets), enabling a unified representation across diverse perturbation types. We encode prompts with a pretrained BERT model. For example to encode the chemical compound butyric acid, a drug affecting cell growth, we use the prompt:

Similarly, for a CRISPR perturbation, the prompt is structured as:

Training Objective

To align image profiles with perturbation text, we combine the standard CLIP loss with a Continuously Weighted Contrastive Loss (CWCL). CWCL softly reweights pairs based on morphological similarity, ensuring that biologically related perturbations remain close in the embedding space while preserving retrieval performance. This objective balances cross-modal alignment with biologically meaningful intra-modal structure.

For the profile-to-perturbation direction, \( \mathcal{U} \rightarrow \mathcal{V} \), our adapted CWCL objective is given by

\[ \mathcal{L}_{\mathrm{CWCL}, \mathcal{U} \rightarrow \mathcal{V}} = \frac{1}{N} \sum_{i=1}^N \frac{1}{\sum_{j \in [N]} w_{ij}^{\mathcal{U}}} \left[ \sum_{j=1}^N w_{ij}^{\mathcal{U}} \cdot \log \frac{\exp(\langle p_i \cdot q_j \rangle / \tau )} {\sum_{k=1}^N \exp(\langle p_i \cdot q_k \rangle / \tau )} \right]. \]

Here \( p_i \) represents the output of CrossChannelFormer applied to images corresponding to perturbation \( i \), \( q_i \) corresponds to our encoding of the natural language description of perturbation \( i \), and \(w_{ij}^{\mathcal{U}}\) denotes the similarity between the pooled profiles \(\mu(U_i)\) and \(\mu(U_j)\), used to reweight the alignment from modality \(\mathcal{U}\) to modality \(\mathcal{V}\)

For the perturbation-to-profile direction, \( \mathcal{V} \rightarrow \mathcal{U} \), we apply the standard CLIP loss. Thus, the final training loss for CellCLIP is

\[ \mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{CWCL}, \mathcal{U} \rightarrow \mathcal{V}} + \mathcal{L}_{\mathrm{CLIP}, \mathcal{V} \rightarrow \mathcal{U}}. \]

Together, these design choices allow CellCLIP to leverage powerful pretrained models, efficiently represent high-dimensional Cell Painting data, and generalize across diverse perturbation types.

Cross-modality retrieval

We present our results for cross-modality retrieval tasks (i.e., perturbation-to-profile and profile-to-perturbation) in Table 1. For our benchmarking we compared CellCLIP against previous state-of-the-art CL methods for Cell Painting perturbation screens: CLOOME and MolPhenix. Overall, we found that CellCLIP demonstrated substantially higher performance at both cross-modal retrieval tasks compared to baseline methods. To understand the source of performance gains in CellCLIP, we conducted a series of ablations to understand the contribution of each component in CellCLIP to retrieval performance. Specifically, starting with CLOOME's proposed encoding scheme, where individual images encoded using ResNet50 are aligned with chemical perturbations encoded using Morgan fingerprints combined with an MLP, we gradually replaced each of CLOOME's components with those of CellCLIP and assessed each change's impact on model performance.

Language can effectively represent perturbations.

We began by replacing CLOOME's chemical structure encoder, which feeds chemicals' Morgan fingerprints through a multi-layer perceptrion (MLP), with the natural language encoder used in CellCLIP while holding all other model components fixed. We found that this change alone yielded significant performance gains for both retrieval tasks. Notably, we found this result continued to hold even after replacing the generic MLP used in CLOOME with an MPNN++ network designed specifically for molecular property prediction

Cross-channel reasoning improves retrieval performance

We next investigated the impact of CrossChannelFormer's ability to reason across global Cell Painting channel information. To do so, we replaced CLOOME's ResNet image encoder with our CrossChannelFormer encoder. To isolate the effects of CrossChannelFormer's ability to reason across channels from the effects of per-perturbation pooling, in this experiment we removed the CrossChannelFormer's pooling function and trained the resulting model (denoted as CrossChannelFormer\(^{\dagger}\) in Table above) on individual image-perturbation pairs. We found that this change led to an additional increase in model performance. To understand how our approach compared to previous channel-aware vision encoding approaches, we ran this same experiment using a channel agnostic MAE (CA-MAE) and ViT (ChannelViT) as vision encoders. We found that CrossChannelFormer consistently outperformed these baselines on both tasks. This demonstrates that image embeddings extracted from models pretrained on natural images can be effectively leveraged with CrossChannelFormer. Additionally, our approach substantially reduced training time, achieving a 3.9 times speedup compared to CLOOME and a 2.2 times speedup compared to other channel-agnostic methods. This efficiency stems from CrossChannelFormer operating in a compact feature space and requiring only \(C+1\) tokens per instance.

Pooling yields improved alignment and computational efficiency

Finally, we evaluated the impact of the attention-based pooling operator within CrossChannelFormer compared to instance-level training (i.e., no pooling). We found that including pooling resulted in yet another increase in model performance. In addition, by reducing the number of pairs for contrastive loss computation, pooling yielding a 6.7 times speedup in training time relative to instance-level training. We also explored the impact of varying our choice of pooling operator, and found that attention-based pooling yielded the best results among the pooling operators we considered.

Overall, our results demonstrate that each of our main design choices in CellCLIP contributes to improved retrieval performance and computational efficiency compared to prior work.

Recovering known biological relationships

We report recall of known biological relationships among genetic perturbations in RxRx3-core. Across all benchmark databases, we found that CellCLIP achieved the best recovery of known gene-gene relationships compared to baseline models. Notably, we found that replacing the CWCL loss used in CellCLIP with the standard CLIP loss leads to worse performance on this task, illustrating the benefits of using soft labeling for alignment in the image profile space \( \mathcal{P} \). Altogether, these results further demonstrate that CellCLIP can recover meaningful relationships between perturbations in its image profile latent space.