PaperlayerView exact revision on arXiv

Source-linked AI summary

Physics-Constrained Flow Matching: Sampling Generative Models with Hard Constraints

Utkarsh Utkarsh, Pengfei Cai, Alan Edelman, Rafael Gomez-Bombarelli, Christopher Vincent Rackauckas

arXiv:2506.04171v2cs.LGcs.AIcs.CEcs.NAmath.NA

TL;DR

Generative models can sample PDE solutions efficiently, but ensuring that generated samples satisfy physical constraints remains challenging. PCFM applies hard constraint enforcement to pretrained flow models at inference time, and it outperforms state-of-the-art baselines across diverse PDEs while achieving exact constraint satisfaction.

  • Problem

    Ensuring that generated PDE samples satisfy governing physical constraints remains challenging, and existing sampling methods may be costly, fail to satisfy constraints exactly, or have limited robustness on sharp PDEs.

    Source: §1 · pp. 1–2

  • Method

    PCFM projects intermediate states of pretrained flow-matching models onto constraint manifolds at inference time to enforce arbitrary nonlinear equality constraints without retraining or architectural changes.

    Source: §1 · p. 2

  • Results

    Across diverse PDEs, PCFM outperforms state-of-the-art baselines with lower error and exact constraint satisfaction; reported improvements reach up to 99.5% in standard metrics.

    Source: §1 · p. 3, §5 · pp. 9–10

  • Takeaways & Limitations

    PCFM offers a principled approach to enforcing physical feasibility in generative models for scientific simulation and design.

    Source: §5 · p. 10

    The reverse-flow step is approximate: backward ODE integration can incur O(1) error and become unstable for high-Lipschitz vector fields, motivating an OT-interpolant approximation.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

Deep generative models have recently been applied to physical systems governed by partial differential equations (PDEs), offering scalable simulation and uncertainty-aware inference. However, enforcing physical constraints, such as conservation laws (linear and nonlinear) and physical consistencies, remains challenging. Existing methods often rely on soft penalties or architectural biases that fail to guarantee hard constraints. In this work, we propose Physics-Constrained Flow Matching (PCFM), a zero-shot inference framework that enforces arbitrary nonlinear constraints in pretrained flow-based generative models. PCFM continuously guides the sampling process through physics-based corrections applied to intermediate solution states, while remaining aligned with the learned flow and satisfying physical constraints. Empirically, PCFM outperforms both unconstrained and constrained baselines on a range of PDEs, including those with shocks, discontinuities, and sharp features, while ensuring exact constraint satisfaction at the final solution. Our method provides a flexible framework for enforcing hard constraints in both scientific and general-purpose generative models, especially in applications where constraint satisfaction is essential.

1 Introduction

Generative models can support efficient sampling and uncertainty quantification for PDE solutions, but ensuring physical constraints remains difficult. PCFM enforces hard constraints at inference time and improves generation on challenging PDEs while preserving exact constraint satisfaction.

1 Introduction

  • Generative models offer efficient sampling and uncertainty quantification for PDE solutions, but their samples must still respect governing physical constraints.Soft constraints and architectural priors do not readily adapt to enforcing physical laws in dynamical systems, and soft constraints may fail when exact satisfaction is essential.Source: §1 · pp. 1–2
  • Existing sampling methods face challenges with exact constraints: gradient-based approaches can be expensive or fail, while ECI is limited on linear constraints and sharp, high-gradient PDEs.Because constraints need only hold exactly in the final solution, stochastic sampling complicates hard enforcement throughout the process.Source: §1 · p. 2
  • PCFM enforces arbitrary and multiple physical constraints as hard requirements at inference time without modifying the training objective.Supported constraints include conservation laws, boundary conditions, and nonlinear constraints; the framework operates on pretrained flow-matching models without constraint gradients during training.Source: §1 · p. 2
  • PCFM improves PDE generation by up to 99.5% in standard metrics such as mean squared error while ensuring zero constraint residual.The method is evaluated on PDEs with shocks, discontinuities, and sharp spikes, improving accuracy at inference time without retraining or fine-tuning.Source: §1 · p. 3

