Confounder-aware Cell Painting latent diffusion model

Cell Painting stains six cellular compartments and reads out how a compound reshapes cell morphology, which makes it the workhorse assay for image-based profiling in early drug discovery. Its weakness is that the morphological signal is entangled with the experiment: which laboratory ran the plate, which batch it belonged to, and where on the plate the well sat all shift the pixels. Standard remedies such as Harmony correct batch effects after the fact and assume the variation is a linear transformation of derived cell profiles — an assumption that has left image-based methods weak at generalizing to compounds they have never seen.

This model, developed by Pfizer with Sinkove, Imperial College London, and King's College London, takes the opposite route: it treats the confounders as causal variables inside the generative process rather than as noise to subtract afterwards. A structural causal model over source laboratory, batch, and well position is embedded in a latent diffusion model that is additionally conditioned on chemical structure, so the network learns to generate a Cell Painting image given a compound and a specified experimental context. A g-estimation-inspired procedure then synthesizes balanced image sets in which confounder combinations are held fixed while the compound varies, isolating the morphological effect attributable to the chemistry.

The authors deliberately do not give the model a product name; the paper refers to it throughout as the confounder-aware foundation model. It was published in npj Imaging after a December 2024 preprint.

#Key Features

• Causal structure inside the generator: Source, batch, and well position enter the latent diffusion architecture as a structural causal model, so generated images reflect the relationships among compound, phenotype, and experimental context.
• Chemistry conditioning: A compound's SMILES string is encoded by a pretrained MolT5 chemical language model, and the resulting embedding conditions the diffusion process — which lets the model be queried with molecules absent from training.
• Synthetic balanced datasets: Rather than correcting real images, the model generates counterfactual image sets across confounder combinations, from which cell profiles are extracted.
• No task-specific heads: The same profiles support mechanism-of-action identification and target prediction through nearest-neighbor comparison of biosimilarity, with no fine-tuning.

#Technical Details

Pretraining used 13,361,250 Cell Painting images spanning 107,289 compounds from the JUMP Cell Painting Consortium dataset, which supplies the multi-laboratory, multi-batch structure the causal model exploits. Evaluation drew on the subset with Drug Repurposing Hub annotations — 525 compounds for mechanism prediction and 465 for target prediction, all at a single 10 µM dose. Two synthetic set sizes were generated: 5,000 images across 10 confounder combinations and 50,000 across 100.

Against real images, batch-corrected real images, and an otherwise identical model without the causal component, the confounder-aware model was strongest where generalization matters most. For compounds seen in training but in unseen batches, it reached 0.66 ROC-AUC on mechanism and 0.65 on target, against 0.59 and 0.57 for real data. For compounds entirely unseen during training, mechanism prediction held at 0.65 while target prediction rose to 0.73 ROC-AUC, with mean average precision of 0.22 — versus 0.59 ROC-AUC for the non-causal model and 0.45 for real data. Mechanism prediction showed no significant gap between the causal and non-causal models, which the authors read as evidence that compound conditioning alone carries mechanism signal while the experimental confounders bear more heavily on target identification.

#Applications

The intended use is hit expansion in early drug discovery: given a compound's structure, generate its expected morphological response, compare the resulting profile against annotated references, and nominate a mechanism and a target before the compound is ever plated. Because generation is conditioned on structure rather than requiring an image, this extends image-based profiling to laboratories without Cell Painting instrumentation, and to regions of chemical space no screen has covered. The counterfactual mode is also useful for assay design, showing what a compound would look like under a different laboratory's imaging conditions.

#Impact

Compound-conditioned Cell Painting generators such as MorphoDiff had already shown that diffusion models can paint plausible perturbation images; this work established that such a model, benchmarked on mechanism and target prediction at scale, can beat the real images it was trained on. Its more transferable claim is methodological: encoding known experimental confounders causally within a generative model outperforms correcting for them downstream, on the task where correction methods are weakest — an argument that extends to any assay whose signal is stratified by site and batch. The caveats are real. The causal graph covers only source, batch, and well position; dose, dose-response shape, and polypharmacology are unmodeled, and the JUMP subset used carries one annotated mechanism per compound. Absolute ROC-AUCs in the 0.65–0.73 range describe a useful triage signal rather than a confident call. Release is partial: the DMJUMP codebase ships as a supplementary zip attached to the paper, and the synthetic images, cell profiles and metric values are browsable from a public S3 bucket, but the trained weights are not distributed and there is no model card, so the model cannot be run on a new compound outside the authors' hands.