🤖 machine learning

Information Geometry of Message Passing

This paper introduces Natural-Gradient Message Passing (NGMP), a variational inference framework on Forney-style factor graphs that improves uncertainty calibration by locally projecting exact beliefs into edge-specific exponential families rather than averaging factors under neighboring beliefs.

Original authors: Mykola Lukashchuk, Kyrylo Yemets, Alex Ledbetter, żsmail Şenöz

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mykola Lukashchuk, Kyrylo Yemets, Alex Ledbetter, żsmail Şenöz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of modern science and engineering, computers are constantly asked to make sense of messy, incomplete information. Whether predicting the weather, diagnosing a disease, or guiding a robot, the goal is to reason under uncertainty. To do this, researchers use a framework called Bayesian inference, which treats knowledge as a set of beliefs that are updated as new evidence arrives. Imagine trying to guess the location of a hidden object; you start with a general idea, then adjust that idea every time you get a new clue. The challenge arises when the clues are complex and the hidden object has many moving parts. In these situations, the math required to calculate the exact answer becomes so heavy that even the fastest supercomputers cannot solve it in a reasonable time. Scientists have therefore developed shortcuts, known as approximate methods, which trade perfect accuracy for speed. These shortcuts work by simplifying the problem, often by assuming that different parts of the system are independent or follow simple, predictable patterns. However, these simplifications can sometimes throw away crucial details about how uncertain the system really is, leading to overconfident and potentially dangerous predictions.

A team of researchers has developed a new way to navigate this trade-off, offering a method that keeps the speed of these shortcuts while recovering the lost details about uncertainty. Their work focuses on a specific type of mathematical map called a factor graph, which breaks a complex problem into small, local pieces that can be solved individually. Traditionally, when these pieces are solved, the information passed between them is either the exact, messy truth (which is too slow to use) or a simplified, averaged version (which is fast but often inaccurate). The researchers discovered a middle path. They found a way to take the exact, messy information and project it onto the simplified shape that the computer can handle, but to do so in a way that preserves the most important part of the original message. They call this approach natural-gradient message passing. Instead of averaging out the complexities or ignoring them, this method carefully extracts the specific component of the uncertainty that fits the simplified model, ensuring that the final answer remains honest about what it does and does not know.

The researchers tested this new method on several real-world problems where uncertainty is a constant companion. In one experiment, they looked at a chain of data points, similar to tracking the sunspot activity over many years. When large gaps appeared in the data, the old methods became dangerously confident, predicting a smooth path even when the data was missing. The new method, however, correctly widened its uncertainty bands, acknowledging that it did not know what was happening in the gaps. In another test involving a stream of data arriving in batches, like a sensor feeding information to a computer in real-time, the old methods accumulated small errors that eventually caused the system to become overconfident and lose track of the true signal. The new method avoided this collapse, maintaining a stable and accurate view of the data even as it processed information piece by piece. These results show that the method is particularly valuable when the system is dealing with non-standard, complex relationships where the usual shortcuts fail to capture the full picture.

The core of this discovery lies in how the researchers handle the "messages" sent between different parts of the model. In the old approach, when a complex piece of data was sent to a simpler part of the system, the system would often average out the details to make the math work, effectively blurring the picture. The new method changes this by looking at the specific shape of the uncertainty at the receiving end. It takes the complex message and finds the best possible fit within the simpler shape, keeping the parts that matter and discarding only what cannot be represented. This is not a matter of guessing or averaging; it is a precise mathematical adjustment that ensures the simplified model retains the essential information about how uncertain the system truly is. The researchers found that when the data is clear and the uncertainty is low, the new method performs just as well as the old ones. But when the data is noisy, incomplete, or arrives in a stream where errors can build up, the new method shines, providing a much more reliable estimate of the true state of the world.

