PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yixuan Sun, Anirban Samaddar, Sandeep Madireddy

arXiv:2608.18004v1cs.LGphysics.comp-ph

TL;DR

Flow matching models lack a direct, broadly explored energy formulation for PDE-field tasks such as OOD detection and inverse problems. This paper derives explicit transport-time energies from standard potential-parameterized flow matching and shows their composition with physical and observational energies supports improved generation, OOD scoring, and inverse inference.

  • Problem

    Flow matching for PDE fields lacks a standard explicit energy formulation for inference tasks such as OOD detection, composition, and inverse problems, while existing hybrids require customized training.

  • Method

    A potential-parameterized flow matching model directly reads out transport-time energies, which are composed with physical or observational energies in EnergyPC and posterior sampling.

  • Results

    Across PDE-field tasks, energy adjustment improved physical feasibility and generally improved inverse reconstruction, while data and physical energies provided effective complementary OOD scores.

  • Takeaways & Limitations

    Explicit energies give one trained flow matching model a flexible inference-time interface for physics-aware generation, OOD detection, and posterior sampling.

  • Takeaways & Limitations

    The framework adds computational complexity, and its read-out energy is singular at terminal transport time t = 1, so the true terminal energy is not directly recoverable.

Abstract

from arXiv · show

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.

Introduction

The introduction motivates combining flow matching’s efficient training and sampling with EBMs’ explicit unnormalized densities for PDE-field generation, OOD detection, and inverse-problem inference. It proposes potential-parameterized flow matching, whose conservative scores enable time-dependent energies and physics-aware inference.

  • Motivation: Flow matching generates high-dimensional scientific data, including accurate PDE solution fields with sampling-based uncertainty quantification.
  • Motivation: EBMs promise one generative model supporting generation, OOD detection, and inverse-problem inference through explicit energy functions.
  • Method: The proposed approach combines flow matching’s training and sampling efficiency with EBMs’ explicit unnormalized density by equating marginal scores with negative energy gradients.
  • Method: Potential parameterization makes converted scores conservative by construction, yielding time-dependent EBMs along the transport path while retaining standard flow-matching training.
  • Contributions: The work augments data energy with known physical structures and introduces EnergyPC for physics-aware unconditional generation and posterior sampling.

Related Work

Related work establishes EBMs as explicit-energy models for density estimation and OOD detection, while score matching and flow matching provide alternatives for learning distributions and sampling. Recent work connects flow or diffusion models with EBMs; this paper extends that connection to time-dependent marginal energies for PDE fields and known physical structures.

  • Energy-based models: EBMs use parameterized energy functions with Boltzmann distributions and provide explicit access to unnormalized densities for density estimation and OOD detection.
  • Energy-based models: Score matching avoids expensive MCMC but faces ill-defined scores and estimation difficulties off low-dimensional data manifolds.
  • Energy-based models: Multilevel EBMs and noise annealing improve EBM learning but still require careful noise schedule design.
  • Flow Matching: Flow matching learns a neural velocity field transporting samples from an easy-to-sample source distribution to a complex target distribution, using numerical ODE or SDE solvers.
  • Energy-based flow/diffusion models: Prior work relates flow matching or diffusion models to EBMs, whereas this work accesses time-dependent marginal energies for PDE fields through a parameterized potential and augments data energy with known physical structures.

Methods

