[Paper Review] Belief Propagation Neural Networks

Jonathan Kuck, Shuvam Chakraborty+5 more|arXiv (Cornell University)|Jul 1, 2020

Machine Learning and Data Classification47 references16 citations

TL;DR

This paper introduces Belief Propagation Neural Networks (BPNNs), a differentiable, parameterized generalization of belief propagation that maintains key theoretical guarantees while learning to estimate partition functions in factor graphs more accurately and efficiently. BPNNs achieve 1.7x faster convergence than standard BP on Ising models and compute model counts up to 100x faster than handcrafted methods with comparable accuracy, generalizing well to out-of-distribution graphs and larger instances.

ABSTRACT

Learned neural solvers have successfully been used to solve combinatorial optimization and decision problems. More general counting variants of these problems, however, are still largely solved with hand-crafted solvers. To bridge this gap, we introduce belief propagation neural networks (BPNNs), a class of parameterized operators that operate on factor graphs and generalize Belief Propagation (BP). In its strictest form, a BPNN layer (BPNN-D) is a learned iterative operator that provably maintains many of the desirable properties of BP for any choice of the parameters. Empirically, we show that by training BPNN-D learns to perform the task better than the original BP: it converges 1.7x faster on Ising models while providing tighter bounds. On challenging model counting problems, BPNNs compute estimates 100's of times faster than state-of-the-art handcrafted methods, while returning an estimate of comparable quality.

Motivation & Objective

• Address the gap in efficient, accurate approximate counting for combinatorial problems, which remain largely reliant on hand-crafted solvers despite advances in learned neural solvers for decision and optimization tasks.
• Develop a neural architecture that generalizes belief propagation (BP) while preserving its desirable theoretical properties, such as convergence on trees and provable lower bounds in loopy graphs.
• Enable generalization to out-of-distribution factor graphs and larger instances than seen during training, improving robustness and scalability.
• Improve estimation accuracy and speed for challenging problems like model counting, where standard BP fails to converge.
• Bridge the gap between differentiable neural networks and probabilistic inference by designing a layer that respects symmetries and maintains BP’s structural invariances.

Proposed method

• Propose BPNN-D, a differentiable, iterative, parameterized operator that generalizes standard belief propagation while preserving convergence guarantees on tree-structured graphs and provable lower bounds on loopy graphs under mild constraints.
• Introduce BPNN-B, a regression-based layer that maps belief trajectories over a fixed number of iterations to the true partition function, trading off some theoretical guarantees for improved flexibility and performance.
• Design both BPNN-D and BPNN-B to be invariant under factor graph isomorphisms, ensuring consistent predictions regardless of node labeling or graph permutation.
• Train the models end-to-end using mean squared error between predicted and true log partition functions on a dataset of factor graphs with known partition functions.
• Use a standard Adam optimizer with a batch size of 8 and a fixed learning rate of 2e-4 for 300 epochs on a single GPU to optimize the parameters.
• Leverage the structure of factor graphs and message-passing dynamics to embed learned message updates that generalize across graph topologies and sizes.

Experimental results

Research questions

• RQ1 Can a neural network architecture be designed to generalize belief propagation while preserving its key theoretical properties, such as convergence on trees and provable lower bounds?
• RQ2 How well can a learned BP variant generalize to factor graphs with different distributions of edge and class probabilities, or with significantly more variables than seen during training?
• RQ3 To what extent can BPNNs outperform standard belief propagation and graph neural networks in estimating partition functions and marginals on dense, fully connected graphs like stochastic block models?
• RQ4 Can BPNNs achieve high-accuracy model counting estimates for SAT problems—where standard BP fails to converge—using only a small number of training examples?
• RQ5 How does the performance of BPNNs compare to handcrafted approximate model counters in terms of speed and accuracy, particularly on hard instances?

Key findings

• On Ising models, BPNN-D converges 1.7x faster than standard belief propagation and consistently finds fixed points that yield tighter lower bounds on the partition function.
• BPNN-D generalizes effectively to Ising models drawn from a different distribution than the training set and to graphs with nearly twice as many variables, achieving significantly better log partition estimates than BP or a standard GNN.
• On community detection problems, BPNN outperforms both BP and a standard GNN, demonstrating improved performance in settings where BP is known to be effective both empirically and theoretically.
• For approximate model counting, BPNN learns accurate estimates from only tens of training problems and computes them 100x faster than state-of-the-art handcrafted solvers, while maintaining comparable estimation quality.
• On out-of-distribution stochastic block models with up to 20 nodes (80% more edges than in training), BPNN achieves lower RMSE on log partition function estimates than both BP and GNN, showing superior generalization.
• While BPNN is nearly 5 orders of magnitude more accurate than BP in marginal probability estimation on 15-node graphs, it is only about an order of magnitude more accurate than GNN, suggesting that marginal estimation is a harder task for BPNN despite its focus on partition functions.