This work does not claim to solve every problem in artificial intelligence or statistics, nor does it suggest that perfect answers are now possible for all complex systems. Instead, it offers a practical improvement for a specific and common class of problems where uncertainty is persistent and difficult to manage. The researchers demonstrated that by changing how information is projected from one part of a model to another, they could prevent the system from becoming falsely confident. This is a significant step forward for applications like forecasting weather, monitoring financial markets, or guiding autonomous vehicles, where knowing the limits of your own knowledge is just as important as making a prediction. The method is computationally efficient, meaning it does not require massive amounts of extra time or power to run, making it a viable option for real-world systems that need to be both fast and honest about their uncertainties. By bridging the gap between exact calculations and practical approximations, this new approach provides a clearer, more reliable way for machines to reason in an uncertain world.

Technical Summary: Information Geometry of Message Passing

Problem Statement

Bayesian inference requires evaluating the posterior distribution $p(z|y)$ , which involves a high-dimensional normalizing constant that is generally intractable. While Variational Inference (VI) addresses this by replacing integration with optimization over a tractable family of distributions, standard approaches face a trade-off between global and local formulations:

Global VI: Optimizes a single global distribution $q_\lambda(z)$ . While this allows for a clean natural-gradient characterization of stationarity (Khan, 2025), the approximation is a monolithic object. For large graphical models, the sufficient statistics and Fisher matrices required for a global family may be computationally prohibitive or poorly matched to the model's factorization.
Local VI (Message Passing): Decomposes the problem into local computations on a factor graph (e.g., Belief Propagation, Variational Message Passing). This preserves modularity but struggles with closure. Exact Belief Propagation (BP) messages through non-conjugate factors often result in functional forms (e.g., non-Gaussian, non-Gamma) that cannot be represented by the finite-dimensional exponential families used in the surrounding graph.
The Gap: Existing local methods like Variational Message Passing (VMP) restore closure by averaging the factor under neighboring beliefs (tilting), which discards specific information about the uncertainty of the incoming messages. The paper asks whether the global natural-gradient stationary condition of VI can be localized into a rule for messages on a factor graph that preserves more of the exact BP information without requiring a global joint approximation.

Methodology: Natural-Gradient Message Passing (NGMP)

The authors propose Natural-Gradient Message Passing (NGMP), a scheme derived by constraining selected edge marginals to exponential families within the Bethe free energy framework.

Theoretical Derivation

Edge-Local Constraints: Starting from the Bethe free energy, the authors impose an exponential-family form constraint $q_i \in \mathcal{E}_i$ on specific edge marginals.
Stationarity Condition: By differentiating the constrained Bethe Lagrangian, they derive a local stationary condition (Theorem 1). At a stationary point, the natural parameter $\lambda_i$ of an edge equals the sum of two projected messages, one from each incident factor.
The Projection Rule:
- Let $\ell_{a \to i}$ be the exact BP log-message from factor $a$ to edge $i$ .
- The message sent by the factor is not the exact $\ell_{a \to i}$ (which may be outside the edge family) but its natural-gradient projection onto the tangent space of the receiving edge's exponential family at the current marginal $q_i$ .
- Mathematically, the projected natural parameter is:
  $\eta_{a \to i} = \nabla_{\mu_i} \mathbb{E}_{q_i}[\ell_{a \to i}]$
  where $\mu_i$ are the mean parameters of the receiving edge.
- The outgoing message is then $\hat{\mu}_{a \to i} \propto \exp(\eta_{a \to i}^\top T_i(z_i))$ .

Key Distinctions from Existing Methods

Vs. VMP: VMP replaces the factor with a "tilted" surrogate by averaging the log-likelihood over the other variables' marginals. This averaging discards the determinant correction terms arising from the coupling between variables. NGMP, conversely, projects the exact BP log-message (computed with the exact cavity distribution) onto the receiving family. It retains the representable component of the exact update while discarding only the orthogonal residual.
Vs. Expectation Propagation (EP): EP forms a cavity, restores the factor to create a tilted distribution, and projects the marginal onto the family. NGMP projects the log-message (the factor-to-edge update) directly.
Vs. Non-Conjugate VMP (NCVMP): NCVMP takes a single projection step of the tilted VMP message. NGMP projects the exact BP message.
Surrogate Interpretation: The authors show that NGMP can be implemented by replacing non-conjugate factors with conjugate surrogate leaves (pseudo-observations) on an auxiliary graph. The parameters of these surrogates are determined by the projected natural parameters. This allows standard BP to run on the surrogate graph while the outer loop updates the projection points.