The method parameterizes flow matching with a scalar potential, enabling direct time-dependent energy read-outs from the learned flow and their use in physics-corrected sampling and OOD scoring. It augments the learned energy with PDE-residual structure and applies energy-based predictor–corrector inference while accounting for terminal-time numerical instability and finite-training inconsistency.

  • Energy and Potential Flow: A scalar potential parameterizes the flow velocity, yielding an explicit energy whose gradient recovers the time-marginal score and whose flow has the correct marginals.The curl-free restriction is not limiting at the Gaussian-path optimum because the drift terms are gradients.
  • Energy and Potential Flow: The energy is read out with c(t)=0, recovering the Gaussian source energy at t=0, while terminal-time evaluation is stabilized by imposing a minimum noise level σmin>0.The integration constant does not affect fixed-time MCMC sampling or OOD scoring, but the t=1 representation is indeterminate and numerically unstable.
  • EBM training through Flow Matching: The learned potential is trained with the standard flow-matching regression objective, after which its potential and energy are directly available for downstream tasks.At the population optimum, the target marginal, induced marginal, and exp{−Eθ}/Zθ coincide; finite training does not enforce this consistency.
  • Sampling: The total inference energy combines the learned data energy with a weighted squared L2 PDE-residual term, whose weight is annealed toward zero at high noise.The intermediate targets form an inference-time annealing family rather than exact interpolant marginals, and the terminal member recovers the learned prior corrected by known physics.
  • Sampling: EnergyPC uses the flow ODE as predictor and energy-based MCMC correctors, with MALA accepted through a Metropolis adjustment that is invariant to exp{−Etot(x,t)} at each noise level.The framework can also use Hamiltonian Monte Carlo, replica exchange, and Jarzynski-reweighted variants; the adjustment removes discretization bias from unadjusted Langevin steps.
  • Energy as Scoring Function: Terminal total energy Etot(x,tmax), with tmax=1−σmin, is used as an inference-time OOD discriminator whose AUROC measures correctly ranked in-distribution and out-of-distribution pairs.The evaluation uses equal numbers of in-distribution test fields and corrupted fields per tier.

Experiments

Experiments evaluate the potential flow-matching framework for unconditional PDE-field generation, OOD detection, and inverse problems. Energy-based correctors improve generation and posterior sampling, while combining learned data energy with PDE residual energy strengthens OOD detection.

  • Unconditional generation: EnergyPC with MALA reduces PDE residual and spectral distance across Poisson, Helmholtz, and Burgers datasets, while ULA outperforms standard PC on Poisson and performs similarly on Helmholtz and Burgers.MALA is superior to the purely score-based ULA corrector.
  • Out-of-distribution detection: The learned data energy and PDE residual are complementary OOD scores, and their total Etot provides a single scoring function that flags cases separated by either term.Most test cases achieve perfect or near-perfect detection with either separate energy term, while Etot addresses their failure modes.
  • Out-of-distribution detection: Etot improves detection for blurred and spatially rolled samples when data energy or physical energy alone struggles because of high-likelihood smoothing or preserved physical feasibility.Blurred samples lose high-frequency details yet remain visually consistent, while Burgers rolls preserve physical feasibility but leave the data distribution.
  • Inverse modeling: Energy-corrector extensions improve reconstructed solution and inferred coefficient fields within both DiffusionPDE and FlowDPS posterior-sampling frameworks.The corrector is applied after each guidance-updated ODE integration step using the posterior energy.
  • Inverse modeling: Energy correctors show enhanced posterior-sampling performance because they are rooted in an explicitly defined posterior energy, with the largest and most consistent gains over the MALA corrector.The inverse tasks reconstruct u and infer coefficient field a from 500 randomly sampled solution-field observations.

Conclusion

The framework provides explicit transport-time energies that combine with physical structure to form augmented priors and support EnergyPC refinement. Its main limitations are added computational complexity and an algebraically singular read-out energy at t = 1.

  • Potential-parameterized flow matching with a Gaussian probability path provides explicit energy functions for transport-time marginal distributions.
  • Combining these energies with physical structure creates an augmented prior that enables EnergyPC for refined sampling with improved sample quality and physical feasibility.
  • Limitations: The potential parameterization adds computational complexity relative to standard flow matching models, affecting training and sampling efficiency.
  • Limitations: At terminal transport time t = 1, the read-out energy becomes algebraically singular and cannot directly recover the true terminal energy.

A Energy formulation

The section derives a linear-Gaussian transport SDE from the interpolation and identifies a transport-marginal energy by setting its additive constant to zero. It also specifies simplified terminal-energy evaluation and noise-floor choices for visualization, AUROC, generation, and inverse problems.

  • A.1 The identification of f and g: The linear interpolant transports x_0 ~ N(0, I) to data through x_t = t x_1 + (1 − t) x_0 as t runs from 0 to 1.
  • A.1 The identification of f and g: Introducing diffusion time s = 1 − t reverses the parameterization, with s running from 1 to 0 in the forward data-to-source direction.
  • A.1 The identification of f and g: The interpolation is represented by a linear forward SDE, dx = f(x, s) ds + g(s) dw, with drift f(x, s) = h(s) x.
  • A.1 The identification of f and g: The interpolation coefficients satisfy α̇ = h(s) α(s) and σ̇^2 = 2h(s)σ^2(s) + g^2(s).
  • A.1 The identification of f and g: Substituting α(s) = 1 − s and σ^2(s) = s^2 yields f(x, s) = −x/(1 − s) and g^2(s) = 2s/(1 − s).
  • A.2 Terminal Energy: Setting the additive constant c(t) identically to zero gives the energy of the transport marginals.
  • A.2 Terminal Energy: For visualization and AUROC, the implementation uses t_max = 0.99999 and E(x) = ∥x∥^2 − 2Φ(x, t_max), omitting a positive constant and the factor t_max.The simplification largely preserves the energy-landscape geometry and minimally affects AUROC values.
  • A.2 Terminal Energy: For EnergyPC in unconditional generation and inverse-problem experiments, the noise floor is set to σ_min = 10^-3.

