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
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 ϕ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 ρX∣Y(x∣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 ρ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∼N(0,1) (left) and Z∼ρ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 ρX∣Y(x∣y=y^) for y^∈{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 ρX as the source distribution
Figure 10: Particles from the prior ρX ( ∙ ) and posterior ρX∣Y ( ∙ ), and observation y^ ( – ) on the (x1,x3) and (x2,x3) 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 ρX (part of the training dataset), and samples from ρ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
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)
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.
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 Nstep
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 X and 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
Mean
Std. dev
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)
U(7.1,19.2)
Distance of the inclusion’s center from the bottom edge (mm)
U(7.1,27.6)
Radius of the inclusion (mm)
U(3.5,7.0)
Ratio between the inclusion’s and background’s shear modulus (mm)
U(1,8)
Table 5: Details of random variables comprising the parametric prior distribution for X in the quasi-static elastography application
Figure 19: Five typical realizations of X and 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 20: Training loss, test loss, and the moving average of the test loss for the velocity network in the quasi-static elastography application with experimental data
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 X and 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
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 ρ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 ρX∣Y estimated using samples generated by the trained velocity network at different stages of training and using more training data
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[x1∣xt,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× fewer network evaluations than competitive flow-based methods. Varying the number of sampling steps also controls the distortion-perception trade-off without retraining.
Shirin Shoushtari, Edward P. Chandler, Xiao Shi +1
Sampling from the posterior is the central task of computational Bayesian inverse problems. The standard workhorse in Bayesian inference - Markov chain Monte Carlo (MCMC) - is sequential, yields correlated samples, and must be rerun for each observation. Conditional flow matching offers an amortized alternative: a transport map, trained once on joint samples of parameter and data, that yields independent approximate posterior samples for any observation at negligible online cost, without new likelihood evaluations. We give a careful, MCMC-literate assessment of flow matching for PDE-based inverse problems with function-valued parameters. Exploiting the flow's tractable density, we derive computable accuracy estimates of the underlying approximate posterior in total-variation distance and Kullback-Leibler divergence and, moreover, propose a hybrid sampler that is asymptotically exact by Metropolization. We validate the accuracy estimates and demonstrate the amortization in several numerical examples, including electrical impedance tomography and a likelihood-free Lotka-Volterra model.
Jan Blechschmidt, Oliver G. Ernst, Moritz Poguntke +1
Data assimilation and scientific inverse problems require reconstructing high-dimensional physical states from sparse and noisy observations, ideally with uncertainty-aware posterior samples that remain faithful to learned priors and governing physics. While training-free conditional generation is well developed for diffusion models, corresponding conditioning and posterior sampling strategies for Flow Matching (FM) priors remain comparatively under-explored, especially on scientific benchmarks where fidelity must be assessed beyond measurement misfit. In this work, we study training-free conditional generation for scientific inverse problems under FM priors and organize existing inference-time strategies by where measurement information is injected: (i) guided transport dynamics that perturb sampling trajectories using likelihood information, and (ii) source-distribution inference that performs posterior inference over the source variable while keeping the learned transport fixed. Building on the latter, we propose D-Flow SGLD, a source-space posterior sampling method that augments differentiable source inference with preconditioned stochastic gradient Langevin dynamics, enabling scalable exploration of the source posterior induced by new measurement operators without retraining the prior or modifying the learned FM dynamics. We benchmark representative methods from both families on a hierarchy of problems: 2D toy posteriors, chaotic Kuramoto-Sivashinsky trajectories, and wall-bounded turbulence reconstruction. Across these settings, we quantify trade-offs among measurement assimilation, posterior diversity, and physics/statistics fidelity, and establish D-Flow SGLD as a practical FM-compatible posterior sampler for scientific inverse problems.
Meet Hemant Parikh, Yaqin Chen, Jian-Xun Wang
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