Elucidating the Design Space of Flow Matching for Cellular Microscopy

Charles Jones; Emmanuel Noutahi; Jason Hartford; Cian Eastwood

cs.CVMar 25, 2026

Elucidating the Design Space of Flow Matching for Cellular Microscopy

Authors: Charles Jones, Emmanuel Noutahi, Jason Hartford, Cian Eastwood

Organizations: Imperial College London · Valence Labs · University of Manchester

Abstract

Flow-matching generative models are increasingly used to simulate cell responses to biological perturbations. However, the design space for building such models is large and underexplored. We systematically analyse the design space of flow matching models for cell-microscopy images, finding that many popular techniques are unnecessary and can even hurt performance. We develop a simple, stable, and scalable recipe which we use to train our foundation model. We scale our model to two orders of magnitude larger than prior methods, achieving a two-fold FID and ten-fold KID improvement over prior methods. We then fine-tune our model with pre-trained molecular embeddings to achieve state-of-the-art performance simulating responses to unseen molecules. Code is available at https://github.com/valence-labs/microscopy-flow-matching

Figures & tables

Figure 1: We ablate, simplify, and scale flow matching for generative modelling of cellular microscopy, achieving state-of-the-art performance on: ( top images ) the RxRx1 genetic intervention benchmark; ( middle images ) the BBBC021 small molecule perturbation benchmark; and ( bottom images ) the BBBC021 unseen molecule virtual screening task. In the bottom plot , we display performance vs. total training compute on the RxRx1 benchmark.

	Model	ExaFLOPs	Flow	OT	FID $\downarrow$	KID $\downarrow$
PhenDiff ( Bourou et al., 2024 )	U-Net	0.96	$\mathcal{N\leftrightarrow D}$	✗	65.9	51.9
IMPA ( Palma et al., 2025 )	StarGAN	0.39	-	-	41.6	29.1
CellFlux ( Zhang et al., 2025 )	ADM	7.83	$\mathcal{C\rightarrow P}$	✗	33.0	23.8
CellFluxV2 ( Zhang et al., 2026 )	DiT-XL/2	6.89	$\mathcal{C\rightarrow P}$	✗	19.0	9.30
Ours – Config A	ADM	4.70	$\mathcal{C\rightarrow P}$	✗	29.5	14.0
Ours – Config B (undertrained)	ADM	3.93	$\mathcal{N\rightarrow D}$	✗	12.3	3.90

Table 1: Systematically improving upon state-of-the-art microscopy generative modelling on RxRx1. Our final configuration achieves over two-fold FID improvements and ten-fold KID improvements over prior and contemporaneous work. ExaFLOPs measure the total training compute, as estimated in Appendix A . KID numbers are scaled up by a factor of

10^{3}

for ease of interpretation.

Figure 2: Samples from our best performing model on RxRx1. Each column represents a unique gene perturbation and each row represents a different biological experiment. The model generates distinct phenotypes for each perturbation whilst capturing biological batch effects. See Figure 8 for side-by-side comparisons of real and generated images.

Figure 3: Conditioning strategy. While (pre)training on seen perturbations (a), we use learnable linear embeddings of the perturbation label. When finetuning for unseen perturbations in Section 5 (b), we freeze the base model and perturbation encoder. We train a small adaptor to align the embeddings to the model’s conditioning.

Figure 4: Generation vs. counterfactual mode sampling. Top row represents the flow trajectory for generation mode, starting with Gaussian noise and ending at a perturbed image. Bottom row represents the trajectory for counterfactual mode. This involves starting with a control image, solving the reverse probability flow ODE to compute its corresponding noise, then generating a perturbed image from the inferred noise in the usual way.

Figure 5: Counterfactual sampling on RxRx1. Top row represents real untreated control images. Middle row is the predicted counterfactual for the relevant perturbations (as labelled by the column headers). Bottom row represents the predicted individual treatment effect (difference of counterfactual to factual, averaged over channels). Our model is able to transform cell morphology whilst preserving features not relevant to the perturbation.

Figure 6: Our microscopy transformer (MiT) architecture. Differences to the standard diffusion transformer include long-range skip connections, RMSNorm, and dropout (all highlighted in orange).

	Conditioning	Seen		Unseen
	Conditioning	FID $\downarrow$	KID $\downarrow$	FID $\downarrow$	KID $\downarrow$
