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

We thank the reviewer for their thoughtful and constructive review. We appreciate your recognition of the method’s ability to enforce both linear and nonlinear physical constraints on flow-based generative models in a zero-shot manner and your general positive inclination towards our method. Below, we address the specific concerns raised.

R4-1: Novelty Relative to the ECI Framework

We appreciate the reviewer’s observation regarding the similarities to the ECI framework. While our method shares a conceptual goal: enforcing constraints at inference time, it differs significantly in both generality and mechanism, and we clarify these distinctions below:

• Generality of Constraints: ECI introduces a novel mixing-based correction strategy but is primarily demonstrated on simple, linear, and decoupled constraints. It does not provide a general framework for handling nonlinear, coupled, or global constraints, nor does it show robust performance on PDEs with shocks or discontinuities. In contrast, PCFM is built to support arbitrary nonlinear equality constraints, including nonlocal mass conservation, and shows strong empirical performance in such regimes.

• Theoretical Connection: As detailed in Appendix F, ECI is a special case of PCFM when constraints are linear, non-overlapping, and the penalty is omitted (λ = 0). PCFM generalizes this setting to support more complex and coupled constraints.

• Extensibility and Differentiability: Although ECI is described as “gradient-free,” it depends on analytical projections that do not extend naturally to nonlinear or global constraints. PCFM uses general-purpose automatic differentiation and Jacobian-vector products, enabling scalable enforcement in high-dimensional and nonlinear settings.

• Inference Strategy: ECI corrects the final output through constrained optimization, which can be expensive in large systems. PCFM instead uses efficient projection steps during generation and applies a single correction step at the end (Eq. 6) to ensure feasibility without high overhead.

We agree that Eq. 6 is a key contribution: it ensures feasibility while promoting alignment with the learned flow field, preserving generative consistency. We validate this in Appendix L.2, where PCFM outperforms ECI even under linear constraints, and more strongly under nonlinear and coupled constraints.

That said, Eq. 6 is part of a broader pipeline. PCFM integrates projection-based solvers, reversible integration, and constraint-guided updates, drawing from numerical PDE solvers to enable general, zero-shot constraint enforcement. Unlike ECI, which solves a constrained optimization problem at every sampling step, PCFM uses lightweight intermediate projections and a final correction. To ensure scalability and efficiency, we implement a batched, differentiable projection solver using reverse-mode Automatic Differentiation and Schur-complement updates, avoiding retraining while supporting complex constraints.

R4-2: Experimental Scope — Small-Scale Systems

We respectfully clarify that while the reviewer refers to our setting as having “only two degrees of freedom,” the Navier–Stokes system is inherently three-dimensional, spanning two spatial dimensions and one temporal dimension. Moreover, we enforce three types of constraints simultaneously—initial conditions, local consistency, and nonlinear global conservation laws—highlighting the method’s applicability to high-dimensional and tightly constrained PDE systems.

To ensure thorough evaluation, we included PDEs of varying difficulty: from relatively simple systems like the Heat equation, to intermediate ones such as Reaction–Diffusion, and finally to nonlinear, shock-prone systems like Inviscid Burgers and Navier–Stokes.

All datasets and spatiotemporal resolutions are sourced directly from prior work (Appendix B). Evaluating on 2D and 3D PDE systems is standard practice in the literature on neural operator learning [1, 2], scientific foundation models [3, 4], and constrained generative modeling [5, 6], making our experimental setup both representative and widely accepted.

[1] Li, Zongyi, et al. "Fourier Neural Operator for Parametric Partial Differential Equations." NeurIPS (2020).

[2] Li, Zongyi, et al. "Physics-Informed Neural Operator for Learning Partial Differential Equations." arXiv:2111.03794 (2021).

[3] Herde, Maximilian, et al. "Poseidon: Efficient Foundation Models for PDEs." arXiv:2405.19101 (2024).

[4] Sun, Jingmin, et al. "Towards a Foundation Model for Partial Differential Equation: Multi-Operator Learning and Extrapolation." arXiv:2404.12355 (2024).