2 Related Work

Related work spans flow-based generation, constraint-guided sampling, and physics-informed learning. These approaches provide useful tools, but limitations in computational cost, exactness, or constraint generality motivate combining constrained numerical integration with flow-based generation.

2 Related Work

  • Flow matching transports samples from a tractable prior to a target distribution through a learned time-dependent vector field, and functional variants model PDE solutions as functions.Functional flow matching and Denoising Diffusion Operator extend the framework to spatiotemporal data and function spaces.Source: §2 · p. 3
  • Constraint-guided PDE generation includes gradient-based methods with high computational cost and approximate satisfaction, while ECI has limited evaluation on nonlinear, coupled, and shock-related constraints.Soft-penalty physics-informed methods also lack guarantees of hard constraint satisfaction; constrained inference and structured architectures offer other approaches.Source: §2 · p. 3
  • Constraint-aware integration methods use projection-based updates to enforce feasibility and motivate principled constrained sampling without retraining.Examples include constrained neural ODEs, differential algebraic equations, and geometric integrators.Source: §2 · p. 4

3 Methodology and Setup

PCFM narrows a pretrained flow model’s solution support to samples satisfying additional constraints at inference time, without retraining. It interleaves flow updates and constraint corrections, using an OT displacement approximation for reverse updates and a relaxed correction to balance solution fidelity and constraint satisfaction.

3.1 Problem Setup

  • PCFM targets solutions satisfying both the governing PDE and an additional constraint, imposing that constraint at inference time without retraining the pretrained model.The assumed model approximates the solution set for a fixed PDE family and parameter set.Source: §3.1 · p. 4

3.2 Generative Models via Flow-Based Dynamics

  • Functional Flow Matching learns flows between PDE solution functions using a neural operator and can be augmented at inference time to enforce additional physical constraints.It approximates the target measure over the PDE solution support using a pretrained model.Source: §3.2 · p. 4

3.2.1 Constraint Types in PDE Systems

  • Table 2 organizes common PDE constraint types by their mathematical form and scope.The accompanying text notes that constraints may not be implicitly satisfied by generated solutions.Source: Table 2 · p. 5, §3.2.1 · p. 5

3.3 PCFM: Physics-Constrained Flow Matching

  • PCFM alternates a forward flow solve with a Gauss–Newton projection, then uses an OT displacement interpolant to propagate the correction back to the intermediate state.The projection shifts the endpoint toward a linearized feasible point and recovers the exact solution when the constraint function is affine.Source: §3.3 · pp. 5–6

3.4 Reversibility and Interpolant Compatibility

  • The OT displacement approximation addresses potential instability in reverse ODE integration, and PCFM can also extend to variance-preserving score-based diffusion models.The paper states that straight-line displacement approximates the true marginal flow in the small-step limit and that VP flows benefit from bounded Lipschitz constants due to Gaussian smoothing.Source: §3.3 · p. 6, §3.4 · p. 6

3.3 PCFM: Physics-Constrained Flow Matching

  • A relaxed penalized correction encourages constraint satisfaction at an extrapolated point while preserving local flow alignment, with a final full projection used if residuals remain high at τ = 1.As the step size approaches zero, the relaxed correction aligns with the OT-based marginal objective and can provide robustness at coarse discretization.Source: §3.3 · p. 6

3.4 Reversibility and Interpolant Compatibility

  • For Poseidon evaluation, square-input requirements led to bilinear interpolation to 128×128, which may introduce artifacts in FPD; D-Flow results were omitted for numerical instability.Source: §3.4 · p. 7

3.4.1 Tradeoff Between Flow Steps and Relaxed Constraint Correction

  • Fewer flow steps speed inference but increase numerical error that can violate intermediate constraints; the relaxed correction compensates, while more steps reduce its importance.The paper describes the tradeoff between step count and relaxed-correction strength.Source: §3.4.1 · p. 7

