Joint Flow Matching Enables Continuous Dose-Conditioned Cell Morphing

Lea Bogensperger; Manuela Merlo; Martin Baumgartner; Michael Krauthammer; Bernard Ciraulo

cs.CVAug 17, 2026

Joint Flow Matching Enables Continuous Dose-Conditioned Cell Morphing

Authors: Lea Bogensperger, Manuela Merlo, Martin Baumgartner, Michael Krauthammer, Bernard Ciraulo

Organizations: University of Zurich · University Children’s Hospital Zurich

Abstract

Generative modeling has shown increasing promise for predicting cellular perturbation effects under chemical compound treatments. Existing approaches either model perturbation as a distribution-to-distribution mapping without explicit concentration handling, or treat concentration as a discrete class label, precluding continuous dose control. We introduce a joint flow matching approach that simultaneously models cell latents and drug concentration via a dual-timestep formulation, enabling dose-conditioned single-cell morphing through the invertibility of flow matching. The joint formulation induces a monotonic dose-response geometry in latent space and additionally supports concentration estimation from cell morphology. As proof of concept, we further demonstrate generalization to an unseen dose held out during training. Empirically, our method achieves competitive or improved per-concentration metrics on two compounds compared with representative baselines, while enabling capabilities structurally unavailable to discrete-class methods.

Figures & tables

Figure 1: Joint Flow Matching for continuous dose-response cell image generation. Given a real control cell, our model encodes it into a VAE latent $z_{1}$ , inverts the flow to noise $z_{0}$ , and re-integrates the forward ordinary differential equation (ODE) at target concentrations $c^{\star}$ to generate morphed cells across a continuous dose range. The dual-timestep formulation jointly models $p(z,c)$ , enabling continuous concentration control.

ODE steps $K$	PSNR $\uparrow$	SSIM $\uparrow$
10	$27.79_{\pm 2.9}$	$0.866_{\pm 0.04}$
25	$31.93_{\pm 3.0}$	$0.894_{\pm 0.04}$
50	$33.81_{\pm 3.0}$	$0.902_{\pm 0.04}$
100	$34.50_{\pm 3.0}$	$0.905_{\pm 0.05}$
200	$34.70_{\pm 3.0}$	$0.906_{\pm 0.05}$
500	$34.77_{\pm 3.0}$	$0.907_{\pm 0.05}$

Table 1: Inversion fidelity for C7 on the test set (

N{=}100

control cells).

Figure 2: Qualitative comparison of dose-response morphing for C7. Each row shows a real control cell morphed across concentrations of C7 (columns: 0.037, 0.11, 0.33, 1.0, 3.0 $\mu$ M). Latent CellFlux and IMPA are shown only at control and morphed dose. Individual samples from concentration-agnostic methods can vary in visual quality across cells.

		Concentration ( $\mu$ M)
Method	Metric	0.037	0.11	0.33	1.0	3.0	Pooled
		( $N{=}892$ )	( $N{=}1028$ )	( $N{=}868$ )	( $N{=}1040$ )	( $N{=}1270$ )	( $N{=}5098$ )
Latent CellFlux [ 42 ] $\dagger$	FID	—	—	—	—	—	25.19
	KID	—	—	—	—	—	1.90
IMPA [ 28 ] $\dagger$	FID	—	—	—	—	—	17.58
	KID	—	—	—	—	—	0.97

Table 2: FID/KID comparison (lower is better) for C7. Per-concentration values shown for methods with explicit concentration modeling; concentration-agnostic baselines (Latent CellFlux, IMPA) are reported at pooled only, as they do not natively support concentration targeting. Best metrics per column are in bold , second-best underlined .

Figure 3: Dose-response geometry and generalization to unseen concentrations for C7. Left: Normalized latent displacement from the reference as a function of drug concentration, comparing our joint model, our conditional-only ablation, and PhenDiff. Right: After retraining our joint model with

0.33\,\mu

M held out during training (train concentrations:

0.037,0.11,1.0,3.0\,\mu

M), forward integration at the held-out dose (highlighted star) yields a monotonic position between its neighbors.

Figure 4: Concentration specificity. KID between generated cells conditioned on concentration $c_{i}$ (rows) and real cells at concentration $c_{j}$ (columns). Diagonals are outlined; red asterisks mark each row’s argmin. Panels (left to right): our joint model, PhenDiff, and the real-vs-real reference ceiling.

VAE training	Metric	0.037	0.11	0.33	1.0	3.0	Pooled
		Concentration ( $\mu$ M)
C7-only	FID	35.26	33.39	34.37	33.89	31.54	24.84
	KID	2.18	2.09	2.14	2.06	2.01	1.94