[5] Huang, Jiahe, et al. "DiffusionPDE: Generative PDE-Solving under Partial Observation." arXiv:2406.17763 (2024).

[6] Négiar, Geoffrey, Michael W. Mahoney, and Aditi S. Krishnapriyan. "Learning Differentiable Solvers for Systems with Hard Constraints." arXiv:2207.08675 (2022).

Thank you for pointing out other minor presentation issues. In the revised manuscript, we will fix those typographical errors.

We thank the reviewer for the thoughtful assessment, noting the generality of our method and clarity of presentation. We appreciate recognition of PCFM’s handling of complex constraints at inference and strong empirical performance on challenging PDE benchmarks.

R3-1: Comparison to Prior Work on Intermediate Projection

While PCFM shares inference-time constraint satisfaction goals, it differs from ECI and Projected Diffusion Models (PDM) in generality, theory, and extensibility.

Comparison to ECI:

1. ECI is a special case of PCFM (Appendix F), limited to linear, uncoupled constraints without correction penalty (λ = 0). PCFM generalizes this to support nonlinear and coupled constraints, such as nonlocal mass conservation.

2. ECI is tested on simple linear cases and lacks scalability for nonlinear/global constraints. Its correction involves a full constrained optimization, which is expensive in large systems. PCFM uses lightweight projections during generation for guidance and final correction for feasibility.

3. Though described as "gradient-free", ECI relies on analytical projections, limiting extensibility as robust optimization methods are gradient-based methods, which makes the extension to nonlinear constraints nontrivial. PCFM uses our custom batched differentiable solver for scalable PDE enforcement.

Comparison with PDM

While PDM also uses projection during sampling, its assumptions and modeling context differ fundamentally from PCFM.

1. PDM enforces constraints at every timestep, which can over-constrain the trajectory, especially under nonlinear or global constraints like ICs or mass conservation that only apply at the final state. PCFM avoids this by guiding the trajectory through reversible integration and applying a single final correction, preserving consistency and flexibility.

2. PDM is built on stochastic diffusion models, while PCFM uses deterministic flow matching. Though related via the probability flow ODE [10], their sampling behaviors and practical implementation differ.

3. PDM has not been evaluated on PDEs with shocks or discontinuities, and its scalability to nonlinear global constraints is unclear. In contrast, PCFM is designed for such regimes and shows strong performance on complex PDEs like Burgers and Navier–Stokes.

To ensure fair comparison, we implemented a PDM-style ablation on our flow models. While it performs comparably on simpler cases like Heat, PCFM clearly outperforms on harder nonlinear dynamics, validating our design choices.

Table A: Comparison between PDM and PCFM

R3-2: Runtime Analysis and Overhead

We agree that runtime is important for assessing PCFM’s practicality. Accordingly, we include detailed runtime and memory comparisons using consistent batch sizes (8 for Navier–Stokes, 32 otherwise) and 100 flow steps on a single NVIDIA V100 GPU.

PCFM incurs a modest overhead relative to unconstrained baselines like vanilla FFM, due to exact constraint enforcement. However, it achieves 4–6× speedups over DFlow [4], which requires costly backpropagation through the full ODE unroll. DiffusionPDE is faster but struggles with hard constraints, especially for nonlinear problems like Burgers and Reaction-Diffusion. ECI’s runtime grows with mixing steps but still fails to satisfy nonlinear constraints.

For linear problems like Heat and Navier–Stokes, projections converge quickly, reflected in PDM’s runtime. Overall, PCFM adds a small, necessary overhead for robust constraint enforcement. We will summarize these findings in the main text.

Table B: Runtime and Memory Comparison

R3-3: Generality Beyond Synthetic PDE Data

We appreciate the reviewer’s suggestion and will clarify the intended scope of our generality claim. Specifically, by "general", we refer to PCFM’s ability to support a wide variety of constraint types—including nonlinear, coupled, and global equality constraints—on a general class of widely studied PDE systems. We chose flow matching models (FFMs) to reflect the functional nature of PDE data [1].

