Direct Message Approximation (DMA): A Consistency-Based Framework for Tractable Approximate Inference on Factor Graphs

Ralf Herbrich; Rainer Schlosser; Jan Lemcke; Johann Ukrow; Anna Kazachkova; Nicolas Alder; Leonhard Hennicke; Theo Bardey

cs.LGSep 24, 2026

Direct Message Approximation (DMA): A Consistency-Based Framework for Tractable Approximate Inference on Factor Graphs

Authors: Ralf Herbrich, Rainer Schlosser, Jan Lemcke, Johann Ukrow, Anna Kazachkova, Nicolas Alder, Leonhard Hennicke, Theo Bardey, +6 more

Organizations: Hasso Plattner Institute University of Potsdam Potsdam, Germany

Abstract

Approximate message passing on factor graphs underlies two dominant families of probabilistic inference algorithms: expectation propagation (EP) and variational message passing (VMP). Both methods approximate the marginal at each factor edge, forcing an iterative round-robin schedule, risking negative-precision messages, and, for VMP, collapsing to point estimates at Dirac-delta factors. We introduce Direct Message Approximation (DMA), which approximates factor-to-variable messages directly rather than the marginal. For normalisable factors, we define a consistency condition (requiring exactness when all other incoming messages are Dirac deltas) to guide message construction. We prove a master theorem (proper messages, any graph) bounding marginal KL from message KL, with three structural corollaries: Dirac-input consistency, no EP-style inner-loop iteration, and no negative-precision messages. Further, we prove a complementary $O(1/r^2)$ guarantee for the inherently improper backward message of the product factor, whose closed-form treatment has resisted prior work. As a concrete instantiation, we derive explicit DMA messages for the product and leaky-ReLU factors and assemble a Bayesian neural network (BNN) inference algorithm with one forward/backward sweep per training example and no gradient learning-rate hyperparameter, validating that the structural guarantees translate to predictive uncertainty that widens in data-sparse regions, including under model mismatch.

Figures & tables

Panel of Figure 1 — Figure 1: Factor graph of a BNN (schematic; two learnable layers shown). Circles are variable nodes; filled squares are factor nodes. Blue-tinted circles represent the $d_{l}\!\times\!d_{l-1}$ independent scalar weight variables $W^{(l)}_{ij}$ (shown collectively per layer for clarity), each with a Gaussian prior factor $\pi$ . Each $\times$ node is a matrix-vector product factor : shorthand for $d_{l}\!\times\!d_{l-1}$ 1D product factors $\delta\!\left({z_{ij}-W^{(l)}_{ij}\,x^{(l-1)}_{j}}\right)$ (Sec. 3.1 ) combined with $d_{l}$ sum factors $z^{(l)}_{i}=\sum_{j}z^{(l)}_{ij}$ , together implementing the inner product (see Sec. 4.1 ). R factors are element-wise leaky-ReLU factors (Sec. 3.3 ); $\ell$ is the Gaussian likelihood factor (conjugate; $\delta=0$ ); the double circle is the observed output $y$ .

Figure 2: Left : Posterior predictive mean (solid) and $\pm 2\sigma$ intervals (shaded) versus the true data-generating function (dashed), where $\sigma{=}\sqrt{\mathrm{Var}[f(x)]+\beta^{2}}$ combines posterior variance with observation noise ( $\beta{=}0.2$ ). Training points shown as crosses ( $N=200$ , $x\in[-2.5,\,1.5]$ ); dashed verticals mark the training boundaries. Right : Hinton diagram of posterior weight beliefs after training. Square area $\propto$ posterior mean magnitude; transparency $\propto$ posterior variance (opaque = certain).

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

Approximates	marginal	marginal	marginal	message
	ADF	EP	VMP	DMA (ours)
Backward weight updates	$\times$	$\checkmark$	$\checkmark$	$\checkmark$
No cavity division	$\checkmark$	$\times$	$\times$	$\checkmark$
No inner-loop iteration	$\checkmark$	$\times$	$\times$	$\checkmark$
Consistency axiom	$\times$	$\times$	$\times$	$\checkmark$
Edge-local KL bound	$\times$	$\times$	$\times$	$\checkmark$

Appendix

Table 1: Comparison of approximate inference frameworks on properties relevant to BNN inference. Below the first row,

\checkmark

