PaperlayerView exact revision on arXiv

Source-linked AI summary

Conditional Neural Optimal Transport for Predicting Cellular Phenotypes from Molecular Structure

Gauthier Avité, Maxime Sanchez-Renauld, Nicolas Bourriez, Auguste Genovesio

arXiv:2608.14293v1cs.CVcs.LG

TL;DR

Predicting cellular phenotypes across chemical space is experimentally infeasible, motivating models that infer image-derived responses from controls and molecular structure. The paper introduces molecule-conditioned Neural Optimal Transport and finds that it outperforms baselines on molecules unseen during training.

  • Problem

    The chemical space of drug-like molecules exceeds imaging capacity, creating a need to predict image-derived phenotypes without acquiring treated cells.

  • Method

    Molecule-conditioned Neural Optimal Transport learns an inductive map from negative-control phenotypes to perturbed phenotypes using molecular structure as conditioning.

  • Results

    On molecules unseen during training, the model outperforms baselines and recovers molecule-specific phenotypic effects with improved cross-batch comparability.

  • Takeaways & Limitations

    The findings support chemically conditioned transport as a framework for phenotype prediction beyond observed compounds, while identifying molecular representation quality as central to generalization.

  • Takeaways & Limitations

    Molecular representation quality remains the main bottleneck, and the study did not evaluate a large pretrained graph encoder for out-of-distribution generalization.

Abstract

from arXiv · show

High-content microscopy enables systematic profiling of cellular responses to chemical perturbations, but the scale of the chemical space makes exhaustive phenotypic characterization experimentally infeasible. This motivates computational models that can predict image-derived phenotypes without acquiring the corresponding treated cells. We formulate molecule-induced phenotype prediction as an inductive conditional transport problem in image representation space. Given a negative-control phenotype and the structure of a molecule, we aim to predict the phenotype induced by the corresponding molecule. We first evaluate classical optimal transport baselines and show that static couplings do not yield useful predictions on large-scale phenotypic image datasets. We then introduce a molecule-conditioned Neural Optimal Transport (NOT) model with a Monge-Gap regularization training objective that learns to transport negative-control unperturbed phenotypes toward perturbed phenotypes using molecular structure as conditioning information. NOT recovers molecule-specific phenotypic effects while reducing microscopy-associated technical variation, thereby facilitating comparisons across experimental batches. On unseen active molecules, the model outperforms baseline approaches, demonstrating that chemically conditioned transport can generalize beyond the molecules observed during training. We identified the molecular encoder as the main limitation to this generalization, while transport in a compressed representation space improves performance and scalability. These results establish NOT as a promising framework for predicting cellular phenotypes from molecular structure and negative-control phenotypes, while highlighting the development of more informative molecular representations as a key direction for improving out-of-distribution performance.

1 Introduction

This section frames molecule-induced phenotype prediction as a conditional transport problem in Cell Painting representation space and introduces molecule-conditioned Neural Optimal Transport to predict perturbed phenotypes from controls and molecular structure. The approach recovers molecule-specific effects, improves cross-batch comparability, and generalizes to unseen active molecules, while exposing molecular representation quality as a key limitation.

  • Motivation: Cell Painting profiles cellular responses at scale, but drug-like chemical space far exceeds the coverage of any imaging campaign.Large-scale initiatives such as JUMP-CP provide standardized profiles for hundreds of thousands of perturbations.
  • Problem formulation: The task predicts perturbed phenotypes from negative-control DMSO representations and molecular structure, matching predicted and observed phenotype sets without paired control-treated wells.Images are encoded with self-supervised DINOv2 and aggregated at the well level.
  • Method: Molecule-conditioned Neural Optimal Transport learns an inductive map that transforms individual negative-control phenotypes while matching predictions to the observed perturbed distribution.The method builds on the Monge-Gap estimator and its conditional extension.
  • Results: Neural OT recovers molecule-specific phenotypic effects and produces profiles more comparable across experimental batches than raw or DMSO-normalized measurements.The model is evaluated under complementary generalization settings involving held-out plates and held-out active molecules.
  • Results: On held-out active molecules, the model outperforms random and identity baselines, supporting generalization beyond compounds observed during training.The molecular encoder is identified as the main limitation to this generalization, while compressed representation-space transport improves performance and scalability.

