Guided Diffusion Sampling on Function Spaces with Applications to PDEs

We thank the reviewer for their time and their thoughtful, constructive feedback.

[Strengths And Weaknesses]

Answer: We thank the reviewer for their valuable feedback and for recognizing the theoretical importance of our work. Due to the special care needed in function spaces, our formulation requires introducing notations from functional analysis. To enhance clarity, we will add a high-level introductory paragraph at the beginning of the Method section. Our goal is to ensure that foundational concepts are presented more simply to make the technical details accessible to a broader audience.

[Q1] Equation 13 is the approximation you use for the velocity when solving the ODE. This approximation is highly inefficient in the context of inverse problems and gradient descent. This is because if c is large (you want to actually fit the data) the ODE becomes highly stiff and you will need an implicit type method. This is why almost every PDE constrained optimization problem is solved using a Gauss-Newton type method rather than SGD. Your comparisons with other highly sub-optimal methods does not strike confidence.

Answer: We agree that for traditional PDE-constrained optimizations, where the goal is to find a single optimal solution, the gradient descent methods are often inefficient due to stiffness. However, our approach operates in a different framework (generative posterior sampling), solving the probability flow ODE directly with an explicit and well-posed definition. Specifically, our goal is not to find a single point estimate but to draw samples from an approximate posterior distribution. This process is governed by a reverse-time ODE that integrates a powerful data-driven prior (diffusion model) with a likelihood guidance term. We would like to further clarify several points about our implementation:

• FunDPS is sampler agnostic: Equation 13 defines a conditional score that provides the velocity for the reverse ODE. Since it does not prescribe any specific numerical solver nor do we use any specific solver throughout the proofs, our framework is general and can be integrated with any suitable ODE solver. The choice of sample is separate from the formulation of the guidance term (i.e. likelihood term).
• We use a higher-order solver, not SGD: We do not use first-order Euler steps (analogous to not SGD). As shown in Algorithm 1, we combine the method with a second-order deterministic sampler inspired by the EDM framework [1]. This higher-order solver mitigates the instability and stiffness issues related to simpler first-order schemes.
• Extensive Comparisons: During this work, we compared deep learning and GenAI baselines, including the SotA works (DiffusionPDE, FNO, PINN), and demonstrated our method's superior speed and accuracy.

[1] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." Advances in neural information processing systems 35 (2022): 26565-26577.

We hope our detailed answers are helpful, and we kindly ask for a re-evaluation of our submission in light of this discussion.

We thank the reviewer for their time and their thoughtful, constructive feedback.

[W1] While the paper's contributions are significant, the core components are not entirely novel in isolation. The concept of function-space diffusion models builds upon prior work (e.g., Score-based Diffusion Models in Function Space), and plug-and-play guidance is a well-established technique in finite-dimensional diffusion models.

Answer: Thank you for your comment. We agree that our work builds upon important prior research. However, our primary contribution is based on a novel integration of these ideas. Most importantly, FunDPS provides essential theoretical and practical contributions to make this combination work and be effective for solving ill-posed inverse problems arising in PDEs. While plug-and-play guidance was established in finite dimensions, a rigorous theoretical framework in the infinite-dimensional setting of function spaces was absent. Our work introduces the first generalization of Tweedie’s formula to Banach spaces. This is a key novelty that enables us to propose principled gradient-based guidance for diffusion models operating directly on a function. It goes beyond a simple combination of existing works, enabling design choices that are crucial for superior accuracy, speed, and multi-resolution strategies.

[Q1] How sensitive is the model's performance to the choice of \zeta? Please provide an ablation study showing performance across a range of \zeta. Moreover, is \zeta manually chosen or is automatically determined under certain theoretical background?

Answer: This is an excellent question regarding the key hyperparameter in our framework. Currently, the guidance weight $\zeta$ is manually chosen for each task by tuning on a small validation set. This is a common practice of guided diffusion models, as optimal weight often depends on the forward operator and prior data distribution. We provide a list of tuned $\zeta$ values used for all five PDE tasks in Appendix F, Table 4. In our experience, the model’s performance is stable in wide ranges. However, values that are too high can cause the sampling to become unstable and diverge, while the values that are too low result in weak guidance and less accurate reconstructions.