To ensure comparability with prior work, we focused on standard benchmark PDE datasets, including smooth equations such as Heat and Reaction–Diffusion, and nonlinear, shock-forming systems such as Burgers and Navier–Stokes. These settings are commonly used in the constrained generative modeling literature [2, 3, 4], neural operator learning [5, 6], and foundation PDE models [7, 8].

We agree with the reviewer and will soften claims about real-world generality in the paper. While our method is broadly applicable to other constrained generation tasks (e.g., control, imaging, or molecular design), we reserve such extensions for future work.

R3-4: Q1

As per the reviewer's request, we also include an out-of-distribution (OOD) constraint experiment to assess PCFM's generalization beyond constraints implicitly present in the data during training.

Specifically, we consider initial condition (IC) shifts in the Burgers equation, where the pretrained model has not seen such constraints during training. We generated a new test dataset where the location of the initial sigmoid profile $u(x,0) = \frac{1}{1 + \exp\left(\frac{x - p_{\text{loc}}}{\epsilon}\right)}$ is sampled from the intervals $[0.0, 0.1]$ and $[0.9, 1.0]$, while the training set used $p_{\text{loc}} \in [0.2, 0.8]$. The pretrained FFM model is reused without retraining, and all sampling methods are evaluated on the same OOD dataset to ensure fair comparison.

We adopt 200 flow matching steps for the Burgers problem and the table below summarizes the results across methods.

Table C: Performance on out-of-distribution dataset

PCFM achieves the lowest MMSE and SMSE while satisfying both ICs and nonlinear mass conservation even in this challenging OOD regime. These results demonstrate the robust generalization of our inference-time constraint enforcement, without requiring retraining or dataset-specific tuning. We will include this OOD evaluation in the final version of the paper.

R3-4: Q2

We thank the reviewer for the insightful question. As shown in Figure 9, increasing the number of flow matching steps improves MMSE by enabling the learned vector field to guide samples more accurately toward high-likelihood regions. However, we observe a mild increase in SMSE, particularly when stronger penalty weights are applied.

This can be attributed to reduced sample diversity: more frequent projection-based corrections during integration may cause samples to collapse toward a narrower region of the constraint manifold, lowering ensemble spread. This behavior reflects a trade-off between sample fidelity (MMSE) and diversity (SMSE).

Flow steps and step size are user-controlled hyperparameters in flow matching frameworks [9]. In our case, the flow integration length may slightly suppress SMSE due to stronger adherence to the constraint manifold. However, this is tunable and can be adapted based on application needs.

[1] Kerrigan et al., arXiv:2305.17209, 2023.

[2] Cheng et al., 2025.

[3] Huang et al., arXiv:2406.17763, 2024.

[4] Ben-Hamu et al., arXiv:2402.14017, 2024.

[5] Li et al., arXiv:2010.08895, 2020.

[6] Li et al., arXiv:2111.03794, 2021.

[7] Herde et al., arXiv:2405.19101, 2024.

[8] Sun et al., arXiv:2404.12355, 2024.

[9] Lipman et al., arXiv:2412.06264, 2024.

[10] Albergo et al., arXiv:2303.08797, 2023.

Thank you for your thoughtful engagement and for reconsidering your assessment in light of our clarifications and additional results.

We appreciate your acknowledgment of the empirical comparisons with ECI and PDM, and we agree that their inclusion in the main paper will help clarify the methodological distinctions and strengthen the positioning of PCFM. Nevertheless, PCFM employs a shoot–project–reverse framework with lightweight continuous guidance, which is inherently different from projecting intermediate states in the flow trajectory. This avoids PDM’s tendency to over-constrain noisy intermediates—especially under nonlinear constraints—and improves upon ECI, which is a special case of PCFM restricted to linear, decoupled constraints without penalty and lacks extensibility to nonlinear or global settings.

We will incorporate the runtime and memory benchmarks into the main text to improve transparency regarding computational overhead. We also thank you for recognizing the practical relevance of the OOD experiments and the clarification regarding SMSE trends. As committed, all these analyses, comparison tables, runtime metrics, and new baselines will appear in the final version.

