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

Information Geometry of Message Passing
Aug 16, 2026Mykola Lukashchuk, Kyrylo Yemets, Alex Ledbetter +1Probabilistic Graphical Models Variational Inference
Eindhoven University of Technology · Lviv Polytechnic National University · Lazy Dynamics, Utrecht, the Netherlands
Composing Non-Conjugate Factor Graphs with Closed-Form Variational Inference
May 28, 2026Mykola Lukashchuk, Kyrylo Yemets, Wouter M. Kouw +4Factor Graph Optimization Mixture-of-Experts Inference
Eindhoven University of Technology, the Netherlands · Lviv Polytechnic National University, Lviv, Ukraine · Lazy Dynamics, Utrecht, the Netherlands +1
Amortized Variational Inference for Joint Posterior and Predictive Distributions in Bayesian Uncertainty Quantification
May 5, 2026Nan Feng, Xun HuanUncertainty Quantification Amortized Inference
Department of Mechanical Engineering, University of Michigan, Ann Arbor, Michigan, United States

Aug 16, 2026, cs.LGJ/K move · Enter open · S save

Information Geometry of Message Passing

Mykola Lukashchuk, Kyrylo Yemets, Alex Ledbetter, İsmail Şenöz

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

Probabilistic Graphical Models Variational Inference Information Geometry

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Information Geometry of Message Passing

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

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