Composing Flow-Matching Energies with Known Physics: Generation, OOD Detection, and Inversion on PDE Fields

Yixuan Sun; Anirban Samaddar; Sandeep Madireddy

cs.LGAug 18, 2026

Composing Flow-Matching Energies with Known Physics: Generation, OOD Detection, and Inversion on PDE Fields

Authors: Yixuan Sun, Anirban Samaddar, Sandeep Madireddy

Organizations: Mathematics and Computer Science, Argonne National Laboratory

Abstract

Probabilistic modeling of physical fields benefits from both a data-driven prior and known physical structure such as the governing equations. Energy-based models (EBMs) are a natural fit since energies compose additively, which enables augmenting physics information during inference. However, EBMs have been difficult to train and sample from due to the intractable partition function. We show in this work that flow matching models with a potential-induced velocity yield an explicit scalar energy at all transport times, whose gradient is exactly the converted learned score and which recovers the marginal negative log-density at the population optimum. The time-dependent energy functions are obtained purely from the matching regression objective on an independent linear Gaussian interpolation, without a variational form or additional MCMC steps, and the sampling retains the flow ODE. Access to the energy function from a trained model serves three roles: energy-corrected data generation, energy as a scoring function for out-of-distribution (OOD) detection, and energy compositional posterior sampling for inverse problems. In particular, we show the explicit energy permits general MCMC samplers in the predictor-corrector sampling framework, reducing PDE residual and spectral distance compared to the flow ODE baseline. Furthermore, we demonstrate utilizing the data energy and physics-based energy (e.g., PDE residuals) as complementary mechanisms to improve detection accuracy for OOD tasks. In addition, we explore the connection to MCMC-based inference for inverse problems by composing the energy with a quadratic observational likelihood that yields a posterior energy, used as an explicitly chosen family of inference-time targets.

Figures & tables

Figure 1: Energy landscapes and flow generated samples of 2D datasets from potential parameterized flow matching training. The flow ODE samplers correctly generate samples resembling true data points. At the same time, the learned energy landscapes share visually the same geometries as the true negative log densities, where a lower value represents higher likelihood.

Dataset	Sampler	NFE	Sliced- $W_{2}\!\downarrow$	Spectral $\downarrow$	PDE resid $\downarrow$	MMSE $\downarrow$	SMSE $\downarrow$
Poisson	ODE120	120	$0.0532{\pm}0.0046$	$1.180{\pm}0.079$	$13.66{\pm}5.46$	$2.66{\pm}3.14$	$16.33{\pm}4.30$
	ODE240	240	$0.0520{\pm}0.0048$	$1.181{\pm}0.079$	$14.59{\pm}5.91$	$2.92{\pm}3.51$	$12.72{\pm}3.75$
	ODE360	360	$0.0518{\pm}0.0048$	$1.181{\pm}0.076$	$14.94{\pm}5.82$	$3.01{\pm}3.71$	$\mathbf{11.58{\pm}3.69}$
	Standard PC	240	$0.0449{\pm}0.0047$	$1.013{\pm}0.034$	$4.54{\pm}0.93$	$3.12{\pm}2.73$	$22.51{\pm}4.91$
	EnergyPC-ULA (ours)	240	$0.0449{\pm}0.0050$	$0.979{\pm}0.024$	$3.65{\pm}0.57$	$3.12{\pm}2.85$	$22.51{\pm}4.87$
	EnergyPC-MALA (ours)	360	$\mathbf{0.0420{\pm}0.0050}$	$\mathbf{0.915{\pm}0.032}$	$\mathbf{2.46{\pm}0.64}$	$\mathbf{1.81{\pm}2.06}$	$18.35{\pm}4.27$

Table 1: Unconditional generation metrics across samplers and datasets, with 95% bootstrap confidence intervals (2000 resamples of the 100-sample ensemble). MMSE/SMSE

{\times}10^{-4}

; Burgers PDE residual

{\times}10^{-4}

. Bold marks the best mean value per column within each dataset.

Figure 2: Visualization of samples of Poisson and Burgers samples generated using the baseline flow ODE and EnergyPC with MALA.

