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

• Flow matching supports generation, but OOD detection and inverse problems require additional inference machinery; EBMs offer explicit energies but are difficult to train and sample.Source: Introduction · p. 1
• A potential parameterization of standard flow matching yields time-dependent energies whose converted scores are conservative by construction and directly readable from the trained model.The energies correspond to unnormalized negative log densities along the transport path.Source: Introduction · pp. 1–2
• For PDE fields, total energy additively combines learned data energy with known physical structure, and EnergyPC enables physics-aware unconditional generation and posterior sampling.Source: Introduction · p. 2
• On Poisson, Helmholtz, and Burgers datasets, the approach improves generation, OOD detection, field reconstruction, and coefficient inference over flow ODE, predictor–corrector, and guidance-based posterior sampling baselines.The introduction reports these gains across high-dimensional PDE datasets.Source: Introduction · p. 2
• Data energy and physical energy provide complementary OOD scoring functions, while their combined total energy is effective across the tested corruptions.Source: Introduction · p. 2

Energy-based models

• EBMs provide explicit unnormalized densities for density estimation, OOD detection, composition, and uncertainty quantification, but training and sampling have required expensive MCMC or related methods.Score matching avoids expensive MCMC but has ill-defined-score and off-manifold estimation challenges; multilevel EBMs improve learning but require careful noise schedules.Source: Related Work · p. 2

Flow Matching

• Flow matching learns transport velocities and marginal distributions, but does not directly expose likelihood or energy; physical-field applications have imposed physics through constraints, guidance, or velocity fine-tuning.Likelihood computation by change of variables is expensive for high-dimensional physical datasets; this work instead uses a potential whose gradient approximates velocity to construct marginal energies for additive physics and direct OOD scoring.Source: Related Work · p. 2

Energy-based flow/diffusion models

• Related energy-based flow and diffusion methods use customized training targets or additional procedures, whereas this work uses standard flow matching with a potential-induced flow to obtain time-dependent marginal energies for PDE fields.The cited alternatives include MCMC-based maximum-likelihood training, a variational Poisson-equation objective, and equilibrium constraints; this work augments the data prior with known physical structures.Source: Related Work · pp. 2–3

Energy and Potential Flow

• The potential-induced velocity yields an explicit marginal energy whose gradient is the converted score; for a Gaussian path, the curl-free restriction is no restriction at the optimum.At the population optimum, the target marginal, the marginal induced by the learned velocity, and the density proportional to exp(−Eθ) coincide; finite-training read-outs need not be transport-consistent. Because the terminal-time energy is numerically unstable, the method uses a minimum noise level as a close approximation to the data energy.Source: Methods · p. 3

EBM training through Flow Matching

• Training parameterizes the flow potential with a neural network and learns it through the standard flow-matching regression objective on a linear Gaussian interpolant.The resulting energy is a direct read-out of the learned potential; under standard neural-network implementations it defines a normalizable density for t < 1, without conventional EBM training’s expensive, unstable MCMC sampling.Source: Methods · p. 3

Sampling

• EnergyPC combines the learned energy with a time-weighted squared PDE residual and uses energy-based MCMC correctors to refine samples toward physical feasibility.The residual weight vanishes at high noise, and the terminal target is the learned prior corrected by known physics. MALA with a Metropolis adjustment is invariant to exp(−Etot) at each noise level, removing the discretization bias of an unadjusted Langevin step.Source: Methods · pp. 3–4

Energy as Scoring Function

• The terminal total energy serves as an OOD score: inputs above a threshold are classified as OOD, and AUROC measures correct ranking of in- and out-of-distribution pairs.The evaluation uses equal numbers of in-distribution test fields and corrupted fields per tier; AUROC of 1 indicates perfect separation, while values below 0.5 indicate score failure.Source: Methods · p. 4

Posterior Energy for Inverse Problem

• For inverse problems, the method composes the learned field energy with a quadratic Gaussian observation-likelihood energy to define a posterior energy for the joint field.The likelihood is evaluated on the current field at every noise level, so the annealing family is accurate only as the field becomes clean; its terminal member coincides with the flow posterior. The posterior energy gives the MCMC corrector an explicit target at each noise level.Source: Methods · pp. 4–5

Unconditional generation

• EnergyPC with MALA reduces PDE residual and spectral distance across Poisson, Helmholtz, and Burgers generation, while its MMSE and SMSE effects are less consistent.EnergyPC with ULA outperforms standard PC on Poisson and performs similarly on Helmholtz and Burgers.Source: Experiments · p. 5

Out-of-distribution detection

• The data energy and PDE residual complement one another: their total Etot improves blur and roll detection, while all energies detect channel-shuffled pairs jointly.Across most OOD cases, either component alone achieves perfect or near-perfect detection; data energy struggles with blur, and physical energy alone struggles with Burgers rolls.Source: Table 2 · p. 6, Experiments · p. 6