indicates a desirable property. Approximates : whether the method targets the factor-to-variable message directly or approximates the marginal at each edge. No cavity division : whether updates avoid dividing out a stored incoming message (eq. 20 ). No inner-loop iteration : whether a single forward/backward sweep per example suffices, with no per-example fixed-point iteration. Backward weight updates : whether weight beliefs are updated via a backward message sweep. Consistency axiom / edge-local KL bound : see Definition 2.1 and Theorem 2.3 .

Factor	Direction	Proper?	Theorem
Gaussian prior	both	Yes	Cor. 2.4 (exact)
Gaussian likelihood	both	Yes	exact (conjugate)
Linear / copy	both	Yes	Cor. 2.4
Product ( $z{=}xy$ )	fwd ( $z$ )	Yes	Thm. 2.3
Product ( $z{=}xy$ )	bwd ( $x,y$ )	No	Thm. 3.3
Leaky-ReLU ( $\alpha{>}0$ )	fwd	Yes	Thm. 2.3

Appendix

Table 2: Theorem coverage for each factor direction in the BNN. “Proper?” refers to whether the DMA message is a valid (positive-precision) Gaussian. All linear and copy factors satisfy Corollary 2.4 exactly; the Gaussian likelihood is conjugate and exact. See Remarks B.2 and B.6 for the improper cases.

Figure 3: Empirical validation of Theorem 2.3 across 192 leaky-ReLU factor configurations. Left : log-log scatter of message error $\delta$ versus marginal error; the dashed line $C\delta$ (with $C=\|m_{{X_{j}}\to{f}}\|_{\infty}/Z$ ) lies above every point, confirming $O(\delta)$ scaling. Right : the marginal KL normalised by its per-configuration bound $\|m_{{X_{j}}\to{f}}\|_{\infty}/Z\cdot\delta$ ; all 192 ratios lie below the bound line at 1 (maximum ratio $0.97$ ), directly certifying the theorem. The bound is tight: the $\|m_{{X_{j}}\to{f}}\|_{\infty}/Z$ prefactor evaluates the Gaussian sup-norm, which is largest where the incoming message is most concentrated and the marginal error is therefore small anyway.

Figure 4: Actual marginal KL (solid, coloured) versus the theoretical bound

\|m_{{X_{j}}\to{f}}\|_{\infty}/Z\cdot\delta

(dashed, black) as each parameter is swept individually (legend shown in panel (a)). (a) Input SNR

r

(

\alpha=0.3

\sigma_{x}=1

fixed): the bound decays monotonically; the actual KL is non-monotone, peaking near

r\approx 2

–

3

where the ReLU kink is hardest to match, and small at both low and high

r

. (b) Leaky slope

\alpha

(

r=2

\sigma_{x}=1

fixed): both decrease as

\alpha\to 1

(linear factor,

\delta=0

) and grow as

\alpha\to 0

(approaching the improper hard-ReLU limit). (c) Input width

\sigma_{x}

(

r=2

\alpha=0.3

fixed): the actual KL increases with

\sigma_{x}

(wider incoming message exposes more backward-message error in the marginal), while the bound decreases through its

1/\sigma_{x}

prefactor; the conservatism is greatest at small

\sigma_{x}

. In all panels the bound lies strictly above the actual KL, with a maximum ratio of

0.68

across all three sweeps (width sweep, panel (c)).

Figure 5: KL divergence between IS reference marginal and DMA marginal as a function of input SNR

r

. Left : product factor backward message to

X

from

Z=XY

; incoming messages

Y\sim\mathcal{N}(\mu_{y},(\mu_{y}/r)^{2})

and

Z\sim\mathcal{N}(\mu_{z},(\mu_{z}/r)^{2})

with

\mu_{y}=4

\mu_{z}=10

, and the incoming message from

X

fixed at

\mathcal{N}(3,1)

. Right : ReLU backward message to

X

from

Y=\mathrm{relu}(X;\,\alpha)

with

\alpha=0.1

; incoming message

X\sim\mathcal{N}(\mu_{x},(\mu_{x}/r)^{2})

with

\mu_{x}=1

. Both panels show empirical

O(1/r^{2})

decay. The product backward (left) is covered by Theorem 3.3 ; the leaky-ReLU backward (right,

\alpha{=}0.1

, proper message) is consistent with Theorem 2.3 applied to the message KL.