Figure 3: Synthetic OOD sample examples on the Poisson $a$ and $u$ pairs.

	Poisson			Helmholtz			Burgers
OOD tier	$E$	$R$	$E_{\rm tot}$	$E$	$R$	$E_{\rm tot}$	$E$	$R$	$E_{\rm tot}$
Gaussian noise	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00
Noise ( $10\%$ field std)	1.00	1.00	1.00	1.00	1.00	1.00	1.00	$0.843{\pm}0.027$	1.00
Noise ( $50\%$ field std)	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00
Channel shuffle	$0.992{\pm}0.005$	1.00	1.00	1.00	1.00	1.00	–	–	–
Cross-PDE	1.00	1.00	1.00	1.00	1.00	1.00	–	–	–

Table 2: OOD detection AUROC with the energy, with 95% bootstrap confidence intervals ( italic =at or near chance). The learned energy

E

and the PDE residual

R

are complementary: each is blind where the other separates, and the total energy

E_{\rm tot}

flags every tier where either does.

Dataset	Method	aRE $\downarrow$	uRE $\downarrow$
Poisson	DiffusionPDE	$0.285{\pm}0.020$	$\mathbf{0.146{\pm}0.017}$
	+ $E_{\rm pos}$ -ULA	$0.304{\pm}0.057$	$0.151{\pm}0.018$
	+ $E_{\rm pos}$ -MALA	$\mathbf{0.255{\pm}0.010}$	$0.150{\pm}0.020$
	FlowDPS	$0.378{\pm}0.008$	$0.039{\pm}0.001$
	+ $E_{\rm pos}$ -ULA	$0.373{\pm}0.007$	$0.038{\pm}0.001$
	+ $E_{\rm pos}$ -MALA	$\mathbf{0.353{\pm}0.012}$	$\mathbf{0.036{\pm}0.001}$

Table 3: Relative

L_{2}

errors of reconstructed

u

and inferred

a

in the inverse problem. Fields are the ensemble mean over 16 samples; energy correctors are added at the base’s own predictor levels.

\pm

is the standard deviation over resamples of the 16 draws.

Figure 4: Poisson solution field reconstruction and coefficient inference comparison using energy-corrector with MALA (+E) extended diffusion posterior samplers.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Unconditional generation of Poisson equation data with various samplers. All energy-based samplers reduce the PDE residuals compared to the baseline flow ODE sampler.

Sampler	physics	NFE	Sliced- $W_{2}\!\downarrow$	PDE resid $\downarrow$
ODE120 †	–	120	0.0532	13.66
ODE + ULA polish	terminal	220	0.0408	10.33
ODE + MALA polish	terminal	320	0.0409	7.17
Annealed + $\sigma$ -sched	gated	1752	0.0880	1.15

Appendix

Table 4: Additional unconditional samplers on Poisson, not reported in the main text.

Figure 6: Unconditional generation of Burgers equation data with various samplers. Energy-based samplers, except the energy annealing sampler, reduce the PDE residual compared to the ODE baseline sampler.

Sampler	physics	NFE	Sliced- $W_{2}\!\downarrow$	PDE resid $\downarrow$
ODE120 †	–	120	0.0521	3.59
ODE + ULA polish	terminal	220	0.0531	2.58
ODE + MALA polish	terminal	320	0.0513	2.54
Annealed + $\sigma$ -sched	gated	1752	0.1928	7.73

Appendix

Table 5: Additional unconditional samplers on Burgers, not reported in the main text.

Physics-Informed Generative Modeling

Physics-informed generative modeling uses physical laws and data-driven generators to create plausible samples and solve scientific inverse problems.

See all 108 papers

May 25, 2026

Guided Flow Matching for Forward and Inverse PDE Problems with Sparse Observations: Algorithm and Theory