Inverse Modeling

• Posterior-energy correctors improve DiffusionPDE and FlowDPS overall, with the largest and most consistent gains in coefficient inference, except for a slight Poisson solution-error increase with DiffusionPDE.Correctors are applied after guidance-updated ODE steps using 500 randomly observed solution-field locations; MALA performs better than ULA on Poisson, while ULA shows marginal improvement over MALA on Helmholtz.Source: Experiments · pp. 6–7, Table 3 · p. 7

• The potential-parameterized flow-matching model with a Gaussian probability path provides explicit energy functions for transport-time marginal distributions.Source: Conclusion · p. 7
• Combining these energies with physical structure creates a total-energy prior for EnergyPC, which improves sample quality and physical feasibility.Source: Conclusion · p. 7
• The total energy combines complementary terms into a single scoring function for OOD detection.Source: Conclusion · p. 7
• Composing the total energy with observational likelihood energy leads to better field reconstruction and coefficient inference in PDE inverse problems.Source: Conclusion · p. 7

Limitations

• The framework adds computational complexity versus standard flow matching, and its read-out energy is algebraically singular at t = 1, preventing direct recovery of the true terminal energy.Future work will focus on addressing these challenges.Source: Conclusion · p. 7

A.1 The identification of f and g

• The interpolation is x_t = t x_1 + (1 − t)x_0, with x_0 ∼ N(0, I) as t runs from 0 to 1.Source: §A.1 · p. 10
• Using diffusion time s = 1 − t makes s run from 1 to 0 in the transport direction from data to source, with α(s) = 1 − s and σ(s) = s.Source: §A.1 · p. 10
• The corresponding linear forward SDE is dx = f(x, s) ds + g(s) dw, with drift f(x, s) = h(s)x.Source: §A.1 · p. 10
• The coefficient relations yield h(s) = −1/(1 − s), f(x, s) = −x/(1 − s), and g^2(s) = 2s/(1 − s).These follow by substituting α(s) = 1 − s and σ^2(s) = s^2 into the stated SDE coefficient equations.Source: §A.1 · p. 10

A.2 Terminal Energy

• Setting c(t) ≡ 0 gives the energy of the transport marginals.Source: §A.2 · p. 10
• For visualization and AUROC, they use t_max = 0.99999 and E(x) = ∥x∥2 − 2Φ(x, t_max); EnergyPC experiments use σ_min = 10^−3 and t_max = 0.999.The simplified energy drops the positive constant 1/(1 − t_max) and factor t_max; the authors report that this largely preserves landscape geometry and minimally affects AUROC.Source: §A.2 · p. 10

Scalar-potential parameterization

• The scalar potential is the input–UNet dot product, whose input gradient trains the model; this form performed better empirically than the denoising-autoencoder objective (Du et al. 2023).The UNet backbone follows Dhariwal and Nichol 2021, and time conditioning uses t and σ = 1 − t.Source: B Potential flow model training details · p. 10

Data normalization

• Data are standardized channel-wise using fixed statistics, which are reused to restore generated samples to physical scale for PDE-residual computation.Distributional metrics are computed in normalized space.Source: B Potential flow model training details · p. 10

Architecture

• Poisson, Helmholtz, and Burgers share one architecture and hyperparameter set, including four residual blocks per resolution and attention at downsample factors 4 and 8.Attention uses 64 channels per head.Source: B Potential flow model training details · p. 11

Model training

• Training runs for 40,000 steps with batch size 4, 5,000 warmup steps, gradient clipping at 3, and weight decay 0; evaluation uses raw weights at step 34,000 for Poisson and step 40,000 for the other datasets.The random seed is 49.Source: B Potential flow model training details · p. 11
• The conditional-OT interpolant uses target velocity v⋆ = x1 − ε, samples t uniformly from [10^-5, 0.99999], and sets αt = t and σt = 1 − t.No additive noise floor is used for σt, and time shifting is disabled.Source: B Potential flow model training details · p. 11

Noise ladder and predictor

• The noise ladder is geometric in σ(t)=1−t, spanning σmin=10^-3 to σmax=1−10^-5; one probability-flow ODE Euler step precedes each corrector.The stated t range is 10^-5 to 0.999, and the predictor is evaluated at t=1−σmin before the corrector runs at the next ladder level.Source: C Additional details on EnergyPC · pp. 11–12

EnergyPC target

• The corrector targets e^-Etot/Z, using the physical-energy leaf gradient for Langevin proposals and Etot directly in Metropolis steps.Etot combines the time-dependent data energy with the discrete PDE residual energy.Source: C Additional details on EnergyPC · p. 12

Adaptive corrector steps

