FUSE: Fast Unified Simulation and Estimation for PDEs

W1. Concerns about the novelty/applicability: We thank the reviewer for this important question. We will exemplify the applicability of the methodology by providing an example in PWP. The goal in this setting is to predict the output function $s$ given a finite-dimensional vector of parameters $\xi^*$. In problems such as personalized medicine, $\xi^*$ is not known a priori, as it represents quantities such as vascular resistance that cannot be measured non-invasively, and it needs to be inferred from available measurements $u$, e.g. PPG. This is a strongly ill-posed problem, thus a probabilistic approach needs to be considered. This is modelled by approximating a posterior probability $\rho(\xi^*|u)$. When the posterior probability is approximated, posterior samples are drawn and used to run a flow solver to estimate continuous quantities such as local vascular pressure $s$, which cannot be measured non-invasively, and uncertainty estimates. More generally, FUSE targets the problem of calibrating vector-valued parameters of PDEs, which is applicable to many areas in the sciences such as climate modeling, computational fluid dynamics, material science, wave scattering etc.

W2. Reasons for Unification under a Neural Operator framework: We understand the reviewer's confusion and thank them for asking for clarification. There are two strategies that can be considered for the inverse problem, which for the sake of brevity we are going to call "classical" and "ML-based". Given measurements $u$ the "classical", (e.g. MCMC) strategies guess the value of $\xi^*$, and iteratively correct their guess. For complex problems such as the ones presented here, these approaches are computationally infeasible due to the high cost of the numerical solver, e.g. see lines 709-710 for ACB, to compute the likelihood probability.

The "ML-based" approaches are built in two stages. They first consider a training stage where given pairs of $(u_i, \xi^*_i), i=1,...,N$, with $u_i =$Sim $(\xi^*_i)$ and Sim a numerical solver, the model learns an approximate distribution $\rho(\xi_i^* | u_i)$ using variational inference or normalizing flows. During the evaluation stage a new input $u$ is given and the ML model provides samples from $\xi^* \sim \rho(\xi_i^* | u_i)$. To get an ensemble of continuous quantities $s$ and uncertainty estimates, following both the "classical" and "ML-based" inversion strategies, a numerical solver needs to be considered, which is very expensive as well.

For this reason, we propose a methodology that uses the same ML architecture to infer $\xi^*$ given $u$ and to act as a solver surrogate to make predictions of $s$ given $\xi^*$ jointly during evaluation. During training we consider all $(u, \xi^*, s)$ as known but during evaluation we consider only $u$ known and predict $(\xi, s)$. Using a neural operator surrogate for the solver, we can consider $u$ to not be equal to $s$, and e.g. get predictions at only specified locations of interest, e.g. the aorta, without need to run the simulation through the whole cardiovascular system.

The joint analysis of the inverse and forward problems with FUSE results to the following important benefits:

1. Propagated parametric uncertainty: The joint formulation of both problems under a rigorous mathematical framework (Eqn. 1) allows assessing the influence parametric uncertainty in the physical model has on the predicted continuous quantities in a very efficient manner.

2. Unified validation: Since both the forward and inverse models are nonlinear it is hard to track the error accumulation in the different components. The joint evaluation is a solution to assessing forward and inverse errors separately. Furthermore, this mathematical framework facilitates the comparison of different combination of inverse/forward models and training objectives. A priori, the errors of an inverse and a forward model trained and validated separately may amplify nonlinearly at concatenation, increasing the chance that samples generated by the former are OOD for the latter.

3. Function space simulation-based inference: Parametric PDEs as described above require a communication between the infinite-dimensional function spaces of $u$ and the finite-dimensional space of $\xi$. In order to allow operating on function spaces, we combine a finite-dimensional flow matching model with neural operators, which is (to the best of our knowledge) an entirely new concept. This approach can generalize to any flow based or diffusion approach.

Answers to the questions:

Q1. We kindly refer to the answer provided above.

Q2. We thank the reviewer for pointing out this ambiguous nomenclature and will adopt the proposed change. The term "discrete" emerged in opposition to "continuous" spatially varying parameters, which is of course not the right wording.

Q3. We are happy to incorporate any literature the reviewer may be pointing at that amplifies our literature review in the proposed direction. Following the reviewer's suggestion, we would like to add several references based on Gaussian processes for forward and inverse problems (arxiv 2204.02583), kriging (2012.11857), and GP based models to explore complex high-dimensional spaces by low dimensional representations (2101.00057), as well as others.