Note that after several trials, we follow ideas from the DPS paper, and we use adaptive guidance weighting based on noise level as highlighted in Appendix F. We agree that developing a method to automatically determine $\zeta$ is an important direction for future research, which could help overcome inherent challenges in this field.

Finally, we followed the suggestion and conducted an ablation study to show the performance of Darcy and Poisson equation forward/inverse problems across a range of $\zeta$ values. We can observe that FunDPS has expected behaviour across different ranges.

We hope our detailed answers are helpful, and we kindly ask for a re-evaluation of our submission in light of this discussion.

We thank the reviewer for their time and their thoughtful, constructive feedback.

[W1] The key contribution lies in the directly extending DiffusionPDE to the function space through function representation. The contribution on the multi-resolution training and inference for acceleration, and conditional sampling algorithm for solving PDE problems may be limited, since they are mostly integrated from previous works such as DiffusionPDE.

Answer: Thank you for this critical question. We would like to clarify the distinction and novelty of our work compared to DiffusionPDE. While we use their datasets for a fair and direct comparison, our methodology is fundamentally different. Shortly speaking, the only design choice we borrowed from them is the joint-embedding mechanism. In other aspects, FunDPS has important contributions as follows:

• A New Discretization-Agnostic Framework: Our core contribution is a framework that is inherently discretization-agnostic. FunDPS achieves this by using a U-shaped Neural Operator (U-NO) and Gaussian Random Field (GRF) noise model, which operates on continuous function representations. This differs sharply from DiffusionPDE, which is based on convolutional networks and is fundamentally tied to a fixed grid resolution (Please check Table 2 in the response to reviewer wgk8). This is not only a change in representation, but a critical shift in the underlying model that leads to inherent resolution-invariance, which was the key limitation of prior works. Our ablation study in Table 9 further confirms the benefits of function-space approach with significant performance boost.
• Novel Theoretical Guarantees in Function Spaces: To provide principled guided sampling on the infinite-dimensional settings, we propose and prove the generalization of Tweedie’s formula to Banach spaces. This theoretical result is a critical contribution that provides a mathematical relation between the score function and the denoising operator in function spaces, which is crucial for many plug-and-play guidance mechanisms, including ours.
• Multi-resolution Capabilities and State-of-the-Art Performance: Our function space formulation inherently leads to powerful multi-resolution training and inference strategies that were impossible with prior works as they are tied to fixed resolutions. For instance, our ReNoise inference pipeline achieves 25x wall-clock speedup over DiffusionPDE. Furthermore, our resulting framework obtains 32% average accuracy improvement with 4x fewer sampling steps.

We will revise the introduction to more explicitly highlight these fundamental distinctions.

[W2] The results with different inverse solvers in Table 5 are interesting. The results show that DAPS and DDNM is far worse than the DPS, which may diverge from the results in previous works. It would be helpful to elaborate what could be the underlying reason, like due to the function space formulation or PDE task?

Answer: We thank the reviewer for highlighting this interesting result. The significant underperformance of DDNM and DAPS in our setting is primarily due to the extreme sparsity of observation data, not the function space formulation, which we found after extensive investigation.

• DDNM relies on projecting the sample onto measurement subspace at each step. With only 3% of observation points, the measurement subspace is extremely low-dimensional compared to the overall function space. As a result, the projection provides very weak guidance and leads to poor reconstruction.
• DAPS uses a decoupled noise annealing schedule. The core issue comes from the localness of guided updates during Langevin dynamics. Only the points with observation are updated and others are just added with noise. This makes the intermediate state after Langevin dynamics discontinuous and out-of-distribution, which hinders performance. If we increase the # of sampling steps to 20,000, the accuracy can match DPS, but with 100x more time. The same trend is also observed in image-domain frameworks as we provide the quantitative results.