• Per-sample adaptive Langevin steps shrink on stiff off-manifold states, retain nonzero acceptance at high noise, and cap moves when the score is small.The rule combines an SNR-adaptive step from Song et al. 2020 with a classic NCSN cap from Song and Ermon 2019, allowing unclipped physics guidance while preserving MALA’s exact target.Source: C Additional details on EnergyPC · p. 12

Physics weight schedule

• A noise-scale gate confines physics guidance to the band where the Tweedie estimate is meaningful; setting λmax=0 recovers Standard PC without physics at every level.At high noise, the estimated state is far off-manifold and the residual and its gradient can dominate the score.Source: C Additional details on EnergyPC · p. 12

Inverse problems

• For inverse problems, EnergyPC retains unmodified DiffusionPDE or FlowDPS predictors and corrects toward the true posterior energy with a squared-residual Gaussian likelihood and exact MH ratio.FlowDPS guides only on observations, whereas DiffusionPDE guides on both observations and the PDE residual.Source: C Additional details on EnergyPC · p. 12

D PDE datasets

• Each DiffusionPDE dataset contains 10,000 samples at 128×128 resolution, split into 9,000 training and 1,000 test samples; evaluations use only the test split.The samples are read from the distributed .mat files.Source: D PDE datasets · p. 13
• For Poisson, the two channels are forcing a and solution u, with residual F = ∆_h u − a for the zero-boundary equation ∆u = a.Source: D PDE datasets · p. 13
• For Helmholtz, the two channels are forcing a and solution u, with residual F = ∆_h u + u − a for ∆u + u = a and zero boundary conditions.Source: D PDE datasets · p. 13
• The Burgers dataset has one channel u(t, x) and residual F = u_t + u u_x − νu_xx, with ν = 0.01.Source: D PDE datasets · p. 13
• Poisson and Helmholtz residuals use a zero-padded 5-point Laplacian with h = 1/(S − 1), S = 128, then zero the boundary edges; Burgers uses central-difference 3-tap stencils.Burgers uses unit grid spacing (dt = dx = 1), zero padding, and a 128×128 space–time grid with rows indexing time and columns space; residual discretizations match DiffusionPDE exactly.Source: D PDE datasets · p. 13
• Fields follow DiffusionPDE fixed per-channel standardization, with mean 0 for every channel.The passage says standard deviations are listed, but their values are not provided here.Source: D PDE datasets · p. 13

E Evaluation Metrics

• Distributional metrics use normalized coordinates to weight channels comparably, while PDE residuals are evaluated in the original physical scale.Source: E Evaluation Metrics · p. 13
• With sample size limit m = 100, sliced Wasserstein distance measures discrepancy between generated and true data samples, estimated using 256 i.i.d. Gaussian directions normalized to the unit sphere and 128 evenly spaced quantiles per projection.Source: E Evaluation Metrics · p. 13
• Radial log-spectrum L1 measures the mean absolute log10 power-spectrum discrepancy per radial frequency bin between generated and real fields.Power is averaged over samples, channels, and directions; on the 128 × 128 grid, radial bins run from 0 to 91.Source: E Evaluation Metrics · p. 13
• The ensemble mean PDE residual reports ∥R(ˆx)∥2, while MMSE and SMSE respectively measure mean/bias error and spread mismatch.Source: E Evaluation Metrics · p. 14
• OOD detection uses the learned prior read-off E = ∥x∥2 − 2Φ(x, tmax) and mean-square residual R, combined with balanced in-distribution spread.The fitted balance values are 1.729×10^11 for Poisson, 1.018×10^11 for Helmholtz, and 5.412×10^4 for Burgers.Source: E Evaluation Metrics · p. 14
• Inverse-problem errors report recovered-channel aRE and observed-channel uRE in physical units, with confidence intervals reflecting resampled posterior draws.Observations comprise 500 uniformly selected channel-u pixels with fixed locations and Gaussian noise; intervals use 5,000 resamples for inverse tables and 2,000 for other metrics.Source: E Evaluation Metrics · p. 14

• ODE-sample refinement applies 100 ULA or MALA steps using the terminal-time energy, initializing chains with baseline ODE samples.The step size is SNR-adaptive.Source: F Additional results · p. 14
• Energy annealing starts from Gaussians and sweeps 120 noise levels from σ = 1 to 10^-3, taking eight Langevin steps per level before Tweedie denoising.The physics weight ramps up as noise falls; sampling switches from ULA to MALA below σ = 0.3.Source: F Additional results · p. 14
• All energy-based samplers reduce PDE residuals relative to the flow ODE sampler for unconditional Poisson data generation.Source: Figure 5 · p. 15
• For Burgers generation, energy-based samplers reduce PDE residuals versus the ODE baseline, except energy annealing.Source: Figure 6 · p. 16