Reconstructing PDE solutions from sparse observations is a core challenge in scientific computing. We present FM4PDE, a flow-matching generative framework that learns the joint distribution of PDE coefficients (or initial states) and solutions (or final states), enabling both forward simulation and inverse recovery with limited paired data. At inference, sampling is guided by a composite loss that enforces agreement with sparse measurements and reduces the PDE residual; we support deterministic, stochastic, and hybrid samplers. We provide error guarantees for these guided procedures. For the deterministic optimizer, a coercivity condition ensures trajectory boundedness and a phase-wise contraction yields logarithmic complexity in the target accuracy. For the stochastic sampler, we introduce adaptive guidance and assume dissipativity of the velocity field to obtain uniform moment bounds independent of the noise-floor parameter. This leads to polynomial-time error bounds, and a matching lower bound shows constant guidance induces an unavoidable positive bias, motivating adaptivity. A hybrid deterministic-stochastic analysis is also provided. Experiments on static and time-dependent benchmark PDEs demonstrate competitive accuracy and faster inference than diffusion-based generative models.

Xifeng Zhang, Jin Zhao

School of Mathematical Science Capital Normal University Beijing, 100048, China · Academy for Multidisciplinary Studies Capital Normal University2026 Beijing, 100048, China

Flow Matching PDE Solving PDE Surrogate Modeling PDE Inverse Problems Inverse Problems Inference-Time Guidance

Ask PDF

Aug 5, 2025

Physics-Constrained Fine-Tuning of Flow-Matching Models for Generation and Inverse Problems

We present a framework for fine-tuning flow-matching generative models to enforce physical constraints and solve inverse problems in scientific systems. Starting from a model trained on low-fidelity or observational data, we apply a differentiable post-training procedure that minimizes weak-form residuals of governing partial differential equations (PDEs), promoting physical consistency and adherence to boundary conditions without distorting the underlying learned distribution. To infer unknown physical inputs, such as source terms, material parameters, or boundary data, we augment the generative process with a learnable latent parameter predictor and propose a joint optimization strategy. The resulting model produces physically valid field solutions alongside plausible estimates of hidden parameters, effectively addressing ill-posed inverse problems in a data-driven yet physicsaware manner. We validate our method on canonical PDE benchmarks, demonstrating improved satisfaction of PDE constraints and accurate recovery of latent coefficients. Our approach bridges generative modelling and scientific inference, opening new avenues for simulation-augmented discovery and data-efficient modelling of physical systems.

Jan Tauberschmidt, Sophie Fellenz, Sebastian J. Vollmer +1

DSA, German Research Center for Artificial Intelligence (DFKI) · Department of Computer Science, University of Kaiserslautern–Landau (RPTU) · Department of Mathematics, Imperial College London

Physics-Informed Generative Modeling Flow Matching Generative Modeling Physics-Informed ML PDE Inverse Problems

Ask PDF

Apr 4, 2026

From Euler to Dormand-Prince: ODE Solvers for Flow Matching Generative Models

Sampling from Flow Matching generative models requires solving an ordinary differential equation (ODE) whose computational cost is dominated by neural network forward passes. We derive four classical ODE solvers -- Euler, Explicit Midpoint, Classical Runge-Kutta (RK4), and Dormand-Prince 5(4) -- from first principles via Taylor expansion, implement them from scratch in PyTorch, and systematically benchmark their efficiency on Conditional Flow Matching tasks ranging from 2D toy distributions to MNIST digits. On the quantitative side, we use sliced Wasserstein distance to construct NFE-quality Pareto frontiers,finding that RK4 at 80 function evaluations achieves sample quality comparable to Euler at 200. Beyond reproducing known convergence rates, we report two empirical observations: (1) the Jacobian eigenvalue spectrum of the learned velocity field stiffens sharply near t=1, explaining why the adaptive Dormand-Prince solver automatically concentrates its step budget at the end of the trajectory; (2) the quality gap between low-order and high-order solvers widens for undertrained and smaller models, indicating that solver choice matters most when the model is imperfect. Code and all experiment scripts are publicly available.

Hao Xiao

ATLAS AI Lab

Numerical Integration Flow Matching Conditional Flow Matching

Ask PDF

Abstract

Physics-Informed Generative Modeling

Explore similar work

Guided Flow Matching for Forward and Inverse PDE Problems with Sparse Observations: Algorithm and Theory

Physics-Constrained Fine-Tuning of Flow-Matching Models for Generation and Inverse Problems

From Euler to Dormand-Prince: ODE Solvers for Flow Matching Generative Models