Figure 6: Copy-factor KL divergence over the 2D parameter space, computed from the closed-form expression in Lemma B.3 (no Monte Carlo). Left : Step 1 error $\mathrm{KL}[\mathrm{LN}(\cdot;\,\mu_{Y},\sigma_{Y}^{2}),\;\hat{m}_{f\to X}(\cdot)]$ swept over $(\mu_{Y},\sigma_{Y})$ . Horizontal bands confirm that the KL depends only on the log-space variance $\sigma_{Y}^{2}$ , not on $\mu_{Y}$ ; the analytical rate is $\frac{3}{4}\sigma_{Y}^{2}+O(\sigma_{Y}^{4})$ . Right : Step 2 error $\mathrm{KL}[\mathrm{LN}_{\mathrm{m}}(\cdot;\,\mu_{X},\sigma_{X}^{2}),\;\mathcal{N}(\cdot;\,\mu_{X},\sigma_{X}^{2})]$ swept over $(\mu_{X},\sigma_{X})$ . Dashed white lines are level curves of $r=\mu_{X}/\sigma_{X}$ ; diagonal banding confirms the KL depends only on $r$ , consistent with the $O(1/r^{2})$ rate.

Figure 7: Product factor

f(X,Y,Z)=\delta(Z-XY)

. Each row shows marginals

X

Y

Z

(left to right). Blue: DMA. Gray: IS. (a) Nominal (

\mu_{X}

= 3,

\sigma^{2}_{X}

= 1,

\mu_{Y}

= 4,

\sigma^{2}_{Y}

= 1,

\mu_{Z}

= 10,

\sigma^{2}_{Z}

= 5): the

Z

forward marginal is well approximated by the DMA Gaussian; the

X

and

Y

backward marginals are already noticeably skewed in the IS reference. (b) Stress (

\mu_{X}

= 1,

\sigma^{2}_{X}

= 4,

\mu_{Y}

= -2.5,

\sigma^{2}_{Y}

= 1,

\mu_{Z}

= -5,

\sigma^{2}_{Z}

= 1):

Z

remains well approximated; the

X

and

Y

backward marginals become severely non-Gaussian.

Figure 8: ReLU / Leaky-ReLU factor (

\alpha

= 0.1), nominal scenario (

\mu_{X}

= 1,

\sigma^{2}_{X}

= 1). Left: backward marginal

X

. Right: forward marginal

Y

. Blue: DMA, Gray: IS reference

Panel of Figure 9 — Figure 9: Left : Training negative log-likelihood (NLL $=-\frac{1}{N}\sum_{i}\log p(y_{i}\mid x_{i})$ , nats/example) versus epoch for DMA (solid black) and Adam at four learning rates (dashed). DMA’s outer updates stop automatically at epoch 17; Adam’s speed and final NLL depend critically on $\eta$ . Values above $8$ are clipped to $8$ for display. Right : Adam ( $\eta=0.1$ ) point prediction outside the training range $[-2.5,\,1.5]$ (dotted verticals); the true function (dashed) can deviate arbitrarily from the point estimate, with no uncertainty quantification available. Compare with the widening DMA posterior in Figure 2 . DMA’s per-epoch cost is $0.80\,\text{ms}$ versus $0.45\,\text{ms}$ for Adam on this architecture ( $1.8\times$ ), but DMA’s early stopping at epoch 17 versus Adam’s ${\sim}80$ gives an overall ${\sim}5\times$ reduction in total training time.

Figure 10: Left : Training NLL vs. epoch for DMA (solid black), Adam

\eta=0.1

(dashed grey), AdamW

\eta=0.1

at four weight decay values (solid coloured), and two poorly-tuned AdamW configurations (

\eta=0.01

and

\eta=1.0

, both

\lambda=1.0

, dotted) illustrating hyperparameter sensitivity. Right : AdamW

\eta=0.1

\lambda=0.1

(best weight decay on this seed) point prediction outside the training range

[-2.5,\,1.5]

(dotted verticals); the true function (dashed) is tracked well on this seed but no uncertainty quantification is available.

Method	Epochs	Extrapolation NLL
		Median	IQR
DMA (posterior predictive)	$\mathbf{16}$	$\mathbf{\phantom{0}0.80}$	$1.95$
Adam $\eta=0.1$	$200$	$12.09$	$11.52$
AdamW $\eta=0.1$ , $\lambda=0.01$	$200$	$\phantom{0}8.05$	$10.73$
AdamW $\eta=0.1$ , $\lambda=0.1$	$200$	$\phantom{0}5.52$	$\phantom{0}9.86$
AdamW $\eta=0.1$ , $\lambda=1.0$	$200$	$\phantom{0}4.07$	$10.47$

