Gradient-Free Generation for Hard-Constrained Systems
Authors: Chaoran Cheng
Boran Han
Danielle C. Maddix
Abdul Fatir Ansari
Andrew Stuart
Michael W. Mahoney
Bernie Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Flow Matching
Generative Model
Constrained Generation
Partial Differential Equations
Conservation Laws
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose ECI sampling, a gradient-free approach for guiding pre-trained generative models for hard-constrained generation.

Abstract: Generative models that satisfy hard constraints are critical in many scientific and engineering applications, where physical laws or system requirements must be strictly respected. Many existing constrained generative models, especially those developed for computer vision, rely heavily on gradient information, which is often sparse or computationally expensive in some fields, e.g., partial differential equations (PDEs). 
In this work, we introduce a novel framework for adapting pre-trained, unconstrained flow-matching models to satisfy constraints exactly in a zero-shot manner without requiring expensive gradient computations or fine-tuning. 
Our framework, *ECI sampling*, alternates between extrapolation (E), correction (C), and interpolation (I) stages during each iterative sampling step of flow matching sampling to ensure accurate integration of constraint information while preserving the validity of the generation. 
We demonstrate the effectiveness of our approach across various PDE systems, showing that ECI-guided generation strictly adheres to physical constraints and accurately captures complex distribution shifts induced by these constraints. 
Empirical results demonstrate that our framework consistently outperforms baseline approaches in various zero-shot constrained generation tasks and also achieves competitive results in the regression tasks without additional fine-tuning.

Published as a conference paper at ICLR 2025
GRADIENT-FREE GENERATION FOR HARD-CONSTRAINED SYSTEMS
Chaoran Cheng'; Boran Han?, Danielle C. Maddix?, Abdul Fatir Ansari?,
Andrew Stuart Michael W. Mahoney", Yuyang Wang?
l University of Illinois Urbana-Champaign
2 Amazon Web Services
3 Stores Foundational AI, Amazon
4 "Amazon SCOT
chaoran7@illinois.o
(boranhan, dnmaddix, ansarnd, andrxstu, zmahmich, yuyawang) amazonn,
ABSTRACT
Generative models that satisfy hard constraints are critical in many scientific and engineering applications, where physical laws or system requirements must be strictly respected. Many existing constrained generative models, especially those developed for computer vision, rely heavily on gradient information, which is often sparse or computationally expensive in some fields, e.g., partial differen- tial equations (PDEs). In this work, we introduce a novel framework for adapt- ing pre-trained, unconstrained flw-matchinn models to satisfy constraints ex- actly in zero-shot manner without requiring expensive gradient computations or fine-tuning. Our framework, ECI sampling, alternates between extrapolation (E), correction (C), and interpolation (I) stages during each iterative sampling step flow matching sampling to ensure accurate integration constraint informa- tion while preserving the validity of the generation. We demonstrate the effec- tiveness of our approach across various PDE systems, showing that ECI-guided generation strictly adheres to physical constraints and accurately captures com- plex distribution shifts induced by these constraints. Empirical results demon- strate that our framework consistently outperforms baseline approaches in vari- ous zero-shot constrained generation tasks and also achieves competitive results in the regression tasks without additional fine-tuning. Our code available https: / github......... com/ amazon weeen-sciieeeeeeeeeeeeeeeeeeeeee Caeance/EEII-I EEII-ampling.
INTRODUCTION
Diffusion and flow matching models have achieved remarkable success in generative tasks of image generation (Ho et al., 2020; Esser et al., 2024a), language modeling (Lou et al., 2023; Gat et al., 2024), time series prediction (Lin al., 2024; Kollovich et al., 2024), and functional data modeling (Lim et al., 2023; Kerrigan et al., 2023). Constrained generation built upon these generative mod- els for solving various inverse problems has also been explored the image domain (Kawar al., 2022: Ben-Hamp et al., 2024). These approaches predominantly rely on gradient information with respect to some cost function as a soft constraint. While such soft-constrained methods have been successful in the image domain, applications in other domains often require the generation to adhere to certain constraints strictly. Such hard-constrained generation, i.e., generative modeling where the natural hard constraints are requirements and not just suggestions, is crucial for tasks in many scientific domains. For example, in scientific computing, numerical simulations often require gener- ated solutions to adhere to specific physical constraints (energy mass conservation (Hansen al., 2023; Mouli et al., 2024)) and satisfy boundary condition constraints on the values or derivatives of the solutions (Saad al., 2023). An intuitive example of the generation of PDE solutions is demon- strated in Figure 1, where the solution set narrows and shifts when the constraint is imposed. While existing approaches for constrained generation in image inverse problems can always be adapted
Work done during internship AWS.

---

Published as a conference paper at ICLR 2025
(Esser et al., 2024b; Ansari et al., 2024), hard-constrained generation for complex systems like PDE solutions presents the following challenges:
IC-Guided Solution
IC
Unconstrained Generative Prior
[/2.8)
k 5
w- ~~U,8 k U-2,,,20
Hard Constraint
BC
Stokes Problem
BC-Guided Solution
U3 = VUrr
PDE System
u(z, 0) Ae ke (kxz)
Initial Condition (IC)
A~ 072,
u(O,t) A cos(wit)
Boundary Condition (BC)
Figure 1: Hard-constrained generation of PDE solutions with the boundary condition (BC) or initial condition (IC) prescribed a posteriori. w, k are PDE parameters that determine the IC and BC.
Scarcity of constraint information. In contrast computer vision (CV) tasks, e.g., image inpaint- ing, which typically assumes considerable amount of context information, common constraints like BC and IC in Figure have measure with respect to the spatiotemporal domain. For exam- ple, standard setting with a spatiotemporal resolution of 100 100 leads to only 1% of known pixel values context. This is significantly less than the context provided in typical CV application.
Exact constraint satisfaction. The exact conservation of mass, energy, or momentum is often es- sential in the simulation of PDEs (Le Veque, 1990) for ensuring physically feasible and consistent solutions. This in contrast with inverse problems CV, e.g., super-resolution and de-blurring, where constraints are often implicitly defined modeling assumptions and evaluation metrics, e.g.. peak signal-to-noise ratio (PSNR) do not directly depend on the exact satisfaction of these con- straints. Existing CV approaches fail to guarantee exact satisfaction of constraints.
Issues with gradient-basee based methods. Existing zero-shot frameworks for constrained generation predominantly rely on gradient guidance from differentiable cost function on the final step of gen- eration. This information can be prohibitively expensive, especially in large 3D PDE spatiotemporal systems. Previous work has also indicated the drawbacks of these soft-constrained gradient- approaches. Well-known limitations include gradient imbalances in the loss terms (Wang et al., 2020; 2021) and ill-conditioning (Krishnapriyan et al., 2021), which can lead to failure modes in scientific machine learning (SciML) tasks.
To address these challenges for hard-constrained generation, we propose general framework that adopts gradient-fres and zero-shot approach for guiding unconstrained pre-trained flow matching models. For the diversity of practical hard constraints for PDE systems, our proposed framework provides unified and efficient approach without the need for expensive fine-tuning gradient backpropagation. We term our framework ECI sampling since it interleaves extrapolation (E), cor- rection (C), and interpolation (I) stages each iterative sampling step. ECI sampling effectively and accurately captures the distribution shift imposed by the constraints and maintains the consistency of the generative prior. We summarize our main contributions as follows:
Unified gradient-free generation framework. We introduce ECI sampling, unified gradient- free sampling framework for guiding an unconstrained pre-trained flow matching model. By in- terleaving extrapolation, correction, and interpolation stages each iterative sampling step. ECI offers fine-grad grained iterative control over the flow sampling that can accurately capture the distri- bution shift imposed by the constraints and maintain the consistency of the system.
2

---

Published as a conference paper at ICLR 2025
Exact and efficient satisfaction of hard constraints. ECI sampling addresses the unique chal- lenges imposed by hard constraints. The memory- and time-efficient gradient- free approach guar- antees the exact satisfaction of constraints and mitigates gradient issues known with existing gradient-based methods in CV domains.
Zero-shot performance on generative and regression tasks. Comprehensive generative met- rics of distributional properties manifest the superior generative performance of our ECI sampling compared with various existing zero-shot guidance methods on various PDE systems. We also show that ECI sampling can be applied zero-shot regression tasks, still achieving competi- tive performance with state-of-the-art Neural Operators (NOs) (Li et al., 2020a) that are directly trained on the regression tasks.
RELATED WORK
Diffusion and Flow Matching Models. We first review existing generative models for functional data that serve as the generative prior for the spatiotemporal solutions in our approach. Diffusion models (Song& & Ermon, 2020; Ho al., 2020; Song et al., 2020) rely on a variational bound for the log-likelihood and learn reverse diffusion process to transform prior noise into meaningful sam- ples. Lim et al. (2023) propose the diffusion denoising operator (DDO) to extend diffusion models to function spaces. Flow matching models (Lipman et al., 2022; Liu et al., 2022) are family of continuous normalizing flows that learn a time-dependent vector field that defines the data dynamics via a flow ordinary differential equations (ODEs). Kerrigan al. (2023) proposes functional flow matching (FFM) as an extension of existing flow matching models for image generation. Compared to DDO, FFM has more concise mathematical formulation and better empirical generation results. We also note the close connection between diffusion and flow models in Lipman et al. (2022); Al- bergo et al. (2023). Therefore, while we focus on guiding flow-based models in this work, we also compare our approach with a wide range of diffusion-bassee methods.
Constrained Generation. Flow-based generative modeling in function spaces has not been ex- plored until recently, and constrained generation for PDE systems remains largely unexplored. The neural process (NP) (Garnelo et al., 2018; Kim et al., 2019; Sitzmann et al., 2020) is one traditional constrained generation approach that learns to map a context set of observed input-output pairs to a distribution over regression functions. Hansen et al. (2023) introduce a generic model-agnostic approach to control the variance of the generation and enjoy the exact satisfaction of constraints. Lippe et al. (2024) use gradient guidance from the constraint to guide each sampling step of the pre-trained diffusion model in a zero-shot fashion, similar to diffusion posterior sampling in Chung et al. (2022). Huang et al. (2024) further incorporate the gradients from physics-informed losses (Shu et al., 2023). Other approaches focus on controlling NOs in regression tasks (Négiar et al., 2023; Saad al., 2023; Li et al., 2024; Mouli et al., 2024). As the sampling for flow diffu- sion models is an iterative procedure, these methods cannot be directly adopted to provide iterative control. They can be, however, incorporated in our correction stage (see Section 3.2).
Inverse Problems. Zero-shot generation for pre-trained diffusion and flow models has been ex- plored in the image domain for solving various inverse problems, including image inpainting, de- blurring, and superresolution (Bai et al., 2020). Existing approaches predominantly rely on gradient guidance. Liu et al. (2023b); Ben-Hamu et al. (2024); Wang et al. (2024) propose to modify the prior noise or vector field by differentiating through the ODE solver, thus being extremely time- and memory-consuming. Pan et al. (2023a;b) use the adjoint sensitivity method to mitigate the high memory assumption. Other works propose gradient-fres control, usually with strong prior assump- tions on the constraints to derive control at each sampling step. Lugmayr et al. (2022) propose to mix forward diffusion steps from the context with the model's 's prediction, and Kawar et al. (2022) similar variational approach to solve linear inverse problems.
Although these existing methods have achieved decent generation results for inverse problems in image generation, the approximate enforcement of physical laws (soft-constrained) (Négiar et al., 2023) failed to address the challenges imposed by hard constraints (Wang et al., 2020; 2021; Krish- napriyan al., 2021). Our proposed zero-shot guidance approach focuses on the challenges of exact constraint satisfaction, using fine-grain iterative control at each sampling step for more consistent generations. The gradient-free approach is also more efficient and provides unified framework for various linear and non-linear constraints. We summarize the major differences between the existing methods and our proposed method in Table 1.
3

---