I thanks the reviewer for their detailed comments.

My understanding from their responses is the following:

1. Data is available in triples of the form $(s, u, \xi)$, so there is no latent variable inference, etc.
2. The training of the inverse and forward losses is performed in parallel and there is no shared parameters / common components etc.
3. The first time the models speak to each other is during inference / prediction, where the models are simply composed in the obvious way.

Based on this, I feel my original challenge about novelty still holds. My understanding remains that this is two separate methods,

(i) Learning a FNO model for the parameter to output problem.

(ii) Learning a conditional density via a flow-based model for the inverse problem from the input to the parameter distribution. Each approach is generally well-established, each trained in isolation (in parallel). These approaches are then combined at inference time to provide probabilistic inference for the associated inverse problem.

The assumption that you would have a data set of the form (parameter, functional input, functional output) would only be reasonable when you're working purely synthetic data from a PDE -- so that this provides utility mostly as a probabilistic surrogate / emulator for the inverse problem. This is different from, say INVAERT which does not prescribe intermediate values at training time. It seems quite unfair to me that it would be used as a baseline for FUSE, as it has to solve a substantially harder problem.

In terms of the benefits:

1. Propagated parametric uncertainty: Yes, this is true.
2. Unified validation: This is good, but there's no really way of back-propagating that unified validation to adjust the individual model.
3. Function space simulation-based inference: This is good and a clear feature of FNO layers within the models.

I will slightly bump up my score to reflect the helpfulness of the clarification, and because the general paper offers useful insight.

We sincerely appreciate the reviewer's commitment to this discussion, and we will gladly incorporate these clarifications in the updated manuscript. We also thank the reviewer for further increasing their score to acceptance.

W1. We thank the reviewer for pointing out possible difficulties in understanding our model and data. An updated version of the model illustration is provided in the 1-page pdf, including the requested details. We are happy to incorporate any further specific suggestions, should we have missed a crucial part. In terms of data visualization, we point to the extensive figures in the appendix, in particular the sensitivity studies in Figs. A.9 and A.16, as well as A.11 showing the data-generating 2D fields for the ACB experiment.

W2. The reviewer raises an excellent point. Flow matching and diffusion are competitive methods, each with their own advantages and disadvantages. As explained in "Flow Matching for Generative Modeling" (ArXiv 2210.02747), flow matching may use a diffusion defined probability path, but an optimal-transport path may also be selected, as in FMPE. The result is that FMPE is able to train on less data, converge faster, and provide faster sampling at inference time. Given that PDE training data often come from expensive numerical simulations, data efficiency is a key component for machine learning approaches in scientific computing. To showcase this, we have performed some initial experiments using a conditional DDPM model, and present the results in the 1-page pdf. Keeping in mind that these are only preliminary results, we observe that the experimental findings support the arguments stated above.

W3. We understand the reviewer's concern and agree that more details regarding FMPE models would benefit this paper. However, given the amount of space available, fully describing these details was not possible within the main text. To clarify, We have followed [23] in this regard (lines 135-144), and we would kindly refer interested readers to their work for a full derivation and explanation of flow matching for simulation-based inference. We would also include details on FMPE in the appendix of a CRV, if accepted.

Q1, Q2. The true parameters are used during training, with training samples being tuples of the form $(u,\xi^*,s)$. The inverse part of the model is trained with $(u,\xi^*)$ to learn a $u$-conditional probability path to $\xi^*$, while the forward map is learned from $(\xi^*,s)$. Only at evaluation time, when merely $u$ is available as input, the parameters $\xi$ sampled from the flow matching are input to the forward neural operator to obtain $s$.

Q3. Training is done in parallel. As the reviewer correctly points out, each loss function only depends on a single set of parameters during training. At evaluation time, we combine these two models to unify the inverse and forward problems, enabling us to study the relationships within the system, such as uncertainty propagation and sensitivity analysis.

Q4. We assume the reviewer is asking about the number of artificial time steps known from diffusion models. Although FMPE is defined along a time-dependent probability path, this pseudo-time variable is not discretized as in a diffusion model, but computed in a single step. Should the reviewer be referring to the number of samples drawn by the FMPE, our results use 100 samples from $\rho(\xi | u)$ at inference time.

Q5. After considering other generative approaches in the initial phase of our project, such as discrete normalizing flows and generative-adversarial models, the flow matching approach presented in this work proved the most competitive in our research. Please refer to W2 for a comparison to diffusion models.