4 Experiments and Results

Across linear and nonlinear PDE constraints, PCFM improves generated-solution fidelity while satisfying constraints, including nonlinear mass conservation and local shock-related constraints. The experiments also examine how adding complementary constraints affects fidelity and existing constraint satisfaction.

4 Experiments and Results

  • The experiments use the same pretrained FFM model and evaluate generated subsets conditioned on held-out initial or boundary conditions, including constraint satisfaction and distributional similarity.The reported distributional metrics are MMSE, SMSE, and FPD; constraint errors cover IC, BC, and global mass conservation.Source: §4 · pp. 7–8

4.1 Linear Hard Constraints

  • On the 2-D Heat and 3-D Navier-Stokes equations, PCFM outperforms compared methods in MMSE, SMSE, and FPD, with machine-level precision for IC and mass constraints.Relaxed constraint corrections can improve fidelity with fewer flow-matching steps; randomized u0 across batches is used for high-dimensional Navier–Stokes.Source: §4.1 · p. 8

4.2.1 Nonlinear Global and Boundary Constraints

  • For Reaction-Diffusion, PCFM achieves the lowest constraint errors and best generation fidelity, satisfying IC and nonlinear mass conservation to machine precision for all t.The passage reports that this also enforces Neumann boundary fluxes and improves solution quality.Source: §4.2.1 · p. 8, Figure 2 · p. 9
  • For Burgers’ equation with Dirichlet and zero-flux Neumann boundary conditions, PCFM satisfies the constraints while matching the fidelity of other methods.Source: §4.2.1 · p. 8

4.2.2 Nonlinear Local Dynamics and Shock Constraints

  • For Burgers’ equation, PCFM outperforms all baselines across metrics while satisfying IC and mass constraints and capturing the localized shock structure.The method uses five unrolled local flux collocation points in the residual.Source: §4.2.2 · p. 8, Figure 3 · p. 9

4.3 Enhancing Fidelity through Additional Physical Constraints

  • Adding constrained collocation points improves Burgers’ solution fidelity without worsening IC or mass-constraint satisfaction; complementary constraints can help, though conflicting constraints require balancing.The study also reports that TVD constraints improve heat-equation smoothness and fidelity while preserving IC and global mass constraints.Source: §4.3 · p. 9

5 Conclusion

PCFM enforces arbitrary nonlinear equality constraints in pretrained flow-based generative models through zero-shot inference. Across diverse PDEs, including systems with shocks and nonlinear dynamics, it achieves lower error and exact constraint satisfaction; adding complementary constraints can improve generation quality.

5 Conclusion

  • PCFM combines flow integration, tangent-space projection, and relaxed penalty corrections to enforce constraints while staying consistent with the learned generative trajectory, without retraining or architectural changes.Source: §5 · p. 9
  • Across diverse PDEs, including systems with shocks and nonlinear dynamics, PCFM achieves lower error and exact constraint satisfaction.Source: §5 · pp. 9–10
  • Adding complementary constraints can improve generation quality, while extending PCFM beyond equality constraints and exploiting constraint-Jacobian structure remain directions for future work.The conclusion identifies inequality constraints and Jacobian structure as promising directions for improving scalability.Source: §5 · p. 10, Figure 4 · p. 10

A Proof of Constraint Satisfaction

The proof establishes final constraint satisfaction under regularity conditions and a sufficiently converged final projection. It relies on a full-row-rank constraint Jacobian near the constraint manifold and an initial guess within the Newton solver’s convergence basin.