Published as a conference paper at ICLR 2025
Table I: Comparison between existing constrained generation methods and our ECI sampling.
Zero-shot
Gradient-free
Exact constraint
Iterative control
Conditional FFM [18]
x
ANP [19]
/
ProbConserv [13]
/
X
DiffusionPDE [16]
/
X
D-Flow [4]
/
ECI (ours)
/
3 METHOD
3.1 PRELIMINARY
Problem Definition. Consider a PDE system Fau(2) 0. R with PDE parameters é E Q. For a fixed PDE family F, let Uy = {u(z) 36 E b Fa((2) 0. E }} denote the set of all the plausible PDE solutions by varying the PDE parameters. Consider some constraint operator Gu( = 0, Xg S X' defined on subset of the PDE domain, and let Ug = {u(r) Gu(x) = 0, ø E Xg3) } denote the solution set for the constraint. We are interested in finding a subset of solutions Uzig
U5 Nug C U in which both the original PDE and the constraint are satisfied. We assume that the prior solution set Uy can be captured by a generative model, e.g., an FFM pre- trained on solution set. We aim to guide the pre- pre-trained model to satisfy the constraint operator in a zero-shot fashion towards the narrowed solution set UFg A more mathematically rigorous definition using measure theoretic terms is provided in Appendix A.1.
Prior Generative Model. FFM (Kerrigan et al., 2023) extends the flow matching framework (Lip- man et al., 2022) to model measures over the Hilbert space of continuous functions. Given a set 21 of well-behaved (see Appendix C.2) functions u : X' + IR, FFM learns a time-dependent vector field operator v, 4 [] + U that defines time-dependent diffeomorphism vi l x [0, U called the flow via the flow ODE:
(eve(u) = (1b(()), vo(u) = 24p
(1)
The flow t/ induces a pushforward mea- sure Ae := (1))) where Ho is the prior noise measure from which ug can be sam- pled. FFM showed that a tractable flow matching objective can be derived when
Algorithm 1 Sampling from FFM (Euler Method)
1: Input: Learned vector field ve, Euler steps N.
2: Sample noise function uo mo(u).
3: for < 0, 1/N, 2/N, (N 1)/N do
the flow is conditioned on the target fune- tion i1 sampled from the target measure interested in guiding pre-
4: Ut+1/N Uy + og(u: t)/N
μ1- As trained generative models, we will assume
5: return
that FFM can well approximate the target measure jú R fn over Uy. We want to guide the pre- trained FFM towards the conditional measure JFF over UFig- The sampling procedure the FFM, similar all diffusion and flow models, is an iterative process in which the initial noise function is iteratively refined into target functions via the learned vector field with the dynamics described in Equation 1. Algorithm describes sampling from the FFM method using the Euler method. In prac- tice, the vector field can be parameterized by some discretization-inaarte al., 2020b;a) such that the whole generative framework is pnteereizzaionnivvariii that FFM can be naturally adapted for zero-shot superresolution (Kerrigan et al., 2023).
NOs (Lu et al., 2019; Li This indicates
3.2 HARD CONSTRAINT INDUCED GUIDANCE
Although various regression NOs have been proposed to satisfy specific constraints their predic- tion (Liu et al., 2023a; Négiar et al., 2023; Saad et al., 2023; Liu et al., 2024; Mouli et al., 2024), the difficulty in applying these existing methods to flow matching models lies the iterative nature of the flow sampling in Algorithm I. For flow matching sampling, intermediate generations noised data instead of final predictions, making the correction algorithm inapplicable as the constraint can only be applied to the last generation step.
4

---

Published as a conference paper at ICLR 2025
Following previous zero-shot guidance frameworks (Lugmayr al., 2022; Kawar al., 2022), each iterative sampling timestept, t, we model the constrained generation process with the conditional generation probability pg) p(ù, (u, G)p(u,, where the unconditional probability p(u,) can be modeled by a pre-trained generative model. We assume correction algorithm C(u,, 9) is read- ily available for final predictions that uses orthogonal/oblique projection (Hansen et al., 2023) to ensure hard constraint satisfaction (see Appendix A.3). To offer iterative control over the sam- pling process for hard constraints, we note that if the final prediction can be extrapolated from intermediate noised generations, the correction can be safely applied. Furthermore, if intermediate noised generations can be interpolated back from the corrected prediction, we can "propagate" the constraint information back to each time step. Indeed, let po(ui (ue) denote the probability of ex- trapolating the final prediction of u1, given the current noise data Uy and the learnable parameter 0; and let g( ( denote the probability of interpolating the intermediate noised data ün given the corrected data ú = C(u,,9), then the constraint-condiiinn variational way by marginalizing over the auxiliary variable U1:
probability can be decomposed in a
p(ue/u,, 9) E., (o1 (g(0 (((1)))())))))))
(2)
We call such a constraint-guided step in each flow sampling step an ECI Step (for extrapolation- correction-intee interooaiiiteroottioo..... With the determinis- tic flow ODE formulation, good properties can be naturally deduced with the flow-matching frame-
Algorithm 2 ECI Step
1: Input: Learned vector field va, constraint
g. current noise data uy timestep L' > I.
u1 u, + t)
Extrapolation
work. The extrapolation probability po (1 |ux) learned by the pre-trained unconditional model by doing a deterministic one-step extrap- olation of the current predicted vector field
3: â1 t C(u1, 9)
> Correction
ue < (1 t'Yuo + t'ú
> Interpolation
5: return u.
u1 = it - (1 -t)u The interpolation probability ((ū)(ū1) also well-defined in the FFM along the OT-path of measures (Kerrigan et al., 2023), where the ground truth data are linearly interpolated with random noises as i, (1 t)ug + tūn, Me Mn(u). this way, the expectation can be discarded, and the conditional probability can be further simplified as
p(ūe (ue 9) = a( C(u + (1 (ue), 9)).
We summarize the ECI step in Al- gorithm 2. in which we also allow for L' 2 & to advance the solver. Furthermore, we note that our ECI steps can be applied recursively to the same timestep in a manner sim- ilar to previous work that applied multiple rounds of guidance (Lug- mayr al., 2022; 2024). With the constraint informa- tion "backpropagating" into the noise data via one ECI step, recursive ap- plication of such steps can promote further information mixing between the constrained and unconstrained re- gions, leading to a more consistent generation. At the last mixing step, we instead interpolate the new time step L' > with sample uy to advance the ODE solver. The number of total mixing steps M is controllable hyperparameter. Algorithm 3 describes the complete sampling process for flow matching models using iterative ECI steps, for which we term ECI sampling. ECI guarantees exact satisfaction of the constraint and facilitates information mixing between the constrained and unconstrained regions to produce a more consistent generation. See Appendix A for details on the sampling and proof of exact satisfaction.
Algorithm 3 ECI Sampling (Euler Method)
Input: Learned vector field to, Euler steps N. mixing iterations M, constraint G.
2: Sample noise function uo po(
3: for + 0, 1/,1 2/N, (N 1)/N do
Wi (0)
4:
+ Ut
5: for m < 0, 1, M do
et al.,
if m (m+1) ECIStep(ve G. u,
M
1 then
(m)
u
8:
else
> Advance the flow ODE solver
U1N1/N t ECIStep(Ve, 9, u.m)) t 1/N)
9:
10: return 21
3.3 CONTROL OVER STOCHASTICITY
ECI sampling provides a unified zero-shot and gradient-free guidance framework for various con- straints. Nonetheless, the varieties in the constraint types may impose challenges. For example, if
5

---

Published as a conference paper at ICLR 2025
the constraint contains little information, the variance is expected to match the unconstrained case. On the other hand, with enough constraints to result in a well-posed PDE with a unique solution, the generation variance should be as small possible (see Section 3.4).
Inspired by the work of the stochastic interpolant that unified flow matching and diffusion (Albergo et al., 2023), we propose control the generation stochasticity with different sampling strategies of un during our ECI sampling. For flow-based models, stochasticity only appears in the sampling process of un ~ mo(u) for interpolating the new üz In the previous context, we assume the initially sampled uo at / = 0 is used throughout the sampling process. This makes the trajectory straighter as the interpolation always starts with the same point. An alternative approach is sample new noises on each step, effectively marginalizing over the prior noise measure Ao and resulting in more concentrated generation. We use hyperparameter R to denote the length of the interval for re-sampling the noise function for interpolation. Intuitively, smaller R leads to more concentrated marginalized generation, whereas larger R leads to larger generational diversity. We quantitatively demonstrate the effect of this hyperparameter in Section 4.3 and also propose heuristic rules for choosing it.
3.4 GENERATIVE FFM AS A REGRESSION MODEL
We illustrate the potential of our ECI sampling framework for various regression tasks. Specifically, the solution set Uyig defined in Section 3.1 will degenerate to unique solution given enough constraints G, leading to a well-posed PDE system. this way, the generative task degenerates into deterministic regression learning task, with the target distribution now shifted to the Dirac measure over the solution 6um. Although our constrained generative framework is not designed for these regression tasks and has intrinsic stochasticity, in our empirical study, show that framework able to generate competitive predictions to standard regression models, e.g., the Fourier Neural Operator (FNO) (Li et al., 2020a).
EMPIRICAL EVALUATION
In this section, we present an extensive evaluation of our ECI sampling and other existing zero-shot guidance models for both generative and regression tasks. Covering wide range of diverse PDE systems, Table summarizes the dataset specifications used our evaluation. See Appendix for more details regarding the PDE systems and the data generation process. For most zero-shot guidance methods, we first pre-train an unconditional flow matching model as the common prior generative model (See Appendix C.3). We test our proposed method using more than one mixing iteration in most cases. See Appendix for the detailed experimental setup and baseline adaptations for flow sampling, and Appendix E for additional visualizations and results.
Table 2: Dataset specifications (IC/ /C/ CL for initial condition / boundary condition / conservation law). For different resolutions, we test models with the zero-shot superresolution setting.
Type
Dataset
Underspecification Spatial Resolution
Temporal Resolution
Constraint
Stokes Problem
IC & BC
IC BCC
Heat Equation
IC & diffusion
100
100
IC
Generative
Darcy Flow
BC & field
101 x 101
N/A
BC
NS Equation
IC & forcing
64 x 64
50
IC CL
Heat Equation
diffusion
100 / 200
100 / 200
PME
diffusion
100 / 200
100 200
CL
Regression
Stefan Problem
shock range
100 / 200
100 200
CL
Stokes Problem
IC & BC
100
100
IC & BC
NS Forward
forcing
first 15 frames
4.1 GENERATIVE TASK
In generative tasks, the solution set is non-degenerate with infinite cardinality after imposing the constraint. We expect the generation distribution to match the ground truth shifted distribution. The prior FFM is pre-trained larger solution set, whereas during the zero-shot constrained

---

Published as a conference paper at ICLR 2025
generation, impose specific PDE constraint narrow down the solution set. Additionally, use a conditionally conditionllly-tranned trained FFM (CondFFM) that takes known constraints as conditions during both training and sampling upper bound for the zero-shot guidance models. We conduct extensive experiments on various 2D and 3D) PDE datasets with different types of constraints.
mean
mean
error
std
error
GT
CondFFM
ECI
ANP
ProbConserv
DiffusionPDE
D-Flow
Figure 2: Generation mean and standard deviation errors for the Stokes problem with IC fixed. Note that grey indicates 0 error, red for positive error, and blue for negative error. Noticeable IC errors (left column) can be observed for the gradient-based methods.
To compare two distributions over continuous functions, we follow Kerrigan et al. (2023) to calculate the pointwise mean and standard deviation (std) of the ground truth solutions and the solutions generated from various sampling methods. Mean squared errors (MSEs) of the mean (MMSE) and of the std (SMSE) are then calculated assess the resemblance in these statistics. To evaluate the satisfaction of the constraint, the constraint error (CE) is defined as the MSE for the constraint operator CE(u)= = (s)(G(u). 0) on the constrained region Xg. Inspired by the generative metric of Fréchet inception distance (FID) in image and PDE generation domains (Lim et al., 2023), we calculate the Fréchet distance between the hidden representations extracted by pre-trained PDE foundation model Poseidon (Herde et al., 2024). We call this metric Fréchet Poseidon distance (FPD). A lower FPD indicates closer match to the ground truth distribution regarding the high- level representations. To alleviate the impact of randomness in sampling, for the 2D datasets and 100 solutions for the 3D datasets for all methods to calculate the metrics.
sample 512 solutions
We carry out experiments ID Stokes problem, ID heat equation, 2D Darcy flow, and 2D Navier-Stokes (NS) equation. The detailed description these PDE systems, as well as the set of PDE parameters for pre-training and constrained generation, are provided in Appendix B. Note that we also ensure at least two degrees of freedom (underspecification) in generative tasks, as demonstrated in Table 2. This ensures that distribution over the constrained generations does collapse to a Dirac distribution as in regression task (see Figure 1).
mcan
mean
error
15
std
std
10
error
82
.03
-1.
GT
CondFFM
BCI
ProbConserv DiffusionPDE
D-Flow
Figure 3: Generation mean and standard deviation errors of the trajectories for the NS equation with IC fixed. 25 downsampled time frames are plotted for each method.
Table 3 summarizes the results for the generative tasks with the best zero-shot method performance highlighted in bold. Figure shows the visualization of pointwise emrors for the generation statistics for the Stokes problem with IC fixed together with the ground truth reference. Figure 3 shows the visualization for the NS equation with 25 downsampled time frames. Our proposed ECI sampling achieves state-of-the-art performance most metrics and also enjoys constraint emrrs. Notice-

---

Published as a conference paper at ICLR 2025
Table 3: Generative metrics various constrained PDDE. The best results for zero-shot methods
(CondFFM is not zero-shot) are highlighted in bold. The unconstrained pre-trained FFM is provided.
Dutaset
Metric
ECI
ANP
ProbConserv
DiffusionPDE
D-Flow
CondFFM
FFM
MMSE/ 2
0.090
13.750
2.014
0.601
1.015
2.068
SMSE 10-2
0.127
1.528
0.016
1.572
Stokes IC
CE/ 10-2
3.711
0.721
1.021
0.017
9.792
FPD
0.076
26.685
13.075
0.202
13.342
MMSE/ / 10 10-2
0.005
3.003
2.938
0.603
4,800
0.032
3,692
SMSE/ / 10
0.003
3.055
2.302
0.341
>-00
0.047
3.320
Stokes BC
CR/ 10-2
29.873
0.007
53.551
0.012
>> 100
FPD
0.010
5.339
2.780
1.162
27.680
0.048
2.824
MMSE/ 10 -2
1.605
5.127
3,744
0,979
1.174
0.008
3.954
SMSE / 10-2
0.157
0.740
4,453
1.291
9.928
4.737
Heat Equation
CR/ 10-2
0.183
11.646
94.551
0.002
49.296
FPD
1.365
2.432
2.639
0,951
3.831
0.006
2.799
MMSE/ 10 -
2.314
14.647
56.608
10.442
1,728
0.043
58.879
SMSE / 10-2
0,592
1.269
65.967
1.097
6.343
0.037
69.432
Darcy Flow
CE/ 10-2
3.152
0.656
14.553
0,071
10
FPD
0.946
5.805
100
8.027
10.043
0.021
100
MMSE/ 10 2
7.961
59.020
19.141
18.749
SMSE 10-2
5.846
14,990
8.306
5,941
0.554
8,423
NS Equation
CE 10
5.527
2,295
0.776
0,079
11.451
FPD
1.131
57.263
2,171
5.750
1.651
0.232
2,185
ably, ECI sampling excels in capturing the higher-order distributional properties of the shifted con- ditional measure, e.g., the standard deviation and FPD, and even surpasses the conditional FFM in several simple 2D cases. Specifically, we observe noticeable artifacts around the IC for the gradient- based methods, i.e., DiffusionPDE and D-Flow, they do not guarantee the exact satisfaction of the constraint. For ProbConserv (Hansen et al., 2023), although the constraint is satisfied exactly, the naive non-iterative correction approach leaves a sudden change around the boundary that vi- olates the physical feasibility. In addition, Table 4 shows that our gradient-free approach enjoys sampling efficiency in both 2D and 3D cases compared with gradient-based methods. Noticeably, ECI-1 achieved x440 acceleration and x310 memory saving compared to gradient-based D-Flow on the Stokes problem. More visualizations for the other PDE systems and the different generation statisties provided in Appendix E.1.
Table 4: Generation setup (batch size Fuler step), time per sample in seconds, and GPU memory for 2D (Stokes IC) and 3D (NS equation) PDE systems. A fraction sample size is available for ANP.
Datasct
Resource
ECI-1
ECI-5
CondFFM
ANP
ProbConsery
DiffusionPDD
D-Flow
Hsample/HEuler
128/200
128/200
128/200
32/A
128/200
128/200
2/200
Stokes IC
Time/sample/s
0.065
0.325
0.057
0.009
0.058
0.131
28.774
Me
5.4
5.4
5.4
7.4
26.4
#samplo/Euller
25/100
25/100
25/100
0.2/NA
25/100
25/100
1/20
NS equation
Timessample/s
0.415
2.067
0.669
2.324
0.675
0.676
8.456
GPU Memory/GB
16.3
16.3
16.3
11.0
16.3
27.0
27.1
4.2 REGRESSION TASK
We additionally experiment with regression scenarios (see Section 3.4), where enough constraints reduce the generative task into the standard regression setting in neural operator learning. As the ground truth targets are available in this case, we calculate the MSE and other related evaluation metrics and compare them with traditional NOs.
Uncertainty Quantification. In the uncertainty quantification task, random context points from the true solution sampled together with the conservation laws pin down the unique solution. Following Hansen et al. (2023), we experiment with our ECI sampling on the Generalized Porous Medium Equation (GPME) family of equations using identical PDE parameter choices. We evaluate
8

---

Published as a conference paper at ICLR 2025
Table 5: Uncertainty quantification metries constrained PDEs from the Generalized Porous Medium Equation (GPME) family of equations. All baseline results are taken from Hansen et al. (2023). The best results are highlighted in bold.
Dataset
Metric
ECI
ANP
SoftC-ANP
HardC-ANP
ProbConsery
CE/ CE/10-3 10-3
0
4.68
3.47
0
0
Heat Equation
LL
1.90
2.72
2.40
3.08
2.74
MSE/ / 10
0.81
1.71
2.24
1.37
1.55
CE 10-3
0
6.67
5.62
0
0
PME
I.I.
2.19
3.49
3.11
3.16
3.56
MSE/ / 10
0.19
0.94
1.11
0.43
0.17
CE 0//10 10-2
0
1.30
1.72
0
0
Stefan Problem
L. MSE/ / 10-3
3.30
3.53
3.57
2.33
3.56
1.89
5.38
6.81
5.18
1.89
Table 6: Neural operator learning metrics two constrained PDEs. We abbreviate 0OM for out- of-memory, The results for the best generative models are highlighted in bold.
Dataset
Metric
ECI
ANP
ProbConserv
DiffusionPDE:
D-Flow
FNO
MMSE / 10-3
0.050
16.470
103.136
5,839
89.514
0.033
Stokes Problem
SMSE / 10-2
0.028
0.033
7.509
0.473
136.936
0
MMSE
0.069
OOM
46.673
19.015
1.310
0.380
NS Equation
SMSE/ / 10
0.002
OOM
28.789
13.914
4.068
0
the generation results on three instances of the GPME, i.c., the heat equation, porous medium equa- tion (PME), and Stefan problem, where the additional constraints various physical conservation laws. We also follow the paper to report the pointwise Gaussian log-likelihood (LL.) based on the sample mean and variance. The results in Table show that our ECI sampling able achieve comparable results with ProbConserv (Hansen et al., 2023), which is specifically tailored for these uncertainty quantification tasks. We also notice that, though being a generative model with intrinsic stochasticity from the prior noise distribution, our method is able to produce quite confident predic- tions with little variance. It surpasses the other baselines terms of MSE but at the cost worse LL. See Appendix E.3 for visualizations and more discussions.
Neural Operator Learning. We further experiment with the zero-shot neural operator learning task on the Stokes problem and Navier-Stokes equation. For the Stokes problem, both BC and IC are prescribed give unique solution. For the NS equation, we follow Li et al. (2020a) to consider the regression task from the first 15 frames to the other 35 frames. In addition, a standard FNO (Li et al., 2020a) was trained as the regression baseline. Table 6 summarizes the results. The large amount of context pixels in the NS equation makes it infeasible for the ANP model. Our ECI sampling significantly outperforms gradient-basse better MSEs than the NOs directly trained on these regression tasks. Similar to the uncertainty quantification task, we observe that our ECI sampling quite confident about its prediction with little variance. We also provide comparison of ECI sampling with different numbers of frames fixed as the constraint in Appendix E,4 to demonstrate how the variance gradually reduces with more constraint information preseribed.
methods and reaches comparable or even
4.3 ABLATION STUDY
Similar to most existing controlled generation methods for diffusion or flow models, the number of mixing iterations is tunable hyperparameter for our ECI sampling control the extent of informa- tion mixing. Intuitively, an insufficient number of mixing iterations may fail to mix the information between the constrained and the unconstrained regions, leading more artifacts. On the other hand, an excessively large number of mixing iterations may discard the prior PDE structure learned the pre-trained generative model, leading inconsistency with the PDE system. We further test with

---

Published as a conference paper at ICLR 2025
the length of noise re-sampling interval R, discussed in Section 3.3, to control stochasticity for different constrained generation tasks.
We test the combination of mixing iterations M (1, 10, 20), with re-sampling interval R E {1, 2, 5, 10). and no re-sampling settings on the Stokes problem with the IC or BC prescribed. Figure illustrates the results of different ECI sampling settings together with other baselines. For the IC task, R None performs the best, whereas for the BC task, R = None performs worse, with R and R achieving lower MSEs. In both cases, the MSEs do not decrease monotonically with increasing mixing iterations. Therefore, the mixing iteration impacts the final performance in a task-specific manner. We hypothesize that the difference in the shifted variance may cause this difference in the performance after imposing specific constraint. For the Stokes problem, the prior ground truth variance is higher for the BC than for the IC. Prescribing BC values provides more information and reduces the uncertainty, making it easier for the model to be guided. In contrast, prescribing IC values provides less information, which results in larger variance in the shifted distribution and makes it harder to guide the model.
Inspired by the aforementioned observations, we propose the following heuristic rules for choosing the appropriate number of mixing iterations M and re-sample interval R: 1) We usually limit M between 1-10, as a larger number of iterations both computation-inten intensive and leads to worse performance; and 2) For easier tasks with lower variance, we choose smaller M and smaller R; and for harder tasks with higher variance, we choose a larger M and larger R or no re-sampling. Indeed, we use mixing iteration for the relatively easy tasks the heat equation and Darcy flow and use 10 mixing iterations for the harder NS equation.
MMSE vs. SMSE for Stokes IC
MMSE vs. SMSE for Stokes BC
AN
Prab
serv
10
ProbConsery
x DiffusionPDD
SMSE
M-1
SMSF
M-1
DiffusionPpE
10-1
M-2
M-5
CondFFM
M+5
D-Flow
M-18
M-10
M-20
M-20
10-1
R=1
R-2
R-2
10-4-
Red
CondFFM
R-10
R-1li
R-None
10-1
10-3
10"1
10-4
10-3
10-2
10-1
MMSE
MMSE
Figure 4: MMSE and SMSE of ECI sampling with mixing iterations M in different colors and re- sampling intervals in different markers the Stokes problem with IC (left) or BC (right) fixed. "None" for no re-sampling (the initial noise is used).
CONCLUSION
In this work, we present ECI sampling as a unified zero-shot and gradient-free framework for guid- ing pre-trained flow- flow-matching models towards hard constraints. Instantiated on PDE systems, our framework obtains superior sampling quality and efficient sampling time compared to existing zero- shot methods, and it also enjoys the additional benefit of a zero-shot regression model with com- parable performance to the traditional regression models. We also highlight the potential of ECI sampling on other domains outside of SciML, including supply chain and times series, which we leave as future work. We further note a limitation of our framework: the number of mixing itera- tions and the re-sampling interval have task-specift impact on the final performance. While have provided empirical evaluations of these hyperparameters and suggested heuristic guidelines for selecting them, theoretical analyses would be beneficial.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions.. arXiv preprint arXiv:2303 08797, 2023, 3,
Abdul Fatir Ansari, Lorenzo Stella, Caner Turkmen, Xiyuan Zhang, Pedro Mercado, Huibin Shen, Oleksandr Shchur, Syama Syndar Rangapuram, Sebastian Pineda Arango, Shubham Kapoor, Jasper Zschiegner, Danielle C. Maddix, Hao Wang, Michael W. Mahoney, Kari Torkkola, An- drew Gordon Wilson, Michael Bohlke-Schneider and Yuyang Wang. Chronos: Learning the language of time series. arXiv preprint arXiv:2403.07815. 2024. 2
Yanna Bai, Wei Chen, Jie Chen, and Weisi Guo. Deep learning methods for solving linear inverse problems: Research directions and paradigms. Signal Processing, 177:107729, 2020.
Heli
Omri Puny, Itai Gat. Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. arXiv preprint 240.::2402.11111., 2024.
1, 3, 4, 23
Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03666 2023. 19
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018. 19
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14687, 2022.
3, 23
J.R. Dormand and P.J. Prince. A family of embedded runge-kutta formulae. Journal of Computa- tional and Applied Mathematics, 6(1):19-26, 1980. ISSN 0377-0427. 21
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024a.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion En- glish, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis, 2024b. URL https:: https://arrivoorr/aass. 2403. 03206.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint
2018. 3
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint arXiv:2407.15599 2024.
Derek Hansen, Danielle C Maddix, Shima Alizadeh, Gaurav Gupta, and Michael W Mahoney. Learning physical models that can respect conservation laws. International Conference on Machine Learning, pp. 124409-1550.. PMLR. 2023. 1, 5, 8, 9, 15, 17, 18, 19, 22, 27
Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes. arXiv preprint arXiv:2405.11 19101, 2024. 21
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840 6851, 2020. 1,3, 19
Jiahe Huang, Guandao Yang, Zichen Wang, and Jeong Joon Park. DiffusionPDE: Generative PDE- solving under partial observation. arXiv preprint arXiv:2406.1776. 2024, 3, 4,
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. volume 35, pp. 2599-233006, 2022. 3,

---

Published as a conference paper at ICLR 2025
Gavin Kerrigan, Giosue Migliorini, and Padhraic Smyth. Functional flow matching. arXiv preprint arXiv:2305.17209, 2023. 1, 1,3, 4, 14, 16, 19, 20, 22
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Gamelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv: 2019.
3, 4, 22
Marcel Kollovieh, Abdul Fatir Ansari, Michael Bohlke-Schneider, Jasper Zschiegner, Hao Wang, and Yuyang Bernie Wang. Predict, refine, synthesize: Self-guiding diffusion models for proba- bilistic time series forecasting, Advances in Neural Information Processing Systems, 36, 2024.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes in physies-informed neural networks. In Advances in Neural Information Processing Systems, volume 34, pp. 26548 265560, 2021. 2,
Randall LeVeque. Numerical Methods for Conservation Laws. Lectures in mathematics ETH
Zürich. Birkhäuser Verlag, 1990. 2
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv: 2020a. 3, 4, 6, 9, 17, 18, 20, 23
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprini arXiv:2003.0344 03485, 2020b. 4
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. ACM/JMS Journal of Data Science, 1(3):1-27, 2024. 23
Jae Hyun Lim, Nikola Kovachki, Ricardo Baptista, Christopher Beckham, Kamyar Azizzade- nesheli, Jean Kossaifi, Vikram Voleti, Jiaming Song, Karsten Kreis, Jan Kautz, et al. Score-based diffusion models in function space. arXiv preprint arXiv:2302.07400, 2023. 1, 3, 7, 20
Lequan Lin, Zhengkun Li, Ruikun Li, Xuliang Li, and Junbin Gao. Diffusion models for time-series applications: survey. Frontiers of Information Technology & Electronic Engineering, 25(1): 19-41, 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXhv:2210.02747, 2022. 3, 4. 15, 19
Phillip Lippe, Bas Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. In Advances in Neural Information Processing Systems, volume 37, 2024. 3
Ning Liu, Yue Yu, Huaiqian You, and Neeraj Tatikola. INO: Invariant neural operators for learning complex physical systems with momentum conservation. In Proceedings of The 26th Interna- tional Conference on Artificial Intelligence and Statistics, volume 206, pp. 6822-222222222 6838. PMLR, 2023a. 4
Ning Liu, Yiming Fan, Xianyi Zeng, Milan Klöwer, Lu Zhang, and Yue Yu. Harnessing the power of neural operators with automatically encoded conservation laws. In International conference on machine learning. PMLR, 2024. 4
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209. 000000,,,0, 2022. 3
Xingchao Liu, Lemeng Wu. Shujian Zhang, Chengyue Gong, Wei Ping, and Qiang Liu. Flow- grad: Controlling the output of generative odes with gradients. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24335 -2334,, 2023b.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310. 10. 2023. 1
12

---

Published as a conference paper at ICLR 2025
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem operators. arXiv preprint ar arXh::1910.03190 2019. 4
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 11461-11471, 2022. 3,
S. Chandra Mouli, Danielle C. Maddix, Shima Alizadeh, Gaurav Gupta, Andrew Stuart, Michael W. Mahoney, and Yuyang Wang. Using uncertainty quantification characterize and improve out- of-domain learning for pdes. In International Conference Machine Learning. PMLR, 2024. 1, 3,
Geoffrey Négiar, Michael W. Mahoney, and Aditi Krishnapriyan. Learning differentiable solvers for systems with hard constraints. In International Conference on Learning Representations, 2023. 3,4 4
Jiachun Pan, Jun Hao Liew, Vincent YF Tan, Jiashi Feng, and Hanshu Yan. Adjointdpm: Adjoint sensitivity method for gradient backpropagation of diffusion probabilistic models. arXiv preprint arXiv:2307.10711 2023a.
Jiachun Pan, Hanshu Yan, Jun Hao Liew, Jiashi Feng, and Vincent YF Tan. Towards accurate guided diffusion sampling through symplectic adjoint method. arXiv preprint vrXiv:2312.12030 2023b.
M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: deep learn- ing framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019. ISSN 0021-9991. 23
Nadim Saad, Gaurav Gupta, Shima Alizadeh, and Danielle C. Maddix. Guiding continuous operator
learning through physics-base based boundary constraints. In International Conference on Learning Representations, 2023. 1, 3, 4,
Dule Shu, Zijie Li, and Amir Barati Farimani. A physics-informed diffusion model for high-fidelity flow field reconstruction. Journal of Computational Physics, 478: 11s:121992,,. 2023. 3
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. In Advances in Neural Information Processing Systems, volume 33, pp. 7462-7473, 2020. 3
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint XXiv:2010.025022 2020. 3
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution.
In Advances in Neural Information Processing Systems, volume 32, 2019. 19
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems, volume 33, pp. 12438-12448, 2020. 3
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-2826, 2016. 21
Ashish Vaswani. Attention is all you need, arXiv preprint arXiv: (7t.::3706...722. 2017. 20
Hengkang Wang, Xu Zhang, Taihui Li, Yuxiang Wan, Tiancong Chen, and Ju Sun. Dmplug: A plug- in method for solving inverse problems with diffusion models. arXiv preprint 4rX1v:2405.16749, 2024. 3
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient pathologies in physics-informed neural networks. arXiv preprint wwt.ivv00010045500 2020. 2, 3
Sifan Wang. Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric par- tial differential equations with informed deeponets. Science advances, 7(40):eabi8605, 2021. 2,3
13

---

Published as a conference paper at ICLR 2025
Supplementary Material
A DETAILS OF CONSTRAINED SAMPLING
In this section, further discuss additional details of our proposed ECI sampling framework. Our proposed ECI sampling interleaves extrapolation, correction, and interpolation stages at each sam- pling step for flow-hased generative models enforce the hard constraint. We will elaborate further on the implementation details below.
Extrapolation
Correction
Interpolation
Figure 5: Intermediate results (trajectory) of the extrapolation, correction, and interpolation stages for the Stokes problem with IC fixed. The trajectories demonstrate how the initial random noise gradually transforms into the controlled generation with decreasing artifacts around the boundary.
A.1 FORMAL DEFINITION OF CONSTRAINED GENERATION
rigorous measure- -heorrett definition of constrained generation for func- tional data as a generative task. Consider a well-posed PDE system Fau(2) 0, E X' with PDE parameters é E Φ then the solution map T fixed distribution pe over the parameter Φ, we can naturally define the measure over the solution set the push-forward of the solution map: μa T.Ps. Now consider some constraint operator Gu(#) = 0, E Xg C X defined on a subset of the PDE domain. Denote the solution set for G as Ug {u(x): Gu(x) 0, x E Kg} and the characteristic function over the solution Xg such that
First, provide
u(x) is well-defined mapping. For
xo(u) = f1.
u E EE
u & Ug
With the non-degenerating assumption that fu Xod1i > O (this is essentially saying that we should have at least one solution), the conditional measure can be obtained as
== := =- J xoduz / xoduT, VACU,
(5)
where U/ the Hilbert space of all "well-behaving" functions (see Kerrigan al. (2023) for regu- larity conditions). We assume faF can be approximated well by generative priors like FFM, and we want to guide it towards the conditional measure μFiG-- In practice, the ground truth solution map T can be obtained by either exact analytical solutions or by using numerical PDE solvers. With pre-defined distribution over PDE parameters (often uniform over some interval, see Table and Appendix B for detailed specification), solution map to obtain the sampled functions for pre- pre-training.
easily sample the PDE parameters and apply the
For constrained sampling, instead of directly calculating the conditional measure μrng, can sample the PDE parameters from the pull-back probability PΦig = TT* In practice, we directly assume such conditional probability distribution Paje over the PDE parameters is known. Indeed, for all the experiments in this paper, we simply choose a subset of PDE parameters and apply the solution map obtain functions as the ground truth constrained solutions.
14

---

Published as a conference paper at ICLR 2025
A.2 EXTRAPOLATION OF SOLUTIONS
The continuous-time and deterministic formulation of the flow-matching sampling makes it possible to apply different step sizes to advance the solver for the flow ODE. Specifically, at timestep t, we can make one-step prediction with step size of 1
that directly advances the ODE solver to t = 1 as W1 = 4 + (1 - t)e. Theoretically, if the optimal transport paths are learned exactly, the flow should be straight (Lipman al., 2022), and any arbitrary discretization should lead the same target. Indeed, in the left plot in Figure 5, the extrapolation gives reasonable prediction even timestep close to O. The constraint, the other hand, is not necessarily satisfied in this stage.
A.3 SOLUTION CORRECTION
Our ECI sampling scheme is unified framework for different constraint operators long as the corresponding correction algorithm at = is readily available. Specifically, we consider constraint operators G of two following categories: 1) value constraints that specify exact values for subset of the PDE domain: u(x) g(x), E Xg S X. and 2) region constraints that specify an exact value for the integration over some subset of the domain: Sx,5x u(x)dx
R. These two classes of constraints encompass a wide range of practical constraints for PDE systems.
Example 1. For X' = Ω2 x [0, 7' where Ω E P-1,, the follow examples are value constraints:
Initial condition (IC): u(x, 0) = g(z), E Ω
Boundary condition (BC): u(x,t) f(t), 02, [O.T].
Example 2. For X' = 0] x [0, T), the following examples are region constraints:
Mass conservation: fo u(x, ()dr = 0,, E [0, T7.
Periodic boundary condition (PBC): u(0,t) u(I.t),t [0,
Furthermore, any linear conservation law can understood as a region constraint, specifying a conserved quantity over region in the PDE domain. As these constraints are linear constraints, we follow Hansen et al. (2023) to apply an oblique projection as a correction. The corresponding correction algorithms are outlined in Algorithm 4 and 5.
Algorithm 4 Value Constraint Correction
1: Input: function U1, constraint function g, region Xg.
2: 01 < 1[z E Xg] @ + [|x & Xg] @ u1
3: return ún
Algorithm Region Constraint Correction
1: Input: function u1 conservation value a, region Ng.
2: g + ūi + (a- Sx. un (xx ( Sx, dr
3: 0 - 1[x E Agl @ + 1 $ Xzl C
4: return ún
is easy to verify that the corrected function ú these correction algorithms indeed satisfies the corresponding constraint exactly. We can also have multiple (even infinitely many) overlaaalap constraints. For example, consider the generalized mass conservation laws fo u(xx f)dr = ft) = 0, E [0,2] for some given function (t), we have different conservation law for each PDE timestep t. After discretization, we have a finite set of non-overlapping constraints that can be corrected according to Algorithm 5. The correction stage is demonstrated in the middle plot in Figure 5. Note how directly applying the correction algorithm creates noticeable artifacts around the boundary and how such artifacts gradually disappea when advancing the ODE solver.
15

---

Published as a conference paper at ICLR 2025
A.4 INTERPOLATION OF SOLUTIONS
The iterative nature of flow sampling requires interpolation back to arbitrary timestep / during each sampling step. This can be understood the forward noising process along the conditional path of measures and can be efficiently calculated as the linear interpolation the FFM formulation (Kerrigan al., 2023). In the right plot in Figure we demonstrate how the initial Gaussian process noise is transformed into the constrained generation. We have also discussed the impact of the re-sampling interval length to control stochasticity in the generation in Section 3.3.
= (1- - tt + + according
A.5 PROOF OF EXACT SATISFACTION OF CONSTRAINTS
Proposition 1. Suppose the corrected algorithm C(u1, g) in Equation 2 satisfies the constraint G exactly, then for any number of mixing steps M 2 1, the ECI sampling scheme described in Algorithm 3 exactly recovers the constraint in the final generation 1.
Proof. As the advancing step for the ODE solver (the last mixing step) will always perform, it suffices to consider M 1. Consider the last Euler step at t = I/N, the linear interpolation procedure a (ú1) 0- ug + 1 the final interpolant at timestep t' = 1. Therefore, the interpolation will deterministically produce q( (ū)) ún C(u + ())) g) which satisfies the constraint exactly.
ú₁ will be deterministic the noise will not contribute to
B DATASET DESCRIPTION AND GENERATION
This section provides a more detailed description the datasets used in this work and their gener- ation procedure. In addition to the statistics in Table 2, we further provide additional information in Table 7 for pre-training the FFM as our generative prior. For dataset types, synthetic indicates that the exact solutions calculated both the training and test datasets. We manually assign the training set with 5k solutions and the test set with 1k solutions. On the other hand, simulated indicates that the solutions pre-generated using numerical PDE solvers and are different for the training and test datasets.
the fly based the randomly sampled PDE parameters for
Table 7: More dataset specifications for pre-training the prior FFM.
Dataset
PDE Parameter
Split
Spatial Domain
Time Domain
Type
U[2, 20]
Stokes Problem
5k / 1k
[0, 1]
[0, 1]
Synthetic
w U[2,8 8
a U[1
Heat Equation
5k / 1k
[0,2m]
[0, 1]
Synthetic
k (see below)
Darcy Flow
10k 1k
[0.
NA
Simulated
C 21-
NS Equation
Wo. / (see below)
10k 1k
[0, 1]
[0,49]
Simulated
PME
U[1,
5k lk
[0. 1]
[0,1]
Synthetic
Stefan Problem
~ U(0.55..077
5k 1k
[0,
[0,0.11
Synthetic
B. STOKES PROBLEM
The ID Stokes problem is given by the heat equation:
Ut = Vilans
x E [0, 10 E [0,
u( 0) = Ae ke cos(k:rr)
x E [0, 1].
(6)
t) = A cos(w:t),
E [0,
with viscosity 2 oscillation frequency w, amplitude A 0. and k /w/(2v). The analytical solution is given by (x, Ae-ka cos(ker wt). Note that & and w independently and
16

---

Published as a conference paper at ICLR 2025
uniquely define the IC and BC, respectively, and fixing both values reduces the PDE to well-posed system with a unique solution. We follow Saad et al. (2023) to construct the dataset by fixing A == test two
and sampling ~ [2, 8), k U[2 20] for pre-training. During constrained sampling, different settings of preseribing IC with k = 5 or BC with w = 6, respectively. In the regression task, provide both the IC and BC values. Figure shows the ground truth solutions for the unconstrained and constrained Stokes problem.
B.2
HEAT EQUATION
The ID heat (diffusion) equation with periodic boundary conditions given
We aM w(r, 0) sin(2 + sp),
ar E [0, 2T), E [0, 1],
E [0, 2m], E [0, 1],
v(0, t) = u (2 t),
where a denotes the diffusion coefficient and yp denotes the phase of the sinusoidal IC. The exact solution is given (2, e sin(2 + (p). We sample U[1, 5] and ip ~ U0, æ] for pre-training. During constrained sampling, we fix the phase Φ = T/4. Figure 6 shows the ground truth solution family.
In the uncertainty quantification task, we follow Hansen et al. (2023) to fix Φ = and vary only the diffusion coefficient in U11, 5] during pre-training. During constrained generation, we use the same
setting to fix 1, = 0.5 for calculating the MSE and LL. The global conservation law for the heat equation in Equation 7 is written as
2m u(x, t)dx O, t [0,1].
(8)
Unnormalized
Normalized
Figure 6: Ground truth solutions for the heat equation with a U[1,5), m).
Figure 7: Ground truth solutions for the Darcy flow with C U,- 2]. Solution values are unnormalized on the left and nor- malized per-solution the right.
B.3 DARCY FLOW
The 2D Darcy flow is time-independentt second-order, elliptic PDE with the following form:
-V (((K())V((z)) = f(x), x E D =0,01
u(x) = C,
aD.
where denotes the permeability field and f denotes the forcing function. We follow Li et al. (2020a) to fix f(x) = 1 and sample & ~ (O,,(0, (-A 9/)-2) with zero Neumann boundary conditions on the Laplacian. The mapping s - R takes the value 12 for positive numbers and 3 for the negative numbers, and the pushforward is defined pointwise. Unlike previous work, the boundary condition sampled from C UI- 2, 2]. It can be verified that, for the same kt if un is the solution for prescribing the boundary condition u(x)
0, E θD, then un + C' is the solution when prescribing the boundary condition u(x) = C, x E åD. Therefore, we simulate 1000 samples with different & and the zero boundary condition and sample C on the fly for both the training and testing datasets. During constrained generation, we fix C = 1. Figure shows the ground truth solution family.
17

---

Published as a conference paper at ICLR 2025
B.4 NAVIER-STOKEE EQUATION
The 2D Navier-Stokes (NS) equation for a viscous, incompressible fluid in the vorticity form with periodic boundary conditions is given as
d,w(x, t) + u(2, t)
Vw(x, t) = vAw( 1) f(x),
x E 1| [O,T],
V V- w(r, 1) = O,
x E 00 1J,, [O,T], z x [0,
(10)
(xx, 0) = (I))
where denotes the velocity field, w V x denotes the vorticity, and wo denotes the initial vorticity. We follow Li et al. (2020a) to sample wo ~ (O, (-/2 + 49/)-6/2) with periodic boundary conditions. The forcing term is defined as f(x) = sin (iT((2((z1 + x2) + 6) where O0, x/ /22, and the viscosity fixed to forces with linearly spaced ő in the interval [0, /2). We mesh-grin gessh the vorticities and forces to build 10000 solutions for the training set. For the testing set, we sample another 10 initial vorticities and mesh-grid them with the same 100 forces, resulting in another 1000 solutions. We use 7 49 and sample 50 snapshots the numerical solutions. We use the same generative NS equation datasets training for the regression task that maps the first 15 frames to the other 35 frames. The regression evaluation performed on the first sample in the test set.
10-3 We sample 100 initial vorticities and 100
B.5 POROUS MEDIUM EQUATION
The nonlinear Porous Medium Equation (PME) with zero initial and time-varying Dirichlet left boundary conditions is given
d V (u
[0, 1],
[0, 1], * E [0, 1], t E [0, 1],
u(r 0) = O,
(11)
u00, t) (mt)l/m
u(1, t) = O,
t e [0, 1].
where m 2 The exact solution is given as Uexae t) (m ReLU(t - z)) We follow
Hansen et al. (2023) to sample m (111,5) 5] for pre-training and fix m 1,f = 0.5 for sampling. The conservation law for the PME in Equation 11 written
L. ((x,t)dx= t)dr
1mt)
m ++ EE [0, 1].
(12)
Note that this conservation law implicitly contains information about m.
B.6 STEFAN PROBLEM
The Stefan problem is a challenging and nonlinear case of the Generalized Porous Medium Equation (GPME). With fixed Dirichlet boundary conditions, is given as:
4 V (k(u)Vu),
x [0, 1|, e [0, T],
u(x, 0) = 0, u(0, f) 1,
x E [0, 1), E [0, T), t E [0,7],
(13)
u(1,t) 0,
where k(u)) denotes the nonlinear step function with respect fixed shock value ut:
1, u 2 n* k(u) = 10, < n²
(14)
The exact solution is given
Mexact (x,, 1) 11 2 w²] 1 - u²) erf(x/(v87)
(15)
crf(a)
18

---

Published as a conference paper at ICLR 2025
where ez) -े So² exp(-t?)dt is the enoor function and a is uniquely determined by u" the
solution for the nonlinear equation (1 w²) ( VT u" erf(a)a exp(a²). We follow Hansen et al. (2023) sample the shock value u" U [0055,,00 0.7] and T 0.1. During sampling, we fix 2" = 0.6.6. t = 0.05. The conservation law for the Stefan problem in Equation 13 is
Lo u(z,
2(1- u²) erf(a)
()dr u(a,()dr=
LE [0.1].
(16)
C FUNCTIONAL FLOW MATCHING AS THE GENERATIVE PRIOR
This section provides additional mathematical backgrounds on flow matching, functional flow matching, and pre-training settings for FFM the generative prior.
C.I FLOW MATCHING
We first provide more details regarding the flow matching framework on common Euclidean data (e.g., pixel values of images) before proceeding to functional flow matching for functional data (e.g., PDE solutions). Flow matching (Lipman et al., 2022) is a generative framework built upon continuous normalizing flows (Chen et al., 2018). This flow-based model can viewed as the continuous generalization of the score matching (diffusion) model that allows for more flexible design of the denoising process following the optimal transport formulation (Lipman et al., 2022). Flow matching tries to learn the time-dependent vector field IR x [0, 1] + E that defines continuous time-dependent diffeomorphism called the flow the Rd x [0, 1] + Rd via the following flow ODE:
0t (c)( vy (v)()), ~ Po(o,
(17)
where po is the initial noise distribution. The flow induces a probability path with the push-forward Pr (vt))PP Po for generative modeling. The vanilla flow matching loss can be written as
FFM = 10- (w) (x,, (x,))1))
(18)
where := (1r) is the noise data, pi is the target data distribution, and wx(xe) is the ground truth data vector field at Tr and timestep The vanilla flow matching objective is generally intractable, as do not know the ground truth vector fields. Lipman et al. (2022) demonstrates key observation that, when considering the conditional probability path ( [1) conditioned on the target data, the conditional vector field (x( (211 can be calculated analytically while sharing the same gradient with the vanilla flow matching objective. The conditional flow matching objective can be written as
fCrM (0.0U00.11. (z) [ll uo (Ie,, [=<=((=1)11))
(19)
Following Chen & Lipman (2023), when further conditioned on the noise xo po(z), the CFM loss can be reparameterized into simple form:
LCRM = Bxw.:/j......oorreeeee...... (1l (a1 xo)1 x =1= (1- 1)20 + tz (20)
The linear interpolation above corresoonds to the aptinal-transpor probability path (OT-path) in Lipman et al. (2022), which enjoys additional theoretical benefits of straighter vector fields over other options, including the variance-preserving path (Ho et al., 2020) or the variance-exploding path (Song & Ermon, 2019). During sampling, the flow ODE in Equation 17 is solved with the learned vector field obtain the final generation during sampling. The deterministic flow ODE makes it potentially easier to guide than the stochastic Langevin dynamics or stochastic differential equations (SDEs) in the diffusion formulation.
C.2 FUNCTIONAL FLOW MATCHING
Functional flow matching (FFM) (Kerrigan et al., 2023) further extends the conditional flow match- ing framework to modeling functional data - intrinsically continuous data like PDE solutions. As mentioned in the original work, the major challenge of extending the CFM framework lies in the fact that probability densities are ill-defined over the dimensoail Hilbert space of contin- uous functions. FFM proposed to generalize the idea of flow matching to define path of measures
19

---

Published as a conference paper at ICLR 2025
using measure-theore theoreeic formulations. Specifically, consider real separable Hilbert space U of functions u : X' + R equipped with the Borel o-algebra B(U). FFM learns a time-dependent vector field operator v, .4 [0, 1] + U that defines time-dependent diffeomorphism 1) :u x [0, 1] - u called the flow via the differential equation
(u) = (v(r)())), Ug ~ po(u),
(21)
where Mo is some fixed noise measure from which random continuous functions be sampled.
Similar to CFM, the flow 2e induces a push-forward measure Ae := (16))= for generative model- ing. FFM demonstrates that the conditional formulation in CFM to deduce a tractable flow matching objective can also be adapted for the path of measures under some regularity conditions. In this way, also relying on the optimal transport path of measures, the FFM objective shares similar format as the CFM loss in Equation 20:
LIFM = (P.0UU/0.11,0 (o mamp pa (w) (ll,, () - (w1 mo))l?), 4r = (1 f)ug + tw1, (22)
where Wt == ( ) is the interpolation along the conditional path of measures and Yn is the larget data measure. The norm is the standard -n-m for square-integrable functions. Similarly, sampling for FFM can be thought of as solving the flow ODE with the learned vector field. We demonstrate the Euler method for FFM sampling in Algorithm 1.
We also noted that Lim et al. (2023) proposed the diffusion denoising operator (DDO) as an exten- sion of diffusion models to function spaces. DDO relies the non-trivial extension of Gaussian measures on function spaces. Compared to DDO, FFM has a more concise mathematical formula- tion and better empirical generation results. Therefore,
FFM as our generative prior.
C.3 PRE-TRAINING FOR FFM
We follow Kerrigan et al. (2023) to use Fourier Neural Operator (FNO) (Li et al., 2020a) as the vector field operator parameterization with additional relative coordinates and sinusoidal time embeddings (Vaswani, 2017) concatenated to the noised function as inputs. Following FFM, for all 2D data (ID PDE with a time dimensional or Darey flow), the prior noises are sampled from the 2D Gaussian process with a Matérn kernel with a kernel length of 0.001 and kernel variance of 1. For 3D) data (2D) NS equation), sampling from the Matérn kernel prohibitively expensive. Though Kerrigan et al. (2023) has indicated that the white noise does not meet the regularity requirement from mathematical considerations, we empirically found that such white noise prior performed well enough. Therefore, for 3D data, we always sample from the standard white noise.
For 2D data, we use four-layer FNO with a frequency cutoff of 32 x 32, time embedding channel of 32, hidden channel of 64, and projection dimension of 256, which gives total of 17.9M trainable parameters. All FFMs are trained on single NVIDIA A100 GPU with a batch size of 256, initial learning rate of 3 10-4, and 20k iterations (approximately 1000 epochs for 5k training dataset).
For 3D data, use two-layer FNO with frequency cutoff of 16 x 16 x 16, time embedding channel of 16, a hidden channel of 32, and a projection dimension of 256, which gives a total of 9.46M trainable parameters for efficiency concern.. This model is trained on NVIDIA A100 GPUs with a batch size of 24 (per GPU) for approximately total number of 2M iterations (or 5000 epochs) with an initial learning rate of 3 x 10-4
It is worth noting that our proposed ECI sampling framework, as a unified zero-shot approach for guiding pre-trained flow-matching models, does not require conditional training like the conditional FFM baseline. Therefore, we only need to train separate model for each different PDE system, but not for each different constraint. For example, the two generative tasks and the regression task on the Stokes equation are based on the same pre-trained FFM on the unconstrained Stokes problem dataset.
To evaluate the pre-training quality, we provide MMSE and SMSE between the generation by the pre-trained unconstrained FFM model and the test dataset in Table as the evaluation metries. It can be seen that most FFMs can achieve a decent approximation of the prior unconstrained distribution over the solution set with small MMSEs and SMSEs. For the 2D Darcy flow, the boundary condition has more significant impact on the final solution all pixel values should be shifted by the same value, leading to larger errors in the distributional properties. The generation, on the other hand, does seem reasonable.
20

---

Published as a conference paper at ICLR 2025
Mean and standard deviation MSEs between the generation by the pre-trained uncon- strained FFM model and the test dataset.
Table 8:
Dataset
Stokes Problem
Heat Equation
Darcy Flow
NS Equation
MMSE / 10-2
0.045
0.130
5,421 4.121
0.992
SMSE/ / 10-2
0.039
0.198
0.371
We also provide the generative metrics of the unconstrained pre-trained FFM without any guidance as a sanity check in Table 3. It can be demonstrated that there is indeed a shift in the distribution from the generative prior learned by the unconstrained FFM, as many errors significant. Our proposed ECI sampling can successfully capture such a shift in the distribution with significantly lower MSEs and closer resemblance the ground truth constrained distribution.
D EXPERIMENTAL SETUP
In this section, we provide further details regarding the experimental setups and evaluation metrics. We also give brief introduction the baselines used in our experiments.
D.1 SAMPLING SETUP
For the baseline models, we use the adaptive Dopris ODE solver (Dormand Prince, 1980) for CondFFM and ProbConserv. For ECI and DiffusionPDE, we used 200 Euler steps for 2D datasets and 100 steps for 3D datasets. For D-Flow, due to the limitation of GPU memory, we use 100 Euler steps for 2D datasets and 20 Fuler steps for 3D datasets. The hyperparameters for sampling summarized in Table 4.
Generative models have intrinsic stochasticity, the sampling procedure has randomness in the prior noise distribution. To provide a more robust evaluation of the generated results, we always generate 512 samples for every 2D dataset and baseline for computing all evaluation metries and the sampling time in Table 4. 100 samples are generated for all the NS equation tasks (3D dataset). In this way, we aim to minimize the impact of randomness in the generation and provide better estimation of the statistics and metrics.
D.2 EVALUATION METRICS
As we have discussed in Section 4, for generative tasks, calculate distributional properties like the mean squared errors mean and standard deviation (MMSE and SMSE) to measure the errors from the ground truth statistics. We also use the constraint error (CE) measure the violation of constraints.
The Fréchet Poseidon distance (FPD) is inspired by the metric of Fréchet Inception distance (FID) widely used in the image generation domain to evaluate the generative quality using the pre-trained Inception V3 model (Szegedy al., 2016). By comparing the Fréchet distance between the hidden activations of the data captured by a pre-trained discriminative model, such a score provides mea- surement of the similarity between data distributions. A smaller Fréchet distance the ground truth data distribution indicates a closer resemblance to the ground truth data and, thus, a better generation quality. Inspired by FID, we propose to use the pre-trained PDE foundation model Poseidon (Herde et al.. 2024) to generate the hidden activations. We used the base version of Poseidon with 157. 85.88M model parameters. Poseidon takes the initial frame, z-velocity field, y-velocity field, pressure field, and timestep inputs and outputs predictions of the frame after the timestep. We always zeroed out the zevelocity field, y-velocity field, and pressure field the inputs. For 2D data, the generated solutions are directly fed into Poseidon as the initial frames, and the timestep is always fixed to O. For 3D data, each time frame of the solution trajectories fed separately into Poseidon. The Fréchet distances are calculated separately for different time frames and averaged to obtain the final FPD score. As we want to leverage the pre- trained Poseidon to extract high-level representations of the generated solutions, extract the last hidden activations of the encoder for Fréchet distance caleu-
21

---

Published as a conference paper at ICLR 2025
lation. The hidden activation size for the Poseidon base model is 784 x 4 x 4, which we mean-pool into a 784-dimensional vector for FPD calculation.
For regression tasks, CE for constraint violation. Following Hansen et al. (2023), the log-likelihoog (LL) is calculated pointwise as
directly calculate the MSE with respect to the ground truth and the
LL = logp (u) (N (21)) - 4i -
log ō log (2T),
(23)
272
2
where u1 is the ground truth data and fr, â are the generation mean and standard deviation, respec- tively.
D.3 CONSTRAINT ENFORCEMENT
For value constraints, we can easily enforce the constraint as deseribed in Algorithm 4. Specifi- cally, the initial condition is defined as the first column values after discretization, and the boundary condition is defined the first row values. For 2D Darcy flow, the boundary condition is defined as the outer-most pixels in all four directions, thus leading to an additional resolution along each dimension. For conservation laws, the integral is approximated with the Riemann sum across the constrained region with uniform discretization along each dimension. The constraint error is also calculated using the same Riemann sum.
In the uncertainty quantification task, we follow Hansen et al. (2023) to experiment with the zero- shot superresolution setting. While the pre-trained FFM is trained on spatiotemporal resolution of 100 x 100, during constrained generation, resolution of 200 x 200 is instead enforced. Following Hansen et al. (2023), 100 random context points (0.25%) of the ground truth solution are provided as the value constraint, with additional conservation laws given in Appendix B.
D.4 BASELINE MODELS
The conditional FFM model (Kerrigan et al., 2023) assumes the family of constraints known a priori. During training, the constraints are fed to the FFM model as additional informa- tion. In this way. CondFFM can only handle value-based constraints but not conservation laws. In practice, we copy the input function and set the pixels in the constrained region to the corresponding values and the pixels in the unconstrained region to zero. This condition is then concatenated to the input channel-wist as additional information. Similarly, this condition is fed into the model during the constrained sampling.
CondFFM.
We use an almost identical FNO-based encoder for CondFFM except for one additional channel for the condition. For fair comparison, the training hyperparameters the same the corresponding generative FFM. As CondFFM has additional information on the constraint in both training and sampling, we naturally expect it to exhibit better performance than other zero-shot (training-free baselines. As a drawback for all models that require training adjustment, different CondFFM models need be trained separately for different constraints even on the same PDE system (e.g., IC and BC for the Stokes problem), making it less flexible to different constraints.
ANP. The Attentive Neural Process (ANP) (Kim et al., 2019) models the conditional distribution of function at a specific set of target points {x, her given another set of context points {a; hec- The ANP uses the attention mechanism and variational approach to maximize the variational lower bound for the data likelihood. As it is special case of neural processes, the ANP can take both target and context points of arbitrary sizes.
In practice, we use 100 random context points and 1000 random target points during training. The prior model is trained on the same amount of data as the corresponding FFM models. During sampling, the context points are fixed be the values in the constraints, and all discretized spa- tiotemporal coordinates are used as the target points in the generation. ANP model is also zero-shot and can be applied different value constraints.
ProbConserv. ProbConsery (Hansen et al., 2023) is general black-box posterior sampling method specifically designed for the exact satisfaction of various PDE constraints. ProbConsery adopts the gradient-fres approach that directly projects the generation to the corresponding solution
22

---

Published as a conference paper at ICLR 2025
space using the least squares method. When applied to value constraints, ProbConserv (the version that ensures the exact satisfaction of constraints) directly modifies the corresponding pixels to the constraint values, often leaving noticeable artifacts between the constrained and the unconstrained regions. For a fair comparison, the ProbConserv model in our experiments built upon the same pre-trained FFM model. As ProbConserv directly operates the final generation, it does not offer fine-grained control over the intermediate steps of the iterative sampling process of flow-matching models.
DiffusionPDE. DiffusionPDE (Huang et al., 2024) uses gradient guidance from the constraint loss and the PINN loss (Raissi ct al., 2019; Li et al., 2024) to modify the vector field at each step to minimize these losses while maintaining the physical structure of the PDE system, This gradient- based approach is also representative of many models for inverse problems, specifically, the diffusion posterior sampling approach (Chung et al., 2022). As our PDE system may have variations in the PDE parameters, the PINN loss is not applicable. Thus, we only use the constraint loss to guide the model each time step. In practice, follow the original paper to test several guidance strengths from 102 to 10 and choose the value with the lowest MSE.
D-Flow. D-Flow (Ben-Hamu et al., 2024) guides the final generation via optimizing the initial noise uo rather than modifying the vector field. Assuming the constraint loss L(ū,) is differen- tiable, D-Flow backpropagates the gradient of loss through the ODE solver to the initial noise as u. Lú1), hoping such a direct optimization can reduce the final constraint loss. Such direct gradient- based approach is extremely memory-demanding and time-consuming, as the simplest Eu- ler solver requires about 100 steps to achieve decent generation results. We follow the original work to use the LBEGS optimizer with a learning rate of and 20 maximum iterations.
FNO. The Fourier Neural Operator (FNO) (Li al., 2020a) used as the additional baseline in regression tasks of neural operator learning. We use an almost identical architecture as the en- coder for FFM, except for the missing time embedding part. Similar CondFFM, the pixels in the constrained region are set to the constraint values, while other pixels are set to zero. For fair comparison, the FNO model tries to predict all the pixel values for the solution domain, including the constrained ones already given the input. As the regression models are easier overfit, apply carly-stopping techniques to choose the best model checkpoint for evaluation. The FNO was trained for 10k iterations for the Stokes problem and 500k iterations for the NS equation when the valid loss plateaued.
E ADDITIONAL RESULTS AND VISUALIZATIONS
In this section, we provide additional experimental results and ablation studies to further demonstrate the effectiveness of our proposed ECI sampling scheme. Additional visualizations of the generations and evaluation metrics are also provided in this section.
E.1 GENERATIVE TASKS
E.1.1 ADDITIONAL VISUALIZATION
mean
06
std
CondFTM
ProbConserv DiffusionPDD
D-Flow
Figure 8: Generation mean and standard deviation for the Stokes problem with fixed.
23

---

Published as a conference paper at ICLR 2025
mean
std
GT
CondFFM
BCI
ANP
ProbConserv DiffusionPDE
D-Flow
Figure Generation mean and standard deviation for the Stokes problem with BC fixed.
mean
mean
error
std error
std
GT
CondFFM
ECI
ANP
ProbConsery DiffusionPDE
D-Flow
Figure 10: Generation mean and standard deviation errors for the Stokes problem with BC fixed.
In addition to the error plots Figure 2 and 3, we further provide visualizations of the generation statistics in Figure (Stokes problem with IC fixed), 9 (Stokes problem with BC fixed), 11 (heat equation), 13 (Darcy flow), and 15 (2D NS equation).
Errors for generation statistics for different systems also provided in Figure 10 (Stokes problem with BC fixed), 12 (heat equation), and 14 (Darcy flow). It can be demonstrated more clearly that gradient-based methods like DiffusionPDE and D-Flow had more artifacts around the boundary and did not satisfy the constraint exactly. For ProbConserv, though its gradient-free approach indeed ensured the exact satisfaction the constraint, the non-iterative control of the prior flow model left noticeable artifacts between the constrained and the unconstrained regions. Similar trends can also be observed in other PDE systems.
mean
std
GT
CondFFM
ECI
ANP
ProhConsery DiffusionPDE
D-Flow
Figure 11: Generation mean and standard deviation for the heat equation with IC fixed.
We also noticed that D-Flow tended to lead high-variance generations. This is probably because large number of Euler steps leads to an exceedingly complex dynamic that is difficult to optimize. Also, note that directly taking the gradient with respect to the initial noise may break the initial structure sampled from the Matérn kernel and may lead to initial noises never seen by the model during training. These two reasons may also account for the better performance of the D-Flow for the NS equation, in which we used fewer Euler steps and the standard white noises as the prior noise functions.
24

---

Published as a conference paper at ICLR 2025
mean
mean
std error
std
GT
CondFFM
ECI
ANP
ProbConsery DiffusionPDE
D-Flow
Figure 12: Generation mean and standard deviation emrors for the heat equation with IC fixed.
mean
std
CondFFM
ANP
ProbConsery DiffusionPDE
D-Flow
Generation mean and standard deviation the Darcy flow with BC fixed. Note that scaling range is different and large variance values cropped for better comparison.
Figure 13:
E.2 STOKES PROBLEM WITH DIFFERENT CONSTRAINTS
We provide additional generative evaluation the Stokes problem with various ICs/BCs in Table in addition to the results in Table 3. Recall that during unconstrained pre-training, we varied the [ 8), k U22 20), and in Table we fixed & 5 with the results Table 9, it can be clearly demonstrated that our proposed ECI sampling can achieve the baselines in zero-shot manner with the flexibility all different IC settings l: 5, 10, 15 and BC settings w = 4, 6,8.
PDE parameter
6. Combined
consistent performance improvement
E.2.1 GENERATIVE SUPERRESOLUTION
In Table 10, provide the evaluation results under the zero-shot superresolution setting, where the generation resolution 200 x 200, four times larger than the training set. is interesting see that the conditional FFM is the most sensitive method in such zero-shot setting, probably due to the hard-coded connection of the masked input as the condition.
mean
mean
std
std eгror
GT
CondFFM
ECI
ANP
ProbConserv DiffusionPDE
D-Flow
Figure 14: Generation mean and standard deviation errors for the Darcy flow with BC fixed. Note that large error values are cropped for better comparison.
25

---

Published as a conference paper at ICLR 2025
mean
-2
1.0
std
0.4
GT
CondFFM
ECI
ProbConserv DiffussonPDE
D-Flow
Figure 15: Generation mean and standard deviation for the 2D NS equation with IC fixed.
Table 9: Generative metrics on the Stokes problem with different constraints. The best results for
zero-shot methods (CondFFM is not zero-shot) are highlighted in bold.
Dataset
Metric
BCI
ANP
ProbConserv
DiffusionPDE
D-Flow
CondFFM
MMSE/ 10 2
0.113
4.616
1.792
0.512
0.774
0,307
Stokes IC k 10
SMSE 10
0.124
12.400
0.356
0.301
59.377
0.060
CE. 10-2
13.032
0
0.068
3.004
0.011
FPD
0.156
6.694
3.619
0.340
11.013
0.454
MMSE/ 10-2
0.190
11,148
0.797
0.216
0.088
0.125
Stokes IC ::
SMSE 10-2
0.183
3.354
2.751
0.470
12.063
0.010
CE/ 10
18.278
0
0.194
13.149
0.010
FPD
0.105
6,509
6,425
0.832
4.745
0.020
MMSE/ / 10 10-2
0.676
9.273
3.034
1.488
1.728
0.031
Stokes BC w 4
SMSE 10-2
0.378
2.777
2.462
2.278
6.343
0.081
CE/ 10-2
115.335
0
39.99
14.553
0.012
FPD
1.671
5.309
1.686
3.954
10.043
0.071
MMSE 10-2
0.042
4.115
9,008
5.794
3.801
0.028
Stokes BC w
SMSE 10
0.026
2.462
2.710
3.288
22.376
0.010
CE/ 10-2
0
73.916
0
160.397
51.056
0.017
FPD
0.109
6.663
0.843
9.356
13.202
0.111
E.2.2 ADDITIONAL METRICS FOR NS EQUATION
We provide the per-frame FPD scores for the NS equation generation results in Figure 16. We noted that ECI sampling tended to have more accurate generations for frames close the initial condition, with the FPD scores increasing with the time frames. In gradient-basee models of DiffusionPDE and D-Flow, we also noticed a peak in the FPD scores an early stage, which should have been
Table 10: Generative metrics on the Stokes problem with the zero-shot superresolution setting to 200 200. The best results are highlighted in bold.
Dataset
Metric
ECI
ANP
ProbConserv
DiffusionPDE
D-Flow
CondFFM
MMSE/ 10-2
0.274
14.120
4.338
5.109
4.242
SMSE 10
0.515
8.720
2.933
4.167
4.152
Stokes IC
00M
CE / 10
3.563
5.392
12.508
FPD
2.647
26.965
15.977
19.654
24.798
MMSE/ -
0.863
2.819
5.223
1.610
4.669
SMSE / 10-2
0.418
2,870
1.491
1.768
2,706
Stokes BC
00M
CE / 10-2
29.168
0
0.068
21.567
FPD
1.903
5.282
2.021
3.830
9.504
26

---

Published as a conference paper at ICLR 2025
easier for the model. As expected, CondFFM outperformed all zero-shot guidance models large margin such a complex task for almost all frames.
Per Frame FPD Score
10
BCI
CondFFM
ANP
ProbCensery
DiffusioPPDE
D-Flew
FPD
Frame
Figure 16: Per frame FPD score for the NS equation with IC fixed.
E.3 REGRESSION TASKS
E.3.1 UNCERTAINTY QUANTIFICATION
For the uncertainty quantification tasks, the specific snapshot at the test time as described in Ap- pendix B (following the same settings in (Hansen al., 2023)) is plotted in Figure 17, with the gen- eration mean and 3 times of the generation standard deviation as the confidence interval. Our results share the similar trend described in (Hansen al., 2023), in which they also observed more uncer- tainty with the increasing difficulty from the heat equation, PME, to the Stefan problem. Specifically, the solution in the Stefan problem has a shock in the function value, which leads higher uncer- tainty around the shock position. For other linear parts of the solution, our ECI sampling generation able achieve quite good predictions with low variance and low uncertainty.
Taut Egarton
PMF
Goald Tral
Crnral
022
-36
322
07.55
3.2%
Figure 17: Uncertainty quantification visualization results for three instances of the Generalized Porous Medium Equation (GPME). The 30 confidence intervals are plotted.
E.3.2 NEURAL OPERATOR LEARNING
For NO learning tasks, we provide additional visualizations the generation statistics and the refer- ence FNO predictions in Figure 18 (Stokes problem) and 19 (NS equation). It can be demonstrated that, though not directly trained on such regression tasks, our ECI sampling naturally produces generation results with low variance in such scenarios. In comparison, other models often have variance scale similar to those the generative tasks, making them less suitable for such zero-shot regression tasks.
27

---

Published as a conference paper at ICLR 2025
mean
GT
0.00
4.6
std
0.00
0.50
FNO
ECI
ANP
ProbConserv DiffusionPDE
D-Flow
Figure 18: Regression mean and standard deviation for the Stokes problem with both IC and BC fixed. Note that the scaling range for the standard deviation is different from that in the generative task.
mean
GT
std
FNO
ECI
ProbConserv
DiffusionPDE
D-Flow
Figure 19: Regression mean and standard deviation for the NS equation with the first 15 frames as the constraint. Note that the scaling range for the standard deviation is different from that in the generative task.
E.4 MORE ABLATION STUDIES
E.4.1 STOKES PROBLEM FPD
thorough investigation into the impact of mixing iterations M and re-sampling interval A on the final performance, we further calculate the FPD scores in Figure 20 for the two settings for the Stokes problem addition Figure The FPD scores more consistent than the MMSE and SMSE. For the Stokes IC problem with higher variance, generations with no re- sampling work the best, with the FPD scores peaking at 10 mixing iterations. For the Stokes BC problem with lower variance, the best performance combination of hyperparameters has diagonal pattern with more mixing iterations pairing best with a smaller re-sampling interval. Such results are also consistent with our heuristic suggestions on choosing the best M and R for different tasks with different variances.
To provide
E.4.2 PRIOR GENERATIVE MODEL.
We further provide the ablation studies on the impact the quality of the prior generative model. Intuitively, stronger generative prior that better captures the underlying information of the PDE system can be easier to guide. The results are demonstrated in Table 11, with checkpoints from 5k, 10k, and 20k training iterations. A clear trend of better -endede-enerttion performance can be observed in both and BC constraints. In this way, the quality of FFM indeed plays an important role in ECI sampling framework.
28

---

Published as a conference paper at ICLR 2025
FPD Score Stokes IC
FPD Score for Stokes BC
16
1-
-8
14
Ke-sampling Interval
Re-sampling Interval
12
5
10
10
6
4
None
None
10
2
20
Mixing Iteration
Mixing Iteration
Figure 20: FPD score of ECI sampling with different mixing iterations and re-sampling intervals in the Stokes problem with IC (left) or BC (right) fixed.
Table II: Generative metrics the Stokes problem with different constraints using the different FFM checkpoints with different training iterations.
Dataset
Stokes IC
Stokes BC
#iter
10k
15k
20k
10k
15k
20k
MMSE/ 10-2
1.273
0.601
0.090
0.888
0.011
0.005
SMSE 10-2
1.693
0.177
0.127
0.868
0.015
0.003
FPD
1.863
1.031
0,076
3.101
0.009
0.010
E.4.3 AMOUNT CONSTRAINT INFORMATION
We also provide additional ablation studies on ECI sampling behavior as a regression model. Specif- ically, we explore the regression losses when more information is provided in the constraint. Intu- itively, serving as a deterministic regression model, the ECI generation should exhibit fewer errors and less variance with more information provided. In experimental setup, fixed the number of mixing iterations to and the re-sampling interval to 1. To control the information the con- straint, varied the number of initial frames fed into the model the value constraint from 5 to 25. As our ECI sampling guarantees the exact satisfaction of the constraint, we instead calculated the MSE and standard deviation MSE based on the unconstrained region for fair comparison between different settings. The losses are plotted in Figure 21, in which a clear and almost linear trend the log scale can be distinguished for both metries increases. Similarly, the visualization of the mean eπтог and the logarithm of the standard deviation are provided in Figure 22. The model made more accurate predictions with less uncertainty as the number of constrained frames increased. The regions with large mean errors also matched those with large variance (uncertainty). this sense, our ECI sampling provides a bonus to transform the pre-trained unconditional FFM model into decent zero-shot uncertainty quantification model with potentially wider applications.
the number of constrained frames
29

---

Published as a conference paper at ICLR 2025
MMSE vs. SMSE for NS Equation
frames
10 2
SMSE
10-3
10 frames
10
frames
10-5
20 frames
25 frames
10-4
10-3
10-2
MMSE
Figure 21: Regression MMSE and SMSE for the NS equation with different numbers of frames fixed as the constraints. Notice clear and almost log-linear trend for the decrease in both metrics as the number of constrained frames increases.
0.
0.1
mean error
0.0
-0.1
-2
log std
5 frames
10 frames
15 frames
20 frames
25 frames
Figure 22: Regression error of the mean and logarithm of the standard deviation for the NS equation with different numbers frames fixed as the constraints. We use the same scale for plotting different settings.
30


Reviews:
[{"review_id": "W21GLO5vPn", "reviewer": "Reviewer_MQJx", "summary": "In this work, the authors propose ECI for gradient-free generation for constraint systems. ECI is zero-shot and doesn't need tuning the model. It works through alternating between extrapolation (E), correction (C), and interpolation (I) in each sample step. The empirical results show the model achieves competitive performance across various PDE systems compared with baseline models.", "questions": "1. How does proposed ECI compared with FFM without ECI in experiments? It would better show the improvement from ECI in sampling. \n2. In section 3.3, the authors mention resampling noise during sampling. Can the authors explain why this helps better generative results?\n3. The model assumes a perfectly trained FFM, however there is inevitable error in the trained generative model. How does such error affect the performance? Will such error break the constraint in sampling?\n4. Empirical results are on low-resolution problems. Can the authors comment on how ECI work at larger scale?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The work is well-motivated, it's important in scientific applications to have constraint generation\n2. The proposed method is effective and does have constrained generation. Also, the method is zero-shot which is a big benefit for constraint sampling method. Although it's based on functional FM not FM in general\n3. The experimental results show advantages over other baseline models.", "weaknesses": "1. The writing is also a bit confusing. Actually comments from reviewer KM5e help me better understand the algorithm 2. \n2. The authors are recommended to better formulate the contribution of this work. In particular, the model works in functional space and applies projection to constraint spaces to guarantee hard-constraint met. \n3. The authors mention supply chain optimization in abstract and introduction as a hard-constrained system. However, it lacks experiments on such problems. It would be better to showcase some applications beyond PDE learning (which I believe is broad).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UQ183j1RRl", "reviewer": "Reviewer_bR4B", "summary": "The paper introduces ECI sampling, a framework that adapts pre-trained, unconstrained generative models to exactly satisfy hard constraints in a zero-shot manner without requiring gradient computations or fine-tuning. This framework alternates between extrapolation, correction, and interpolation stages during each iterative sampling step to integrate constraint information while preserving the validity of generated outputs. Empirical results demonstrate that ECI sampling strictly adheres to physical constraints across various PDE systems and outperforms baseline approaches in zero-shot constrained generative and regression tasks.", "questions": "1. This is the most critical question. How does ECI guarantee that the achieved solution satisfies the PDE system? In computer vision tasks where differences from ground truth images are permissible.  In this work, boundary conditions are guaranteed, but how does ECI ensure that the generated data truly represents solutions to the PDE? More specifically, why the linear interpolation between two PDE solutions is still a valid PDE solution?\n\n2. Can you elaborate more on the stochasticity of generated solutions, as is discussed in Sec 3.3. Does more stocahsticity mean the generated distribution is distribution-wise closer to the uniform distribution over the space of constrained PDE solutions?\n3. For the extrapolation step, will the algorithm performance be better if more a refined scheme is used? For example, mutiple step extrapolation when time $t$ is far from $1$ and one step when $t$ is close to $1$. From my understanding, the quality of predicted $\\hat u_1$ is crucial to the algorithm performance. An additional ablation study could be performed here.\n4. Can you comment on the computational efficiency of ECI compared with existing approaches?\n\nThe current score assumes a positive answer to the first question. My score will be adjusted accordingly based on the authors' response.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is easy to follow with high readability. The problem setting is well motivated and important.\n2. The paper provides quite comphrenseive numerical results with comparison to relevant benchmarks. The extension to regression setting is also quite impressive.\n3. The paper also provides a quite detailed ablation study on the choice of algorithm hyperparameters.", "weaknesses": "*1. Clarification of Problem Setup*\n\nThe problem setup is not sufficiently explained, which may lead to confusion. Although the authors repeatedly emphasize that ECI is intended for the generative modeling of constrained PDE solutions, it remains unclear what distribution of constrained solutions ECI aims to recover. Specifically, is it targeting a uniform distribution over the space of constrained solutions $\\mathcal{U}_{|\\mathcal{G}}$? The authors should invest more effort in formally clarifying the problem setup.\n\n*2. Discussion of ECI’s Advantages and Comparisons*\n\nThe paper appears to lack a section that discusses the benefits and potential sources of ECI’s superior performance, as well as comparisons with other existing approaches mentioned in the work. Including such a discussion would help readers better understand both the relevant literature and the proposed method.\n\n*3.Validation of PDE System Satisfaction*\n\nWhile ECI ensures the exact satisfaction of the constrained operator by applying a correction operator after each one-step extrapolation, it is unclear why the solution obtained at time 0 still satisfies the PDE system $\\mathcal{F}_{\\phi} u(x) = 0$. In Algorithm 2, the next-step solution is generated by forward noising starting from the predicted solution at time 1. Does this noising process preserve the satisfaction of the PDE system? Additionally, is the linear interpolation of two PDE solutions a valid PDE solution, implying that the PDE system is linear? Further explanations are needed to clarify these points.\n\n*4.Missing References on Gradient-Free Guidance*\n\nSome references on gradient-free guidance are missing, such as [1] and [2]. It appears that the proposed method is related to [2], involving functional FM and additional projection steps. Providing more clarification on these connections would be appreciated.\n\n[1]Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli S Sastry, Siddharth Gururani, Sageev Oore, and Yisong Yue. Symbolic music generation with nondifferentiable rule guided diffusion. arXiv preprint arXiv:2402.14285, 2024.\n\n[2]Chung, Hyungjin, et al. \"Diffusion posterior sampling for general noisy inverse problems.\" arXiv preprint arXiv:2209.14687 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KtSYaBkHu7", "reviewer": "Reviewer_rhkG", "summary": "This paper tackles the problem of conditional generation of physical systems with physical constraints (such as physical conservation laws). They propose “a framework for adapting pretrained zero-shot, unconstrainted generative models to satisfy exacts constraints” in a computationally efficient manner. Subsequently, the apply the framework to Partial Differential Equation (PDE) systems and demonstrate empirical results on a variety of 1D and 2D physical problems.", "questions": "* Is it possible to address the conservation laws in a more principled way? For example, minimal local perturbations to achieve the desired conservation (rather than a global correction), or monitoring which cells are causing the violation and thus correct the violations in a more targeted manner?\n* Can we apply ECI to a 3D problem/higher resolution/more complex problem to see if the correction step is much worse? It would also be good to see an experiment at a much higher resolution and with many more time-steps.\n* Why are we using Frechet Poseidon Distance? For these problems, we know the physics and so would it not be better to simply compute the ground truth for evaluation?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 1, "strengths": "* The method does seem to operate as advertised. The generated results have 0 constraint error, compared to > 0 constraint error of other soft-constrained methods. At the very least, the claim that ECI guarantees hard-constraint satisfaction is well supported\n* The ablation study that shows the effect of the resampling intervals and the mixing iterations is quite good.\n* Reporting both generation quality as well as runtime for various experiments helps objectively evaluate the proposed methodology.\n* Overall, the presentation of the paper if quite good. The problem statement and context are very clearly outlined at the start of the paper. The advantages of the methodology are plain described, and the research problem is clearly state and motivated within existing literature. Thorough and well written description of experiments and results. Nice figures though could benefit from more detailed captions describing what we are looking at.", "weaknesses": "__Theoretical Concerns__\n* While the authors do a great job of citing the existing context for SciML applications and flow-matching models, they do not cite the large existing bodies of work related to constraint satisfaction/constrained optimization. In particular, many methods share many similarities with ECI in that they perform unconditional update steps, and then project back into the feasibility region. As such, the method lacks theoretical motivation and context of existing works.\n\n__Experimental Concerns__\n* Many of the experiments performed are quite simple, with either low spatial resolution, low dimensionality, or a small number of time-steps. In order to validate that the correction step is not overly destructive, it would be beneficial to perform experiments with higher resolution, and a greater number of time steps. Additionally, there is a concern that ECI performs poorly in complex settings, as it is out-performed by one model or another on all 2D experiments (in terms of MMSE), and thus it would be good to see a 3D fluid setting to evaluate how ECI performs in higher-dimension problems.\n* Additionally, fluid flow can exhibit quite complex behaviour. It is not clear why an increased variance in the solutions is a sign of a bad solution, as opposed to the system exhibiting complex dynamics (such as turbulence). Thus, the justification for SMSE as an evaluation metric is not quite clear.\n\n__Contribution Concerns__\n* First, the concept of perform Extrapolate-Correct-Interpolate steps in order to solve constraint satisfaction is not a novel idea. In fact, and is well studied in constraint satisfaction settings. For example, projected gradient descent functions very similarly by first taking a step to minimize the objective (Extrapolate) and then taking a step to project back into the feasible region (Correct).\n* Second, as mentioned in the paper, the extrapolate step assumes a pre-trained generative model (and thus is not a novel contribution), the interpolate step simply adds noise to the sample (and thus no novel contribution contribution), so the only real contribution is the correction step. However, the correction step is extremely trivial. In the case of the boundary conditions, the method simply sets the boundary values to the desired constraint. In the case of conservation, the method subtracts the deviation from the conserved amount from all cells equally. Both approaches are extremely simple, and don’t really capture realistic physics. We can construct pathological examples where this type of correction leads to systems which satisfy the constraints but are completely implausible and useless for any sort of SciML application. In fact, this correction step is more or less the most trivial way to project back into the feasibility region.\n* Furthermore, to address new problems, you need to implement your own correction steps, which may be difficult, computationally intensive or simply intractable for more interesting problems. For example, there is no treatment for conservation of a vector quantity (i.e. momentum, angular momentum), which would occur in many settings of interest.\n* Additionally, the paper deals with a pretty niche problem. Notably hard-constrained generation using flow-matching models for physical system simulation. However, when digging into the ECI framework, we note that the E and I steps come from pre-existing generative models, and the C step is incredibly simple. Thus, the overall contribution is very limited in scope of application, as being quite simple from a theoretical and practical point of view.\n* Overall, given a pre-trained flow matching model, ECI is a somewhat simple extension that allows for hard constraints to be satisfied. This is guaranteed by the correction step, which is quite simplistic. In particular, we would expect that such coase correction steps would lead to poor solution quality. In fact, from the experimental results, we can see that ECI performs well in 1D settings, but under performs compared to other models in the more complex 2D problems. In fact, we would expect it to perform even worse with high spatio-temporal resolution or in 3-dimensional problems.\n\n\nOverall, the paper is very well written and presented, but the ECI framework (in particular the correction step) do not present a significant theoretical or practical contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wyo1UxxXkk", "reviewer": "Reviewer_KM5e", "summary": "This work proposes an ECI sampling framework that strictly satisfies hard constraints such as IC/BCs and conservation laws through extrapolation, correction, and interpolation stages. The method has gradient-free generation and zero-shot inference for parametric PDE and dynamical systems.", "questions": "For pre-trained operator networks, fusing BC/ICs to enforce hard constraints can indeed improve predictions at initial time and boundaries, but can ECI sampling directly correct or improve the prediction values in the domain? And how effective is this ability under zero-shot inference?\nIn addition, u_1 in line 2 of Algorithm 4 should be in brackets and needs to be corrected.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. There are many experimental results, sufficient comparison with baselines, and complete test metrics.\n2. The analysis presented in line 260 is basically confirmed by experiments: when the physical constraints are strong or contain a lot of information, the generation variance is small, and conversely the variance is large.\n3. Compared with Boundary enforcing Operator Network (BOON), this work implements hard constraints more concisely. The ECI sampling framework satisfy physical laws and system requirements strictly.", "weaknesses": "1. Although the results in Table 3 is excellent, Figures 6-8 may not clearly reflect the advantages of ECI sampling, and further explanation of the figures and settings is needed.\n2. This method can achieve zero-shot performance and exact satisfaction of different IC/BCs. However, Figures 5-7 are shown with IC or BC fixed. It is necessary to supplement the visualization of final predictions under different ICs. For the Stokes problem in Appendix B.1, we recommend testing more k and w values to demonstrate the generalization and applicability of ECI sampling.\n3. There is a lack of an overall framework diagram, and the entire training and generating process is not clear or intuitive enough. The PDE problems considered in the experiment are limited to 2-3 dimensions. And those experimental descriptions, data formats and prediction tasks should have a clearer and unified presentation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rqAmvbW60D", "reviewer": "Reviewer_11X9", "summary": "The authors introduce a novel framework for adapting pre-trained, unconstrained\ngenerative models to exactly satisfy constraints of PDEs in a zero-shot manner.\nSampling which alternates between extrapolation (E), correction (C), and interpolation\n(I) stages during each iterative sampling step to ensure accurate integration of constraint information of PDEs\nis proposed \nand called ECI sampling.\nThe new method is evaluated against various existing approaches and for different PDEs with both value constraints \n(initial and boundary conditions) and conservation conditions.\nThe topic is interesting, but the paper requires a complete rewriting focusing on the main contribution and towards a sound mathematical notation.", "questions": "Minor remarks:\n\n- p. 3: PDE family is $\\mathcal F_\\phi$, but then $\\mathcal U_{|\\mathcal F}$ without $\\phi$ \n(skip the $x$ in $u(x)$ in definition of $\\mathcal U$ since you mean the function and not the function value);\nthe family could be $\\mathcal F:=\\{\\mathcal F_\\phi: \\phi \\in \\Phi\\}$\n\n- in the definition of $\\mathcal U_{|\\mathcal G}$ I am missing the domain of $x$ in $\\mathcal F_\\phi u(x) = 0$\n\n- what is the reason for the clumsy notation, why not calling it just $\\mathcal U_{\\mathcal G}$ and $\\mathcal U_{\\mathcal F}$\n\n- there are errors in English, e.g. and of p. 3 ''denote'' instead of ''denotes''\n\n- when the authors use $:=$ (push-forward measure) and when not (most of all other definitions)\n\n- formula (2) is not readable since $q$ is not defined here; later the authors try to explain the formula\n\n- put important formulas on an extra line\n\n- While Alg 1 is superfluous, how Alg 3 and 4 are embedded in Alg 2", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- interesting and relevant research question\n\n- extensive evaluation of the new ECI sampling method versus other zero-shot guidance models for both generative and regression tasks\nand experiments on various 1d/2D and one 3D PDE;\nhowever I am not an expert in all these PDE learning methods", "weaknesses": "The writing of the paper is  unsatisfactory, everything boils down to Alg. 2 together with the\nbeginning of p. 5 (Extrapolation, Correction, Interpolation), \nbut the authors \ntry to explain something in a non understandable, fuzzy way.\nIn particular, often the notation is not introduced when it is needed.\nFor example in formula (2) it is not clear how $q$ is defined.\n\n The paper could be considerable shortened or better some parts from the appendix could be moved to the main part\nto make the methods better understandable.\\\\\nProp 1 is straightforward and holds just by construction; it could go into the main part \n\n Alg. 2  is obscured by notational\noverload. Basically what is happening in Algorithm 2, line 6 is the\nfollowing:\n\n- drawing $u_1\\sim p_\\theta(u_1|u_t^{m-1})$ amounts to computing\n$$u_1=u_t^{m-1}+ (1-t)v_{t,\\theta}(u_t^{m-1})$$ In particular\n$p_\\theta(u_1|u_t^{m-1})$ is a delta measure resulting from the one step\napplication of the vector field.\n- then compute $\\hat{u}_1:=C(u_1,\\mathcal{G})$ in order to project to\na function $\\hat{u}_1$ that satisfies the constraints\n\n- sample from $q(u_t^m|\\hat{u}_1)$ by sampling from $u_0\\sim\\mu_0$ and\ncalculating $u_t^m=(1-t)u_0+t\\hat{u}_1$. This means that one assumes that\n$q(u_t^m|\\hat{u}_1)$ is distributed as $(1-t)U_0+ \\hat{u}_1$ where\n$U_0\\sim \\mu_0$.\n\nIn line 7 the only difference is that in the last step one computes\n$u_{t+1\\backslash N}=(1-(t+1\\backslash N))u_0+(t+1\\backslash N)\\hat{u}_1$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)