2 Related Work

Prior work aligns molecular structures with microscopy phenotypes, compares distributions using optimal transport, and learns parametric transport maps for response prediction. These approaches motivate molecule-conditioned transport while leaving gaps in native morphology-space shift modeling and stability beyond fixed couplings.

  • Cross-modal and generative phenotype modeling: Contrastive methods align molecular structures and microscopy phenotypes in a shared space for cross-modal retrieval, but do not model shifts from negative-control morphology.CLOOME, MoCoP, and MolPhenix can generalize molecular representations to unseen molecules, yet score against acquired images rather than phenotype distributions in native morphology space.
  • Optimal transport: Classical Wasserstein, Gromov–Wasserstein, and Fused Gromov–Wasserstein methods compare or align distributions using sample, relational, and combined feature-structural costs.GW supports alignment across different metric spaces, while FGW is suited to structured and multimodal alignment when feature information is partial.
  • Neural Optimal Transport: Neural Optimal Transport learns parametric maps beyond fixed-sample couplings, including perturbation-conditioned Brenier-potential models and CellOT for unpaired single-cell response prediction.The Monge Gap regularizes an unconstrained map toward optimality, avoiding convexity constraints and min–max instabilities associated with earlier estimators.

3 Method

The method learns an inductive molecule-conditioned transport map from negative-control phenotypes using a Monge-Gap regularizer and optional unbalanced resampling. Molecular structure conditions a neural architecture that transports phenotypic representations toward perturbed distributions.

  • Inductive conditional transport: The model learns a parametric transport map Tθ that can be evaluated on new points, rather than a coupling tied to a fixed sample.This provides the inductive formulation needed for prediction on unseen inputs.
  • Monge Gap: The Monge-Gap regularizer guides a generic neural map toward ε-optimal transport without imposing convex architectural constraints.It measures excess transport cost relative to the entropic OT optimum and vanishes when the map is ε-optimal for its source and pushforward.
  • Unbalanced resampling: Unbalanced resampling relaxes exact marginal matching to address outliers and over- or under-represented regions across Cell Painting plates and batches.The heuristic replaces source and target batches with OT-based resampled counterparts before evaluating the training objective.
  • Architecture and conditioning mechanism: The conditioned map Tθ(x, fm) combines a negative-control phenotype with a molecular representation derived from chemical structure.The molecular representation is projected through a trainable encoder before conditioning the phenotype transformation.
  • Architecture and conditioning mechanism: Multi-head attention followed by a feed-forward residual block produces the transported phenotype, while the objective uses separate entropic scales εfit and εreg and a cosine ground cost.The cosine cost matches the geometry of DINOv2 phenotypic representations, and both objective terms are differentiable in θ.

4 Data and Evaluation Protocol

The evaluation uses JUMP-CP Cell Painting phenotypes with DMSO source distributions and repeated positive controls. It tests plate-level generalization across held-out batches and molecule-level generalization to unseen structures through cosine-similarity retrieval.

  • Dataset: JUMP-CP provides standardized, large-scale Cell Painting perturbation experiments across sites and batches in U2OS cells.Each five-channel microscopy image is converted into a single image-level phenotypic representation.
  • Dataset and controls: DMSO wells define source phenotype distributions, while eight repeated positive-control molecules benchmark generalization across held-out plates with molecule identity fixed.The controls recur across plates, laboratories, and experimental batches.
  • Plate-level split: Plate-level splitting holds out experimental plates and batch effects while retaining all eight molecules in training, using 5-fold cross-validation and three random seeds.Each training split averages 260,262 DMSO points and 142,196 non-DMSO positive-control points; validation averages 65,065 and 35,549, respectively.
  • Molecule-level split: Molecule-level splitting evaluates structure-driven generalization because training and evaluation molecules are disjoint, using 5-fold cross-validation and 3 random seeds.The benchmark is restricted to the 12,400 molecules in the top 10% of phenotypic activity.
  • Retrieval evaluation: Predicted phenotypes are ranked against measured candidates by cosine similarity, with success defined by retrieving the correct molecule among the top k.Plate-level evaluation reports Retrieval@1 for eight positive controls, whereas molecule-level evaluation reports R@10rep and R@10plate; plate-level retrieval reduces replicate noise and is emphasized as more stable.

