Efficient Message Passing for Partial Differential Equation Priors

Anna Kazachkova; Leonhard Hennicke; Rainer Schlosser; Ralf Herbrich

cs.LGSep 26, 2026

Efficient Message Passing for Partial Differential Equation Priors

Authors: Anna Kazachkova, Leonhard Hennicke, Rainer Schlosser, Ralf Herbrich

Organizations: Hasso Plattner Institute University of Potsdam Potsdam, Germany

Abstract

Prior information for real-world physical quantities is most elegantly expressed via partial differential equations (PDEs). In this paper, we propose a novel way to solve PDEs using probabilistic inference on a factor graph. In general, factor graphs provide a natural way to encode prior knowledge into a model as explicit factors; here, this knowledge is provided by a governing PDE, which narrows the solution space, while observed data further shape the posterior over the parameters. The approximate parameter posterior is inferred using message passing based on moment matching, without posterior sampling or global gradient-based optimization. We demonstrate our approach on the first-order advection and the second-order semi-linear Fisher-KPP equations, where it achieves predictive accuracy comparable to a standard baseline while providing structured predictive uncertainty. Moreover, the inferred posterior marginal means and uncertainty structure match more closely those obtained using Hamiltonian Monte Carlo than the evaluated mean-field variational inference baseline, while requiring up to 10x less training time in our experiments, with inference speed comparable to variational inference.

Figures & tables

Figure 1: A generic factor graph for PDEs. White circle nodes refer to the hidden variables, the gray node to the observed data

{\bm{x}}

, and black squares to the functions (factors). Left subgraph (blue rectangle) : data subgraph, where

f_{U}({\bm{x}},{\bm{W}})

models a neural network that is trained to match the label data - only initial condition in our case. Right subgraph (red rectangle) : PDE subgraph, that enforces a PDE on a subset of collocation points

\mathcal{C}

with

f_{R}

	Advection				Fisher-KPP
	Ours	VI	HMC	SGD	Ours	VI	HMC	SGD
RMSE ( $\times 10^{-2}$ )	4.89	19.5	4.85	4.86	2.72	72.7	2.73	13.7
Cov. 90 % (%)	14.5	11.1	27.1	–	50.0	0.0	68.3	–
CRPS ( $\times 10^{-3}$ )	39.0	160	37.9	39.2	3.78	718	3.65	112

Table 1: Predictive accuracy and uncertainty diagnostics against the exact solutions, aggregated over all runs. For the metric descriptions, refer to Appendix A.6 . RMSE : mean per-run root mean square error of the predictive mean. Cov. 90 % : fraction of grid points for which the 90% predictive interval contains the exact solution. CRPS : continuous ranked probability score. HMC (grey) serves as the reference. Stochastic gradient descent (SGD) refers to the non-Bayesian PINN baseline (Section 4.1 ); its CRPS reduces to the mean absolute error. Best results among our method, VI, and SGD are highlighted in bold.

Figure 2: Comparison of the message passing approach (ours) and the baseline solution from PDEBench for the advection equation across different wave speeds, $b\in\{0.1,0.4,1,2,7\}$ . All heatmaps are plotted in the coordinates of ${\bm{x}}=(s,t)$ , where $s$ denotes space and $t$ denotes time. (a) the exact solution; (b) predicted mean values of our message passing solution, which preserve the propagation pattern across the evaluated wave speeds; (c) predicted variance values of our message passing solution, showing increasing uncertainty further in time from the observed initial condition.

Figure 3: Comparison of our approach (ours) to the HMC implemented based on Yang et al. (2021) for the advection equation. We map posterior weights of HMC (x-axis) to our approach (y-axis) - points in (a) refer to the mean values, points in (b) refer to the standard deviations; our approximate approach converges to posterior weights close to the ones from HMC. For readability, 25 random weights per trajectory are shown.

Figure 4: Computational cost of our approach (ours), HMC and VI for the advection equation across wave speeds

b

