Learning Graphical Model Parameters with Approximate Marginal Inference

Likelihood based-learning of graphical models faces challenges of computational-complexity and robustness to model mis-specification. This paper studies methods that fit parameters directly to maximize a measure of the accuracy of predicted marginals, taking into account both model and inference approximations at training time. Experiments on imaging problems suggest marginalization-based learning performs better than likelihood-based approximations on difficult problems where the model being fit is approximate in nature.

💡 Research Summary

The paper addresses a fundamental difficulty in learning parameters of high‑dimensional graphical models such as Markov random fields (MRFs) and conditional random fields (CRFs). Traditional maximum‑likelihood (ML) learning requires exact computation of the log‑partition function and the model marginals, which is intractable for graphs with large tree‑width. Approximate ML approaches (e.g., pseudo‑likelihood, contrastive divergence) replace the exact marginals with those obtained from an approximate inference algorithm, but they still suffer from a mismatch: the inference algorithm used at test time is not taken into account during training, and the learning objective does not directly reflect the quality of the predictions that will actually be used.

The authors propose a “marginalization‑based” learning framework that directly optimizes a loss defined on the approximate marginals produced by the inference algorithm that will be employed at test time. In other words, instead of maximizing a surrogate likelihood, the training objective measures how close the inferred marginal distributions are to the ground‑truth labels (e.g., using cross‑entropy or squared error). This approach has two major benefits: (1) it incorporates the inference approximation into the learning process, allowing the parameters to compensate for systematic inference errors; (2) it is more robust when the model is misspecified, because the loss directly reflects the quality of the final predictions rather than the fidelity of the underlying probabilistic model.

To make this framework practical, the paper develops two key technical contributions for computing gradients of the marginal‑based loss with respect to the model parameters:

1. Perturbation‑based gradient estimation – The authors observe that the marginal vector is the gradient of an (approximate) log‑partition function. By running the approximate inference algorithm twice—once with the current parameters θ and once with a slightly perturbed parameter vector θ + Δθ—and measuring the change in the loss, they obtain a finite‑difference estimate of the gradient. This “perturbation” method is simple to implement, works with any black‑box inference routine, and can be integrated with automatic‑differentiation frameworks.

2. Truncated fitting – Conventional variational inference iterates until convergence before computing the loss, which is computationally expensive during training. The authors propose to stop the inference after a fixed number of updates (e.g., a few mean‑field sweeps or message‑passing iterations) and use the resulting intermediate marginals in the loss. This truncated approach dramatically reduces training time while still providing useful gradient information, because the loss is evaluated on the same approximate marginals that will be used at test time.

The paper reviews two popular approximate inference algorithms within this framework:

• Mean‑field (MF) – A fully factorized variational approximation that replaces the marginal polytope with a tractable subset. The MF updates are derived as block‑coordinate ascent steps, and the approximate log‑partition function is a lower bound on the true value.

• Tree‑reweighted belief propagation (TRW) – A convex relaxation that expands the feasible set to the local polytope and replaces the true entropy with a tractable upper bound involving singleton entropies and mutual informations weighted by edge appearance probabilities ρ_c. When TRW messages converge, the resulting marginals are the maximizers of an upper‑bounded variational objective.

Both algorithms fit naturally into the marginal‑based learning scheme: the loss is computed on the MF or TRW marginals, and gradients are obtained via perturbation or truncated fitting.

Experimental evaluation focuses on image segmentation tasks, a domain where accurate marginal predictions (pixel‑wise class probabilities) are crucial. The authors train CRF models on standard datasets using MF and TRW as the inference back‑ends. They compare three training regimes:

• Exact or approximate ML (using pseudo‑likelihood or contrastive divergence),
• Marginal‑based learning with full convergence of the inference algorithm,
• Marginal‑based learning with truncated inference (few updates).

Results show that marginal‑based learning consistently outperforms ML‑based baselines in terms of pixel accuracy and Intersection‑over‑Union (IoU). The advantage is most pronounced when the model is deliberately misspecified (e.g., using a simplified graph structure) or when the inference algorithm is far from exact. Moreover, truncated fitting reduces training time by 30‑50 % without sacrificing accuracy, confirming the practical benefit of the proposed approach.

Significance and limitations – By aligning the learning objective with the inference algorithm that will be used at deployment, the paper introduces a principled way to handle inference approximation and model misspecification. This “inference‑aware” learning can be applied to any differentiable approximate inference method, making it broadly relevant for computer vision, natural language processing, and other fields that rely on large graphical models. However, the method assumes that the chosen inference algorithm is stable; if message‑passing diverges or mean‑field updates become trapped in poor local optima, gradient estimates may become noisy. Additionally, the design of the marginal loss (choice of divergence, weighting) remains problem‑specific and may require domain expertise.

In conclusion, the work presents a compelling alternative to likelihood‑based training for graphical models. By directly optimizing the quality of approximate marginals, it achieves better predictive performance, greater robustness to model errors, and faster training through truncated inference. Future directions include extending the framework to more sophisticated variational families (e.g., structured mean‑field, amortized inference) and automating the selection of loss functions that best reflect downstream task metrics.