We have figures available to back our findings. We will include them in the camera-ready version and hope that it can provide insights into the advantages and disadvantages of different inverse solvers.

[W3] A key baseline compared in the paper is the DiffusionPDE, with the main difference in finite v.s. infinite representation space for training diffusion models, with claimed key contribution in extended Tweedie’s Formula. But another possible simpler approach is to represent measurement in the function space and inject to each diffusion sampling step to avoid estimating the Tweedie’s Formula. This was introduced in the diffusion inverse solvers before the introduction of DPS. Would this be possible to solve the target problem?

Answer: Thank you for this insightful question about alternative approaches. To ensure we understand accurately, we would be grateful if the reviewer could clarify which specific prior works they have in mind. This may refer to resampling-based or projection-based methods. Proceeding with that assumption, we can explain why such an approach, while powerful in some domains, is ill-suited for the PDE inverse problems here.

• Instability of forward operators on noisy states: Representing measurement would often require applying the forward operator $A$ (i.e. masking or PDE solver) to a noisy intermediate state $a_t$ to enforce consistency. PDE solvers are often ill-posed or numerically unstable when their input is a noisy, non-physical function, which is the case in huge noise spaces. FunDPS prevents this by first using denoiser to predict a clean state $\hat{a_0} = D_\theta(a_i, t_i)$, and then applying the forward operator to this physically plausible estimate $A(\hat{a_0})$. This leads to more stable gradient guidance.
• From a practical standpoint, our generalized Tweedie's formula is computationally inexpensive. Plus, both methods involve approximations, so there is no clear advantage of one over another.

With that said, our generalization of Tweedie’s formula bridges the denoising target and the score function required for guidance. We believe this contribution is not just crucial for our method but can also empower a wide range of future guided diffusion solvers that operate in function spaces.

[W4] The paper claims the sampling from posterior distribution with DPS guidance, while DPS is shown not to provably sample from posterior, such as in [1]. Can the author elaborate more on how to guarantee the posterior sampling as claimed?

Answer: Thank you for this important point and for providing the reference. The reviewer is correct that our method, like other guided diffusion techniques, provides an approximation of true posterior distribution. The approximation arises because the time-dependent log-likelihood term in the conditional score (Equation 6) is intractable, and we approximate using the conditional expectation of the clean sample (Equation 9 & 10). We want to highlight that even in standard finite dimensions, there is no inference-time inverse solver that has guarantees for true posterior (w/o training). Methods, such as Sequential Monte Carlo (SMC, TDS, etc.), have stronger asymptotic guarantees, but they are computationally more intense. It would be interesting to extend such methods to function spaces in future work. We will revise our manuscript to claim “approximate posterior sampling”. We will further add a sentence in Section 3.3 to explicitly state the source of this approximation by referring to the suggested paper [1].

[1] Cardoso, G., Janati El Idrissi, Y., Le Corff, S., & Moulines, É. (2024). Monte Carlo guided Denoising Diffusion Models for Bayesian Linear Inverse Problems. In Proceedings of ICLR 2024

We hope our detailed answers are helpful, and we kindly ask for a re-evaluation of our submission in light of this discussion.

We thank the reviewer for their time and their thoughtful, constructive feedback.

W1

Problem: In order to obtain a closed form expression for the likelihood, the conditional measure of the reverse diffusion process is approximated with its expectation. Can you comment on the error induced by this approximation?

Answer: This is an important question and we thank the reviewer for pointing it out. In the Euclidean setting, the error of this kind of approximation is described precisely by the Jensen gap of the conditional density, i.e., $E_x [p(y|x)] - p(y|E[x])$. Depending on the form of the noise, an upper bound on this gap can be given [1]. We believe a similar result will hold in our setting, where the notion of a density is now replaced by an appropriate notion of a Radon-Nikodym derivative. However, further research is needed for a formal proof, and we plan to spend more effort in this direction.

[1] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).

W2