MorphoDiff ( Navidi et al., 2025 )	RDKit2D	65.8	79.9	-	-
PhenDiff ( Bourou et al., 2024 )	One-Hot	49.5	31.0	67.7	34.5
IMPA ( Palma et al., 2025 )	Morgan Fingerprints	33.7	26.0	44.5	30.7
CellFlux ( Zhang et al., 2025 )	Morgan Fingerprints	18.7	16.2	42.0	13.1
CellFluxV2 ( Zhang et al., 2026 )	Morgan Fingerprints	7.90	2.30	19.2	8.00

Table 2: Simulating effects of seen and unseen molecule perturbations on BBBC021. KID numbers are scaled up by a factor of

10^{3}

for ease of interpretation. Note that, since FID and KID are biased metrics, they cannot directly be compared across the seen and unseen columns because they are computed over different numbers of samples.

Figure 7: Randomly sampled real and generated images for virtual screening on BBBC021 – the AZ841 molecule was not seen by the model during training. Samples are generated using our model (with Morgan fingerprints).

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Optimiser	Adam { lr : $10^{-4}$ , $\beta_{1}$ : $0.9$ , $\beta_{2}$ : $0.999$ , $\epsilon$ : $10^{-8}$ }
LR schedule	Constant
Gradient norm clipping	0.5
Global batch size	256
Augmentation	Uniform dequantisation $\sim\mathcal{U}(0,1/256)$
	Random flipping (horizontal and vertical)
Solver	DoPri5 { $\texttt{rtol}=\texttt{atol}=10^{-5}$ }

Appendix

Table 3: Training and inference hyperparameters for all RxRx1 models.

	NFE	FDD $\downarrow$	KDD $\downarrow$	FID $\downarrow$	KID $\downarrow$
Config A ( $\mathcal{C\rightarrow P}$ )
guidance strength 1.0	381	272	14.9	34.0	16.7
guidance strength 1.5	826	266	14.1	31.2	14.5
guidance strength 2.0*	918	261	13.8	29.5	14.0
Config B ( $\mathcal{N\rightarrow D}$ )
guidance strength 1.0*	299	51.0	1.75	9.90	1.52

Appendix

Table 4: Extended results of RxRx1 experiments. For each configuration, we include a coarse-grained sweep over guidance strength

\{1.0,1.5,2.0\}

. In addition to conventional FID and KID, we report DinoV2 alternatives, as proposed by Stein et al. (2023) (FDD and KDD respectively). We include the mean number of function evaluations (NFE) used per batch for sampling with each configuration. Notice how using guidance strength

\not=1

approximately doubles NFE, since it requires two forward passes per step. Similarly, using counterfactual generation doubles it again since it requires solving two ODEs. Rows marked with * represent the configurations chosen for Table 1 . KID and KDD numbers are scaled up by factors of

10^{3}

and 10, respectively.

Interpolant		FID $\downarrow$	KID $\downarrow$
Linear	$tx_{1}+(1-t)x_{0}$	9.90	1.52
Variance preserving	$x_{1}\mathrm{sin}(\frac{1}{2}\pi t)+x_{0}\mathrm{cos}(\frac{1}{2}\pi t)$	9.99	1.42
Brownian bridge, $k=0.1$	$tx_{1}+(1-t)x_{0}+\epsilon k\sqrt{2t(1-t)},\quad\epsilon\sim\mathcal{N}(0,1)$	14.5	6.58
Brownian bridge, $k=0.5$	$tx_{1}+(1-t)x_{0}+\epsilon k\sqrt{2t(1-t)},\quad\epsilon\sim\mathcal{N}(0,1)$	103	90.2
Brownian bridge, $k=1.0$	$tx_{1}+(1-t)x_{0}+\epsilon k\sqrt{2t(1-t)},\quad\epsilon\sim\mathcal{N}(0,1)$	180	172

Appendix

Table 5: Interpolant ablations. All runs use the setup from Config B (

\mathcal{N\rightarrow D}

), with guidance strength 1.0.