, ten trajectories each (points: trajectories; lines: medians). (a) Training time, cold start; (b) inference time of the predictive mean and variance on the full grid, steady state. Our training is up to

6.6\times

faster than HMC and up to

13\times

faster than VI; our closed-form inference is as fast as VI’s and up to

150\times

faster than HMC’s sampling-based inference.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Solution spaces for the advection equation ( $b=0.1$ ) for prior and posterior parameters shown as predictive mean (colorful) $\pm\sigma$ (gray shade). (a) prior space before any training; (b) posterior space with parameters trained only on initial condition - the function fits perfectly for $t=0$ , but does not learn any time dependencies and already narrows the space too much; (c) posterior space with parameters trained with a mixed schedule of initial condition data passes and PDE passes - the function learns both spatial and time patterns correctly; (d) exact solution of the PDE.

Figure 6: Comparison of our message passing approach (ours) and the SGD baseline for the advection equation; each point refers to one trajectory. (a) mean squared error (MSE) of the predictive mean against the exact solution; (b) root mean square (RMS) of the PDE residual

R

. Both methods reach errors of the same order at every wave speed.

Figure 7: Calibration and uncertainty structure for the advection equation, ten trajectories per wave speed

b

. (a) CRPS against the exact solution; (b) empirical coverage of the exact solution by the nominal 90 % (solid) and 95 % (dashed) predictive intervals; (c) our mean predictive standard deviation over time

t

(band: minimum to maximum over trajectories), growing with the distance from the observed initial condition at a rate that depends on

b

; (d) per-trajectory median ratio of our predictive standard deviation to that of HMC (bar: median over trajectories), roughly constant within each

b

. Ours and HMC share accuracy and the uncertainty pattern, whereas VI degrades for

b\geq 1

Figure 8: Comparison of our approach (ours) to VI, implemented based on Yang et al. (2021) , for the advection equation. We map posterior weights of VI (x-axis) to our approach (y-axis) - points in (a) refer to the mean values, points in (b) refer to the standard deviations; the means largely agree, while the VI standard deviations deviate from ours increasingly with the wave speed

b

. For readability, 25 random weights per trajectory are shown.

Figure 9: Comparison of VI to HMC, both implemented based on Yang et al. (2021) , for the advection equation. We map posterior weights of HMC (x-axis) to VI (y-axis) - points in (a) refer to the mean values, points in (b) refer to the standard deviations; VI underestimates the standard deviations at

b\leq 0.4

and overestimates them at

b\geq 2

, so it is further from HMC than our approach (Figure 3 ). For readability, 25 random weights per trajectory are shown.

Figure 10: Neural layer width ablation for the advection equation:

H\in\{1,2,3,4\}

neurons in the hidden layer. (a) Relative

L_{2}

error of the predictive mean (points: trajectories; lines: medians); (b) empirical coverage of the exact solution by the nominal 90 % (solid) and 95 % (dashed) predictive intervals; (c) predictive negative log-likelihood (NLL) of the exact solution, mean over trajectories; (d) training time (points: trajectories; lines: medians). HMC (grey) refers to the

H=1

model. Wider layers reduce the error and improve calibration for

b\leq 1

at a higher training cost; at larger

b

H\geq 3

degrades, and

H=4

fails at

b=7

Figure 11: Our message passing approach for the Fisher-KPP equation on one PDEBench trajectory at $\nu=1$ across reaction strengths $\rho\in\{1,2,5,10\}$ , in the coordinates ${\bm{x}}=(s,t)$ . (a) the exact solution; (b) predicted mean values of our message passing solution, which reproduce the reaction-driven saturation; (c) predicted variance values of our message passing solution.

Figure 12: Our message passing approach for the Fisher-KPP equation with a Gaussian initial condition, $f(s,0)=\exp\!\big(-(s-\pi)^{2}/(2(\pi/4)^{2})\big)$ on $s\in[0,2\pi)$ , at $\nu=1$ across reaction strengths $\rho\in\{1,2,5,10\}$ , in the coordinates ${\bm{x}}=(s,t)$ . (a) the exact solution; (b) predicted mean values of our message passing solution; (c) predicted variance values of our message passing solution.

