Variational Message Passing Neural Network for Maximum-A-Posteriori (MAP) Inference
Abstract
Maximum-A-Posteriori (MAP) inference is a fundamental task in probabilistic inference and belief propagation (BP) is a widely used algorithm for MAP inference.
Though BP has been applied successfully to many different ﬁelds, it offers no performance guarantee and often performs poorly on loopy graphs.
To improve the performance on loopy graphs and to scale up to large graphs, we propose a variational message passing neural network (V-MPNN), where we leverage both the power of neural networks in modeling complex functions and the well-established algorithmic theories on variational belief propagation.
Instead of relying on a hand-crafted variational assumption, we propose a neural-augmented free energy where a general variational distribution is parameterized through a neural network.
A message passing neural network is utilized for the minimization of neural-augmented free energy.
Training of the MPNN is thus guided by neural-augmented free energy, without requiring exact MAP conﬁgurations as annotations.
We empirically demonstrate the effectiveness of the proposed V-MPNN by comparing against both state-of-the-art trainingfree methods and training-based methods.
INTRODUCTION
Given a probability distribution of a set of random variables, a Maximum-A-Posteriori (MAP) inference problem involves identifying the most probable conﬁguration of a subset of unobserved random variables with observed evidence for the rest of the variables.
MAP inference problem has been studied in different communities, such as discrete energy minimization [Kappes et al., 2013] where optimization solvers are designed to directly solve for the optimal solution (i e., the most probable conﬁguration).
Solving the
MAP problem exactly is NP-hard, even with binary variables [Kolmogorov and Zabin, 2004, Cooper, 1990].
MAP inference on a probabilistic graphical model (PGM) is a fundamental task in probabilistic inference, where the joint probability distribution of a set of random variables is captured by a PGM.
Such task has lots of real-world applications such as image semantic segmentation in computer vision [Knobelreiter et al., 2020] and protein structure prediction in biochemistry [Soni et al., 2010].
In this work, we focus on MAP inference inside PGM context.
Different probabilistic inference algorithms have been proposed leveraging underlying structures of graphs, with belief propagation (BP) via message passing [Murphy et al., 2013] being a popular and widely used one.
Besides, for efﬁcient approximate inference, variational methods have been widely considered whereby probabilistic inference is reformulated as an optimization problem.
Variational assumptions are introduced over variational distributions such as mean ﬁeld assumption [Barabási et al., 1999] and Bethe assumption [Yedidia et al., 2001a].
Under mean ﬁeld assumption, a variational distribution can be fully factorized which in general does not hold on an arbitrary graph.
Bethe assumption is relaxed and is true on loop-free graphs.
Variational BP is to perform variational inference through message passing and is theoretically grounded on the well-established connection between BP and Bethe free energy [Tatikonda and Jordan, 2002, Yedidia et al., 2003, 2000, 2001a, Heskes, 2004].
Variational BP under Bethe assumption is exact on loop-free graphs, but its performance on an arbitrary loopy graph remains inaccurate without performance guarantee [Cannings et al., 1976, Shenoy and Shafer, 2008].
Different works based on variational BP have been proposed to improve the performance on loopy graphs, all of which rely on speciﬁc variational assumptions, resulting in speciﬁc families of variational distributions.
In this work, we propose a variational message passing neural network (V-MPNN) for improved MAP inference performance on loopy graphs.
V-MPNN leverages both the power of neural networks in modeling complex functions and the well-established algorithmic theories on variational
BP.
In particular, a neural-augmented free energy is proposed where variational distribution is parameterized via a neural network.
An optimal variational condition is explored during training.
Minimization of neural-augmented free energy is achieved through a message passing neural network (MPNN), which performs probabilistic inference through message passing.
The training of the MPNN is guided by neural-augmented free energy, which is different from existing neural-network-based inference methods that require exact inference results as annotations.
Without requiring labeled training data, our proposed V-MPNN is data efﬁcient.
More importantly, our model can scale up to large graphs where exact inference results are unobtainable.
RELATED WORKS
MAP inference.
MAP inference can be directly solved as an integer optimization problem [Wu et al., 2020] or can be relaxed to be a linear optimization problem (LP).
With the constraints on marginals enforcing global consistency, i.e., marginal polytope, exact MAP inference can be achieved under LP relaxation
[Wainwright and Jordan, 2008].
Marginal polytope is in general intractable.
Instead, constraints enforcing local consistency (e.g., pairwise consistency) are considered, that is, local polytope [Sherali and
Adams, 1990].
Local polytope yields pseudo-marginals that are local consistent but is not guaranteed to be exact.
Unfortunately, MAP inference under LP relaxation with local polytope remains computational prohibitive, particularly on large graphs [Yanover et al., 2006].
Variational BP for MAP inference.
Variational BP is to perform variational inference through message passing.
Variational BP is based on the connection between BP and Bethe free energy [Yedidia et al., 2001b].
Since Bethe free energy can exactly capture only loop-free graphs, BP is guaranteed to be exact on loop-free graphs and is only an approximate inference on loopy graphs.
Different techniques have been proposed to improve the performance of BP on loopy graphs, including initialization strategies [Koehler, 2019, Knoll et al., 2018], message update scheduling [Elidan et al., 2012, Knoll et al., 2015, Aksenov et al., 2020] and damping [Murphy et al., 2013, Pretti, 2005].
In addition to these practical techniques, more sophisticated hand-crafted variational distributions are proposed, leading to different variational BP algorithms [Hazan and Shashua, 2010, Riegler et al., 2012].
For example, max-product tree-reweighted message passing (TRW-MP)
[Wainwright et al., 2005a] decomposed the original joint distribution into a convex combination of treestructured distributions.
A tree-reweighted variational free energy is correspondingly derived.
TRW-MP is guaranteed to produce exact MAP conﬁgurations under a certain condition but it suffers from convergence issues.
Existing studies show that the entropy term within a variational free energy heavily affects the algorithm performance
[Ravikumar et al., 2010, Meshi et al., 2012, Lee et al., 2020, Savchynskyy et al., 2011, Hazan and Shashua, 2012].
More speciﬁcally, when the entropy is concave and the variational free energy is thus convex, a class of message passing algorithms is obtained with convergence guarantee [Savchynskyy et al., 2011, 2012, Hazan and Shashua, 2012, Weiss et al., 2012, Meshi et al., 2015].
MAP inference error bound with convex free energy can also be derived.
In this work, we propose to further reduce the MAP inference error bound by leveraging neural networks.
Neural networks for probabilistic inference.
Neural networks have been considered for probabilistic inference tasks.
Yoon et al.
[2019] empirically demonstrated the usage of MPNN
[Gilmer et al., 2017] for probabilistic inference, including MAP inference and marginal inference.
The architecture of MPNNs follows a message passing scheme.
Messages and beliefs are parameterized by neural networks and are learned from observed probabilistic graphs annotated with corresponding exact inference results.
Though inspired by belief propagation, MPNN is solely learned from data.
Different works have been proposed along this line, the majority of which are for marginal inference.
Satorras and Welling [2020] proposed to reﬁne messages from belief propagation via messages learned in MPNN.
Kuck et al.
[2020] proposed a belief propagation neural network (BPNN) where beliefs are regularized by minimizing a Bethe free energy.
Zhang et al.
[2019] proposed a factor graph neural network (FGNN) that can perform MAP inference.
FGNN is proved to be equivalent to BP and thus can perform well only when ordinary BP does well.
Hence, FGNN does not explicitly address the poor inference performance issue of BP on loopy graphs.
All the neural-network-based methods mentioned above require either exact MAP conﬁgurations or exact partition functions as annotations for fully supervised training.
As a result, these methods are limited to small graphs where exact inference results are obtainable.
PROPOSED METHOD
We propose a variational message passing neural network (V-MPNN) for improving inference performance on loopy graphs and scaling up to large graphs.
V-MPNN leverages both the power of neural networks in modeling complex functions and the algorithmic theories on variational BP.
We begin with preliminaries that are necessary for later discussions.
We then introduce our proposed V-MPNN.
Towards the end of this section, we summarize the training objectives of the proposed V-MPNN.
3.1 PRELIMINARIES
In this work, we focus on MAP inference on discrete pairwise markov random ﬁelds (MRFs).
We ﬁrst deﬁne MAP inference on MRFs and then introduce the variational free energy.
We discuss different families of variational distributions and introduce the minimization of a variational free energy through message passing.
In the end, we show the connection between the optimality of minimizing a variational free energy and the exactness of MAP inference.
3.1.1 MAP Inference on Markov Random Field
Given a set of N random variables x = { x 1 , x 2 , ..., x N } in discrete space χ = χ 1 × χ 2 × ...
× χ M , their joint probability distribution is captured by an MRF G =
( V , E ) where | χ
i | =
k i is the number of possible states of each variable x i , |V| = N , |E| =
M with M being the total number of edges in the graph.
The joint probability distribution of x is deﬁned as,
[Equation]
where E refers to the set { ( i, j ) : i ∈ V , j ∈ N ( i ) , i < j } .
θ deﬁnes probability parameters of the graph G .
θ i
( x i ) is the unary potential of variable x
i and θ ij ( x i , x j ) is the pairwise potential of two neighboring variables x i and x j connected via edge ( i, j ) .
Given a graph G and its probability parameters θ , the MAP inference task is formulated as
[Equation]
3.1.2 Variational Free Energy
Variational method converts a probabilistic inference problem to an optimziation problem, solving for a variational distribution by minimizing a variational free energy [Blei et al., 2017].
Given a target joint distribution p ( x ) , Gibbs free energy as a function of a variational distribution q ( x ) is deﬁned as
[Equation]
U ( q ) =
P x q ( x ) E ( x ) is the average energy and the energy function E ( x ) is speciﬁed by p ( x ) .
H ( q ) = − P x q ( x ) ln q ( x ) is the entropy.
T ◦ is the temperature.
For MAP inference, temperature is speciﬁed to be a sufﬁciently small value ǫ ( T ◦ = ǫ ).
An optimal variational distribution is obtained as
[Equation]
Marginal polytope M ( G ) enforces global consistency as M ( G ) =
{ q : q ≥ 0; P x q ( x ) = 1 } .
This constrained optimization is strictly convex and q ∗ achieves zero KL divergence w.r.t.
the target distribution, that is, KL ( q ∗ || p ) = 0 .
Exact inference can be performed with q ∗ .
However, minimizing the Gibbs free energy over marginal polytope is in general computational prohibitive.
Variational assumptions are introduced for tractable variational distribution.
On pairwise MRF with the joint distribution deﬁned in Eq.
1, we have E ( x ) =
− P i ∈V
θ i ( x i ) − P ( i [Equation] j ) ∈E θ ij ( x i , x j ) and the average energy is computed as
[Equation]
The average energy becomes a function of local marginals { q i } i ∈V and { q ij } ( i,j ) ∈E with q i ( x i ) =
P x \ x i q ( x ) and q ij ( x i , x j ) =
P x \ ( x i ∪ x j ) q ( x ) .
We thus assume a variational distribution q ( x ) is a function of { q i ( x i ) } i ∈V and { q ij ( x i , x j ) } ( i,j ) ∈E , referred to as pairwise assumption .
Pairwise assumption is widely used on pairwise MRF and there exist various families of variational distributions under pairwise assumption as introduced below.
(BP)
[Murphy et al., 2013] and TRW-MP [Wainwright et al., 2005a] are the two most representative families of variational distributions under pairwise assumption.
In BP, the family of variational distribution is deﬁned as:
[Equation]
Correspondingly, we obtain a variational free energy (i.e,
Bethe free energy):
[Equation]
N ( i ) denotes the set of neighboring nodes of i -th node.
H ( q i ) =
− P x i q i ( x i )
ln q
i ( x i ) .
H ( q i , q j ) =
− P x i ,x j q ij ( x i , x j ) ln q ij ( x i , x j ) .
In TRW-MP, a convex combination of tree-structured distributions via spanning trees is employed for approximating probability distribution.
The family of variational distribution is deﬁned as
[Equation]
which is closely related to BP but differs in terms of an edge appearance probability ρ ij ∈ (0 , 1] .
Edge appearance probability ρ ij measures the probability of an edge ( i, j ) in a graph G being present in a randomly chosen spanning tree.
A variational free energy is correspondingly obtained as
[Equation]
TRW-MP is guaranteed to perform exact MAP inference under a certain post-checking condition [Wainwright et al., 2005a,b].
In summary, under the pairwise assumption, a variational free energy is of a general form:
[Equation]
Each of the variational BP algorithms (e.g., BP and TRWMP) is speciﬁc to a family of variational distributions, leading to an entropy approximation (i.e., a set of c
i and c ij in
Eq. 10).
The performance of a variational BP algorithm is hence limited by the corresponding variational assumption.
Differently, we propose to leverage the power of a neural network to automatically explore the optimal variational distribution family under the pairwise assumption.
Minimization of a variational free energy.
Given a variational free energy in Eq.
10, the optimal solution set { q ∗ i , q ∗ ij } i ∈V , ( i,j ) ∈E is obtained as:
[Equation]
with the local polytope constraint set L ( G )
= {{ q i , q ij } : q i ≥ [Equation] ; q ij
[Equation]
; P
[Equation] i q
[Equation]
∀ i ∈ V ; q i ( x i ) =
P x j q ij ( x i , x j ) , ∀ ( i, j ) ∈ E} .
This constrained optimization is in general not convex.
Its convexity depends on the concavity of the entropy term, which varies with different variational distribution families.
Solving for optimal solution can be implemented through message passing.
After convergence, ﬁxed-point solutions are guaranteed to be local optimal in minimizing G pairwise .
However, a variational gap usually exists between q ∗ and the target distribution p (i.e., KL ( q ∗ || p ) > 0 ), where q ∗ is computed from { q ∗ i , q ∗ ij } i ∈V , ( i,j ) ∈E .
MAP inference is performed as [Equation] = arg max x i q [Equation]
i ( x i ) .
MAP inference is exact if there does not exist a variational gap.
Otherwise, the inference remains approximate and is prone to errors.
3.2 VARIATIONAL MESSAGE PASSING NEURAL NETWORK
We now introduce the proposed variational message passing neural network (V-MPNN).
We ﬁrst introduce the proposed convex neural-augmented free energy whereby we parameterize variational distribution families via a neural network.
The proposed neural-augmented free energy is provable convex.
The minimal MAP inference error with the proposed neural-augmented free energy is upper bounded by an optimal entropy approximation.
We then introduce the minimization of the proposed convex neuralaugmented free energy through a message passing neural network (MPNN).
The MPNN performs inference through message passing with messages parameterized via neural network parameters.
In the end, we summarize the training objectives together with training procedures.
The overview of V-MPNN is shown in Figure 1.
3.2.1 Convex Neural-augmented Free Energy
Under the pairwise assumption, we introduce the proposed neural-augmented free energy G neural , where we parameterize variational distribution families through neural network parameters Φ .
Such parameterization is implicitly achieved via a neural-network-parameterized entropy approximation:
[Equation]
with input tensors q node = { q i } i ∈V ∈ R N × k and q edge = { q ij } ( i,j ) ∈E ∈ R M × k 2 .
The calculation of U ( q node , q edge ) directly follows the deﬁnition of the average energy and requires no free parameters to be learned.
Neural-networkparameterized entropy approximation is realized through a neural network with three sets of free parameters φ node ∈
R 1 × N , φ edge ∈ R 1 × M , φ ∆ ∈ R N × N .
In particular, a row-wise entropy calculation w.r.t.
each input tensor is ﬁrstly performed, producing intermediate values: h node =
{ H ( q i ) } i ∈V ∈ R N × 1 and h edge =
{ H ( q i , q j ) } ( i,j ) ∈E ∈
R M ×. The approximate entropy is then computed as
[Equation]
where ∆ h ∈ R N × N with ∆ h
( i, j ) =
H ( q i , q j ) − H ( q i ) if ( i, j ) ∈ E .
Otherwise, ∆ h ( i, j ) = 0 .
⊙ denotes element-wise product.
Neural network parameters Φ = { φ node , φ edge , φ ∆ } are unknown and are to be learned.
We theoretically prove the convexity of the proposed neuralaugmented free energy and the minimal MAP inference error bound through the following propositions.
Proposition 1.
Neural-augmented free energy G neural is provable convex with a strictly concave neural-networkparameterized entropy approximation H ( q node , q edge ; Φ) .
Proof: We prove this proposition by showing the neuralnetwork-parameterized entropy approximation is strictly concave.
We ﬁrst introduce the deﬁnition of concave entropy approximation [Heskes, 2004, Weiss et al., 2012]:
Deﬁnition (Concave Entropy Approximation).
An approximate entropy of Eq.
10 is strictly concave over local polytope L ( G ) if there exist ˆ c ij > 0 , ˆ α ij ≥ 0 and ˆ c i such that c i
= ˆ c i − P j ∈N ( i ) ˆ α ij and c ij = ˆ c ij + ˆ α ij
+ ˆ α ji .
The approximate entropy becomes
H ( { q i } , { q ij } ) =
X
[Equation]
[Equation]
With any set of parameters ˆ c ij > 0 , ˆ α ij ≥ 0 and ˆ c i , the approximate entropy of Eq.
14 is strictly concave.
Tensor operation deﬁned in neural-augmented free energy (Eq. 13) is equivalent to Eq.
14, with φ node , φ edge and φ ∆ corresponding to { ˆ c i } i ∈V , { ˆ c ij } ( i,j ) ∈E and { ˆ α ij } i ∈V ,j ∈N ( i ) , respectively.
exp ( · ) ensures the satisfaction of the constraint ˆ c ij > 0 .
ReLU ( · ) ensures the satisfaction of the constraint ˆ α ij ≥ 0 .
By deﬁnition of concave entropy approximation, the neural-network-parameterized entropy approximation H ( q node , q edge ; Φ) is strictly concave.
The neuralaugmented free energy G neural is thus convex over local polytope L ( G ) .
We now show the minimal MAP inference error with the proposed neural-augmented free energy is upper bounded by an optimal entropy approximation.
We ﬁrst deﬁne the
MAP inference error and then present the proposition with its proof.
Let q ∗ Φ denote the optimal solution set { q ∗ Φ ,i , q ∗ Φ ,ij } i ∈V , ( i,j ) ∈E minimizing the neural-augmented free energy G neural parameterized by Φ .
Given a target probability distribution p , the MAP inference error ∆ map ( q ∗ Φ , p ) is deﬁned as
[Equation] [Equation]
with p i
= P x \ x i p ( x ) and p ij = P x \ ( x i ∪ x j ) p ( x ) .
By deﬁnition, ∆ map ( q ∗ Φ , p ) ≥
[Hazan and Shashua, 2010].
Proposition 2.
MAP inference error is upper bounded by an entropy approximation scaled by ǫ , i.e.,
[Equation]
The minimal MAP inference error is hence upper bounded by an optimal entropy approximation with Φ ∗ = arg min Φ H ( q ∗ Φ ; Φ) .
Proof: Given the optimal solution set q ∗ Φ minimizing the neural-augmented free energy G neural parameterized by Φ , we have
[Equation]
By reorganizing the above equation, we have
[Equation]
Given the fact that H ( { p i } , { p ij } ) ≥ 0 , we can have
[Equation]
With an optimal set of neural network parameters Φ ∗ = arg min Φ H ( q ∗ Φ ; Φ) , the error bound becomes
[Equation] [Equation]
We thus show that the minimal MAP inference error is upper bounded by an optimal entropy approximation.
Proposition 2 is the basis for the proposed V-MPNN and it motivates the training of our neural-network-parameterized entropy approximation as we will introduce in Section 3.2.3.
In the end, we provide a brief comparison between the proposed neural-augmented free energy and existing variational BP algorithms:
Proposition 3.
Neural-augmented free energy subsumes existing variational distribution families (e.g., BP and TRW-MP) as a strict generalization.
The optimal MAP inference performance achieved with neural-augmented free energy is superior or comparable to existing variational distribution families, i.e., ∆ map ( q [Equation] Φ [Equation]
p
[Equation] ≤ ∆ [Equation] Proof: By manipulating neural network parameters, different existing variational distribution families can be realized with neural-augmented free energy.
For example, neuralaugmented free energy with Φ speciﬁed as φ node = 1 − |N
( i ) | , φ edge = 1 and φ ∆ = 0 is equivalent to BP.
Furthermore, given the fact that q ∗ Φ ∗ = arg min q G neural ( q ; Φ ∗ ) , we have
[Equation]
with q ∗ Φ fix denotes the optimal variational distribution minimizing neural-augmented free energy speciﬁed with ﬁxed parameters Φ fix .
By subtracting U ( { p i } , { p ij } ) on both sides of Eq. 21 and a re-organization, we have
[Equation]
with ∆ =
H ( q ∗ Φ ∗ ; Φ ∗ ) − H ( q ∗ Φ fix ; Φ ∗ ) .
If ∆ ≤ 0 , it is clear that ∆ map ( q ∗ Φ ∗ , p ) ≤ ∆ map ( q ∗ Φ fix , p ) .
If ∆ > 0 , we can have ∆ map ( q ∗ Φ ∗ , p ) ≤ ∆ map ( q ∗ Φ fix , p ) with a sufﬁciently small coefﬁcient ( ǫ → 0 ). To show the latter point, we ﬁrst note that ∆ = H ( q ∗ Φ ∗ ; Φ ∗ ) − H ( q ∗ Φ fix ; Φ ∗ ) ≤ H ( q ∗ Φ ∗ ; Φ ∗ ) . We then show that H ( q ∗ Φ ∗ ; Φ ∗ ) is upper bounded by a constant and ﬁnite value δ . For clear derivation, we use notation q and Φ and derive for H ( q ; Φ) in the following.
The derivation applies to arbitrary q and Φ .
To derive δ , we ﬁrstly re-organize the entropy approximation H ( q ; Φ) deﬁned in Eq. 13 as [Equation] We now show that both H ( q i ) and H ( q i , q j ) are bounded by a constant value. For H ( q i ) , applying the Jensen’s inequality yields,
[Equation]
where k i indicates the number of states of variable x i .
Similarly, we have H ( q i , q j ) ≤ log k ij where k ij indicates the number of joint conﬁgurations of variables x i and x j .
Given the bounds for H ( q i ) and H ( q i , q j ) , we can conclude
H
[Equation]
δ is only a function of underlying graph and parameters Φ .
Thus, we have ∆ ≤ H ( q ∗ Φ ∗ ; Φ ∗ ) ≤ δ .
With ∆ ≤ δ , we now can show [Equation] δ is not a function of ǫ . Furthermore, given the mild assumption that neural parameters Φ are of ﬁnite values, δ is always ﬁnite.
With a constant and ﬁnite upper bound δ , there always exists a sufﬁciently small ǫ such that [Equation]
[Equation] , p ) . Theoretically, we show that the optimal MAP inference performance achieved with neural-augmented free energy is superior or comparable to existing variational distribution families.
3.2.2 Minimization of Neural-augmented Free Energy with MPNN
To minimize the neural-augmented free energy, we employ a message passing neural network (MPNN). In particular, q node = { q i } i ∈V and q edge = { q ij } ( i,j ) ∈E are parameterized by an MPNN, leading to unary marginal estimate q node (Ψ) and pairwise marginal estimate q edge (Ψ) .
We detail the MPNN module in the following.
Unary marginal estimation.
We map each node in MPNN to a variable in MRF with hidden feature h i ∈
R k i .
k i is the number of possible states of variable x i .
In total, we have node features h
= { h 1 , h 2 , ..., h N } and N is the total number of nodes.
Node feature h i corresponds to the unary marginal estimation in logarithmic space, up to a scale factor z i .
At every iteration t , each node receives a message from each of its neighboring nodes as
[Equation]
[Equation] M is a message function realized via a multi-layer perceptron (MLP).
The messages are then aggregated through summation, i.e., m t
+1 i
[Equation]
Each node then updates its hidden state with the aggregated message as:
[Equation]
[Equation]
U is a node update function and is customized based on
BP’s belief equation, instead of employing a gated recurrent unit (GRU) as a standard MPNN. Scale factor z i is calculated as z t +1 i
[Equation]
( x i )) .
The update process is repeated until convergence.
Estimated marginal probability of variable x
i (i.e., q i ) is obtained as
[Equation]
where h ( T ) i is the hidden feature from the last iteration.
Pairwise marginal estimation.
The pairwise marginal estimation is obtained as
[Equation]
where h ( T ) i and h ( T ) j are the respective hidden features for i -th node and j -th node from the last iteration of MPNN.
Eq. 30 is deﬁned based on BP’s pairwise belief equation.
We customize our MPNN based on BP with only message function M containing free parameters that need to be learned. The free parameters of MPNN Ψ hence refers to parameters within the message function M .
3.2.3 Training Objectives
In summary, we have two sets of parameters to be learned: Φ and Ψ .
The total training objective is based on neuralaugmented free energy, i.e.,
[Equation]
under the local polytope constraint L ( G ) .
To effectively perform the training with the neural-augmented free energy, we consider a two-phase alternative update.
For each iteration r , we ﬁrst update Ψ given the neural-augmented free energy speciﬁed with current Φ r , i.e.,
[Equation]
The constraints within local polytope L ( G ) are naturally satisﬁed by adopting BP belief equations for customizing our MPNN.
We then update Φ .
By deﬁnition of G neural in Eq.
12, we have max Φ G neural (Φ)
=
min Φ H (Φ) and the Φ is updated as
[Equation]
Following proposition 2, we theoretically prove that the entropy is the upper bound of the MAP inference error and hence updating Φ by minimizing the entropy is equivalent to minimizing the MAP inference error.
We update two sets of parameters alternatively until convergence.
After training, only the MPNN module with the optimal parameters Ψ ∗ is required for MAP inference.
MAP conﬁguration is obtained as
[Equation]
EXPERIMENTS
Datasets.
We consider 13 classic graphs for evaluation – these are the most representative graphs of real world models, and are employed widely in related works [Yoon et al., 2019].
Their structures are illustrated in Figure 2.
There are three loop-free graphs, i.e., STAR , TREE and PATH .
The other 10 graphs are loopy graphs, with the COMPLETE graph being the most complex one.
To simulate graphical models with different parameters, we randomly sample from uniform distributions [Wainwright et al., 2005a].
Particularly, we assume θ i ( x i ) =
b i x i and θ ij ( x i , x j ) =
J ij x i x j with x i = {− 1 , 1 } .
Pairwise parameters J ij are sampled from a uniform distribution, i.e., J ij =
J ji ∼ U [ − 1 , 1] .
Unary parameters b i are sampled from a uniform distribution as b i ∼ U [ − 0 .
For each type of graph, we simulate 1000 graphs for training and 100 graphs for testing.
GT MAP conﬁguration of each simulated graph is computed by enumeration.
Since enumeration is a computationally expensive process, we limit the sizes of the graphs.
Particularly, we consider two graph sizes: N=9 and N=15.
Evaluation metrics.
We employ the accuracy of estimated
MAP conﬁguration as the evaluation metric [Yoon et al., 2019].
Given a GT MAP conﬁguration x ∗ = { x ∗ 1 , ..., x ∗ N } , and an estimated MAP conﬁguration ˆ x = { ˆ x 1 , ..., ˆ x N } , the accuracy of ˆ x is calculated as #( x ∗ i =
ˆ x i ) N .
We report the average accuracy over testing graphs.
Experiment settings.
ADAM optimizer is employed for training with a learning rate 1 e − 4 .
In Eq 12, ǫ = 0 . 0001 .
For the message function M , a ﬁve-layer MLP is adopted and hidden features are of dimension 256 .
Messages propagate for T = 10 iterations.
MPNN is pre-trained at a message level, where a mean squared error between messages from M and messages from BP is used as the loss function.
4.1 COMPARISON TO STATE-OF-THE-ART METHODS
We compare the proposed V-MPNN to different state-ofthe-art methods for approximate MAP inference.
Speciﬁcally, we consider both training-free methods and trainingbased methods for comparison.
Training-free methods refer to optimization algorithms that do not contain neural network components and thus require no training procedure, such as the belief propagation algorithm.
In this work, we limit our comparisons to message-passing-based optimization approaches.
Training-based methods refer to neuralnetwork-based methods for probabilistic inference tasks.
4.1.1 Comparison to Training-free Methods
We consider three training-free methods: BP [Murphy et al., 2013], TRW-MP [Wainwright et al., 2005a] and max product linear programming (MPLP)
[Globerson and Jaakkola, 2007].
For all these three methods, we apply the same stopping criterion: if the maximum number of iteration t is larger than 200 or the average difference between beliefs from two consecutive iterations is sufﬁciently small, i.e.
[Equation]
− b t i | 2 < 1 e − 7 , we break the algorithm and obtain the estimated inference results 1 .
Following [Wainwright et al., 2005a], for both BP and TRW-MP, we apply message damping in log-space with damping parameter set to be 0 .
The edge appearance probability in
TRW-MP is set as ρ ij = |V|− 1 |E| .
Results are presented in Table 1.
As shown, we can see that V-MPNN achieves the best average accuracy with both sizes of graphs.
On each type of graph, V-MPNN achieves overall better performance than the other three baselines.
On loopy graphs, though the performance of all the algorithms decreases as the complexity of the graph increases, V-MPNN achieves better accuracy compared to the other three baselines. On CIRCULAR LADDER with
[Table] nodes, V-MPNN achieves 73% accuracy, which is 20% higher than the accuracy achieved by TRW-BP.
On loopfree graphs, such as STAR , TREE , and PATH , BP is guaranteed to produce the exact MAP conﬁguration, and thus always achieves 100% accuracy.
Though the proposed VMPNN is theoretically shown to be a strict generalization of BP, training of the MPNN is not guaranteed to ﬁnd the global optimal, leading to MAP inference errors on loopfree graphs.
4.1.2 Comparison to Training-based Methods
We compare the proposed V-MPNN to a training-based method: node-GNN [Yoon et al., 2019] for MAP inference.
Node-GNN 2 is the state-of-the-art method that employs neural networks for probabilistic inference tasks.
We employ the suggested hyper-parameter settings stated in the paper to perform the experiments.
Results are presented in Table 2.
C - LADDER denotes CIRCULAR LADDER .
As shown, V-MPNN achieves signiﬁcant graph_ [Table] better average accuracy with both sizes of graphs without requiring exact MAP conﬁgurations for training.
Across different types of graphs, V-MPNN achieves overall better performance than node-GNN and signiﬁcantly outperforms node-GNN on loop-free graphs.
On TREE with 9 nodes,
V-MPNN achieves 96% accuracy, which is 19% higher than the accuracy achieved by node-GNN.
These results show that, under the guidance of well-established algorithmic knowledge, the proposed V-MPNN can be trained to achieve outstanding performance, without requiring exact
MAP conﬁgurations as annotations.
4.2 ABLATION STUDY
In our experiments [Table] MPNN is pre-trained with BP’s message equation. In practice
[Table] we ﬁnd this pre-training step is crucial since directly using the neural-augmented free energy objective G neural without pre-training can easily make training diverge.
We thus adopt the pre-training step and then ﬁne tune the model with G neural . To better understand the effectiveness of the neural-augmented free energy objective G neural , we perform an ablation study.
Particularly, we compare the performance of V-MPNN to the performance of V-MPNN with pre-training only. We considerclassic graphs with 9 nodes. Results are shown in Table 3.
[Table] As shown
[Table]
ﬁne-tuning through G neural improves V-MPNN’s average performance compared to V-MPNN with pretraining only.
On CYCLE , the performance of V-MPNN is improved by 5% with G neural through ﬁne tuning.
From the results, we can see that the neural-augmented free energy G neural introduces important effect on the inference performance of V-MPNN, particularly on loopy graphs.
CONCLUSION
In this work, we proposed a variational message passing neural network for MAP inference.
Instead of relying on a speciﬁc family of variational distributions, we proposed a neural-augmented free energy where variational assumptions are parameterized via a neural network.
An optimal family of variational distributions is learned through training.
An MPNN is employed for efﬁcient inference through message passing.
Training of the MPNN is performed under the guidance of neural-augmented free energy, without requiring exact MAP conﬁgurations as annotations.
In our experiments, the proposed V-MPNN outperforms both state-of-the-art training-free and training-based methods for MAP inference, demonstrating the effectiveness of the proposed method. ACKNOWLEDGMENTS
This work is supported by the Rensselaer-IBM AI Research
Collaboration (http://airc.rpi.edu), part of the IBM AI Horizons Network (http://ibm.biz/AIHorizons).