Key Contributions

Localization of Natural Gradients: The paper establishes that the global natural-gradient stationarity condition of VI (Khan, 2025) has an edge-local form on Forney-style factor graphs.
The NGMP Rule: A derivation showing that the optimal message is the Fisher-metric projection of the exact BP log-message onto the tangent space of the receiving marginal.
Surrogate Model Equivalence: Demonstrating that NGMP updates are equivalent to running BP on a surrogate graph with conjugate leaves, providing a practical implementation path.
Theoretical Unification: Placing BP, VMP, EP, and NGMP within a single local variational language, distinguishing them by their information flow and projection points.

Experimental Results

The authors validate NGMP through ablation studies and larger-scale experiments, focusing on scenarios where uncertainty persists and is repeatedly composed.

1. Ablation Studies (Section 7)

Vanishing Uncertainty: In a Normal-Precision model with many observations, edge uncertainty vanishes. Here, NGMP and VMP converge, as the determinant corrections become negligible.
Poisson Smoothing (Sunspot Data): In a state-space model with non-conjugate Poisson observations, long stretches of missing data create persistent uncertainty.
- Result: VMP (and its projective variant PVMP) fails to propagate uncertainty correctly, keeping posterior variance artificially low ( $\approx \sigma^2/2$ ) regardless of distance from data. NGMP correctly propagates the cavity uncertainty, leading to wider credible intervals and significantly better Negative Log-Likelihood (NLL) on held-out data (e.g., 50% mask: PVMP NLL $\approx$ 5.4 vs. NGMP $\approx$ 4.9).
Heteroscedastic Filtering: In an online regression setting with 10 batches of data, VMP accumulates local under-dispersion errors, causing the posterior over noise weights to collapse (overconfidence).
- Result: NGMP remains nearly invariant to batching, maintaining a posterior variance close to the full-batch fit. The "batching penalty" (increase in NLL due to sequential processing) is negligible for NGMP ($0.009$ nats) but significant for VMP ($0.235$ nats).

2. Large-Scale Experiments (Section 8)

UCI Regression: On six datasets using a depth-three heteroscedastic hierarchy, NGMP achieves the lowest NLL on four datasets (Concrete, Energy, Power, Yacht) and competitive results on others, outperforming Bayes by Backprop, deterministic VI, and IVON.
ETTh Forecasting: In a probabilistic ensemble of frozen forecasters, NGMP improves calibration. While point forecasts (RMSE) are competitive across methods, NGMP significantly reduces NLL compared to PVMP and neural gate baselines, preventing the collapse of predictive scale observed in other methods.

Significance and Claims

The paper claims that NGMP offers a principled improvement over classical VMP in scenarios where edge uncertainty cannot be reduced and approximate updates are composed (e.g., smoothing through latent chains or online filtering).

Calibration over Point Accuracy: The primary gain of NGMP is in uncertainty calibration. While point estimates (RMSE) may remain similar, NGMP preserves the "determinant correction" terms that VMP averages out. This prevents the model from becoming overconfident when processing data sequentially or through non-conjugate factors.
Modularity: Unlike global VI, NGMP does not require a single global exponential family. It allows each edge to carry its own exponential family, preserving the modularity of the factor graph.
Practicality: The method is computationally efficient, requiring only a single natural-gradient step per edge update (compared to the inner optimization loops required by Projective VMP).
Future Directions: The authors suggest that this framework enables the compilation of factor graphs into surrogate programs (e.g., for RxInfer.jl) and offers a path toward continuous-state active inference by projecting messages in planning problems where exact factors are intractable.

In summary, the paper argues that by respecting the information geometry of the receiving marginal, NGMP retains the "locally relevant" part of the exact BP update that other local methods discard, leading to superior calibration in complex, non-conjugate, and sequential inference tasks.