Figure 13: Comparison of our message passing approach (ours) and the SGD baseline for the Fisher-KPP equation across reaction strengths

\rho\in\{1,2,5,10\}

(log scale), with the diffusion coefficient

\nu\in\{0.5,1,2,5\}

shown by the marker, five trajectories per setting; each point refers to one trajectory. (a) MSE of the predictive mean against the exact solution; (b) root mean square (RMS) of the PDE residual

R

. Our errors are of the same order in all settings, while the baseline collapses to the trivial solution

f\equiv 0

on most runs at

\rho=10

\nu\geq 1

Figure 14: Computational cost of our approach (ours), HMC and VI for the Fisher-KPP equation across reaction strengths

\rho

, with the diffusion coefficient

\nu

shown by the marker (points: runs; lines: medians). (a) Training time, cold start; (b) inference time of the predictive mean and variance on the full grid, steady state. Our training is about

4\times

faster than HMC and

9\times

faster than VI; our closed-form inference is as fast as VI’s and about

80\times

faster than HMC’s sampling-based inference.

Figure 15: Comparison of our approach (ours) to HMC, implemented based on Yang et al. (2021) , for the Fisher-KPP equation (color:

\rho

; marker:

\nu

). We map posterior weights of HMC (x-axis) to our approach (y-axis) - points in (a) refer to the mean values, points in (b) refer to the standard deviations; the means agree closely, and our standard deviations are slightly smaller than, but proportional to, those of HMC. For readability, 25 random weights per trajectory are shown.

Figure 16: Comparison of our approach (ours) to VI, implemented based on Yang et al. (2021) , for the Fisher-KPP equation (colour:

\rho

; marker:

\nu

). We map posterior weights of VI (x-axis) to our approach (y-axis) - points in (a) refer to the mean values, points in (b) refer to the standard deviations; VI reaches neither our means nor our standard deviations, since it does not converge to the solution within the training budget. For readability, 25 random weights per trajectory are shown.

Figure 17: Comparison of VI to HMC, both implemented based on Yang et al. (2021) , for the Fisher-KPP equation (colour:

\rho

; marker:

\nu

). We map posterior weights of HMC (x-axis) to VI (y-axis) - points in (a) refer to the mean values, points in (b) refer to the standard deviations; VI neither recovers the HMC means nor its standard deviations, which it overestimates by up to two orders of magnitude. For readability, 25 random weights per trajectory are shown.

Figure 18: Calibration and uncertainty structure for the Fisher-KPP equation across reaction strengths

\rho

, pooled over the diffusion coefficients

\nu

t

(band: minimum to maximum over runs); (d) per-run median ratio of our predictive standard deviation to that of HMC (bar: median over runs; marker:

\nu

). Ours closely follows HMC in CRPS but is under-covered relative to it, with median ratios between 0.53 and 0.96; VI fails in all settings.

Figure 19: Effect of the PDE tolerance

\epsilon_{C}

for the advection equation, ten trajectories per wave speed

b

(lines: medians; bands: interquartile ranges; dotted: the default

\epsilon_{C}=0.05

). (a) MSE of the predictive mean against the exact solution; (b) mean squared PDE residual

R

. Tightening

\epsilon

down to about

0.1

lowers both; below that the error no longer improves, and at

b=7

with

\epsilon=0.005

the model collapses to a constant solution, i.e., exhibits a failure mode.

Explore similar work

Aug 16, 2026cs.LG

Information Geometry of Message Passing

