Building a virtual cell capable of accurately simulating cellular behaviors in silico has long been a dream in computational biology. We introduce CellFlux, an image-generative model that simulates cellular morphology changes induced by chemical and genetic perturbations using flow matching. Unlike prior methods, CellFlux models distribution-wise transformations from unperturbed to perturbed cell states, effectively distinguishing actual perturbation effects from experimental artifacts such as batch effects—a major challenge in biological data.

When evaluated on chemical (BBBC021), genetic (RxRx1), and combined perturbation (JUMP) datasets, CellFlux generates biologically meaningful cell images that faithfully capture perturbation-specific morphological changes, achieving a 35% improvement in FID scores and a 12% increase in mode-of-action prediction accuracy over existing methods. Additionally, CellFlux enables continuous interpolation between cellular states, offering a promising tool for studying perturbation dynamics. These capabilities represent a substantial step toward advancing virtual cell modeling in biomedical research.

Our method makes a meaningful contribution to computational biology, providing researchers with new tools for virtual cell modeling and the study of in-silico drug screening and personalized therapy development. The open-source implementation enables the community to build upon this work and accelerate progress in this field.

🔮 Method Overview

Our method aims to predict changes in cell morphology induced by chemical or gene perturbations in silico. The approach addresses the fundamental challenge of distinguishing true biological effects from experimental artifacts through innovative distribution-to-distribution modeling.

CellFlux Overview — **Figure 1:** Method overview. **(a)** Objective: Predict morphological changes from perturbations. **(b)** Data: High-content screening with control and perturbed wells. **(c)** Problem formulation: Distribution-to-distribution transformation. **(d)** Flow matching: Learning velocity fields for continuous transformation. **(e)** Results: Superior performance in generation quality and biological accuracy.

Key components of CellFlux include:

Objective. Predict changes in cell morphology induced by chemical or gene perturbations in silico. In this example, the perturbation effect reduces the nuclear size.
Data. The dataset includes images from high-content screening experiments, where chemical or genetic perturbations are applied to target wells, alongside control wells without perturbations. Control wells provide prior information that contrasts with target images, enabling the identification of true perturbation effects (e.g., reduced nucleus size) while calibrating non-perturbation artifacts such as batch effects—systematic biases unrelated to the perturbation (e.g., variations in color intensity).
Problem formulation. We formulate the task as a distribution-to-distribution problem (many-to-many mapping), where the source distribution consists of control images, and the target distribution contains perturbed images within the same batch.
Flow matching. Our method employs flow matching, a state-of-the-art generative approach for distribution-to-distribution problems. It learns a neural network to approximate a velocity field, continuously transforming the source distribution into the target by solving an ordinary differential equation (ODE).
Results. The method significantly outperforms baselines in image generation quality, achieving lower Fréchet Inception Distance (FID) and higher classification accuracy for mode-of-action (MoA) predictions.

💡 Key Innovation

Our key innovation lies in formulating cellular morphology prediction as a distribution-to-distribution learning problem. Traditional methods often ignore control cells and treat perturbation prediction as a noise-to-image generation problem. In contrast, our approach recognizes that:

Batch effects are systematic biases. Variations in experimental conditions across different runs introduce consistent biases unrelated to the perturbation itself.
Control cells provide crucial context. They serve as a reference to distinguish true perturbation effects from experimental artifacts like batch effects.
Distribution-wise modeling is more robust. Learning transformations between control and perturbed distributions captures true perturbation effects while filtering out batch effects.

This formulation enables our method to generate more accurate cellular responses while maintaining robustness across diverse experimental conditions.

🛠️ Algorithm Details

Our method leverages flow matching—a principled generative modeling approach designed for distribution-to-distribution tasks—to learn continuous transformations between cellular states.

CellFlux Algorithm Overview — **Figure 2:** Algorithm overview. Our method leverages flow matching to learn continuous transformations between cellular states.

Training Process

During training, our model learns a velocity field by fitting trajectories between control cell images (x₀ ~ p₀) and perturbed cell images (x₁ ~ p₁). At each training step, intermediate states are sampled along linear interpolations, and the network minimizes the difference between predicted and true velocities.

Inference Process

At inference, the trained velocity field guides the transformation of control cell states into perturbed states by solving an ordinary differential equation iteratively using numerical integration steps.

Key Technical Components include:

Conditional Flow Matching: Extends flow matching to handle perturbation conditions.
Classifier-Free Guidance: Improves generation fidelity through conditional/unconditional interpolation.
Noise Augmentation: Prevents overfitting and encourages smooth velocity fields.
U-Net Architecture: Captures multi-scale features for accurate cellular morphology modeling.

💎 Results & Capabilities

State-of-the-Art Performance

Our method achieves strong performance across multiple cellular imaging datasets, demonstrating significant improvements in both image generation quality and biological fidelity.

CellFlux Performance Results — **Figure 3:** State-of-the-art performance across multiple datasets. Our method consistently outperforms baseline methods.

Detailed Performance Metrics — **Figure 4:** Qualitative results showing our method generates the most biologically accurate and faithful images under various perturbation conditions.

Our results demonstrate significant improvements:

35% improvement in FID scores compared to existing methods.
12% increase in mode-of-action prediction accuracy.
Strong out-of-distribution generalization to unseen perturbations.

Novel Capabilities

Our method enables new capabilities of batch effect correction and trajectory modeling, advancing the field of virtual cell modeling.

Continuous Interpolation Capabilities — **Figure 5:** Novel capabilities include batch effect correction and modeling bidirectional cellular morphological change trajectories.

Key capabilities include:

Batch Effect Correction: By conditioning on control cells, our method distinguishes true biological effects from experimental artifacts.
Trajectory Modeling: The continuous velocity field learned via flow matching captures cellular morphological change trajectories, offering insights into the dynamics of cellular responses.

These capabilities advance our understanding of cellular behavior and provide powerful tools for drug discovery and personalized therapy development.

Conclusion

We introduce CellFlux, a method that leverages flow matching to generate cell images under various perturbations while capturing their trajectories, paving the way for the development of a virtual cell framework for biomedical research. In future work, we plan to scale up CellFlux to process terabytes of imaging data encompassing diverse cell types and a wide range of perturbations, enabling the full potential of virtual cell modeling.

CellFlux:

Simulating Cellular Morphology Changes via Flow Matching