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

AI summary

Overview

Research area: Approximate probabilistic inference on factor graphs, with a concrete application to Bayesian neural network (BNN) inference.

Technical level: Advanced. The paper is written for readers comfortable with factor graphs, the sum-product algorithm, exponential families, KL divergence, and Gaussian message passing.

Scope: The paper introduces Direct Message Approximation (DMA), a consistency-based alternative to expectation propagation (EP) and variational message passing (VMP) that approximates factor-to-variable messages directly rather than approximating the marginal at each factor edge, and it instantiates this framework as a BNN inference algorithm.

What This Paper Is About

EP and VMP both work by approximating the marginal distribution at each factor edge and then recovering the outgoing factor-to-variable message through a division by the incoming message. That single design choice, the authors argue, causes three problems: it forces an iterative round-robin schedule until a fixed point is reached, it can produce Gaussian messages with negative precision when the projected marginal is wider than the incoming message, and in VMP it collapses to a point mass at Dirac-delta factors, which means VMP cannot maintain a non-trivial weight posterior in a Bayesian neural network.

The paper's goal is to replace that marginal-first design with a message-first one. DMA approximates the factor-to-variable message directly, guided by a "consistency condition" requiring the approximation to be exact whenever all other incoming messages to the factor are Dirac deltas. In that limit the factor-to-variable integral reduces to a deterministic function evaluation, so the exact message is always known and tractable.

Key Contributions

1. The DMA framework and a master theorem. The paper defines a concentration consistency condition as the design criterion for direct message approximation (Definition 2.1), gives a construction recipe using moment matching, and proves a master theorem (Theorem 2.3) bounding the marginal KL at any edge by a quantity linear in the message KL, delta. The bound is edge-local and requires no graph contraction condition, in contrast to the closest prior result (Ihler et al., 2005), which bounds perturbations in loopy BP via graph-level contraction rates. Three structural corollaries follow: asymptotic Dirac-input consistency (Corollary 2.4), no EP-style inner-loop iteration (Corollary 2.5), and no negative-precision messages (Corollary 2.6).

2. Explicit DMA messages for the product factor. The paper derives forward and backward messages for the product factor delta(z − xy) (Propositions 3.1 and 3.2). The forward moments are exact moments of Z = XY under independent Gaussian inputs, so the only approximation is the Gaussian projection. The backward message is the technically demanding case: the distribution of X = Z/Y has no closed-form Gaussian representation, the standard EP backward message is structurally improper, and the ratio of two Gaussians has no finite mean. The authors resolve this with a log-normal intermediate technique and prove a complementary O(1/r²) accuracy bound in the concentrated-input regime (Theorem 3.3), with the probability of the excluded input region bounded by Φ(−r_y/2) = O(e^(−cr²)).

3. Explicit DMA messages for the leaky-ReLU factor. The activation factor delta(y − ReLU_α(x)) is handled by truncated-Gaussian moment matching using the Mills ratio, with an explicit normalisation correction depending on α. The backward message is proper for all α > 0; for α = 0 the backward message is improper (Remark B.6). Both factors are shown to verify the consistency condition analytically.

4. A complete BNN inference algorithm and its validation. The two factors suffice to assemble a BNN inference algorithm with one forward/backward sweep per training example, no EP-style inner-loop fixed-point iteration, and no gradient learning-rate hyperparameter. The algorithm runs in O(sum_l d_l d_(l−1)) time per example, matching the asymptotic complexity of a standard neural network forward-backward pass with a modest constant-factor overhead (1.8× per epoch), and requires only two scalar parameters (tau, rho) per weight belief plus one stored factor message per mini-batch per weight. A 23× larger network in Appendix G confirms the procedure remains computationally viable without algorithmic changes.

Main Findings

• The master theorem bound holds in every tested configuration. Across 192 leaky-ReLU factor configurations spanning a range of slopes, input widths, and SNR values (Appendix C.1), the Theorem 2.3 bound holds in every case with a maximum normalised ratio of 0.97, directly certifying the theorem. The bound is KL[p/Z || p̂/Ẑ] ≤ (||m_{X_j→f}||_∞ / Z) · delta.

• All three pathologies are structurally eliminated. By Corollary 2.5, all factor-to-variable messages for a single training example can be computed in one forward sweep followed by one backward sweep with no EP-style inner-loop fixed-point iteration. By Corollary 2.6, no approximate message can have negative precision because DMA produces a valid member of the exponential family by construction. By Corollary 2.4, the upper bound vanishes in the Dirac limit for factors of the form delta(x_j − g(x_(−j))) with g twice continuously differentiable near the point and nonzero gradient, regular in the tails.

