Official Review · ICLR 2026

This paper introduces a new architecture, Probabilistic DiffusionNet (PDN), for learning uncertainty-aware surrogate models for Partial Differential Equations (PDEs) where the solutions depend on surface geometry. The work extends the existing DiffusionNet (DN) architecture which is a powerful model for learning surfaces, by reformulating the core deterministic diffusion mechanism of DN to be stochastic in nature. This stochastic diffusion operator is derived from the spectral solution of a stochastic heat equation (SPDE) and through this reformulation, the authors are compelled to inject spatially-correlated noise directly into the message-passing layers of their model. The architecture is constructed as a hierarchical probabilistic generative model similar to a Variational Auto-Encoder (VAE) where latent random variables $z$ are introduced at each layer. The model is fitted via amortized variational inference by maximizing the Evidence Lower Bound (ELBO). The formulation of PDN as a latent variable model is sound. The authors clearly define the generative process $p_\theta(u|v, M^h)$, the observational model $p(y_s|u)$, and the amortized variational encoder $q_\phi(Z|Y,V,\mathcal{M})$. The use of the ELBO (Eq. 14) for optimization is a standard and appropriate choice for this framework. Central to the method—the stochastic diffusion operator $\mathcal{SP}_{t,\eta}$—there is nothing arbitrary. The operator is carefully derived from the spectral solution of the stochastic heat equation (Eq. (5)). This derivation is in Sec. 3.1 and proved in App. B which gives a rigorous mathematical justification for the construction of the model. 1. In Appendix C, the authors ablate two types of amortized encoders: "Partial Diffusion (PD)" and "DiffusionNet+ (DN+)". They state that the DN+ encoder is immune to the "basis ambiguity problem", whereas the PD encoder (used in the main paper) suffers from it. The authors choose the PD encoder because it has ~4x fewer parameters and achieves similar performance. This is a slightly unsatisfying trade-off. It is surprising that the theoretically less-robust encoder performs just as well. The paper would be stronger if it either used the more principled DN+ encoder or provided a deeper analysis of why the basis ambiguity issue, which they acknowledge, does not seem to degrade performance in this specific VAE framework. 2. The stochastic operator $\mathcal{SP}_{t,\eta}$ introduces a new and apparently important set of parameters, $\eta$, which shape the noise spectrum through the relation $q_k = \lambda_k^{2\eta}$. These parameters determine how strong and what kind of stochastic perturbations occur at each layer. However, the paper never clearly explains how the $\eta$ parameters are actually chosen or managed. Are these hyperparameters fixed or learned parameters that are optimized during training? If they are learned, an analysis of their converged values would be valuable. If they are fixed, an ablation study on th