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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.
 ModelExaFLOPsFlowOTFID ↓\downarrowKID ↓\downarrow
PhenDiff ( Bourou et al., 2024 )U-Net0.96N↔D\mathcal{N\leftrightarrow D}✗65.951.9
IMPA ( Palma et al., 2025 )StarGAN0.39--41.629.1
CellFlux ( Zhang et al., 2025 )ADM7.83C→P\mathcal{C\rightarrow P}✗33.023.8
CellFluxV2 ( Zhang et al., 2026 )DiT-XL/26.89C→P\mathcal{C\rightarrow P}✗19.09.30
Ours – Config AADM4.70C→P\mathcal{C\rightarrow P}✗29.514.0
Ours – Config B (undertrained)ADM3.93N→D\mathcal{N\rightarrow D}✗12.33.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 10310^{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).
 ConditioningSeenUnseen
FID ↓\downarrowKID ↓\downarrowFID ↓\downarrowKID ↓\downarrow
MorphoDiff ( Navidi et al., 2025 )RDKit2D65.879.9--
PhenDiff ( Bourou et al., 2024 )One-Hot49.531.067.734.5
IMPA ( Palma et al., 2025 )Morgan Fingerprints33.726.044.530.7
CellFlux ( Zhang et al., 2025 )Morgan Fingerprints18.716.242.013.1
CellFluxV2 ( Zhang et al., 2026 )Morgan Fingerprints7.902.3019.28.00
Table 2: Simulating effects of seen and unseen molecule perturbations on BBBC021. KID numbers are scaled up by a factor of 10310^{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

OptimiserAdam { lr : 10−410^{-4} , β1\beta_{1} : 0.90.9 , β2\beta_{2} : 0.9990.999 , ϵ\epsilon : 10−810^{-8} }
LR scheduleConstant
Gradient norm clipping0.5
Global batch size256
AugmentationUniform dequantisation ∼U(0,1/256)\sim\mathcal{U}(0,1/256)
 Random flipping (horizontal and vertical)
SolverDoPri5 { rtol=atol=10−5\texttt{rtol}=\texttt{atol}=10^{-5} }
Appendix
Table 3: Training and inference hyperparameters for all RxRx1 models.
 NFEFDD ↓\downarrowKDD ↓\downarrowFID ↓\downarrowKID ↓\downarrow
Config A ( C→P\mathcal{C\rightarrow P} )     
guidance strength 1.038127214.934.016.7
guidance strength 1.582626614.131.214.5
guidance strength 2.0*91826113.829.514.0
Config B ( N→D\mathcal{N\rightarrow D} )     
guidance strength 1.0*29951.01.759.901.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}\{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 ≠1\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 10310^{3} and 10, respectively.
Interpolant FID ↓\downarrowKID ↓\downarrow
Lineartx1+(1−t)x0tx_{1}+(1-t)x_{0}9.901.52
Variance preservingx1sin(12πt)+x0cos(12πt)x_{1}\mathrm{sin}(\frac{1}{2}\pi t)+x_{0}\mathrm{cos}(\frac{1}{2}\pi t)9.991.42
Brownian bridge, k=0.1k=0.1tx1+(1−t)x0+ϵk2t(1−t),ϵ∼N(0,1)tx_{1}+(1-t)x_{0}+\epsilon k\sqrt{2t(1-t)},\quad\epsilon\sim\mathcal{N}(0,1)14.56.58
Brownian bridge, k=0.5k=0.5tx1+(1−t)x0+ϵk2t(1−t),ϵ∼N(0,1)tx_{1}+(1-t)x_{0}+\epsilon k\sqrt{2t(1-t)},\quad\epsilon\sim\mathcal{N}(0,1)10390.2
Brownian bridge, k=1.0k=1.0tx1+(1−t)x0+ϵk2t(1−t),ϵ∼N(0,1)tx_{1}+(1-t)x_{0}+\epsilon k\sqrt{2t(1-t)},\quad\epsilon\sim\mathcal{N}(0,1)180172
Appendix
Table 5: Interpolant ablations. All runs use the setup from Config B ( N→D\mathcal{N\rightarrow D} ), with guidance strength 1.0.
ModelStable (✓/✗)Notes
DiT-B/2✗Larger variants are stable for longer.
w/ 10−510^{-5} learning rate.✗ 
w/ 0.010.01 AdamW weight decay.✗Loshchilov & Hutter (2019) .
w/ Warmup + inv sqrt decay.✗Adapted from Karras et al. (2024) .
w/ Adam β2=0.995\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 N→D\mathcal{N\rightarrow D} flow.
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.
 NFEFDD ↓\downarrowKDD ↓\downarrowFID ↓\downarrowKID ↓\downarrow
Seen molecules     
Ours (one-hot)23923.20.634.030.31
Ours (uncond.)24225.00.584.630.68
Ours (Morgan)24123.90.504.160.18
Ours (MolGPS)24023.20.464.120.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 10310^{3} and 10, respectively.
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.
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.
More in this topic

Flow Matching

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

See all 639 papers

Explore similar work

CardsList
Apr 17, 2026

The Amazing Stability of Flow Matching

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.

Rania Briq, Michael Kamp, Ohad Fried +2

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

Flow MatchingGenerative ModelingFlow-Based Generative Modeling
AskPDF
Oct 8, 2026

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

Co-denoising pretrained representations such as DINO can substantially improve the training speed and quality of flow matching models, but it introduces a second denoising trajectory and requires carefully designed schedules. We propose a simpler alternative: predict the pretrained representation directly, then condition the model on its own prediction. This removes the need for a second ODE and any representation-specific denoising schedules, while retaining the benefits of representation guidance. Our approach converges substantially faster and achieves better generation quality as measured by FID score. On ImageNet, it outperforms the state of the art in latent space at 2x fewer epochs than prior methods; in pixel space, it improves FID over comparable prior methods by more than 20%. These results support a simple principle: do not denoise what you can predict. Our code is openly available at https://github.com/arijit-hub/dino_forcing.

Arijit Ghosh, Lucas Degeorge, Paul Couairon +3

ENPC, IP Paris · Ecole Polytechnique, IP Paris · AMIAD +1

Self-ConditioningFlow-Based Generative ModelingFlow MatchingRepresentation Steering
AskPDF
Feb 25, 2026

Flow Matching is Adaptive to Manifold Structures

Flow matching has emerged as a simulation-free alternative to diffusion-based generative modeling, producing samples by solving an ODE whose time-dependent velocity field is learned along an interpolation between a simple source distribution (e.g., a standard normal) and a target data distribution. Flow-based methods often exhibit greater training stability and have achieved strong empirical performance in high-dimensional settings where data concentrate near a low-dimensional manifold, such as text-to-image synthesis, video generation, and molecular structure generation. Despite this success, existing theoretical analyses of flow matching assume target distributions with smooth, full-dimensional densities, leaving its effectiveness in manifold-supported settings largely unexplained. To this end, we theoretically analyze flow matching with linear interpolation when the target distribution is supported on a smooth manifold. We establish a non-asymptotic convergence guarantee for the learned velocity field, and then propagate this estimation error through the ODE to obtain statistical consistency of the implicit density estimator induced by the flow-matching objective. The resulting convergence rate is near minimax-optimal, depends only on the intrinsic dimension, and reflects the smoothness of both the manifold and the target distribution. Together, these results provide a principled explanation for how flow matching adapts to intrinsic data geometry and circumvents the curse of dimensionality.

Shivam Kumar, Yixin Wang, Lizhen Lin

Booth School of Business, University of Chicago · Department of Statistics, University of Michigan · Department of Mathematics, University of Maryland, College Park

Flow MatchingGenerative ModelingFlow-Based Generative ModelingStatistical Learning Theory
AskPDF