• Forward product moments are exact, not approximate. For the factor delta(z − xy), the DMA message to Z is N(z; m_z, s_z²) with m_z = mu_x·mu_y and s_z² = sigma_x²·sigma_y² + mu_x²·sigma_y² + mu_y²·sigma_x², which are the exact moments of Z = XY. In the Dirac limit both moments converge to those of a point mass at mu_x·mu_y, so delta = 0.

• The improper product backward message is characterised despite being improper. The DMA backward message in natural parameters is N(tau_hat_x / rho_hat_x, 1/rho_hat_x) with tau_hat_x = tau_y · r · tau_z / D and rho_hat_x = r²/D, where r := tau_y⁴ rho_z / (rho_y (tau_y² + rho_y)) and D := tau_z² rho_y + tau_y² rho_z + rho_y rho_z. Theorem 3.3 bounds the KL to a proper truncated reference by C_kappa · r^(−2) for r := min(|mu_y|/sigma_y, |mu_z|/sigma_z) ≥ 2, assuming the ratio r_y/r_z is bounded above and below by a constant kappa ≥ 1.

• The 1D regression experiment shows structurally widening uncertainty. With N = 200 noise-corrupted observations at inputs drawn uniformly from [−2.5, 1.5], a 7-component standardised feature map, and a two-hidden-layer leaky-ReLU network (d = 6, 5) with Gaussian likelihood (beta = 0.2) trained for up to 200 epochs with mini-batches of 20, the predictive mean tracks the true function closely within the training region and the ±2σ intervals widen in the extrapolation regions where the data provides no information.

• No learning-rate hyperparameter and few epochs are needed on the correctly specified 1D task. Compared with Adam, AdamW at four weight-decay values, and the diagonal Laplace approximation (Appendix E), DMA requires few epochs (median epoch 16) with no gradient learning-rate hyperparameter. Over 20 seeds, DMA achieves a median extrapolation NLL of 0.80 versus 4.05 for AdamW (Appendix E.3).

• Calibration is comparable and not systematically overconfident. Over 20 seeds, DMA has a median calibration error delta = −0.09 and diagonal Laplace has delta = −0.10; neither is systematically overconfident (Appendix E.7).

• Under model mismatch, DMA uncertainty still widens. When the data are generated by a wider network, DMA's posterior predictive continues to widen outside the training range, while Adam provides no epistemic uncertainty (Appendix F).

• Larger networks remain viable. Appendix G evaluates a 1,932-weight, four-output network (6 → 6 → 12 → 48 → 24 → 4, N = 1500). DMA obtains competitive results in 3 epochs (0.72 s total), compared with 1.7 s for Adam (eta = 0.01, requiring approximately 100 epochs), using 100 mini-batches with the default settings in Appendix E.1.

• DMA differs from assumed density filtering (ADF) by adding a formal consistency axiom. ADF also avoids cavity division, but for a structural reason: it is a forward-only sequential algorithm that absorbs each observation into a running prior and never performs a backward sweep, so weight beliefs are never updated via backward messages — the step DMA's product backward message is designed to handle.

Methodology in Plain English

The authors start from an observation about the standard approach. EP and VMP both approximate the marginal distribution at each edge of a factor graph, then divide out the incoming message to get the outgoing message. That division is the root cause of the iterative scheduling, the negative-precision messages, and the point-estimate collapse, so the authors remove it.

Instead, they ask what property a good direct message approximation should have. Their answer is a consistency condition: when all the other inputs to a factor are concentrated at points (Dirac deltas), the factor-to-variable message becomes a simple function evaluation, which is exactly computable. So the approximation should recover that exact answer as the inputs concentrate.

To build messages that satisfy this condition, they use a two-step recipe. First, treat the incoming Gaussian messages as genuine independent random variables and compute the first two moments of the factor's output under that joint. Second, project the result onto the Gaussian family by moment matching, which (via Theorem A.1) minimises the forward KL divergence to the Gaussian family. This is applied to the two non-conjugate factors needed for a feedforward network: the product factor (weight times activation) and the leaky-ReLU activation factor. The conjugate pieces — the Gaussian prior on weights and the Gaussian likelihood — already produce exact messages, so no approximation is needed there.

For the product factor's backward message, direct integration fails because the ratio of two Gaussians has a 1/|y| singularity and no finite mean. The authors route around this by working in log-absolute-value space, where log|W| is approximately normal for |mu_w|/sigma_w much greater than 1 (a first-order delta-method approximation); then log|Z/Y| = log|Z| − log|Y| is the difference of two independent normals, whose parameters propagate in closed form. Converting back to natural parameters gives the explicit message formulas.

