Conditional flow matching for physics-constrained inverse problems with finite training data

Agnimitra Dasgupta; Ali Fardisi; Mehrnegar Aminy; Brianna Binder; Bryan Shaddy; Saeed Moazami; Assad Oberai

stat.MLMar 14, 2026

Conditional flow matching for physics-constrained inverse problems with finite training data

Authors: Agnimitra Dasgupta, Ali Fardisi, Mehrnegar Aminy, Brianna Binder, Bryan Shaddy, Saeed Moazami, Assad Oberai

Organizations: Department of Aerospace & Mechanical Engineering, University of Southern California, Los Angeles, 90089, California, USA · Optimization and Uncertainty Quantification, Sandia National Laboratories, Albuquerque, 87123, New Mexico, USA

Abstract

This study presents a conditional flow matching framework for solving physics-constrained Bayesian inverse problems. In this setting, samples from the joint distribution of inferred variables and measurements are assumed available, while explicit evaluation of the prior and likelihood densities is not required. We derive a simple and self-contained formulation of both the unconditional and conditional flow matching algorithms, tailored specifically to inverse problems. In the conditional setting, a neural network is trained to learn the velocity field of a probability flow ordinary differential equation that transports samples from a chosen source distribution directly to the posterior distribution conditioned on observed measurements. This black-box formulation accommodates nonlinear, high-dimensional, and potentially non-differentiable forward models without restrictive assumptions on the noise model. We further analyze the behavior of the learned velocity field in the regime of finite training data. Under mild architectural assumptions, we show that overtraining can induce degenerate behavior in the generated conditional distributions, including variance collapse and a phenomenon termed selective memorization, wherein generated samples concentrate around training data points associated with similar observations. A simplified theoretical analysis explains this behavior, and numerical experiments confirm it in practice. We demonstrate that standard early-stopping criteria based on monitoring test loss effectively mitigate such degeneracy. The proposed method is evaluated on several physics-based inverse problems. We investigate the impact of different choices of source distributions, including Gaussian and data-informed priors. Across these examples, conditional flow matching accurately captures complex, multimodal posterior distributions while maintaining computational efficiency.

Figures & tables

Figure 1: Visualizing Case 1 corresponding to Eq. 32

Figure 2: Visualizing Case 2 where

\phi_{j}

-s are non-overlapping indicator functions

Figure 3: (a) Train and test data for the toy example used to illustrate effects of overfitting. (b) Train and test loss curve for the velocity network trained using the training data shown in (a). The blue curve above shows the moving average (MA) of the test loss. The MA is computed over a window of size 500

Figure 4: Kernel density estimates of the conditional distribution $\rho_{X|Y}(x\mid y=0.6)$ estimated from samples generated using the trained velocity network at different stages of training

Figure 5: Mean and one-standard-deviation interval of the conditional distribution $\rho_{X|Y}$ estimated using samples generated by the trained velocity network at different stages of training

Figure 6: Training dataset for the spiral problem

Figure 7: Training loss, test loss, and the moving average of the test loss for the velocity network trained on the spiral dataset with $Z\sim\mathcal{N}(0,1)$ (left) and $Z\sim\rho_{X}$ (right) as the source distributions, respectively

Figure 8: Histograms of samples generated using the trained velocity field compared to the samples from the true conditional distribution $\rho_{X|Y}(x|y=\hat{y})$ for $\hat{y}\in\{0.5,0,0.5,1\}$ with different source distributions. The kernel density estimates of the conditional density are also shown

Source distribution	Avg. Sinkhorn distance	Avg. number of sampling steps
Gaussian	0.051	17
Prior	0.056	19

Table 1: Comparison of the two different source distributions for the spiral dataset

Figure 9: Training loss, test loss, and the moving average of the test loss for the velocity network trained on the one-step data assimilation example with (a) Gaussian and (b) the prior $\rho_{X}$ as the source distribution

Figure 10: Particles from the prior $\rho_{{\mathbf{X}}}$ ( $\bullet$ ) and posterior $\rho_{{\mathbf{X}}|Y}$ ( $\bullet$ ), and observation $\hat{y}$ ( – ) on the $(x_{1},x_{3})$ and $(x_{2},x_{3})$ planes for the one-step data assimilation problem. First row: reference solution obtained using the SIR filter with 100,000 particles. Second and third row: 1,000 particles from $\rho_{{\mathbf{X}}}$ (part of the training dataset), and samples from $\rho_{{\mathbf{X}}|Y}$ generated by the conditional flow matching model with a Gaussian and prior source distributions, respectively

Source distribution	Sinkhorn distance	Average number of sampling steps
Gaussian	0.045	17
Prior	0.062	28

Table 2: Comparison of the two different source distributions on the one-step data assimilation example

Panel of Figure 11 — Figure 11: A realization from the advection-diffusion-reaction dataset (a) Piecewise-constant top and bottom wall fluxes. (b) Corresponding concentration field obtained from solving Eq. 52 and sensor measurements (black dots)