Problem: The paper introduces the conditional score of \nu_t^u as the object of interest, however it’s not clear to me exactly how it is used in the given setup for sampling. I think it would be beneficial to make this explicit, for example, by giving the time-reversal of (4) in terms of the conditional score.

Answer: The conditional score of $\nu_t^u$ can be informally thought of as $\nabla_{a} p_t(a_t \mid u)$, i.e., the usual notion of a score function conditioned on an additional variable $u$. One of the main contributions of our work is an approximation of this conditional score in function spaces using an unconditional score and an additional guidance term (Eqn 13). To sample using this score, one starts by initializing $a_T$ as a Gaussian random field, and iteratively moving in the direction of this conditional score (Eqn 14 and L3-11 in Alg 1).

We agree with you that this could be explained more precisely in our submission. Your suggestion of writing the time reversal of (4) in terms of this conditional score will make this more clear to readers, and we will update our paper to include this.

W3

Problem: The submission introduces a lot of different notations for various measures and conditional distributions. I think this is necessary, however I would find a paragraph on the notation useful to refer back to, containing the main measures e.g. \mu, \eta, \gamma, \nu.

Answer: We thank the reviewer for their constructive feedback. In the final manuscript, we will include a more comprehensive notation paragraph to enhance the paper’s readability.

W4

Problem:

• In line 239 “firt” should be “first”.
• In Equation (13) should the scalar be c^2/2 or c/2? Equation (11) uses C_\eta^{-1/2}, which according to line 277 should be given by c^{1/2}I.

Answer: We thank the reviewer for reading the paper thoroughly and finding these typos. Indeed, both points are valid and fixed in the paper. For the second point, we will fix Equation (13) to be $\frac{c}{2}$. As we note in lines 286-288, the constant term is $\frac{c}{2}$ (not $\frac{c^2}{2}$) and absorbed into the final guidance weight hyperparameter $\zeta$. This hyperparameter has been tuned for each specific problem based on a small validation set (Appendix F).

W*

Problem: In summary, I find the paper gives an important and novel contribution in extending diffusion models for Bayesian inverse problems into function space. Experimentally, the method is shown to work well in comparison to existing methods and I recommend accepting this paper. However, I think there are some presentational issues which if fixed would improve the paper. Most importantly, making the infinite dimensional sampling scheme clearer.

Answer: Thank you very much for the positive feedback and constructive suggestions. We will revise the final manuscript based on all valuable reviews we get to make the main idea clearer. The infinite-dimensional sampling scheme is one of our core contributions, and we will enhance the content by rethinking definitions, better organization, and additional discussions.

---

Q1

Problem: For Tweedie’s formula, is there anything that can be said about the assumption that the score of \nu is Fréchet differentiable? When will this hold?

Answer: This is a good question that touches on a subtle theoretical point. The answer depends on the regularity of the data distribution $\mu$. The Fréchet differentiability of the score of $\nu$ is fundamentally tied to the smoothness of the Radon-Nikodym derivative $\frac{d\nu}{d\gamma}$, where $\gamma$ is the base Gaussian measure.

For a Gaussian measure $\gamma$, the natural space of directions for differentiation is its Cameron-Martin space. Since our perturbed measure $\nu$ is equivalent to $\gamma$ (i.e., they are mutually absolutely continuous, per Prop B.2), $\nu$ inherits much of the structure of $\gamma$. Therefore, assuming the score of $\nu$ is Fréchet differentiable with respect to perturbations from the Cameron-Martin space is a reasonable starting point, provided the data distribution $\mu$ is sufficiently regular. Without further assumptions on $\mu$, however, a more general statement is difficult. For a complete technical discussion, we refer to [2, Chapter 5].

One special case, though, where it is possible to give a definitive answer is when the data distribution $\mu$ is itself Gaussian. In this scenario, the perturbed measure $\nu$, being the convolution of two Gaussians, is also a Gaussian measure, and thus its Frechet differentiability is fully characterized by its Cameron-Martin space. See [3, Example 1] for a full proof.

[2] Bogachev, Vladimir Igorevich. Gaussian measures. No. 62. American Mathematical Soc., 1998.