With both factors in hand, they assemble a full BNN factor graph. Each matrix-vector product decomposes into d_l × d_(l−1) 1D product factors plus d_l sum factors implementing the inner product, and each activation vector splits into d_l independent scalar leaky-ReLU factors. Messages are then computed in a single forward sweep followed by a single backward sweep per training example, mirroring the computation graph of an ordinary neural network. They validate on a scalar 1D regression task under a correctly specified model (the ground truth is itself a draw from the BNN prior, a standard simulation-study design that isolates inference quality from misspecification), and compare against Adam, AdamW at four weight-decay values, and the diagonal Laplace approximation. They also run a larger network in the appendix to check that the procedure stays computationally viable.

Why This Matters

Impact on research. The paper reframes a long-standing design choice in approximate message passing. Rather than treating the marginal approximation as primary and the message as derived, DMA treats the message as primary and derives a formal design criterion for it. The master theorem's edge-local, graph-agnostic bound is a different kind of guarantee from the graph-level contraction results that dominate the loopy BP literature (Ihler et al., 2005), and the O(1/r²) result for the improper product backward message addresses a case — the backward message of the product factor — whose closed-form treatment the authors state has resisted prior work. For Bayesian deep learning specifically, the paper offers an alternative to the marginal-first machinery used by EP-based BNNs (Hernández-Lobato and Adams, 2015), VMP (Winn and Bishop, 2005), Bayes by Backprop (Blundell et al., 2015), MC Dropout (Gal and Ghahramani, 2016), SWAG (Maddox et al., 2019), diagonal Laplace (MacKay, 1992), and IVON (Shen et al., 2024).

Real-world applications. The paper evaluates only on scalar regression tasks, so the following are natural connections of the method's properties rather than validated deployments:

• Uncertainty-aware regression where the model must signal that it does not know, such as scientific or engineering surrogate models queried outside the range of training data.
• Safety-relevant prediction settings where systematically overconfident extrapolation is costly, given that DMA's intervals widen in data-sparse regions, including under model mismatch.
• Bayesian neural networks with no gradient learning-rate hyperparameter to tune, since the algorithm requires none.
• Pipeline settings where training time matters and the complexity must match a standard forward-backward pass, given the reported 1.8× constant-factor overhead per epoch.

Industry relevance. The algorithm's cost profile is the main practical argument: O(sum_l d_l d_(l−1)) time per example matches a standard neural network forward-backward pass, the storage requirement is two scalar parameters per weight belief plus one stored factor message per mini-batch per weight, and there is no learning-rate hyperparameter. The reported larger-network experiment (1,932 weights, 0.72 s in 3 epochs versus 1.7 s for Adam at approximately 100 epochs) is a signal that the method is computationally plausible, though it is a single small-scale comparison rather than a broad benchmark.

Future Directions

• Scaling beyond the tested architectures. The largest network reported is 1,932 weights with four outputs and the architecture 6 → 6 → 12 → 48 → 24 → 4. Whether DMA remains competitive on modern deep architectures, convolutional layers, or sequence models is not established, and the paper's experiments are restricted to feedforward networks built from Gaussian, product, sum, and leaky-ReLU factors.

• Other non-conjugate factors. DMA is derived explicitly for two factors: the product factor and the leaky-ReLU factor. Extending the consistency-condition-plus-moment-matching recipe to other nonlinearities and other factor types is the obvious next step, and the paper notes that the standard ReLU (α = 0) has an improper backward message, leaving a gap for the most common activation in practice.

• Sharpening the improper-message guarantees. Theorem 3.3 bounds the KL to a truncated proper reference by C_kappa · r^(−2) with a constant depending only on kappa, and applies for r ≥ 2 with the ratio r_y/r_z bounded. How tight this bound is in practice, and whether the concentrated-input restriction can be relaxed, is left open.

• Evaluation on real data and broader baselines. The validation is a correctly specified 1D simulation plus an appendix extension, compared with Adam, AdamW at four weight-decay values, and diagonal Laplace. Performance under real-world model mismatch, against a wider set of Bayesian deep learning and natural-gradient baselines, and on tasks beyond regression is not reported.

Target Audience

This paper is for researchers and graduate students in probabilistic machine learning and Bayesian deep learning who already know factor graphs, the sum-product algorithm, exponential families, and the mechanics of EP and VMP. Readers working on approximate inference methods, message-passing systems in the tradition of TrueSkill (Herbrich et al., 2006), or uncertainty quantification for neural networks will find the most value. Practitioners looking for a drop-in uncertainty-aware training method should note that the paper reports a 1D regression study plus one larger-network comparison, so the evidence base is a proof-of-concept rather than a production benchmark. Readers without background in message passing should start with Appendix A, which the paper provides for that purpose.

Authors’ 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.