B Potential flow model training details

The potential-flow model uses a scalar-valued network parameterized by a UNet–input dot product, with gradients obtained by automatic differentiation for training. Across Poisson, Helmholtz, and Burgers, training uses a shared architecture and conditional-OT flow-matching schedule with fixed optimization settings.

  • Scalar-potential parameterization: The scalar potential Φθ is formed from the dot product between an ADM/guided-diffusion UNet backbone and input x, with ∇xΦθ computed by automatic differentiation.Time conditioning appends t and σ = 1 − t to the flattened state, while t also uses sinusoidal timestep embedding.
  • Scalar-potential parameterization: The dot-product parameterization outperformed the denoising autoencoder objective empirically in this work and prior work.The comparison is reported against the denoising autoencoder objective form in Salimans and Ho (2021).
  • Architecture: Poisson, Helmholtz, and Burgers use the same model architecture and hyperparameters, including four residual blocks per resolution, SiLU activation, FiLM-style scale-shift normalization, and attention at resolutions [4, 8].The configuration also uses 64 channels per attention head and 19 total attention blocks.
  • Model training: Training uses a conditional-OT flow-matching path with target v⋆ = x1 − ε, αt = t, σt = 1 − t, and uniform timestep sampling over [10^-5, 0.99999].The schedule has no additive noise floor and appends t and σt to the flat input.
  • Model training: The models train for 40,000 iterations with batch size 4, 5,000 warmup steps, zero weight decay, gradient clipping at 3, and seed 49.Evaluation uses raw non-EMA weights at step 34,000 for Poisson and final step 40,000 for Helmholtz and Burgers.

C Additional details on EnergyPC

EnergyPC combines a geometrically spaced noise ladder with an Euler probability-flow predictor and physics-augmented Langevin–Metropolis correction. Its adaptive, gated physics guidance preserves the exact total-energy target and supports posterior sampling for inverse problems.

  • Noise ladder and predictor: The noise ladder is geometrically spaced in σ(t) = 1 − t, with σmax = 1 − 10^-5 and σmin = 10^-3.Thus, t runs from 10^-5 to 0.999.
  • Noise ladder and predictor: The predictor takes one explicit Euler step of the probability-flow ODE before the corrector runs at the next noise level.The update is x ← x + σmin ∇xΦθ(x, 1 − σmin).
  • EnergyPC target: EnergyPC uses the physical-energy gradient for Langevin proposals while Metropolis steps target e^-Etot/Z through the physics-augmented total energy.The physical term uses the leaf gradient of the PDE residual energy, whereas acceptance uses Etot directly.
  • Adaptive corrector: Per-sample, SNR-adaptive step sizes shrink stiff off-manifold physics updates, avoiding drift clipping that would compromise MALA’s exact target.The adaptive rule combines a drift-relative term with a classic NCSN cap and uses d ≈ E∥z∥.
  • Physics weight schedule: The physics-weight schedule activates guidance only where the Tweedie estimate is meaningful and makes Standard PC identical to the physics-free code path when λmax = 0.At high noise, large residual gradients would otherwise dominate the score estimate.
  • Inverse problems: For inverse problems, the corrector targets the true posterior energy with an exact Metropolis–Hastings ratio from the Gaussian observation likelihood.The predictor remains an unmodified DiffusionPDE or FlowDPS guidance step.

D PDE datasets

