Conditional Neural Optimal Transport for Predicting Cellular Phenotypes from Molecular Structure

2026-08-14

AI Summary

This study addresses the challenge of cellular phenotypic characterization arising from the vastness of chemical space by proposing a Molecule-Conditioned Neural Optimal Transport model. Conditioning on molecular structures, this approach integrates Monge-gap regularization with compressed representation space transport to learn mappings from negative controls to perturbed phenotypes, effectively overcoming static coupling failures in large-scale datasets. Experimental results demonstrate that the proposed framework significantly outperforms baseline methods on unseen active molecules by accurately recovering specific phenotypes and mitigating technical variation. Furthermore, it enables cross-batch phenotypic comparison and zero-shot generalization, validating its efficacy for virtual screening applications in drug discovery.

Abstract

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.

Problem

• Cellular Phenotype Prediction
• Molecular Structure
• High-content Microscopy
• Conditional Optimal Transport
• Out-of-distribution Generalization

Innovation

• Neural Optimal Transport
• Monge-Gap Regularization
• Conditional Transport
• Phenotype Prediction
• Compressed Representation Space