Appendix

Table 3: Extrapolation NLL (

\downarrow

better) on the 83-weight 1D regression task, averaged over 20 seeds (median [IQR]). Bayes-optimal

\approx-0.69

nats.

Figure 11: Hyperparameter sweep: minimum training NLL achieved over 500 epochs for each of 160 IVON configurations ( $8\times 4\times 5$ grid over $\eta$ , $\beta_{2}$ , $\delta$ ) on the 83-weight BNN ( $N{=}200$ , $\beta{=}0.2$ ). Only 26 of 160 configurations reach a final NLL below 1.0, concentrated in the $\delta{=}0.1$ and $\delta{=}0.5$ columns. DMA training NLL at convergence: ${-}0.16$ (epoch 17, no hyperparameter search); best IVON: ${-}0.31$ (epoch 500, requires $\delta{=}0.1$ and extensive tuning).

Figure 12: Training NLL minus Bayes-optimal (log scale) vs. epoch for IVON at three learning rates (best

\beta_{2}

\delta

per

\eta

from the sweep) and DMA, on the 83-weight BNN (

N{=}200

\beta{=}0.2

). DMA converges at epoch 17;

\eta{=}0.001

descends slowly but does not converge within 500 epochs;

\eta{=}0.01

converges to NLL

\approx\!0.26

with

\delta{=}0.5

;

\eta{=}0.1

oscillates near the Bayes-optimal line.

Panel of Figure 13 — Figure 13: Left : DMA posterior predictive (identical to Figure 2 ). Right : IVON predictive with the best sweep configuration ( $\eta{=}0.6$ , $\beta_{2}{=}0.9999$ , $\delta{=}0.1$ ) after 2000 epochs. Training range $[-2.5,\,1.5]$ marked by dotted verticals. DMA extrapolation NLL: $0.54$ ; IVON on this seed: $-0.40$ . Across 20 seeds IVON diverges on 14/20; median extrap NLL on finite seeds is $2.11$ (IQR $5.27$ ) vs. DMA median $0.79$ (IQR $1.28$ ).

Figure 14: Left : DMA posterior predictive (same as Figure 2 ). Right : Diagonal Laplace predictive with He-init prior (MAP via Adam with per-layer

L_{2}

decay, diagonal of exact Hessian,

K=1000

MC samples). Training range

[-2.5,\,1.5]

marked by dotted verticals.

Figure 15: Extrapolation calibration curves for the 83-weight network, test points in

[-4,\,-2.5]\cup[1.5,\,3]

\Delta=\overline{\text{coverage}-\alpha}

0

= perfect,

>0

= conservative,

<0

= overconfident. On this seed, DMA (blue,

\Delta=+0.01

) and diagonal Laplace (red,

\Delta=+0.05

);

\Delta

is the mean signed deviation — DMA’s smaller

|\Delta|