[3] Lim, Jae Hyun, et al. "Score-based diffusion models in function space." arXiv preprint arXiv:2302.07400 (2023).

We thank the reviewer for their time and their thoughtful, constructive feedback.

[W1] Concern about incremental innovation compared to DiffPDE. Plus, is PDE loss a negative contribution to DiffPDE and FunDPS?

A: Thanks for the question. Inside the sampling loop, our work offers three key improvements over DiffPDE:

• Architecture: Our core contrib. is a discretization-agnostic generation framework, by using U-NO and GRF noise model operating on continuous representations. However, DiffPDE, which is tied to a fixed grid resolution, cannot generalize to different resolutions well. This generalization is important because in scientific domains, data are often collected in different resolutions. Our ablation study (Tab 9) further confirms the benefits of function-space approach.
• Theory: We introduce and prove the first generalization of Tweedie’s formula to Banach spaces, a critical theoretical guarantee for guided sampling in function spaces, to provide principled guided sampling when using GRF noise.
• Multi-resolution Capabilities and SotA Performance: These contributions unlock new multi-resolution strategies, achieving a 25x inference speedup and a 32% accuracy improvement.

As for the second question, PDE loss is included. It is expressed in a uniform way with the obs. loss. Contrary to the concern, both our work (App F) and DiffPDE (Sec 4.7) demonstrate that PDE loss improves accuracy.

[W2] Highlight the intuition and motivation to introducing the functional space, especially for this PDE problem.

A: Thank you for suggestion. Our main intuition is that physical systems described by PDEs are inherently continuous, and numerical solutions are discretizations of these dynamics. One major challenge in scientific computing is that a model trained on one specific grid often fails to generalize to another. Our function-space formulation addresses this by learning the resolution-agnostic mapping between the parameter and solution functions. This approach is more natural and well-suited as the real-world data inherently comes at varying resolutions. It is also becoming standard in AI4Sci community. Furthermore, this unlocks significant efficiency gains from our multi-resolution training and inference strategies.

[W3] It will be good to add more ablations on the influence of hyperparameters.

A: We agree that selecting hyperparameters is important. In our submission, we analyze various factors: # of observations (App G.5); different inverse solvers (App G.1); different training schemes (App G.3); ratio of low resolution sampling steps (App G.4); and varying design choices (App G.7). Notably, our model has a wide stable range for guidance weights, shown in the rebuttal to mb5h.

[W4] It will be more fair to perform careful hyperparameter on the baseline, especially the error you report here is much higher than the original numbers.

A: Thanks for the important question. To ensure a fair comparison, we used the official implementation and weights from the DiffPDE authors. We were in direct communication with them, and they confirmed that the results are the same as they could reproduce. In fact, for Helmholtz and NS1 cases, the reproduced baseline is even stronger than the original. We're confident this provides the most rigorous foundation for comparison.

[W5] It will be good to include more diffusion-based PDE solver. Also the model 54M is relatively small, not sure if it can generalize well on more general datasets.

A: Thank you for your suggestion. We have included an exp. in App G.1 that evaluates three plug-and-play inverse solvers. Regarding model scale, we chose a 54M parameter model to make a fair comparison with DiffPDE. In ablation, the largest model has 180M parameters, comparable to the typical diffusion model size of 300M for ImageNet64.

To further confirm generalizability, we provide additional exp. on real-world geophysical Darcy dataset, where we outperform all tuned baselines. The dataset is generated following [1].

Finally, we agree that scaling to larger models and more complex datasets is an exciting future direction.

[1] Wen, Gege, et al. "U-FNO—An enhanced Fourier neural operator-based deep-learning model for multiphase flow."

[W6] How the method along with DiffPDE can handle fewer observation and what are the lower limit?

A: Good question! We have investigated this in App G.5. We further conducted exp. with DiffPDE, shown below. Our method consistently outperforms across varying # of observations. As for the lower limit, it really depends on the use case; it maintains under 10% error with as little as 0.5% obs.