Q6. In our work, we did also investigate a fully probabilistic approach. However, given that these data sets come from numerical simulators of PDEs whose continuous outputs are fully deterministic given the input parameters, fully probabilistic models in the way proposed by the reviewer offer no advantage in capturing the parametric uncertainty. As the reviewer suggests, this fully probabilistic approach would be interesting to investigate for instances where the observations are not deterministic given the parameters, such as with SDEs. We leave this interesting line of research to future work.

Q7. We provide the training and inference times in the 1-page pdf, and thank the reviewer for requesting this important information.

W1. The reviewer rightly points out that a reference and comparison to classical numerical methods is indispensable when justifying the use of ML-based methods. Since the application field "parametric PDEs" for our method is rather broad, we are happy to adopt the suggestion to reference a standard text book on the numerical solution of PDEs, as provided by LeVeque in 1992.

W2. We thank the reviewer for highlighting the multiwavelet neural operator (MWT), which is an interesting addition to our extensive review of neural operator approaches in lines 43/44. Since it proves superior performance to FNO in the given reference, it would be an interesting future study to replace the FNO by MWT within FUSE, in particular when it comes to test cases with larger scale separation.

W3, W4. In order to avoid any confusion on terminology up front, we would also like to point out that when Magnani et al. 2022 refer to operator learning for "parametric PDEs", the parameters taken as input to the learning problem are functions in space, $\lambda(x)$. Our setting, in contrast, assumes the parameters $\xi$ to be vector-valued constants, as an intermediate step between the functions $u$ and $s$. We assume that when suggesting to use uncertainty quantification (UQ) for neural operators (NO), they refer to the forward model part $\Xi \mapsto \mathcal{S}$ only.

We thank the reviewer for the suggestion to add the references on UQ for NOs. However, we respectfully disagree on the applicability of UQ for NOs to the problems considered in our paper. Let us explain why: UQ for NOs, in the sense indicated by the given references, quantifies the approximation error by the NO model, instead of the parametric uncertainty in the physical model targeted by FUSE. Thus, benchmarking against the proposed methods, which is listed as the main limitation of our paper, is not possible. As Mouli et al. 2024 point out in their review in the appendix, it is common practice, e.g. in weather forecasting, to perturb a physical model input and parameters to quantify the uncertainty related to the physical state and the model itself, respectively. We deem it important to distinguish these types of uncertainty from the additional approximation error when fitting an ML model. They are also easily distinguished by the probabilistic mappings they model: While an ensemble approach maps an ensemble of parameters (as provided by the FMPE) onto an ensemble of predictions (pushforward), the UQ approaches for NOs equip a prediction on a single input with an uncertainty range. In order to clearly interpret the uncertainties given by FUSE as parametric physical model uncertainty, and given the good performance of the model in our id and ood validation, we consider the PDE parameters as the main source of uncertainty.

W5. The reviewer raises an excellent point. Flow matching and diffusion are competitive methods each with their own advantages and disadvantages. As explained in "Flow Matching for Generative Modeling" (ArXiv 2210.02747), flow matching may use a diffusion defined probability path, but an optimal-transport path may also be selected, as in FMPE. The result is that FMPE is able to train on less data, converge faster, as well as provide faster sampling at inference time. Given that PDE training data often come from expensive numerical simulations, data efficiency is a key component for machine learning approaches in scientific computing. To showcase this, we have run experiments using a conditional DDPM model and present the results in the 1-page pdf. Even though this is initial work, we observe that the preliminary experimental findings support the above arguments. We will include this discussion and results in future versions of the paper.

Q1. We thank the reviewer for the question, which is crucial to the applicability of our method. Indeed, the "parameters" can stem from model properties (ACB: $v$ and $d$, Table 7), or parameterizations of the initial (the remaining ACB parameters) or also boundary conditions (PWP: pulse wave from the heart, lines 644-646 and original reference), or any other model component that is parameterized.

Concerning the nomenclature "discrete", we acknowledge that it is misleading, and we will change it to "finite-dimensional" parameters $\xi$, as opposed to the infinite-dimensional $u(t)$ and $s(t)$. We mean to point out that the inferred parameters are vector-valued constants that are not space or time-dependent.

Q2. The two loss functions are fully decoupled, see line 153 of the experiments section. So, there is no need for hyperparameter tuning.

W1. We thank the reviewer for pointing this out. We include an updated Figure in the one page document, which we believe is much more explanatory.

