🤖 machine learning

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

This paper demonstrates that closed-form variational inference can be preserved in deep probabilistic architectures by composing five specific factor-graph primitives, enabling the construction of universal function approximators like decision trees and Bayesian mixture of experts with calibrated uncertainty without requiring learned gating parameters.

Original authors: Mykola Lukashchuk, Kyrylo Yemets, Wouter M. Kouw, Dmitry Bagaev, żsmail Şenöz, Jeff Beck, Bert de Vries

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mykola Lukashchuk, Kyrylo Yemets, Wouter M. Kouw, Dmitry Bagaev, żsmail Şenöz, Jeff Beck, Bert de Vries

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

Imagine you are trying to predict the weather. You have a team of seven different forecasters: one is great at spotting rain, another is a wizard at predicting heatwaves, and a third is excellent at guessing wind speeds.

In the old way of doing this (standard machine learning), you would ask all seven for their opinions, give them fixed weights (e.g., "Rain Guy gets 20% of the vote, Heat Guy gets 10%"), and average their answers. The problem? Sometimes "Rain Guy" is actually terrible at predicting heatwaves, but the system doesn't know to stop listening to him. It treats his confidence as a fixed fact, not a feeling that changes with the situation.

This paper proposes a new way to build a "super-forecaster" team that is smarter, more flexible, and—most importantly—knows when it is guessing.

The Problem: The "Black Box" Trap

Usually, when you stack layers of smart components to make a deeper, more complex system, you break the math. The equations become so messy that computers can't solve them exactly. They have to guess the answer using trial-and-error methods (like sampling or black-box optimization). This is fast but often inaccurate, and it doesn't tell you how sure the system is about its answer.

The Solution: A Lego Set for Probabilistic Models

The authors discovered a special set of five Lego bricks (mathematical building blocks) that can be snapped together in any order to build deep, complex models. The magic trick is that no matter how you stack them, the math stays simple enough to solve exactly.

Here are the five bricks:

The Softdot: A basic calculator that mixes inputs (like mixing ingredients).
The Exponential Link: A switch that turns a number into a "confidence score" (it ensures the score is always positive).
The Gamma Prior: A rule that says, "We expect this confidence score to be somewhere in this range."
The Gaussian Likelihood: A standard bell-curve rule for how likely an observation is.
The Equality Node: A glue that says, "These two different wires must carry the exact same value."

How It Works: The "Smart Gating" System

The paper shows how to use these bricks to build a system that acts like a traffic controller.

Depth 0 (Static): Imagine a committee where everyone has a fixed seat. The system learns who is generally good, but it doesn't change based on the weather.
Depth 1 (Dynamic): Now, the system looks at the current input (e.g., "It's raining heavily"). It has a "gate" that says, "Okay, for this specific situation, let's trust the Rain Guy 90% and ignore the others." Crucially, the system doesn't just pick a winner; it calculates a probability distribution for who should be trusted. It knows how sure it is about that decision.
Depth 2 (Split-Branch Routing): This is the deep magic. The system builds a decision tree. It asks, "Is it raining?" If yes, go left. "Is it windy?" If yes, go right. It can create complex, branching paths to handle tricky situations (like an "XOR" problem, where the answer depends on a specific combination of factors).

The "Compiler" Analogy

Think of this framework like a programming language:

The Alphabet: The five Lego bricks.
The Grammar: The rules for how you can snap them together.
The Runtime: The computer engine that automatically figures out the math.

In most probabilistic programming, if you write a complex model, you have to manually derive the math equations for how to solve it. It's like writing a program and then having to hand-write the compiler for it every time.

In this paper, the authors built a universal compiler. You just snap the bricks together, and the "Bethe Free Energy" (a fancy mathematical objective function) automatically generates the exact equations needed to solve the model. You don't need to be a math genius to derive the updates; the system does it for you.

The Result: Calibrated Uncertainty

The biggest win is uncertainty.

Old Way: A neural network might say, "I predict 25°C," but it doesn't know if it's guessing or if it's 100% sure.
This Way: The system says, "I predict 25°C, but I'm only 60% sure because the data is weird." It gives you a "confidence interval" that is mathematically guaranteed to be correct based on the model structure.

Real-World Test: Time Series Forecasting

The authors tested this on predicting time-series data (like electricity usage or stock exchange rates). They combined seven different AI models (some good at trends, some at seasonality).

Their system learned to dynamically switch between experts based on the data.
It provided better accuracy than standard "Mixture of Experts" models.
Most importantly, it provided reliable uncertainty estimates. While standard models often became "overconfident" (claiming to be sure when they were wrong), this system correctly signaled when it was unsure.

Summary

This paper gives us a new way to build deep, complex AI models that are:

Composable: You can stack them as high as you want.
Exact: The math is solved exactly, not guessed.
Self-Aware: The model knows when it is uncertain, providing a "calibrated" sense of confidence.

It's like upgrading from a rigid, rule-based robot to a flexible, self-reflecting team of experts that knows exactly when to trust whom.

Technical Summary: Composing Non-Conjugate Factor Graphs with Closed-Form Variational Inference

Problem Statement
Probabilistic modeling often requires hierarchical architectures to capture complex dependencies, such as input-dependent expert reliability or routing mechanisms in ensemble forecasting. However, stacking probabilistic building blocks into deeper architectures typically breaks the property of closed-form inference. Standard approaches to handle non-conjugacy rely on black-box inference methods (e.g., Markov chain Monte Carlo, amortized variational autoencoders, or normalizing flows), which learn or sample the inference procedure but sacrifice analytical tractability. The authors ask whether probabilistic models can be composed hierarchically while retaining closed-form inference, avoiding the need for approximate or sampling-based inference procedures.