reflects cancellation of over- and under-coverage rather than a uniformly tighter fit to the diagonal. Over 20 seeds both methods are comparable (median

-0.09

vs.

-0.10

Method	Extrap NLL	$\Delta$ (extrap)
DMA	$0.82\ [1.96]$	$-0.09\ [0.44]$
Diagonal Laplace	$0.77\ [1.91]$	$-0.10\ [0.35]$
Bayes-optimal	$-0.69$	$0$

Appendix

Table 4: Extrapolation NLL and calibration error

\Delta

on the 83-weight 1D regression task, over 20 independently drawn datasets and true functions (median [IQR]).

\Delta>0

: conservative;

\Delta<0

: overconfident;

\Delta=0

: perfect.

Figure 16: Model mismatch experiment (

N=200

, training range

[-5,5]

\beta=0.2

). The data-generating network has two wider hidden layers (

d=12,10

); both methods learn with the two-hidden-layer model network (

d=6,5

). Left : DMA posterior predictive mean (solid) and

\pm 2\sigma

intervals (shaded) versus the true function (dashed). Training points are shown as crosses; dotted vertical lines mark the training boundaries

x=\pm 5

. Right : Adam (

\eta=0.1

, red) and AdamW (

\eta=0.1

\lambda=0.1

, purple) point predictions; neither carries uncertainty quantification.

Method	Median	Mean	IQR
DMA (posterior predictive)	$\phantom{-}0.01$	$\phantom{-}0.29$	$0.44$
Adam ( $\eta=0.1$ )	$\phantom{-}0.50$	$\phantom{-}2.21$	$2.23$
AdamW ( $\eta=0.1$ , $\lambda=0.1$ )	$-0.36$	$-0.11$	$0.73$
Bayes-optimal	$-0.69$	–	–

Appendix

Table 5: Extrapolation NLL over 20 independent seeds (

\downarrow

better; IQR = interquartile range). Bayes-optimal

\approx-0.69

nats. AdamW uses

\eta=0.1

\lambda=0.1

(best of

\lambda\in\{0.01,0.1,1.0,10.0\}

by median).

Architecture	Weights	$N$	Out	Epochs	NLL	Per-epoch
$8{\to}6{\to}5{\to}1$	83	200	1	17	$-0.16$	$0.8\,\text{ms}$
$6{\to}6{\to}12{\to}48{\to}24{\to}4$	1932	1500	4	3	$-0.73$	$67\,\text{ms}$

Appendix

Table 6: DMA training summary: small baseline vs. large network. NLL is the per-example per-output training log-likelihood

\frac{1}{NK}\sum_{i,k}\log p(y_{ik}\mid x_{i})

at the final epoch (lower / more negative is better). Per-epoch time excludes the first epoch (JIT warm-up).

Figure 17: DMA posterior predictive (

6{\to}6{\to}12{\to}48{\to}24{\to}4

1932

weights,

N=1500

, 3 epochs,

0.72\,\text{s}

total). Each panel shows one of the four output channels: posterior predictive mean (solid) with

\pm 2\sigma

bands against the true function (dashed); training points as faint dots.

Figure 18: Adam (

\eta=0.01

, 100 epochs,

1.7\,\text{s}

total) point predictions with fixed

\pm 2\beta

bands on the same four output channels. The confidence interval is constant across the input range because a point estimate carries no epistemic uncertainty.

Figure 19: Hinton diagram of DMA posterior weight means after 3 epochs. Each block corresponds to one weight matrix; square size encodes

|\mu_{w}|

, colour encodes sign.

Figure 20: Training NLL (per example per output) vs. epoch. DMA (solid black) converges at epoch 3 with no learning-rate tuning. Adam with

\eta=0.01

reaches NLL

-0.79

after 100 epochs;

\eta=0.001

has not converged (NLL

+0.84

);

\eta=1.0

diverges. Values above

8

are clipped to

8

for display. AdamW

\eta=0.1

\lambda=0.1

(solid purple) reaches NLL

-0.47

: weight decay regularises the trajectory but the chosen

\eta

does not reach the same final NLL as the best-tuned Adam.

Figure 21: Extrapolation calibration curve for DMA on the 1932-weight network (

N=1500

), averaged across all four output channels. Test points in

[-6,\,-4]\cup[4,\,6]

\Delta=+0.026

(slightly conservative); the 83-weight result has a 20-seed median

\Delta=-0.09

(Appendix E.7 ).

Explore similar work

Composing Non-Conjugate Factor Graphs with Closed-Form Variational Inference
May 28, 2026Mykola Lukashchuk, Kyrylo Yemets, Wouter M. Kouw +4Variational Inference Factor Graphs
Eindhoven University of Technology, the Netherlands · Lviv Polytechnic National University, Lviv, Ukraine · Lazy Dynamics, Utrecht, the Netherlands +1
Amortized Factor Inference Networks for Posterior Inference
May 26, 2026Joohwan Ko, Justin DomkeBayesian Inference Posterior
Manning College of Information and Computer Sciences University of Massachusetts Amherst
Delta-AI: Local objectives for amortized inference in sparse graphical models
Oct 3, 2023Jean-Pierre Falet, Hae Beom Lee, Esmeralda S. Whitammer +6Probabilistic Inference Probabilistic Graphical Models
Mila – Qu´ebec AI Institute, Universit´e de Montr´eal · Mila – Québec AI Institute, Université de Montréal

May 28, 2026, cs.LGJ/K move · Enter open · S save

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

Mykola Lukashchuk, Kyrylo Yemets, Wouter M. Kouw, Dmitry Bagaev+3

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

Variational Inference Factor Graphs Probabilistic Model Discontinuous Gating Mean-Field Limit

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Amortized Factor Inference Networks for Posterior Inference

Delta-AI: Local objectives for amortized inference in sparse graphical models