We show that the natural-gradient stationary condition of variational inference has an edge-local form on a Forney-style factor graph. We start from the Bethe free energy and constrain a selected edge marginal to an exponential family. At a stationary point, the natural parameter of that edge equals the sum of two projected messages, one from each incident factor. Each projected message is the natural-gradient projection of the exact belief-propagation log-message at the current receiving marginal, or equivalently, the gradient of its expectation in the so-called mean coordinates. We call the resulting scheme natural-gradient message passing (NGMP). The rule is local; each edge may carry its own exponential family, and the message a factor sends depends on the marginal that receives it. Compared with variational message passing, NGMP keeps the part of the exact message that the receiving family can represent instead of averaging the factor under the neighboring beliefs. The two coincide when the uncertainty on the edges entering a non-conjugate factor vanishes, and NGMP is more accurate when that uncertainty persists, for example, along a partially observed latent chain or when parameters are filtered through successive data batches. Experiments on Poisson smoothing, heteroskedastic regression, and hourly ETTh forecasting confirm this and show that the gain appears mainly in uncertainty calibration.

Mykola Lukashchuk, Kyrylo Yemets, Alex Ledbetter +1

Eindhoven University of Technology · Lviv Polytechnic National University · Lazy Dynamics, Utrecht, the Netherlands

Probabilistic Graphical Models Variational Inference Information Geometry

Ask PDF

May 28, 2026cs.LG

Composing Non-Conjugate Factor Graphs with Closed-Form Variational Inference

Stacking probabilistic building blocks into deeper architectures typically breaks closed-form inference. We show that closed-form inference can be preserved. We identify five factor-graph primitives: a bilinear factor, an exponential link, a Gamma prior, a Gaussian likelihood, and an equality node, and prove that any model composed from them admits closed-form variational message passing. The construction works because each primitive preserves a small set of message families: under mean-field factorization, messages on Gaussian variables remain Gaussian and messages on precision variables remain Gamma, while the only non-conjugate interface, the exponential link, remains tractable through the Gaussian moment-generating function and the sufficient statistics of the Gamma family. We demonstrate composition at increasing depth, from static ensembles through input-dependent gating to split-branch routing, and show that stacking routing layers encodes arbitrary decision trees, establishing universal function approximation with closed-form inference. Applied to ensemble time-series forecasting, the framework yields a Bayesian mixture of experts in which gating functions are inferred rather than learned, providing calibrated uncertainty over expert selection across five benchmark datasets.

Mykola Lukashchuk, Kyrylo Yemets, Wouter M. Kouw +4

Eindhoven University of Technology, the Netherlands · Lviv Polytechnic National University, Lviv, Ukraine · Lazy Dynamics, Utrecht, the Netherlands +1

Factor Graph Optimization Mixture-of-Experts Inference Probabilistic Graphical Models Variational Inference

Ask PDF

May 5, 2026stat.ML

Amortized Variational Inference for Joint Posterior and Predictive Distributions in Bayesian Uncertainty Quantification

Bayesian predictive inference propagates parameter uncertainty to quantities of interest through the posterior-predictive distribution. In practice, this is typically performed using a two-stage procedure: first approximating the posterior distribution of model parameters, and then propagating posterior samples through the predictive model via Monte Carlo simulation. This sequential workflow can be computationally demanding, particularly for high-fidelity models such as those governed by partial differential equations. We propose a variational Bayesian framework that directly targets the posterior-predictive distribution and jointly learns variational approximations of both the posterior and the corresponding predictive distribution. The formulation introduces a variational upper bound on the Kullback--Leibler divergence together with moment-based regularization terms. The variational distributions are trained in an amortized manner, shifting computational effort to an offline stage and enabling efficient online inference. Numerical experiments ranging from analytical benchmarks to a finite-element solid mechanics problem demonstrate that the proposed method achieves more accurate predictive distributions than conventional two-stage variational inference, while substantially reducing the cost of online predictive inference.

Nan Feng, Xun Huan

Department of Mechanical Engineering, University of Michigan, Ann Arbor, Michigan, United States

Uncertainty Quantification Amortized Inference Bayesian Inference Variational Inference

Ask PDF