W2. Even though referencing pre-prints is a common practice in ML, the reviewer is right in that, if available, the respective journal or conference of publication should be provided. We will update all references accordingly.

W3. The inference times for FUSE are in the order of milliseconds per sample, details are provided in the one page document. As for the classical solvers, the computational complexity for the ACB case is about half an hour on eight cores per sample (lines 709-710). As cited in line 192, the PWP data set was not created by the authors, but by (Ref. [35]). The authors do not report the computational time, but simulating 1 sample of the system using the OpenBF open source Julia code [openbf-hub, A. Melis, 2018; Melis 2018, EthosID uk.bl.ethos.731549], takes about 80 seconds on one core. As the reviewer points out, this comparison is crucial to justify the use of ML models, and will be emphasized in future versions of the paper.

W4. See Q5 and Q6.

Minor weaknesses. We thank the reviewer for pointing out the typo, which we will of course correct. Regarding the notation, however, we would like to emphasize that l110 defines $\mathcal{G}$, while line 87 defines $\tilde{\mathcal{G}}$. We chose this notation to distinguish the unified operator $\mathcal{G}: \mathcal{U} \mapsto \mathcal{S}$ from the parametric forward model $\tilde{\mathcal{G}}: \Xi \mapsto \mathcal{S}$.

Q1. FUSE stands for Fast Unified Simulation and Estimation, as shown in the title. To facilitate the reading, we will repeat the full name where pointed out by the reviewer.

Q2. We thank the reviewer for asking for clarification on this notation, which we are convinced is correct. The mathematical problem formulation merely presents the distance $d$ that corresponds to the problem on a theoretical level, involving both the true unknown $\tilde{\mathcal{G}}$ and $\mu^*$. Following common practice in ML research, $d$ is subsequently upper bound by loss functions involving $\tilde{\mathcal{G}}$ and $\mu^*$, which are approximated by finite sampling when it comes to implementation, as suggested by the reviewer.

Q3. The lifting function is a commonly used term in ML research, also heavily used in the operator learning literature [3, 9, 10, 31], that describes an affine (linear) transform that increases the dimensionality of a function $u \in \mathcal{U} \subset \mathcal{C} (\mathcal{X}, \mathbb{R}^{d_u})$ defined as $\mathcal{P} : \mathcal{C} (\mathcal{X}, \mathbb{R}^{d_u}) \rightarrow \mathcal{C} (\mathcal{X}, \mathbb{R}^{\tilde{d}_u})$, where $\tilde{d}_u >> d_u$. In practice, a lifting function is implemented by a one-layer fully connected neural network.

Q4. As discussed in lines 601-604, 180 GPU hours is the total computational time for 64 runs of hyperparameter sweep for all models and all experiments contained in the paper, while training a FUSE model only takes about 1 GPU hour. Accounting for the wall-clock times of one simulation of ACB and PWP problems (W3), it is infeasible to repeatedly calibrate the numerical solvers using traditional methods for varying function input $u$ (e.g., MCMC requires about 1,000 to 10,000 simulations for each $u$). However, once trained, the low inference times of the FUSE model (W3) and its operator properties make this calibration possible.

Q5. "Learnable map" denotes a parametric function (for instance a neural network or an affine map) whose parameters are learnt during training. We followed standard practice in ML and did not define this terminology, while believing that an interested reader can check the code for exact architecture of our maps. However, we agree with the reviewer that defining these maps is more reader-friendly and will do so in a camera-ready version, if accepted.

Q6. As explained in the introduction, the method proposed in this paper is applicable to any parametric PDE. It will not work for PDEs that are not a priori parametric, e.g. for a stochastic PDE with a Brownian motion forcing, as explained in the limitations section. An estimate of the sample complexity for different PDEs would require a statistical analysis on a range of problems, which is outside of the scope of this paper. However, we provide an analysis of error scaling with the number of samples for the ACB case in the 1-page pdf. We observe that even if we consider 4096 training samples, half of what we used originally, FUSE would still be more accurate than any of the baselines reported in the main text.

L. The reviewer raises a point about the reliability of the FUSE model if it were applied in clinical practice for the Pulse Wave Propagation case. We believe it is very commendable of the reviewer to think about potential negative societal impact of this method.

Negative Societal Impact, as per NeurIPS author guidelines, is considered for cases where the presented methodology can be immediately used without a prior evaluation protocol, e.g., deep fakes. We present a general purpose framework and do not think that FUSE can be used as a clinical decision making tool as is. For this to happen, a very detailed protocol issued by a public safety organization needs to be considered that contains multiple steps of in-silico, in-vitro and in-vivo validation.