[W7] Will multi-resolution training influence sampling steps and accuracy? Does using coarser grid break the 3% observation assumption? Are you applying it to the diffusion-based baselines as well?

A: i) The number of inference steps and guidance weight are not affected by the strategy because our model can generalize.

ii) The observation is properly handled. 3% observation means 125 and 500 observed pts for 64 and 128 resolution.

iii) The multi-resolution inference cannot be applied to UNet as it cannot process different resolutions without training. The training strategy can be applied to DiffPDE. As shown below, FunDPS still outperforms baseline by a large margin with any training strategy. Notably, the UNet trained on 64x64 performs very badly on 128x128.

[W8] Compare FunDPS with the functional-space diffusion with CFG?

A: Thanks for this insightful comment. However, we believe such comparison is out of scope for the current work because

• Adapting CFG to function space is a non-trivial task and could be a new paper in itself. Our focus is on a principled posterior sampling framework with a plug-and-play solver.
• We compare primarily against DiffPDE, as it was already shown to outperform CFG models in its original paper. This provides a more informative evaluation against a stronger baseline.

[Q1] The main intuition of generalizing the Tweedie's formula to the function space?

A: At a high level, Tweedie’s formula states that if $Y = X + Z$ is a noisy observation of a random variable $X$ corrupted by Gaussian noise $Z \sim N(0, C)$, then the best guess (i.e. cond. expectation) for the value of $X$ given a noisy obs $Y=y$ is given by $E[X|Y=y]=y+C\nabla_y\log p(y)$. In other words, we take the observation $y$ and correct it by adding the vector $\nabla_y\log p(y)$, moving $y$ in the direction which most quickly increases its likelihood. The additional $C$ term accounts for correlations in the Gaussian noise.

When generalizing this to function spaces, the intuition remains the same, with the key difference being $X,Y$ are now random functions, and $Z$ corresponds to a Gaussian random field. This means that the gradient must be replaced with Frechet derivative, which is no longer a vector but rather a function itself. Moreover, $C$ is no longer a covariance matrix, but rather a covariance operator which acts on functions. The shape of this corrective function is determined by the functional derivative of the log-probability density, smoothed by the covariance operator of the noise process.

However, a key theoretical difference in function spaces is that probability densities only exist under certain strict conditions. To ensure a well-defined density exists in function space, the Cameron-Martin theorem (B.1) gives us the precise conditions, effectively requiring our data to live in a specific subspace (the Cameron-Martin space) where the signal and noise measures “agree” with each other.

We will update Sec 3.2 for the sake of more clarity.

[Q2] Compared with DiffPDE, are you using the same neural network arch?

A: No, FunDPS uses a U-shaped Neural Operator, which is designed to be resolution-agnostic. This is the one aspect of our contribution.

[Q3] What are the challenges to generalize this algorithm to more practical diffusion usecase like image generation?

A: Thanks for great question. PDE problems and image generation are important usecases. Generalizing function-space methods to image domain is an active research area. A key challenge is that NOs can oversmooth the sharp features in natural images. In principle, our framework could be adapted, but it is beyond the scope of this paper.

[Q4] Do we need to apply any assumption to this inverse problem formulation if we want to extend Tweedie's formulation to infinite-dimensional Banach spaces?

A: Yes, the core assumption is that the data is supported on the Cameron-Martin space of the noise process. We refer to our Thm B.6. This assumption is the infinite-dimensional equivalent of ensuring a well-defined probability density exists, which is a prerequisite for defining the score function at the heart of Tweedie's formula. As we noted, it is a standard requirement in function-space diffusion literature. While verifying this assumption theoretically can be difficult, a common workaround is to pre-smooth the data, for instance by applying the operator $C^{1/2}$ (the square root of the noise covariance). This effectively projects the data onto the Cameron-Martin space, meeting the assumption. For more details on this technique, we point to [1, Sec 4.2].

[1] Lim, Jae Hyun, et al. "Score-based diffusion models in function space." (2023)

We hope our detailed answers are helpful, and we kindly ask for a re-evaluation after this discussion.