Multi-drug	FID	30.02	27.20	27.50	27.85	25.02	18.98
	KID	1.58	1.40	1.37	1.39	1.28	1.25

Table 3: VAE training-data ablation on C7. Best values in bold .

Conc. transform	Metric	0.037	0.11	0.33	1.0	3.0	Pooled
		Concentration ( $\mu$ M)
Raw ( $c$ )	FID	30.27	27.72	29.94	28.50	27.02	20.31
	KID	1.61	1.46	1.67	1.46	1.48	1.35
$\log_{10}(c)$	FID	30.02	27.20	27.50	27.85	25.02	18.98
	KID	1.58	1.40	1.37	1.39	1.28	1.25

Table 4: Concentration-transform ablation on C7. Best values in bold .

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Qualitative training set samples for ONS76 cell line under compound treatments. Representative real microscopy images of ONS76 cells treated with C7 (top row, 3 $\mu$ M) and CK-666 (bottom row, 10 $\mu$ M) from the training split.

ODE steps $K$	PSNR $\uparrow$	SSIM $\uparrow$
10	$22.27_{\pm 3.2}$	$0.815_{\pm 0.06}$
25	$22.93_{\pm 4.0}$	$0.830_{\pm 0.06}$
50	$23.48_{\pm 4.2}$	$0.838_{\pm 0.06}$
100	$23.50_{\pm 4.3}$	$0.839_{\pm 0.06}$
200	$23.45_{\pm 4.3}$	$0.839_{\pm 0.06}$
500	$23.39_{\pm 4.3}$	$0.838_{\pm 0.06}$

Appendix

Table 5: Inversion fidelity on CK-666 (

N{=}100

control cells).

		Concentration ( $\mu$ M)
Method	Metric	0.625	1.25	2.5	5.0	10.0	Pooled
		( $N{=}1015$ )	( $N{=}1033$ )	( $N{=}1062$ )	( $N{=}1012$ )	( $N{=}956$ )	( $N{=}5078$ )
Latent CellFlux [ 42 ] $\dagger$	FID	—	—	—	—	—	21.48
	KID	—	—	—	—	—	1.60
IMPA [ 28 ] $\dagger$	FID	—	—	—	—	—	20.93
	KID	—	—	—	—	—	1.23

Appendix

Table 6: FID/KID comparison (lower is better) for CK-666. Best metrics per column are in bold , second-best underlined .

C7
Concentration ( $\mu$ M)	0.037	0.11	0.33	1.0	3.0	Pooled
FID	27.09	26.34	27.27	28.71	29.28	18.79
KID	1.32	1.24	1.34	1.26	1.36	1.25
CK-666
Concentration ( $\mu$ M)	0.625	1.25	2.5	5.0	10.0	Pooled
FID	26.61	27.09	27.53	30.30	27.01	19.34

Appendix

Table 7: FID and KID for UW228 (our joint model). Metrics are within a narrow range per compound and pooled values match those reported on ONS76 in the main text.

Figure 6: Qualitative dose-response morphs on the UW228 cell line for two compounds (C7 and CK-666), generated by our joint model. Each row shows a real control cell (leftmost column) morphed to successive concentrations.

Figure 7: Per-cell dose-response distributions. Violin plots of normalized latent displacement (pixel-space for PhenDiff) from the reference at each concentration; dashed lines connect the concentration means. Normalization per method by the maximum concentration-averaged displacement.

Figure 8: Dose-response geometry for unseen concentrations in C7. After retraining our joint model with

0.11

0.33

, or

1\,\mu\text{M}

held out during training (panels left to right, respectively), forward integration at the held-out dose (highlighted star) generally interpolates between neighboring concentrations, closely following the expected monotonic trend.

Figure 9: Confusion matrices for 5-class concentration prediction in C7. Left: Confusion matrix of our joint flow matching model $p(c\mid z)$ , showing that lower and intermediate concentrations in particular are harder to discriminate, confirming the findings in Sec. 4.4 . Right: The corresponding supervised ResNet-18 baseline yields better results at $0.11\,\mu\mathrm{M}$ , $1\,\mu\mathrm{M}$ , and $3\,\mu\mathrm{M}$ , but performs worse at $0.037\,\mu\mathrm{M}$ and $0.33\,\mu\mathrm{M}$ .

Method	Top-1	$\pm 1$ neighbor
Joint FM $p(c\mid z)$	34.0%	68.0%
ResNet-18, real $\to$ real	48.0%	74.8%

Appendix

Table 8: Concentration classification accuracy on C7. Best values in bold .

Concentration ( $\mu$ M)	0.037	0.11	0.33	1.0	3.0
Joint FM $p(c\mid z)$	44.3%	24.1%	21.5%	20.7%	54.6%
ResNet-18 (real $\to$ real)	31.8%	53.9%	14.5%	48.8%	79.2%