We would first like to thank the reviewer for acknowledging the novel framework, the rigorous problem definition, and the integrated approach we take constructing a robust framework.

W1 Regarding the concern with respect to the lack of novelty in the FMPE application, we would like to highlight key differences and contributions in this work.

1. This is not merely an application of FMPE: The metric in the first term in Eq. (1) is bounded by a metric as shown in line 127. Because this metric is hard to compute in practice, our implementation matches $\rho$ and $\rho^*$ with FMPE; however, it may be substituted by other choices such as NPE or diffusion models (see 1-page pdf). In other words, FMPE is not essential to the methodology and any suitable model for matching posterior distributions would suffice.

2. Inverse Problem: Additionally, this paper introduces function space simulation based inference, to infer parameters from functions as opposed to vectors. Therefore, we provide an inverse model that is grid resolution-invariant, bridging a gap between neural operator architectures and simulation-based inference. This is a critical concept for many practical applications and, to the best of our knowledge, has not yet appeared in the literature.

3. Forward Problem: We present a formulation of the supervised operator learning problem for a map between a finite-dimensional parameter and an infinite-dimensional space of output functions. This is accomplished by a novel Lifting Operator, transforming finite-dimensional vectors into the space of band-limited functions, as presented in lines 115-116.

4. Unification: We propose a novel framework in the operator learning literature that unifies forward model simulation and statistical parameter estimation under the same rigorous mathematical framework in Eq. (1). The joint formulation of both problems using the triangle inequality allows to assess what we refer to as the propagated uncertainty. This approach enables a more comprehensive understanding within many scientific problems by introducing explainable model uncertainties, which tie uncertainties in complex, infinite-dimensional spaces to uncertainties in simple parameters which characterize the system.

W2 Uncommon problems: We recognize that the two problems presented may not be widely known to the broader ML community; however, ML-based model calibration is of vital interest in many applied communities. In particular, cardiovascular models have been subject to both classical and ML-based inversion due to their relevance [arXiv 2307.13918]. Climate science is also a major topic within the field of operator learning (arxiv: 2208.05419, 2111.13587, 2306.03838). Due to their high complexity, there is less work on the calibration of small-scale atmospheric LES models, even though their parametric uncertainties are well-known [doi 10.1007/s10546-020-00556-3]. However, the combination of numerical methods and ML presents an opportunity to greatly improve the state of the art [paper ref. 48].

Q To answer the questions, let us first clarify that for the Pulse Wave Propagation problem, we present three scenarios with different levels of input information. In cases 2 and 3, most of the measurement locations on the human body and quantities of interest are not available, and hence masked (lines 196-200). When the reviewer refers to the "worst case", we assume they mean case 3 with least information provided.

Q1-2. First, there are indeed some variables with low sensitivity to the parameters (e.g. Fig. A.9: pressure is insensitive to age), which generally have wider posterior distributions. The wider distributions in Fig. A.1-2.b-c are a result of larger uncertainties as the number of input channels is reduced from 39 (case 1) to 3 (case 2) and 1 (case 3). It would be reasonable to interpret parameters with large uncertainty as a result of low sensitivity to the limited set of inputs in these cases. The uncertainty arising from the ill-posed nature of these cases, as opposed to unknown physics, results in pressure predictions which are quite different from the true values. Essentially, the parameters with a strong sensitivity to pressure may not be determined because they have a weak sensitivity to PPG at the fingertip (case 3). The result is that a more "average" pressure is predicted, with large uncertainty ranges.

Q3. Yes, in fact, the Pulse Wave Propagation experiment has boundary conditions encoded as parameters. The left boundary conditions (i.e. the pulse wave from the heart) is parameterized by Heart Rate, Stroke Volume, Pulse Transit Time, Residual Filling Volume, and Left Ventricular Ejection Time. We kindly refer the reviewer to see original reference [35] for more detailed information.

Dear all,

Thank you for your efforts so far in reviewing the paper and in answering reviewers' questions.

The evaluation is currently borderline, with an average score of 5.2 and a spread from 4 to 7.

The authors-reviewers discussion will close today on August 13. Please take this last opportunity to discuss the paper and the reviews and to reach a consensus, or at least to gather all the necessary information to make a fair and informed decision (in favor of acceptance or rejection).

Thank you for your attention and cooperation. The AC