Panel of Figure 12 — Figure 12: Training and test loss for the advection-diffusion-reaction problem with 400 training samples

Figure 13: Posterior mean and standard deviation of the inferred flux and the true flux for a test case with data size 400 and 1% noise. Left: Gaussian source distribution. Right: Prior source distribution.

Panel of Figure 14 — Figure 14: Training and test loss for the advection-diffusion-reaction problem with 4,000 training samples

Source distribution		Gaussian		Prior
Size of training data		400	4000	400	4000
Metric	Average $N_{\text{step}}$	13.0	14.7	9.0	9.4
	Average error	0.144	0.139	0.173	0.142
	Average std	0.061	0.055	0.066	0.053

Table 3: Comparison of conditional flow matching models with different types of source distributions and trained using different amounts of training data on the advection-diffusion-reaction problem

Figure 15: Five realizations of $\mathbf{X}$ and $\mathbf{Y}$ sampled from the joint distribution of the training dataset for the synthetic quasi-static elastography application. The first row shows the shear modulus field, and the second row shows the corresponding noisy vertical displacement measurements

Figure 16: Training loss, test loss, and moving average of the test loss for the synthetic quasi-static elastography application

Figure 17: Posterior statistics estimated using the MCS and trained velocity network on two test samples for the synthetic quasi-static elastography application

Test sample	RMSE in posterior statistics		Average number of sampling steps
Test sample	Mean	Std. dev	Average number of sampling steps
1	0.014	0.012	5
2	0.012	0.011	7

Table 4: RMSE between statistics estimated using tconditional flow matching and MCS, and the number of sampling steps for the test samples in Fig. 17

Figure 18: Posterior statistics estimated using a severely trained velocity network ( $60{,}000$ iterations) on two test samples for the quasi-static elastography application

Random variable	Distribution
Distance of the inclusion’s center from the left edge (mm)	$\mathcal{U}(7.1,19.2)$
Distance of the inclusion’s center from the bottom edge (mm)	$\mathcal{U}(7.1,27.6)$
Radius of the inclusion (mm)	$\mathcal{U}(3.5,7.0)$
Ratio between the inclusion’s and background’s shear modulus (mm)	$\mathcal{U}(1,8)$

Table 5: Details of random variables comprising the parametric prior distribution for

{\mathbf{X}}

in the quasi-static elastography application

Figure 19: Five typical realizations of ${\mathbf{X}}$ and ${\mathbf{Y}}$ sampled from the training dataset for the quasi-static elastography application with experimental data. The first row shows the spatial distribution of the shear modulus field, and the second row shows the corresponding measurements of the noisy vertical displacement field

Figure 21: Posterior statistics estimated using the trained velocity network on two test samples for the quasi-static elastography application with experimental data

Test sample	RMSE	Avg. number of sampling steps
1	0.300	9
2	0.563	16

Table 6: RMSE between the posterior mean and the ground truth and average number of sampling steps for two test samples for the quasi-static elastography application with experimental data

Figure 22: Posterior statistics estimated using a severely over-trained velocity network on two test samples shown in Fig. 21

Figure 23: Posterior statistics estimated using the trained velocity network on experimental data for the quasi-static elastography application

Figure 24: Experimental setup in the tumor spheroid application (adapted from [ 14 ] )

Figure 25: Five realizations of ${\mathbf{X}}$ and ${\mathbf{Y}}$ sampled from the joint distribution forming the training dataset for the tumor spheroid application. In the first row are instances of the log-normalized Young’s modulus fields, and in the second row are corresponding instances of the noisy measurements. All values have been normalized to [-1,1]

Figure 26: Training and test losses for the velocity network used in the tumor spheroid application, with moving average for the testing loss

Figure 27: Posterior statistics estimated using the trained velocity network for select synthetic and experimental cases for the tumor spheroid application.

Measurement Type	Avg. number of sampling steps
Measurement Type	Test sample 1	Test sample 2
Synthetic	12	12
Experimental	13	12

Table 7: Average number of sampling steps for the tumor spheroid application

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset/ Inverse Problem	Width of hidden layers	Number of hidden layers	Activation function
Spiral ( Section 4.1 )	32	3	ReLU
One step data assimilation ( Section 4.2.1 )	256	4	ReLU
Advection-diffusion-reaction ( Section 4.2.2 )	256	5	ReLU

Appendix

Table A1 : Details of the architecture of the velocity network for the numerical examples in Sections 4.1 , 4.2.1 and 4.2.2

Inverse Problem/ Application	Number of residual blocks	Number of model channels	Channel multiplier	Attention resolution
Quasi-static elastography ( Section 4.2.3 )	2	128	[1, 2, 3, 4]	16
Quasi-static elastography ( Section 4.3.1 )	2	128	[1, 2, 3, 4]	8
Tumor spheroids ( Section 4.3.2 )	2	128	[1, 1, 2, 2, 4, 4]	16