5 Results

Neural OT predicts molecule-specific cellular phenotypes across held-out plates and unseen compounds, while generated profiles improve cross-plate comparability. Out-of-distribution performance depends strongly on molecular encoding and benefits from compressing phenotype representations.

  • In-distribution prediction: Neural OT improves Retrieval@1 for all eight molecules and substantially increases the mean relative to unchanged DMSO controls.It reaches or exceeds independently held-out replicate performance, though it remains slightly below the plate-mean replicate reference.
  • In-distribution prediction: UMAP visualizations show Neural OT transporting controls toward molecule-specific target regions and recovering clusters for the strongest phenotypic responses.The projection was fitted jointly on source and target embeddings.
  • Cross-plate comparability: 0.65 mean average precision was achieved by generated phenotypes in cross-plate retrieval, versus 0.36 for raw measured phenotypes and 0.47 after DMSO normalization.The shared transport map produces standardized surrogate profiles from plate-specific DMSO controls.
  • Out-of-distribution prediction: 0.095 ± 0.007 R@10plate was achieved by NOT on unseen compounds, versus 0.059 ± 0.003 for identity and 0.004 for plate-mean prediction.At replicate level, the ordering was 0.035 vs. 0.018 vs. 0.004.
  • Molecular encoder ablation: 0.095 R@10plate was obtained with the combined morganc+rdkc descriptor, rising from 0.029 with the weakest molecular encoder.Retrieval improved markedly from unimol2 and molformer to morgan and was best with morganc+rdkc.
  • Representation and loss ablations: 0.101 plate R@10 was achieved with a 50D autoencoder for phenotype embeddings, compared with 0.075 for PCA and 0.066 without reduction.Reducing the conditioning embedding instead lowered plate R@10 from 0.095 to 0.036–0.037; cosine and squared-Euclidean costs performed comparably.

6 Conclusion

The study introduces molecule-conditioned Neural OT for predicting Cell Painting phenotypes from plate-specific negative controls and molecular structure, recovering molecule-specific effects and generalizing to unseen molecules. Ablations identify phenotype-space compression and molecular representation quality as key determinants, while limitations motivate stronger encoders and extensions beyond molecules.

  • Core contribution: Molecule-conditioned Neural OT predicts individual Cell Painting phenotypes from plate-specific negative controls and molecular structure through distribution-level training.The formulation treats phenotype prediction as an inductive conditional transport problem.
  • Core contribution: On held-out plates, the model recovers molecule-specific effects and produces surrogate profiles more comparable across batches than raw or normalized real profiles.
  • Generalization: On unseen molecules, the model outperforms baselines, indicating that structure-conditioned phenotype prediction can generalize beyond compounds observed during training.
  • Ablations: Multi-head attention improves conditioning, DINOv2 phenotype-space compression improves retrieval, and molecular compression substantially degrades performance, making molecular representation quality a major OOD bottleneck.Monge-Gap strength controls a trade-off between replicate-level and plate-level retrieval.
  • Limitations: Absolute retrieval remains modest, and the molecule-identity split does not guarantee generalization to structurally distant chemical series.The study also did not evaluate potentially strong closed-source pretrained graph encoders and found static optimal-transport couplings ineffective.
  • Future work: Future work should improve molecular graph encoders or train them during Neural-OT optimization and extend conditional transport to genetic perturbations using gene or protein encoders.Suggested encoders include Geneformer and ESM-2 for predicting knockout or related perturbation phenotypes.

Software and Data

The study will release its source code for reproducibility, and all source data are publicly available through the JUMP Cell Painting Consortium dataset.

  • Software and Data: Source code will be released to ensure reproducibility, while the study data are publicly available in the JUMP Cell Painting Consortium dataset (cpg0016).Data access instructions and documentation are available at the JUMP Hub website.

Use of generative AI. · A Optimal Transport Preliminaries