A Proof of Constraint Satisfaction

  • If h is twice continuously differentiable and its Jacobian has full row rank near the constraint manifold, the final sample satisfies h(u1) = 0 to machine precision when the final projection uses sufficiently many Newton–Schur iterations.Source: A Proof of Constraint Satisfaction · p. 15
  • Under the stated full-rank assumption, the implicit function theorem gives a locally unique nearby root, and sufficiently small initial residual supports quadratic convergence of the Newton–Schur iterates.Source: A Proof of Constraint Satisfaction · p. 15
  • PCFM’s preceding Gauss–Newton projection, OT-based backward solve, and relaxed penalty correction make the incoming residual typically less than 10^-3, supporting the final projection’s convergence-basin condition.The final iteration is continued until the residual falls below a tolerance typically set to machine precision.Source: A Proof of Constraint Satisfaction · p. 15

B PDE Dataset Details

The benchmarks span four PDE systems with varied dynamics and boundary conditions, while PCFM enforces initial, boundary, and conservation constraints tailored to each problem. The pretrained flow models use functional flow matching with a Fourier neural operator backbone.

B PDE Dataset Details

  • The four benchmarks cover linear and nonlinear dynamics, smooth and discontinuous solutions, and varied initial and boundary conditions.They comprise heat, Navier–Stokes, reaction-diffusion, and inviscid Burgers equations.Source: B PDE Dataset Details · p. 15

C FFM Model Pretraining

  • The pretrained FFM models use an FNO backbone, with training data and model configurations tailored to each PDE benchmark.Heat uses analytical solutions; Navier–Stokes, reaction-diffusion, and Burgers use numerical solutions.Source: C FFM Model Pretraining · p. 17

D.2 Nonlinear Equality Constraints

  • PCFM represents constraints with residual operators and enforces nonlinear conservation at inference by iterative projection onto the constraint manifold.The nonlinear density may represent quantities such as energy or entropy; equality constraints can also encode prescribed values or fluxes.Source: §D.2 · p. 18, §D.1 · p. 18

D.3 Heat Equation (IC and Mass Conservation)

  • Heat and Navier–Stokes use initial-condition and global-mass constraints, with periodic boundaries supporting linear conservation constraints.For Navier–Stokes, periodicity and incompressibility make total vorticity constant over time.Source: §B.1 · p. 16, §B.2 · p. 16, §D.1 · p. 18

D.5 Reaction-Diffusion (IC and Nonlinear Mass Conservation)

  • Reaction-diffusion sampling enforces the initial condition and nonlinear mass conservation accounting for reaction and boundary fluxes.Its conservation depends on the nonlinear state-dependent reaction term and boundary-driven transport.Source: §B.3 · p. 16, §D.5 · p. 19

D.7 Burgers (IC, Mass Conservation, and Local Flux)

  • Burgers sampling enforces boundary conditions and nonlinear mass conservation, with an additional configuration enforcing initial conditions and local flux updates.The conservation constraint is sensitive to boundary-induced discontinuities associated with shock formation.Source: §B.4 · p. 17, §D.7 · p. 19, §D.6 · p. 19

E.1 Runtime Analysis and Overhead

PCFM adds projection overhead but is substantially faster than D-Flow and retains competitive runtime across the reported comparisons. Its projection has a tangent-space interpretation, and the framework encompasses ECI as a special case.

E.1 Runtime Analysis and Overhead

  • PCFM achieves a 4–6× speedup over D-Flow, with modest overhead versus vanilla FFM and comparable runtimes to PDM on linear systems.DiffusionPDE runs faster but struggles to enforce nonlinear constraints robustly; runtime comparisons use 100 flow-matching steps on one V100 GPU.Source: §E.1 · pp. 19–20, Table 4 · p. 20

E.2 Tangent-Space Interpretation of the Projection Step

  • The projection is an orthogonal projection onto the tangent space of the constraint manifold’s linearization, and is exact for affine constraints.For h(u) = Au − b, the update projects onto the feasible set Au = b.Source: §E.2 · pp. 20–21

F Connection to ECI as a Special Case

  • With λ = 0 and intermediate corrections skipped, PCFM reduces exactly to ECI, which is tailored to simple linear constraints.The frameworks share the ECI correction step in this linear-constraint regime.Source: F Connection to ECI as a Special Case · p. 21