Model	Stable (✓/✗)	Notes
DiT-B/2	✗	Larger variants are stable for longer.
w/ $10^{-5}$ learning rate.	✗
w/ $0.01$ AdamW weight decay.	✗	Loshchilov & Hutter (2019) .
w/ Warmup + inv sqrt decay.	✗	Adapted from Karras et al. (2024) .
w/ Adam $\beta_{2}=0.995$	✗	Berrada et al. (2024) , Appendix C.
w/ MLP Dropout	✗

Appendix

Table 6: Stabilisation ablations for diffusion transformers. We consider a setup stable if it does not diverge within the first 100k steps when training from scratch on RxRx1. All models are trained using the

\mathcal{N\rightarrow D}

flow.

Panel of Figure 8 — Figure 8: Randomly sampled real and generated images for three SiRNA gene perturbations in RxRx1. Samples are generated using our model (no pretraining), matching the composition of perturbations and biological experiments to the randomly sampled real images.

	NFE	FDD $\downarrow$	KDD $\downarrow$	FID $\downarrow$	KID $\downarrow$
Seen molecules
Ours (one-hot)	239	23.2	0.63	4.03	0.31
Ours (uncond.)	242	25.0	0.58	4.63	0.68
Ours (Morgan)	241	23.9	0.50	4.16	0.18
Ours (MolGPS)	240	23.2	0.46	4.12	0.16
Unseen molecules

Appendix

Table 7: Extended results of BBBC021 experiments. We use guidance strength 1.0 for all conditional models (the unconditional variant uses 0.0). KID and KDD numbers are scaled up by factors of

10^{3}

and 10, respectively.

Panel of Figure 9 — Figure 9: Randomly sampled real and generated images for three seen molecule perturbations on BBBC021. Samples are generated using our model (with Morgan fingerprints), matching the composition of perturbations and experiments to the randomly sampled real images.

Panel of Figure 10 — Figure 10: Randomly sampled real and generated samples for three unseen molecule perturbations on BBBC021. Samples are generated using our model (with Morgan fingerprints), matching the composition of perturbations and experiments to the randomly sampled real images.

Flow Matching

Flow matching learns continuous-time vector fields that transport simple noise distributions into data for efficient generative modeling.

See all 639 papers

The Amazing Stability of Flow Matching
Apr 17, 2026Rania Briq, Michael Kamp, Ohad Fried +2Flow Matching Generative Modeling
Forschungszentrum Jülich · TU Dortmund · Lamarr Institute +4
Dino Forcing Flow Models: Do not denoise what you can predict
Oct 8, 2026Arijit Ghosh, Lucas Degeorge, Paul Couairon +3Self-Conditioning Flow-Based Generative Modeling
ENPC, IP Paris · Ecole Polytechnique, IP Paris · AMIAD +1
Flow Matching is Adaptive to Manifold Structures
Feb 25, 2026Shivam Kumar, Yixin Wang, Lizhen LinFlow Matching Generative Modeling
Booth School of Business, University of Chicago · Department of Statistics, University of Michigan · Department of Mathematics, University of Maryland, College Park

Apr 17, 2026J/K move · Enter open · S save

The Amazing Stability of Flow Matching

Rania Briq, Michael Kamp, Ohad Fried, Sarel Cohen+1

Forschungszentrum Jülich · TU Dortmund · Lamarr Institute+4

Flow Matching Generative Modeling Flow-Based Generative Modeling

Open paper AskPDF arXiv

The success of deep generative models in generating high-quality and diverse samples is often attributed to particular architectures and large training datasets. In this paper, we investigate the impact of these factors on the quality and diversity of samples generated by flow-matching models. Surprisingly, in our experiments on CelebA-HQ dataset, flow matching remains stable even when pruning 50% of the dataset. That is, the quality and diversity of generated samples are preserved. Moreover, pruning impacts the latent representation only slightly, that is, samples generated by models trained on the full and pruned dataset map to visually similar outputs for a given seed. We observe similar stability when changing the architecture or training configuration, such that the latent representation is maintained under these changes as well. Our results quantify just how strong this stability can be in practice, and help explain the reliability of flow-matching models under various perturbations.

Abstract

Flow Matching

Explore similar work

The Amazing Stability of Flow Matching

Dino Forcing Flow Models: Do not denoise what you can predict

Flow Matching is Adaptive to Manifold Structures