The authors used generative AI for manuscript clarity, code assistance, and figure preparation, with all outputs validated by the authors. The optimal-transport preliminaries distinguish Monge maps from Kantorovich couplings and emphasize that standard empirical couplings are transductive rather than predictive.

  • Use of generative AI.: Generative AI assisted manuscript language, code generation and review, and figure preparation.The authors inspected, tested, and validated all generated or modified text, code, and figures.
  • Use of generative AI.: The authors retain sole responsibility for the paper’s scientific content, analyses, interpretations, and conclusions.This responsibility applies despite the use of generative AI tools.
  • A Optimal Transport Preliminaries: The framework treats negative-control DMSO representations as distribution µ and molecule-perturbed representations as distribution ν over metric spaces.The preliminaries define metric spaces (X, d_X) and (Y, d_Y) equipped with probability measures µ and ν.
  • A Optimal Transport Preliminaries: The Monge problem seeks a transport map T: X → Y that pushes µ onto ν while minimizing total displacement under ground cost c(x, y).The Monge formulation is non-convex and may lack a solution, including for discrete measures with unequal support sizes.
  • A Optimal Transport Preliminaries: The Kantorovich relaxation optimizes ground cost over couplings π ∈ Π(µ, ν) with prescribed marginals.Unlike the Monge formulation, this is a convex linear program that always admits a solution.
  • A Optimal Transport Preliminaries: The optimal empirical coupling aligns mass between observed samples but is transductive, so it does not provide predictions for unseen molecules.Its definition is restricted to points available during optimization.

B Cell Painting data … C.1 Implementation Details

The paper uses Cell Painting phenotypes from the large-scale JUMP-CP dataset and defines an out-of-distribution benchmark from the most phenotypically active molecules. Neural Optimal Transport is trained on DINOv2-derived well-level phenotypes, usually compressed to 50 dimensions, with four molecular conditioning representations.

  • B.1 Cell Painting Assay: Cell Painting captures diverse single-cell visual characteristics using fixed fluorescent markers for cellular compartments and organelles.The assay measures structures including the nucleus, cytoskeleton, mitochondria, endoplasmic reticulum, and plasma membrane.
  • B.1 Cell Painting Assay: Cell Painting assumes perturbations affecting similar biological pathways induce similar phenotypes detectable through changes in cell morphology and organization.The datasets are high-dimensional, multichannel, and contain substantial variation unrelated to biological signals.
  • B.2 JUMP-CP: JUMP-Cell Painting is a large-scale dataset of U2OS Cell Painting images spanning chemical treatments, gene overexpression, and CRISPR-Cas9 knockouts.The dataset includes over 116,750 molecules and 12,602 gene ...
  • B.3 Selection of the Most Phenotypically Active Molecules: The out-of-distribution benchmark retains the top (10%) most phenotypically active molecules, ranked by distances between molecule representations and corresponding DMSO negative controls.The threshold was motivated by an estimate that approximately (10%) of screened molecules induce detectable phenotypic responses.
  • B.3 Selection of the Most Phenotypically Active Molecules: Activity scores and rankings were computed independently of model predictions and were withheld from the model during both training and inference.Measured target phenotypes were also not provided to the model.
  • C.1 Implementation Details: Image-level phenotypes are extracted with DINOv2-Giant and aggregated across experimental wells into single well-level representations.These representations serve as the phenotype inputs for Neural Optimal Transport training.
  • C.1 Implementation Details: Unless otherwise stated, phenotype embeddings are compressed to 50 dimensions with an autoencoder before entering the Neural Optimal Transport network.The study also evaluates PCA-based reduction and transport in unreduced DINOv2 space.
  • C.1 Implementation Details: The molecular conditioning comparison includes binary Morgan fingerprints, morganc+rdkc descriptors, MolFormer embeddings, and Uni-Mol2 embeddings.MolFormer uses molecular strings, whereas Uni-Mol2 incorporates molecular geometric information.

C.2 Neural OT Training Details and Hyperparameters · C.3 Out-of-Distribution Ablation Setup and Hyperparameters