G.1 Flow Matching

  • Flow matching learns a time-dependent velocity field transporting prior samples to target samples; functional flow matching extends this objective to function-valued solutions.The functional objective serves as PCFM’s pretraining objective.Source: §G.1 · p. 21, §G.2 · p. 22

H Sampling Setups for PCFM and Other Baselines

PCFM corrects intermediate flow states and applies a final projection, while comparator methods use distinct constraint-guidance strategies. Its nonlinear projection solver uses batched Schur-complement updates designed for efficient inference.

H Sampling Setups for PCFM and Other Baselines

  • PCFM combines explicit Euler flow integration with one Newton projection update per step and a final projection to enforce constraints.Interpolation guidance is optional and used for heat results; λ = 0 is the default for computational efficiency.Source: H Sampling Setups for PCFM and Other… · p. 22
  • The baselines differ in their guidance: ECI mixes and resamples noise, D-Flow optimizes initial noise through the unrolled trajectory, and DiffusionPDE guides each step with a composite loss.PDM also projects during sampling, but its assumptions and modeling framework differ from PCFM’s.Source: H Sampling Setups for PCFM and Other… · pp. 22–23

I Batched Differentiable Solver for Nonlinear Constraints

  • The nonlinear solver iterates Gauss–Newton-style corrections using an m × m Schur system, with cost scaling as O(n) per sample when m ≪ n.The batched implementation parallelizes Jacobian, residual, and Schur-system calculations across samples.Source: I Batched Differentiable Solver for… · pp. 23–24
  • The solver restarts the dual variable at zero each iteration to avoid stale gradient accumulation and improve numerical stability.It supports reverse-mode differentiation and batched execution.Source: I Batched Differentiable Solver for… · p. 23

J Evaluation Metrics

The evaluation measures generation fidelity, physical consistency, and distributional similarity using complementary metrics. These metrics compare generated PDE solutions with ground truth and quantify constraint residuals.

J Evaluation Metrics

  • MMSE and SMSE summarize generation fidelity using the mean and standard deviation of MSE across generated and true PDE solutions.The reported quantities compare the means and standard deviations of the two solution sets.Source: J Evaluation Metrics · p. 24
  • Constraint error is the average ℓ2 norm of constraint-function residuals across generated samples.Residuals measure constraint or conservation violation.Source: J Evaluation Metrics · p. 24
  • FPD compares generated and true solution distributions using Fréchet distance between embeddings from the pretrained Poseidon model.It is computed per frame and averaged for 2D PDEs, or over the full spatiotemporal solution for 1D PDEs.Source: J Evaluation Metrics · p. 24

K Further Results

PCFM’s results vary across PDEs and boundary conditions: some profiles are visually closest to ground truth, while fixed-boundary Burgers results trade slightly worse pointwise errors for constraint satisfaction. Intermediate projections progressively guide samples toward feasible solutions and improve final accuracy and stability.

K Further Results

  • For fixed-boundary Inviscid Burgers, PCFM ensures global mass conservation and low constraint errors over time, despite slightly worse MMSE and SMSE and better FPD.Generated samples satisfy both Dirichlet and zero-flux Neumann boundary conditions while retaining fidelity similar to other methods.Source: Figure 5 · p. 25
  • PCFM’s mean profiles are visually most similar to ground truth for the Heat equation with fixed IC and Reaction-Diffusion and Inviscid Burgers with fixed IC.The figures compare pointwise means and standard deviations across generated samples for the latter two equations.Source: Figure 6 · p. 25, Figure 7 · p. 26, Figure 8 · p. 26

L.1 Effect of Intermediate Projections

  • Final-step-only projections let trajectories drift farther from the constraint manifold, requiring larger corrections and producing slower convergence and lower accuracy.Intermediate projections instead continuously steer samples toward the feasible manifold while preserving diversity, improving final-solution accuracy and stability.Source: §L.1 · p. 26