Methodology
The paper proposes a compositional framework based on Forney-style factor graphs (FFG) and variational message passing (VMP). The core contribution is the identification of a finite "alphabet" of five probabilistic factors that, when composed, preserve closed-form updates under mean-field factorization constraints.

The Alphabet: The authors define five factor-graph primitives:
- Softdot (Bilinear): $f_*(z | w, \phi, \tau) = \mathcal{N}(z | w^\top\phi, \tau^{-1})$
- Exponential Link: $f_{exp}(\gamma | z) = \delta(\gamma - \exp(z))$
- Gamma Prior: $f_G(\gamma | \alpha, \beta) = \mathcal{G}(\gamma | \alpha, \beta)$
- Gaussian Likelihood: $f_N(y | \mu, \tau) = \mathcal{N}(y | \mu, \tau^{-1})$
- Equality Node: $f_=(x_1, x_2, x_3) = \delta(x_1 - x_2)\delta(x_1 - x_3)$
Composition and Grammar:
- Depth 0 (Static Ensemble): A simple combination of experts with fixed precisions.
- Depth 1 (Precision-Gated Experts): Introduces a reusable "word" (the precision word $\pi$ ) composed of a softdot and an exponential link. This allows the precision of an expert to be a learned function of input features, rather than a fixed constant.
- Depth 2 (Split-Branch Routing): Composes multiple precision words with opposing switches to create piecewise-linear precision functions. Stacking these layers allows the model to encode arbitrary binary decision trees, partitioning the input space into regions where different experts are active.
Runtime and Inference:
- The inference algorithm is not prescribed but emerges from the model specification via the minimization of the Bethe Free Energy (BFE).
- Message Passing Types: The framework utilizes a hybrid approach. For non-deterministic factors (softdot, Gamma, Normal), Variational Message Passing (VMP) is used under mean-field constraints. For deterministic factors (Equality and Exponential Link), Belief Propagation (BP) is required because mean-field factorization would lead to degeneracy.
- Q-Conjugacy: A key theoretical result is that despite the non-conjugate nature of the exponential link, the messages remain tractable. The exponential link transforms a Gaussian message on the latent variable $z$ into a log-normal message on the precision $\gamma$ , and vice versa (log-gamma). By imposing form constraints (Gaussian for $z$ , Gamma for $\gamma$ ) and utilizing the moment-generating function of the Gaussian and sufficient statistics of the Gamma family, the authors derive closed-form fixed-point equations for the marginal posteriors. This ensures that all message computations and marginal updates are analytical, requiring no numerical integration or sampling.

Key Contributions

Closed-Form Compositional Inference: The authors prove that any model composed from the five-letter alphabet admits closed-form variational message passing. The set is closed under Q-conjugate updates, meaning the parametric form of messages (Gaussian or Gamma) is preserved throughout the graph, even across the non-conjugate exponential link.
Universal Function Approximation: The paper demonstrates that stacking split-branch routing layers encodes arbitrary decision trees. In the limit of sharp routing (high precision), the posterior mean of the output approximates any continuous function on a compact domain, establishing universal approximation capabilities while maintaining closed-form inference.
Bayesian Mixture of Experts: The framework is applied to ensemble time-series forecasting, yielding a Bayesian mixture of experts where gating functions are inferred rather than learned via gradient descent. This provides calibrated uncertainty over expert selection.

Results
The framework was evaluated on five benchmark time-series datasets (ETTh1, ETTh2, Exchange Rate, Electricity, Traffic) using an ensemble of seven diverse forecasters (including CNN, LSTM, DLinear, and quantile predictors).

Performance: The "Noisy Diagonal" variant (which separates forecaster disagreement from irreducible noise and uses diagonal covariance constraints) achieved the best trade-off between accuracy (MSE), calibration (Negative Log-Likelihood, NLL), and parameter efficiency.
Calibration: Unlike softmax-based Mixture-of-Experts (MoE) baselines, which often produced non-finite NLL values (indicating poor calibration or overconfidence), the proposed framework provided well-calibrated predictive distributions. The uncertainty estimates naturally captured routing ambiguity, with posterior standard deviations concentrating at decision boundaries in sharp-routing regimes.
Efficiency: The diagonal variants significantly reduced parameter counts (e.g., from ~15,000 to ~950 for the Exchange Rate dataset) while maintaining competitive or superior performance compared to full-covariance and deep neural baselines.

Significance and Claims
The paper claims to narrow the gap between the modularity of deep learning and the analytical rigor of probabilistic inference. By treating the alphabet of factors as a language and the Bethe free energy as a compiler, the authors enable the construction of deep, hierarchical probabilistic models without requiring new derivations for each architecture.

Modesty: The authors acknowledge that the alphabet is a specific subset of operations (multiplication, exponentiation, addition) and that purely linear-Gaussian alphabets cannot achieve universal approximation; the exponential link is essential for this expressiveness.
Implication: The work establishes that closed-form inference is not mutually exclusive with deep, non-conjugate architectures. It offers a principled alternative to black-box inference methods, providing marginal posterior distributions and calibrated uncertainty "for free" as a direct output of the inference algorithm, rather than as an add-on or approximation.