The in-distribution Neural OT experiments used an attention-based conditional model trained with plate-level cross-validation, while the out-of-distribution study used a separate molecule-level sweep. The two settings therefore differed primarily in held-out units and evaluation configuration.

  • C.2 Neural OT Training Details and Hyperparameters: In-distribution experiments used 5-fold cross-validation over plates, with three random seeds per fold and 80% training versus 20% validation plates.Each fold contained about 1285 training and 322 validation plates.
  • C.2 Neural OT Training Details and Hyperparameters: The NOT architecture combines attention with feed-forward residual connections in a conditional design.
  • C.2 Neural OT Training Details and Hyperparameters: The model trained for 50 epochs with Adam, an initial learning rate of 10^-4, and warmup, constant, and cosine-decay phases.
  • C.2 Neural OT Training Details and Hyperparameters: Training batches were grouped by perturbation and plate, using a dataloader batch size of one grouped molecule item.
  • C.3 Out-of-Distribution Ablation Setup and Hyperparameters: The molecule-level out-of-distribution study used a separate hyperparameter sweep summarized in Table 7.Unless explicitly ablated, components used the listed baseline values.
  • C.3 Out-of-Distribution Ablation Setup and Hyperparameters: Unlike the in-distribution setup, the out-of-distribution study held out molecules rather than plates and estimated performance with 5-fold cross-validation over three seeds.
  • C.3 Out-of-Distribution Ablation Setup and Hyperparameters: All Section 5.2 results were means ± standard deviation over fifteen runs from the molecule-level cross-validation design.

C.4 Additional Out-of-Distribution Ablation Tables … D.3 Semi-supervised FGW

The paper reports ablations of structure encoding, ground cost, and resampling unbalancedness, then examines static OT, GW, FGW, and semi-supervised FGW baselines. These static couplings remain transductive, and semi-supervised FGW provides only marginal gains on JUMP-CP, with matching near chance even under substantial supervision.

  • C.4 Additional Out-of-Distribution Ablation Tables: The C.4 ablations cover the structure encoder, ground cost, and resampling-heuristic unbalancedness τ, with results reported as means ± standard deviation over 5 folds and 3 seeds.Only finished runs are included.
  • D Static Optimal Transport: Classical static OT couplings serve as baselines but are transductive, returning couplings over observed samples without a prediction rule for unseen molecules.The couplings use negative-control and perturbed morphology distributions as their marginals.
  • D.1 GW: GW compares intra-domain dissimilarities rather than samples directly and is invariant to isometries of each space, but its optimization is a non-convex quadratic assignment.These properties make GW appealing for aligning heterogeneous biological modalities.
  • D.2 FGW: FGW interpolates between feature-level Wasserstein and structural GW costs through α ∈ [0, 1], allowing both feature and relational information to influence transport.Partial supervision can bias the feature cost at known correspondences.
  • D.2 FGW: GW and FGW are typically optimized with conditional-gradient schemes that solve an entropic OT problem at each iteration, giving an overall O(n^3) cost.Low-rank couplings and costs can reduce this toward linear scaling in the number of samples.
  • D.3 Semi-supervised FGW: Semi-supervised FGW rewards known structure–phenotype correspondences through negative feature-cost entries, while structural transport propagates supervision to unlabeled pairs.The parameter λ controls prior strength, α balances supervision against geometry, and S = ∅ recovers unsupervised FGW.
  • D.3 Semi-supervised FGW: On JUMP-CP, even a large fraction of supervised correspondences yields only marginal gains, while the recovered coupling remains close to chance because of fingerprint–morphology manifold mismatch.The method remains transductive and matches only observed samples, providing no map for unseen molecules.
  • D.3 Semi-supervised FGW: Even 80% of supervised correspondences leaves static GW/FGW matching at chance level on the 1125-molecule most-active JUMP-CP subset.The table reports the number of correct structure→phenotype matches.

D.4 Empirical results on JUMP-CP

On the most-active 1% of JUMP-CP, static optimal-transport couplings failed to recover molecule–phenotype matches, even with substantial supervision, consistent with a mismatch between molecular-fingerprint and morphology geometries.

  • Static coupling results: Plain GW recovered essentially no correct structure→phenotype matches on 1125 active JUMP-CP molecules, performing on par with a random coupling.Injecting supervision into 80% of pairs increased recovery only from 1 to 3 matches out of 1125.
  • Interpretation: The near-chance JUMP-CP coupling is consistent with a strong manifold mismatch between fingerprint and morphology geometries, rather than a general limitation of semi-supervised OT.Semi-supervised OT can sharpen alignment when the two domains share recoverable geometry, but that condition is not met here.
Loading 2608.14293v1…