Elucidating the Design Space of Flow Matching for Cellular Microscopy

Cellular microscopy has evolved from traditional imaging to advanced deep learning-based generative models, including GANs, VAEs, diffusion, and flow matching techniques. These models enable high-fidelity image synthesis, aiding phenotypic analysis, drug screening, and disease modeling. Despite progress, challenges remain in model stability, scalability, and generalization, especially for unseen conditions. Prior works like CellFlux and Diffusion Transformers have demonstrated promising results but face limitations in performance and robustness. The need for systematic analysis of model design choices and large-scale pretraining is critical to push the field forward.

Existing models often suffer from overfitting, instability, and limited scalability, hindering their ability to generate realistic images under novel perturbations. The complex design space, including conditioning, interpolants, and coupling strategies, lacks comprehensive exploration, leading to suboptimal configurations. Moreover, current models struggle to generalize to unseen molecules, limiting their utility in virtual screening. Addressing these issues requires a systematic approach to model simplification, training stability, and incorporation of rich biological information, such as molecular embeddings.

The paper introduces a streamlined training recipe that simplifies conditioning and coupling strategies, replacing complex methods with one-hot or multi-modal embeddings. It adopts a diffusion transformer architecture (MiT), scaled via large-scale pretraining on over a billion images, to improve stability and performance. The systematic analysis of design choices reveals that many traditional techniques are unnecessary, leading to a more efficient and robust model. Additionally, integrating molecular embeddings for fine-tuning enables high-quality simulation of unseen molecular perturbations, a novel contribution in bioimage generation.

• �� Construct a simplified conditioning framework using one-hot and multi-modal embeddings, enabling flexible control and generalization;
• �� Explore various interpolants, including Gaussian noise and control image flows, to optimize the trade-off between realism and diversity;
• �� Test different coupling strategies (random vs. optimal transport), assessing their impact on training stability and sampling efficiency;
• �� Replace U-Net with a diffusion transformer (MiT), incorporating long-range skip connections, RMSNorm, and dropout for stability;
• �� Perform large-scale pretraining on over 600 million images from Phenoprints dataset, followed by fine-tuning on RxRx1;
• �� Incorporate molecular embeddings (MolGPS) for conditional fine-tuning, enabling simulation of unseen molecules.

Experiments evaluate performance on RxRx1 and BBBC021 datasets, using FID and KID metrics. Ablation studies compare conditioning methods, interpolants, and coupling strategies, confirming their influence on quality and stability. The training involves multi-stage pretraining and fine-tuning, with model sizes reaching 700 million parameters. The experiments also include testing the model's ability to generate images conditioned on unseen molecules, using different molecular embeddings. Results demonstrate that the proposed approach achieves state-of-the-art performance, with significant improvements over prior methods, validating the design choices.

The model reduces FID from approximately 33 to 10 and KID from 23 to 1.6 on RxRx1, outperforming previous models by over two times. Fine-tuning with MolGPS embeddings yields an FID of 4.12 and KID of 9.95 on BBBC021 for unseen molecules, surpassing prior approaches. Ablation results show that simplified conditioning, multi-modal embeddings, and model scaling are critical for performance. The large-scale pretraining significantly enhances the model's stability and generalization, establishing new benchmarks in bioimage generation.

This model can be used for virtual drug screening, phenotypic response prediction, and disease modeling. Its ability to generate high-fidelity images conditioned on various biological factors accelerates research and reduces experimental costs. The approach supports the development of personalized medicine by simulating individual cellular responses, facilitating drug discovery pipelines, and enabling in silico testing of therapeutic interventions.

Despite advancements, the model's generalization to highly novel or extreme conditions remains imperfect, limited by the expressiveness of molecular embeddings. Training large models is computationally expensive, restricting widespread adoption. Current models primarily generate static images; dynamic cellular processes and multi-modal data integration are future challenges. Further research is needed to improve interpretability and biological relevance of generated images.