Appendix

Table A2 : Details of the architecture of the velocity network for the numerical examples in Sections 4.2.3 , 4.3.1 and 4.3.2

Dataset/Inverse Problem/Application	Learning Rate	Batch Size	Number of iterations	EMA coefficient
Spiral ( Section 4.1 )	0.001	1000	100,000	0.9999
One step data assimilation ( Section 4.2.1 )	0.001	500	100,000	0.9000
Advection-diffusion-reaction ( Section 4.2.2 )	0.001	1000	20,000	0.9999
Quasi-static elastography ( Section 4.2.3 )	0.0001	256	60,000	0.9999
Quasi-static elastography ( Section 4.3.1 )	0.0001	256	150,000	0.9999
Tumor spheroids ( Section 4.3.2 )	0.0001	16	160,000	0.9999

Appendix

Table A3 : Training hyper-parameters for training velocity networks for conditional flow matching

Dataset/Inverse problem/Application	Approximate Wall time	Resource Type (NVIDIA GPU)
Spiral ( Section 4.1 )	0.4 s	Quadro RTX 8000
One step data assimilation ( Section 4.2.1 )	0.4 s	V100
Advection-diffusion-reaction ( Section 4.2.2 )	0.04 s	L40s
Quasi-static elastography ( Section 4.2.3 )	215 s	A100-80GB
Quasi-static elastography ( Section 4.3.1 )	160 s	A100-80GB
Tumor spheroids ( Section 4.3.2 )	100 s	A100-80GB

Appendix

Table A4 : Approximate computational cost of training velocity networks for 100 iterations

Figure B1 : Train and test loss curve for the velocity network with 3 hidden layers each of width 64 trained using the training data shown in Fig. 3 (a). The blue curve above shows the moving average (MA) of the test loss

Figure B2 : Mean and one-standard-deviation interval of the conditional distribution $\rho_{X|Y}$ estimated using samples generated by the trained velocity network with 3 hidden layers each of width 64 at different stages of training

Figure B3 : (a) Train and test data for the toy example used to illustrate the effects of overfitting with more data. (b) Train and test loss curve for the velocity network trained using the training data shown in Fig. B3 (a). The blue curve above shows the moving average (MA) of the test loss

Figure B4 : Mean and one-standard-deviation interval of the conditional distribution $\rho_{X|Y}$ estimated using samples generated by the trained velocity network at different stages of training and using more training data

Conditional Flow Matching

Conditional Flow Matching trains vector fields that transport simple source distributions into target distributions conditioned on observed data.

See all 61 papers

FIRM: Flow-based Imaging via Regularized Minimization
Sep 11, 2026Shirin Shoushtari, Edward P. Chandler, Xiao Shi +1Conditional Flow Matching Image Reconstruction
WashU · UW-Madison
Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems
Oct 1, 2026Jan Blechschmidt, Oliver G. Ernst, Moritz Poguntke +1Posterior Sampling Bayesian Inverse Problems
D-Flow SGLD: Source-Space Posterior Sampling for Scientific Inverse Problems with Flow Matching
Feb 25, 2026Meet Hemant Parikh, Yaqin Chen, Jian-Xun WangPosterior Sampling Bayesian Inverse Problems
Sibley School of Mechanical and Aerospace Engineering, Cornell University, Ithaca, NY, USA · Department of Aerospace and Mechanical Engineering, University of Notre Dame, Notre Dame, IN, USA

Sep 11, 2026J/K move · Enter open · S save

FIRM: Flow-based Imaging via Regularized Minimization

Shirin Shoushtari, Edward P. Chandler, Xiao Shi, Ulugbek S. Kamilov

WashU · UW-Madison

Conditional Flow Matching Image Reconstruction Inverse Problems Image Inverse Problems

Open paper AskPDF arXiv

Flow matching methods for imaging inverse problems typically incorporate measurements through network conditioning or guidance during sampling. Neither approach explicitly applies the forward operator within the learned conditional velocity field. We develop a principled measurement-conditional velocity parameterization that does. For a linear interpolation path, we express the optimal velocity through the posterior mean

E[x_1|x_t, y]

and show that this mean is the unique minimizer of a variational objective with an explicit data-consistency term. The velocity defined by this minimizer provably transports the source distribution to the measurement-conditioned posterior. This result leads to a forward operator-aware velocity field that is trained end-to-end and requires no separate guidance during sampling. Across five imaging tasks, our method achieves leading reconstruction quality with up to

50\times

fewer network evaluations than competitive flow-based methods. Varying the number of sampling steps also controls the distortion-perception trade-off without retraining.

Abstract

Conditional Flow Matching

Explore similar work

FIRM: Flow-based Imaging via Regularized Minimization

Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems

D-Flow SGLD: Source-Space Posterior Sampling for Scientific Inverse Problems with Flow Matching