Appendix

Table 9: Per-concentration accuracy on C7. Best values in bold .

$\lambda_{c}$	Metric	0.037	0.11	0.33	1.0	3.0	Pooled
		Concentration ( $\mu$ M)
0.1	FID	29.83	28.52	30.15	29.95	26.16	20.42
0.1	KID	1.56	1.51	1.62	1.58	1.41	1.37
0.5	FID	30.02	27.20	27.50	27.85	25.02	18.98
0.5	KID	1.58	1.40	1.37	1.39	1.28	1.25
1.0	FID	31.48	28.21	28.91	30.41	25.66	20.32

Appendix

Table 10: Ablation over

\lambda_{c}

on C7. Best values in bold .

	PSNR $\uparrow$			SSIM $\uparrow$
Steps $K$	Euler	Heun	RK4	Euler	Heun	RK4
10	$27.79_{\pm 2.9}$	$27.31_{\pm 2.1}$	$27.14_{\pm 2.1}$	$0.866_{\pm 0.04}$	$0.873_{\pm 0.04}$	$0.872_{\pm 0.04}$
25	$31.93_{\pm 3.0}$	$32.09_{\pm 2.6}$	$32.00_{\pm 2.6}$	$0.894_{\pm 0.04}$	$0.900_{\pm 0.04}$	$0.900_{\pm 0.04}$
50	$33.81_{\pm 3.0}$	$33.99_{\pm 2.8}$	$33.97_{\pm 2.8}$	$0.902_{\pm 0.04}$	$0.906_{\pm 0.05}$	$0.906_{\pm 0.05}$
100	$34.50_{\pm 3.0}$	$34.65_{\pm 2.9}$	$34.65_{\pm 2.9}$	$0.905_{\pm 0.05}$	$0.907_{\pm 0.05}$	$0.907_{\pm 0.05}$
200	$34.70_{\pm 3.0}$	$34.81_{\pm 3.0}$	$34.81_{\pm 3.0}$	$0.906_{\pm 0.05}$	$0.907_{\pm 0.05}$	$0.907_{\pm 0.05}$

Appendix

Table 11: Inversion fidelity comparison across ODE integration schemes on the C7 test set (

N{=}100

control cells).

Method	Batch Latency (ms) $\downarrow$	Per-Cell (ms) $\downarrow$
Latent CellFlux	$4196.5_{\pm 8.2}$	$524.6$
IMPA	$176.2_{\pm 0.3}$	$22.0$
PhenDiff	$26189.5_{\pm 31.9}$	$3273.7$
Ours (cond. only)	$6367.3_{\pm 51.2}$	$795.9$
Ours (Joint Latent FM)	$3237.7_{\pm 31.4}$	$404.7$

Appendix

Table 12: Sampling latency comparison across generative methods and baseline models.

Dose-Response Estimation

See all 8 papers

Elucidating the Design Space of Flow Matching for Cellular Microscopy
Mar 25, 2026Charles Jones, Emmanuel Noutahi, Jason Hartford +1Flow Matching Image Generation
Imperial College London · Valence Labs · University of Manchester
Generalizable single-cell perturbation response prediction using energy-guided flow matching
Sep 27, 2026Jianan Wei, Jiajun Hong, Guikun Chen +3Single-Cell Perturbation Modeling Single-Cell Transcriptomics
Zhejiang University, Hangzhou, Zhejiang, China · Beijing Institute for General Artificial Intelligence, Beijing, China · Peking University Chengdu Academy for Advanced Interdisciplinary Biotechnologies, Chengdu, Sichuan, China
TempoBridge: Source-Conditioned Flow Matching with Optimal Transport Couplings for Single-Cell Population Transitions
Oct 4, 2026Bowen Han, Lingbei Meng, Shihuan Luo +7Single-Cell Perturbation Modeling Conditional Flow Matching
Universiti Sains Malaysia · The Chinese University of Hong Kong Shenzhen · Shenzhen Loop Area Institute +3

Mar 25, 2026J/K move · Enter open · S save

Elucidating the Design Space of Flow Matching for Cellular Microscopy

Charles Jones, Emmanuel Noutahi, Jason Hartford, Cian Eastwood

Imperial College London · Valence Labs · University of Manchester

Flow Matching Image Generation Single-Cell Perturbation Modeling Microscopy Image Analysis

Open paper AskPDF arXiv

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

Abstract

Dose-Response Estimation

Explore similar work

Elucidating the Design Space of Flow Matching for Cellular Microscopy

Generalizable single-cell perturbation response prediction using energy-guided flow matching

TempoBridge: Source-Conditioned Flow Matching with Optimal Transport Couplings for Single-Cell Population Transitions