The study uses three 128×128 DiffusionPDE datasets—Poisson, Helmholtz, and Burgers—with 10,000 samples each, split into 9,000 training and 1,000 test samples. PDE residuals follow the DiffusionPDE discretization, using finite-difference stencils and specified boundary, spacing, and padding conventions.

  • Dataset composition: Poisson, Helmholtz, and Burgers fields come from the DiffusionPDE release at 128×128 resolution, with 9,000 contiguous training samples and 1,000 test samples per dataset.All evaluations use only the test split, indexed [9000, 10000).
  • Governing equations: Poisson satisfies ∆u = a with zero Dirichlet boundaries, Helmholtz satisfies ∆u + u = a, and Burgers satisfies u_t + u u_x − νu_xx = 0 with ν = 0.01.Poisson and Helmholtz contain channels [a, u], while Burgers contains [u(t, x)].
  • Residual discretization: Poisson and Helmholtz residuals use a 5-point Laplacian with zero padding, h = 1/(S −1), S = 128, and zeroed boundary edges.These calculations match the DiffusionPDE codebase exactly.
  • Residual discretization: Burgers residuals use central-difference 3-tap stencils in both axes with unit grid spacing, dt = dx = 1, and zero padding.Rows index time and columns index space on a 128×128 grid.

E Evaluation Metrics

The evaluation uses distributional, spectral, physical-residual, and ensemble-statistics metrics, with normalization chosen to make differently scaled channels comparable. OOD, corruption, inverse-problem, and uncertainty evaluations use explicitly defined score, error, corruption, observation, and bootstrap procedures.

  • Core metrics: Distributional metrics use normalized coordinates, while PDE residuals use the original physical scale.This weighting makes channels with very different physical scales comparably represented in distributional measures.
  • Core metrics: Sliced-W2 measures generated-to-reference distributional discrepancy using 256 random directions and 128 evenly spaced quantiles per projection.The sample-size limit is m = 100, motivating the sliced Wasserstein estimate.
  • Core metrics: Radial log-spectrum L1 reports the mean absolute log10 power-spectrum discrepancy per radial bin, measuring energy mismatch across spatial scales.On the 128 × 128 grid, radial frequencies run from r = 0 to R = 91.
  • Core metrics: PDE residual is the ensemble mean of ∥R(ˆx)∥2, while MMSE and SMSE respectively detect mean/bias error and spread mismatch.MMSE and SMSE compare per-pixel generated ensemble means and standard deviations with those of the test fields.
  • OOD and inverse evaluation: OOD scoring combines the learned prior read-off E and mean-square residual R through scale balancing, with Etot flagging a field if either term does.The fitted λbal values are 1.729×10^11 for Poisson, 1.018×10^11 for Helmholtz, and 5.412×10^4 for Burgers.
  • OOD and inverse evaluation: Inverse-problem errors use ensemble means from M = 16 posterior draws, while confidence intervals use 2000 bootstrap samples and 5000 inverse-table resamples.Observations comprise 500 uniformly sampled pixels of channel u with Gaussian noise near 2% of field standard deviation.

F Additional results

Additional experiments use the learned time-dependent energies for pure MCMC sampling, either refining ODE samples or generating samples through annealing from Gaussians. Across Poisson and Burgers fields, energy-based samplers generally reduce PDE residuals relative to the baseline flow ODE sampler, except Burgers energy annealing.

  • Additional samplers: The learned energy functions enable pure MCMC sampling from either ODE samples or Gaussian initializations.The section introduces two additional sampler families based on access to the explicit energy functions.
  • Energy-based refinement of ODE samples: Energy-based refinement initializes chains with baseline ODE samples and runs 100 steps of either ULA or MALA using the terminal-time energy.The refinement uses Etot(x, 1 −σmin) and an SNR-adaptive step size.
  • Energy annealing MCMC: Energy annealing starts from x ∼N(0, I), traverses 120 noise levels from σ = 1 →10−3, takes K = 8 Langevin steps per level, and applies Tweedie denoising.The step size is δ = c σ2 with c = 0.5, while the physics weight is gated by noise level.
  • Poisson equation generation: All energy-based samplers reduce PDE residuals for unconditional Poisson equation generation compared to the baseline flow ODE sampler.This result is reported in Figure 5 for various samplers.
  • Burgers equation generation: For unconditional Burgers generation, energy-based samplers except the energy annealing sampler reduce PDE residual compared to the ODE baseline sampler.The exception is explicitly reported for the energy annealing sampler.
Loading 2608.18004v1…