We are grateful for your constructive review and support. We also truly appreciate your decision to increase the score; it means a lot to us. Thank you!

We thank the reviewer for their constructive feedback and for appreciating the clarity of the writing, the motivation behind enforcing physical constraints in generative models, and the strength of the shock front results in Figure 1. Below, we address each of the key concerns in detail.

R2-1: Use of the Term “Hard Constraint”

We appreciate the reviewer’s attention to terminological precision. In our work, we enforce equality constraints via Gauss–Newton projection, with residual norms consistently reduced to within machine epsilon. In Float32, this corresponds to violations on the order of $10^{-7}$, and with Float64, the residuals drop to $\sim 10^{-15}$. This level of constraint satisfaction aligns with the definition of hard constraints in the numerical optimization literature, where constraints are considered “hard” if they are satisfied up to solver and machine precision [1]. We will clarify this usage in the revised manuscript and explicitly state that our method achieves hard constraint satisfaction to the limits of numerical precision.

R2-2: Necessity and Role of Intermediate Projections

While constraints in physical systems (e.g., mass conservation or initial conditions) are often imposed only on the final state, we find that intermediate projections offer clear benefits for flow consistency and numerical stability. This aligns with practices in diffusion and flow matching literature, where intermediate states—often initialized as noise—are not expected to satisfy constraints or carry semantic meaning (e.g., intermediate samples in Stable Diffusion are noisy and not valid images).

To assess this, we conducted an ablation in which projections were applied only at the final step, omitting intermediate corrections. Using the same setup (100 flow steps for Heat; 200 for Burgers, RD, and NS), we observed that trajectories drift farther from the constraint manifold, leading to larger corrective steps, slower convergence, and reduced accuracy. This confirms that relying solely on final correction is suboptimal, especially for nonlinear or global constraints.

In contrast, intermediate projections act as lightweight continuous guidance, gradually steering samples toward the feasible set while preserving diversity. This avoids overly rigid enforcement on noisy intermediate states and improves final performance. Our results validate this approach, and we will include the ablation and clarify the motivation in the revision.

Table A: Comparisons between only projecting final step and PCFM

R2-3: Computational Cost of Projection

The reviewer correctly notes that constraint projections can introduce computational overhead. In PCFM, the projection step solves a Gauss–Newton system using iterative solvers (e.g., CG, GMRES) that require only Jacobian-vector products (Jv), computed efficiently via reverse-mode autodiff. We further accelerate this using batched vectorized maps (e.g., vmap in PyTorch) for GPU parallelism.

Our implementation also supports Jacobian-Free Newton–Krylov (JFNK) methods, which approximate directional derivatives via finite differences, eliminating the need for explicit Jacobian construction and scaling well to high-dimensional and nonlocal constraints.

For fair benchmarking, we use consistent batch sizes and 100 flow steps across all methods and datasets, evaluating average per-sample runtime and peak memory on a single NVIDIA V100 GPU. DFlow, which backpropagates through the entire ODE solve, is significantly slower and more memory-intensive; even with adjoint optimization, PCFM is 4–6× faster on average. In contrast, unconstrained baselines like vanilla FFM are faster but cannot guarantee constraint satisfaction. DiffusionPDE is faster than PCFM but sacrifices constraint enforcement and quality on nonlinear tasks. ECI’s runtime varies with mixing steps but fails to enforce nonlinear constraints robustly.

Profiling (Table C) shows the final projection step contributes only 1–3% of total time; most cost is due to sampling itself. Overall, PCFM introduces a modest and practical overhead that is necessary to enforce hard nonlinear constraints, and we will summarize this analysis in the main paper.

Table B: Runtime and Memory Comparison

Table C: PCFM Runtime Breakdown

R2-4: Generalizability and Experimental Scope