L.2 Total Variation (TV) Constraints on the Heat Equation

The heat-equation TVD constraint limits growth in total variation between initial and final times. In the reported test, adding it improves structural fidelity, including Power MSE, without constraint violations.

L.2 Total Variation (TV) Constraints on the Heat Equation

  • The hard constraint relates final-time total variation to initial-time total variation, with a decay factor γ between 0 and 1, and is applied at selected time slices.The paper describes TVD as the principle that diffusive systems should not increase total variation over time.Source: §L.2 · pp. 26–27
  • The TVD constraint is compatible with mass conservation and initial-condition enforcement and can be combined with energy-decay or curvature-regularity conditions.The paper reports reduced spurious oscillations and improved structural fidelity without significantly altering statistical accuracy.Source: §L.2 · p. 27
  • Power MSE measures differences in spatial frequency content and is sensitive to unphysical sharpness or over-smoothing, complementing pointwise metrics.It evaluates structural fidelity in the spectral domain.Source: §L.2 · p. 27
  • 488.42 vs. 500.12 for Power MSE (PCFM vs. ECI); adding the TVD constraint also improves MMSE and SMSE while CE remains zero.The comparison is on the Heat Equation dataset with TVD constraints; lower is better for all metrics.Source: §L.2 · p. 27, Table 7 · p. 28

L.3 Effect of Relaxed Constraint Correction with Flow Matching Steps

The penalty weight’s effect depends on the number of flow steps: it offers greater benefit with coarse integration but only marginal benefit with more steps. Excessively large weights can distort the reverse update and break alignment with the OT interpolant.

L.3 Effect of Relaxed Constraint Correction with Flow Matching Steps

  • With 10 flow steps, increasing λ significantly improves MMSE and SMSE; with 100 steps, the additional penalty yields only marginal benefit.The comparison is for the Reaction-Diffusion dataset.Source: §L.3 · p. 28, Figure 9 · p. 28
  • Overly large λ can distort the reverse update and break alignment with the OT interpolant, so its choice balances constraint enforcement against generative-path consistency.Source: §L.3 · p. 28

L.4 Evaluation on Generating Out-of-Distribution Samples

This section covers a zero-shot Burgers evaluation with out-of-range initial conditions and checklist declarations about human-subjects involvement and LLM use.

L.4 Evaluation on Generating Out-of-Distribution Samples

  • The OOD test shifts sigmoid-profile centers from the training range [0.4, 0.6] to non-overlapping ranges [0.0, 0.2] and [0.8, 1.0].All methods reuse the pretrained FFM backbone without retraining, use 200 flow-matching steps, and enforce the same initial-condition and conservation constraints.Source: §L.4 · pp. 28–29
  • PCFM achieves the lowest MMSE and SMSE while exactly satisfying initial-condition and conservation-law constraints on the OOD test set.The passage reports that unconstrained or partially constrained baselines exhibit significant constraint violations; Table 8 summarizes the comparison.Source: Table 8 · p. 29, §L.4 · p. 29

NeurIPS Paper Checklist

  • The checklist states that the abstract and introduction reflect the paper’s contributions, and that limitations are discussed in the conclusion.Source: NeurIPS Paper Checklist · p. 30
  • The checklist reports access to data and code and points to sections covering baselines, experimental setup, metrics, and compute resources.It also states that SMSE and variance in mass-residual mean plots are reported for statistical significance assessment.Source: NeurIPS Paper Checklist · pp. 31–33
  • The paper marks both participant-related disclosures and IRB review as not applicable because the research involved no crowdsourcing or human subjects.The checklist also states that no participant risks or human-subject review apply.Source: NeurIPS Paper Checklist · pp. 35–36
  • LLM usage is marked not applicable to the core method because LLMs were used only for writing, editing, or formatting.The checklist says LLMs were not an important, original, or non-standard component of the method.Source: NeurIPS Paper Checklist · p. 36

Related papers

Loading 2506.04171v2…