We agree that demonstrating applicability beyond PDE-based systems would broaden the paper’s scope. Our primary focus is on scientific generative modeling, where physical constraints play a central role and quantitative evaluations are well-defined. To ensure comparability with prior work, we focused on standard benchmark PDE datasets, including smooth equations such as Heat and Reaction–Diffusion, and nonlinear, shock-forming systems such as Burgers and Navier–Stokes. These settings are commonly used in the constrained generative modeling literature $2, 3, 4], neural operator learning \[5, 6], and foundation PDE models \[7, 8]. In addition, we also provide runtimes to all our datasets and add an additional baseline (PDM) \[9$ , including on the three-dimensional Navier-Stokes dataset as well. The underlying FNO architecture in the PCFM can also be naturally adopted zero-shot super-resolution $5$ . Nevertheless, we recognize the value of extending to more complex or real-world settings.

[1] Nocedal, Jorge, and Stephen J. Wright. Numerical optimization. New York, NY: Springer New York, 2006.

[2] Cheng, Chaoran, et al. "Gradient-free generation for hard-constrained systems." (2025).

[3] Huang, Jiahe, et al. DiffusionPDE: Generative PDE-solving under partial observation. arXiv:2406.17763, 2024.

[4] Ben-Hamu, Heli, et al. D-flow: Differentiating through flows for controlled generation. arXiv:2402.14017, 2024.

[5] Li, Zongyi, et al. Fourier neural operator for parametric partial differential equations. arXiv:2010.08895, 2020.

[6] Li, Zongyi, et al. Physics-informed neural operator for learning partial differential equations. arXiv:2111.03794, 2021.

[7] Herde, Maximilian, et al. Poseidon: Efficient Foundation Models for PDEs. arXiv:2405.19101, 2024.

[8] Sun, Jingmin, et al. Towards a Foundation Model for Partial Differential Equation: Multi-Operator Learning and Extrapolation. arXiv:2404.12355, 2024.

[9] Christopher, Jacob K., Stephen Baek, and Nando Fioretto. "Constrained synthesis with projected diffusion models." Advances in Neural Information Processing Systems 37 (2024): 89307-89333.

Dear Reviewer aA1E,

Thank you again for your thoughtful comments and questions, which have been very helpful in improving our paper. As the end of the discussion period approaches, we would like to gently remind you to consider our rebuttal and the updated version of the manuscript in your final evaluation.

We hope that our responses have adequately addressed your concerns. If that is the case, we kindly ask you to consider raising your score to reflect this. Of course, if you have any remaining questions or would like further clarification, we would be happy to respond promptly.

Thank you again for your time and engagement.

Best regards,

The authors

Dear Reviewer aA1E,

Thank you again for your follow-up. For elucidated clarity and in the spirit of aligning with your suggestion, we will revise the manuscript to avoid the term "hard constraint" and instead adopt your recommended phrasing of "approximate hard constraint" or "approximately enforcing constraints", i.e., referring to constraint satisfaction up to machine precision.

We appreciate this distinction and agree that it more accurately reflects the practical guarantees achieved by our method. While we are effectively in agreement, we also wish to clarify that our original terminology was chosen to remain consistent with recent literature in scientific machine learning and optimization, where the term "hard constraint" is commonly used for constraints enforced up to numerical precision. As you requested earlier, we have also explicitly cited representative papers that adopt this convention, including recent works in scientific machine learning [1, 2, 3, 6] and widely used numerical optimization toolkits [4, 5].

We hope this resolves the remaining concern. We would greatly appreciate your consideration in increasing the paper’s overall score in light of its contributions.

Warm regards,

The Authors

[1] Donti, P. L., Rolnick, D., & Kolter, J. Z. (2021). DC3: A Learning Method for Optimization with Hard Constraints. In International Conference on Learning Representations (ICLR).

[2] Négiar, G., Mahoney, M. W., & Krishnapriyan, A. (2023). Learning Differentiable Solvers for Systems with Hard Constraints. In The Eleventh International Conference on Learning Representations (ICLR).

[3] Lu, Lu, et al. "Physics-informed neural networks with hard constraints for inverse design." SIAM Journal on Scientific Computing 43.6 (2021): B1105-B1132.

[4] Wächter, A., & Biegler, L. T. (2006). On the Implementation of an Interior-Point Filter Line-Search Algorithm for Large-Scale Nonlinear Programming. Mathematical Programming, 106(1), 25–57.

[5] MATLAB Optimization Toolbox, fmincon Documentation. The MathWorks, Inc.

[6] Hansen, Derek, et al. "Learning physical models that can respect conservation laws." International Conference on Machine Learning. PMLR, 2023.

We thank the reviewer for their thoughtful assessment and for highlighting both the novelty of our shoot–project–reverse cycle and the empirical strength of our method across challenging physical domains. We are encouraged that the proposed framework was found effective and promising for enforcing strict physical constraints at inference time.

R1.1: Computational Cost of Jacobian Construction

We appreciate the reviewer raising this important question regarding scalability. In practice, the constraint Jacobian Jh(x) is not constructed explicitly. Instead, we compute it using batched reverse-mode automatic differentiation, enabled via vectorized mapping (e.g., vmap in PyTorch). This allows us to efficiently compute the full Jacobian across multiple constraint functions in parallel. Since each row of the Jacobian corresponds to a reverse-mode gradient computation, this formulation is highly parallelizable on GPUs and scales well even when the state dimension is large.

For solving the projection step, employing iterative solvers (e.g., conjugate gradient, GMRES) on the Gauss–Newton system will require only Jacobian-vector products (Jv) rather than the full matrix. These products are efficiently handled by automatic differentiation and avoid the memory costs associated with storing large dense Jacobians.

Additionally, our approach is compatible with Jacobian-Free Newton–Krylov (JFNK) methods, where even the Jv products are approximated via finite-difference evaluations. This further reduces memory footprint and makes the solver suitable for extremely high-dimensional or non-local constraint settings.

In summary, the projection step does not pose a computational bottleneck, and we find it to be practical and scalable in our experiments.

R1.2: Wall-clock inference times

We provide additional results, measuring our inference times compared with the existing baselines. To compare fairly on time and peak memory, we use the same batch size for all methods and loop through the test set to record the per sample time and peak memory. We use 100 flow matching steps across all methods and datasets for benchmarking here. We use a single Nvidia V100 GPU with 32GB ram for all experiments. Batch sizes are 32 for heat, reaction diffusion, and Burgers problems and 8 for Navier-Stokes. We note that DFlow runs out of memory even for a batch size of 1 on the V100 chip, thus we improved the algorithm by including an adjoint ODE solve. For PDM, we adopt our constraint projection to every intermediate state along the flow matching trajectory. We report the average per-sample inference time and peak memory usage.

PCFM incurs a modest runtime overhead compared to unconstrained baselines like vanilla FFM, due to projections to satisfy nonlinear hard constraints exactly. However, it achieves 4-6x speedups over D-Flow, which optimizes noise by backpropagating through the full ODE unroll - which is highly memory-intensive and slow. For Heat and Navier-Stokes datasets, the constraints are linear with well-defined structure, the projections converge in a single step, evident also in PDM's runtime. DiffusionPDE is faster than PCFM but it compromises on generative quality and struggles to enforce hard constraint, particularly for nonlinear problems like Burgers and Reaction-Diffusion. Similarly, ECI's mixing iterations can increase runtime overhead when tuned for more complex problems but it fails to achieve hard constraint satisfation for nonlinear problems.

Profiling in Table B reveals 96-99% of PCFM's time is in intermediate projection steps, with final projection at 1-3% of the total runtime.

Table A: Runtime and Memory Comparison

---

Table B: PCFM Runtime Breakdown

Dear Reviewer 8V7Y,

Thank you again for your helpful comments and questions, which have contributed to improving our paper. As the end of the discussion period approaches, we would like to gently remind you to consider our rebuttal and the updated version of the manuscript in your final evaluation.

We hope that we have satisfactorily addressed your concerns. If you have any further questions or would like additional clarification, we’d be happy to respond.

Thank you once again for your time and consideration.

Best regards,

The Authors