Policy Message Passing: Modeling Trajectories for
Probabilistic Graph Inference

Zhiwei Deng 1 Xingguo Li 1 Greg Mori 2

Abstract
Learning to perform flexible reasoning over multiple variables is of fundamental importance forvarious tasks in machine learning. Graph neuralnetwork is an effective framework for building
inference processes among variables. A powerful graph-structured neural network architecture
should operate on graphs through two core components: (1) complex message functions designed
to model relations between nodes; (2) a flexible information aggregation process that executes
the reasoning through passing messages. However, despite the efforts on message designs, existing graph neural networks have limited power ofsystematically modeling flexible reasoning overgraphs. In this paper, we propose the Policy
Message Passing (PMP) algorithm, which takes
a probabilistic perspective and reformulates thewhole information aggregation as stochastic sequential processes. PMP is built upon the Variational Inference framework and defines a set of
automatic agents that observe the states and history to perform actions over the graphs. Theoretical interpretations are provided to show that
our algorithm can achieve improved optimization
efficiency, as well as a more effective learning
process. Experiments show that our algorithm
outperforms baselines by up to 40% on complex
reasoning tasks with parameters reduced up to
80% and is more robust to noisy edges on large
scale graphs.

1. Introduction
The world is constructed through complex schemes and processes involving interactions between elements and objects.
Being able to build models that can perform reasoning based
1
Princeton University 2 Borealis AI. Correspondence to: Zhiwei
Deng <zhiweid@cs.princeton.edu>.

Main part of the paper was done when Zhiwei Deng was interning
at Borealis AI
Copyright 2020 by the author(s).

on observations through establishing connections among elements and passing information, and make conclusions thatcan be generalized to new instances, has been a coveted
guerdon of machine learning research. To perform such flexible reasoning among entities or variables, graph-structured
models have been one of the default options to resort to (Sha& Pereira, 2003; Lafferty et al., 2001; Taskar et al., 2004).
Graph Neural Networks (GNNs) emerged in recent years
as the most popular deep learning models for processing
relational data (Li et al., 2015; Deng et al., 2016), learning
representation in graphs (You et al., 2019; Xu et al., 2018)
and performing complex reasoning based on observations(Teney et al., 2017; Palm et al., 2018; Santoro et al., 2017).
Existing Graph neural networks encompass two key components: (1) message functions Fθ which captures the relations and interactions among nodes; (2) a fixed information
aggregation process τ that passes information around thegraph through message functions. Given input features x
and targets y, the model maximizes a general objective as
follows with two components working cooperatively:
max p(y|x; Fθ , τ ).
θ(1)

Copious works are developed towards building powerful
message functions. Linear weight matrices are used forsemi-supervised node classification in (Kipf & Welling,2016). More complex non-linear message functions are
adopted for high-dimensional relational reasoning (Santoro
et al., 2017; Gilmer et al., 2017). To modulate information
flows in GNNs, gated functions (Li et al., 2015; Deng et al.,2016) and attention schemes (Veličković et al., 2017) areapplied to yield more flexible models. However, building
an algorithm that systematically discusses and models theinformation aggregation process τ has been under-explored.
In our paper, we propose a novel algorithm, termed Policy
Message Passing (PMP) that generalizes existing GNNs byrepresenting the inference process as a learnable component modeled by a neural network πφ . Instead of relying
on a repetitive and fixed aggregation procedure, our algorithm tackles the high-dimensional search space inducedby all possible passing combinations across the sequential
procedure, and perform inference in an active “reasoning”
manner. Our algorithm treats the embeddings in graphs asModeling Trajectories for Probabilistic Graph Inference

states, actively queries information, and performs reasoning
actions. Starting from an initial graph states s0 , the algorithm builds towards a distribution-based perspective thatconsiders all possible reasoning process and learns to generate more effective samples, resulting in a new objective tobe optimized:
max Eτ ∼πφ p(y|s0 , x, τ ; Fθ ) .
(2)through a sequence of local message passing on node states
V, have been developed for varied applications and tasks(Defferrard et al., 2016; Li et al., 2015; Kipf & Welling,
2016; Henaff et al., 2015; Duvenaud et al., 2015). Thesemethods share two common structures: propagation networks and output networks. In our paper, without loss ofgenerality, we follow a standard method that propagates
local messages between every pair of nodes.

More specifically, compared with Eqn. (1), we introduce
a neural network πφ in Eqn. (2) with additional flexibility to represent the distribution over all possible inference
trajectories in GNNs.

Propagation
networks: Assume that a graph state s0 =
{v0i } is initialized from corresponding local node features.
The propagation network, composed of a fundamental unit
termed as “graph-to-graph” module, consecutively generates
a sequence of graph states (s0 , s1 , ..., sT ), t ∈ {0, 1, ..., T },where each states corresponds to updated node embeddings.
At each timestep t, the “graph-to-graph” module takes thecurrent states st and performs message passing locally between node pairs. The functions that generate messages are
implemented as neural networks. The process is executedin the following order:φ,θ

By taking the distribution perspective of message passing
and reformulating the whole information aggregation process as stochastic sequential processes, our model is empowered with abilities to capture more aspects of the graphstructure reasoning problem: (i) Distribution nature - thepossible inference trajectories in graph goes beyond simple synchronous aggregation, and reside in a vast spacewith various possibilities, where each of them can lead to adifferent result; (ii) Consistent inference - reasoning steps
between nodes and across steps are not isolated, maintaining
the memory of reasoning history over time could lead tomore effective and consistent inference, or possible message
function compositions; (iii) Uncertainty - inference itself
encodes a prior knowledge on how to perform certain tasks,
representing the inference from a distributional perspective
maintains the uncertainty of the results of tasks.
Our primary contribution includes: (I) We reformulate thewhole inference aggregation from the Bayesian perspective
and introduce an algorithm that conducts function-set based
message passing with history-aware reasoning over graphs;(II) We propose a Variational Inference based framework
for learning the inference models, and provide theoreticalinterpretation of our model’s superiority from perspectivesof coding theory and optimization; (III) We explore multiple designs of our model and empirically demonstrate thathaving a powerful inference machine can enhance the prediction in graph-structured models while achieving betterparameter efficiency, on complex reasoning tasks and large
scale graphs with noisy edges.

2. Preliminaries: Graph Neural Networks
We start with introducing some notations. Let G ={V, E, X } denotes a graph with local features X = {xi },where V = {v 1 , v 2 , ..., v |V| }, vi ∈ Rd is the set of nodestates, E ∈ E is the graph structure given a priori and eachnode state v i is associated with a local feature xi . We use
Y = {y i } to denote the output set.
A large volume of Graph Neural Network (GNN) models,
which map the input features X to the desired output Y

(i,j)

mt



X

(i,j)
i
= fk(i,j) vti , vtj ; vt+1
=g
mt , vti , v0i ,j∈Niwhere fk(i,j) (·, ·) : R2d → Rd is the message function andk(i, j) : N20 → N0 is the function index determined byindices of node pairs. Function g(·) : R3d → Rd is theaggregator function that collects messages back to node vectors. Ni represents the neighbours of node i in the graph.
The whole aggregation process is fixed and defined as synchronous processes in existing GNNs.
Output networks: The output network of the model maps
the information in the graph states to the target output
through function fo : V → Y. The mapping function is generally represented by a neural network. Note that the target
variable y ∈ Y can take various forms, such as node-level
labels, graph-level labels or even sentences, depending on
the tasks. Without the loss of generality, we use node-level
labels in our following sections.
Limitations: The standard GNN described above represents a constrained class of models that focus on designing
expressive functions capturing interactions between nodes,
but ignore the modeling of the execution processes of passing messages, which is crucial in controlling and editing thenode vector values. This prevents the model from being amore general class of approaches, which learn to actively
reason over graph states through modeling both function
designs and the aggregation process. Potential benefits ofthis class of model include performing edge discovery or
removal dynamically, making compositional usage of message functions across reasoning steps, or learning the orderof solving problems (Chen et al., 2018).

Modeling Trajectories for Probabilistic Graph Inference

Blank message:
Node states

Universal messages 𝑚
01:

)
𝑣".#
#
𝑣".#

Agent

ℎ"

ℎ".#

Message set

Actions𝑣")
𝑣"#

Feature-based messages:

Agent

𝑣"$

0↦1

Message set

𝑣"%
MLP

𝑣"&

4↦3

(#,))

Messages
(#,))

𝒛"

𝑚"

(),#)
𝒛"

(),#)
𝑚"

(#,$)

𝒛"

(#,$)

𝑚"

…

…
(%,$)

Node states
)
𝑣".#
#
𝑣".#

𝑚"

(&,%)
𝒛"

(&,%)
𝑚"

…

…

: Concatenation
Global agent:
)
𝑣".#

4
ℎ"5#

Agent

…

ℎ"4

|3|

$
𝑣".#

(%,$)

𝒛"

: Attention based aggregation

%
𝑣".#
&
𝑣".#

𝑣".#

Local agent:

(),#)

)
𝑣".#

ℎ"5#

#
𝑣".#

Agent
(),#)

ℎ"

Figure 1. PMP algorithm. Our PMP algorithm consists of two components jointly trained under VI framework: (1) Function-set message
passing based on actions proposed by agents; (2) an agent that propose and executes the actions on the message function sets. Left: twopossible designs of messages. Right: two possible designs for agents. Middle: one step reasoning over the graph.

3. Our Approach: Policy Message Passing
In this section, we describe our proposed Policy Message
Passing algorithm, under the same problem setup. Given thegraph defined as G = {V, E, X } and the graph initial statess0 , we build our algorithm to model the full information
aggregation process (referred to as reasoning process), to
generate desired output Y. In the following, we first discuss
the general algorithm framework, then describe our model
designs and show how our proposed model broadens theclass of standard GNNs. We also theoretically show thebenefits of building a more expressive reasoning process.
Framework We formulate our algorithm from a Bayesian
perspective, and treat the reasoning process as a latent variable τ . The reasoning process is defined to be a sequenceof actions τ = (a0 , a1 , ..., aT ), where T is the number ofreasoning steps and at is the action defined over the wholegraph. As the message passing operation is defined locally(i,j)
between pairs of nodes, denoted as zt∈ Z, the graphlevel action can then be defined by the collection of local(i,j)
operations at = {zt }. As a consequence, if consider z
as a dicrete variable with Z denoting the set of possible val2
ues, the search space for the model is |Z|T |V | , where |V |is the number of nodes. Besides the large search space, theinformation passing process also has a distribution nature:
there are multiple possibilities of transmitting informationand manipulating node contents.
To model the distribution over the full reasoning process τ ,we define three key components: (1) the set of message func(i,j)
tions Fθ
= {fk (·, ·; θk )}(i,j) used for passing between
each pair of nodes i and j; (2) a reasoning agent πφ (·; φ),parameterized by φ, which determines the reasoning processes through message function selections; (3) a Variational
Information framework which performs inference and learning of the reasoning processes. Although we are able tospecify varied function sets based on node indices (i, j), we

adopt a shared function set Fθ = {fk (·, ·; θk )} across allnode pair (i, j) for simplicity. Detailed discussions for eachcomponent is provided in the following sections.
3.1. Function set-based message passing
Standard Graph Neural Networks typically assign a uniquemessage function fk (·, ·) to each node pairs, where the kis the index of function type. The message function needs
to account for node states and generate a message as an
output from node i to j. To build a powerful and flexible
reasoning process, we instead define a message function
set {fk (·, ·; θk )} to accommodate multiple possibilities between node pairs. The definition of message functions canbe either driven by the prior knowledge (e.g. potential
semantic relations between nodes), or separate randominitialized neural networks where message types emerge
during learning process. In the following, we present two
designs of message function set Fθ = {fk (·, ·; θk )}.
Feature-based messages Similar to standard GNNs,
feature-based messages are derived from node features orstates through a direct function mapping fk (vti , vtj ; θk ).
The function takes in node states at time t, and produce(i,j)
a message mt,k = fk (vti , vtj ; θk ). The produced mes(i,j)

sage mt,k is feature aware and could vary due to subtle
changes in node states. θk is the parameters for function k.
In our model, the function is represented by a Multi-Layer
Perceptron (MLP).
Universal messages To build models that can execute thepassing message operations to control node content, it might
be beneficial to learn a set of universal messages, which
can unload the modeling from typical non-linear functions
to the agent. Inspired by the concept of Vector Quantisation technique used in distribution approximation (van den
Oord et al., 2017; Louizos et al., 2017), we use a set of“universal” message vectors {m̄k } that represent possible

Modeling Trajectories for Probabilistic Graph Inference

prototypes of information. The prototypical messages arelater selected and transformed linearly (scale and shift):
(i,j)
(i,j)
(i,j)
mt,k = αt,k m̄k + βt,k , in which the selection, scaleand shift coefficients are actions generated by the reasoning
agents. The message function k of Fθ = {fk (·, ·; θk )} andcorresponding learnable parameters, in this case, can be
expressed as fk (i, j; θk = m̄k ) = I · m̄k .

We define the agent to be π(τ ; φ) = π(τ |s0 , s1 , ..., sT ; φ),where τ = (a0 , a1 , ..., aT ) is the sequence of actions thatthe agent takes. The distributionQover action trajectoryTτ can be further decomposed as t=0 q(at |s≤t ), whereq(at |s≤t ) is per-step action distributions. In the following,
we discuss two designs to show how the agents can collectinformation from the graph and propose actions.

Set-based message passing In both cases, we have a setof functions Fθ = {fk (·, ·; θk )}, k ∈ {1, ..., K} defined
between each pair of nodes i and j. Note that in general
the sets assigned for edges can vary or be manually specified. Given the message function set between node i and j,(i,j)
model’s reasoning agents choose to perform an action zt(i,j)
over the elements in the set. The final message mt
attime t passed from node i to node j is derived under the

(i,j)
(i,j)
(i,j)
operation defined as: mt= zt
Fθ , where zt (·)
is the function defined by the action, e.g. indicator function based selection or linear transformation. We define the
(i,j)
(i,j)
(i,j)
action to be a vector of operations [zt,0 , zt,1 , ..., zt,K ]on each message function type, the final messages between
PK
(i,j)
(i,j)
(i,j) node i and j are derived by: mt= k=1 zt,k mt,k .
For example, for feature-based messages, the action is a selection operation, while for universal messages, an extra
linear transformation should also be applied. For time stept, the full update of the node states will be:

Global agent A global reasoning agent collects informationfrom the whole graph. Specifically, it has access to all nodestates {vti } at step t. We use a variant of recurrent neuralnetwork (RNN), namely the Gated Recurrent Unit (GRU)(Chung et al., 2014), to maintain the hidden memory overreasoning history. Denoting the hidden state at time t as ht ,the agent first use a one-hop soft attention (Bahdanau et al.,2014; Deng et al., 2016) to check all the node state vectors,
G
based on its hidden memory state ht . The attention weightPiii
ct is calculated by ct = exp(c̄t )/ j exp(c̄jt ), where c̄it =Ti(Whc hGt ) (Wvc vt ). The attended node information is used
for updating hidden memory:
P i i G hG
t = GRU G
i ct vt , ht−1 .

i
vt+1
=g

K
XX

(i,j)

zt,k



fk (·, ·; θk ) , vti , v0i ,j∈Ni k=1

Action distributions are modeled and generated through an
Action head network. Note that action at is the set of all
local actions for message between each pair of nodes. The
(i,j)
(i,j)
distribution q(zt |vti , vtj , hGaret ) over local action zt
parameterized through the output of an Multi-Layer Perceptron: MLP vti , vtj , hGt , which could generate probability
values for Categorical distributions, or mean and varianceof Gaussian distributions.

where g(·) is a parameterized aggregation function that updates the node state by messages, previous node states andinitial states. fk (·, ·; θk ) is the k-th function that generates(i,j)
the message mt,k at each step t according to the inputs
from node i and j. Ni is the set of neighbors of node iin the graph. Note that we also define an “null” function
f0 (·, ·) = 0 as 0-th function for all sets to allow for the
choice of removing connectivity between nodes.

Local agent The local agent-based modeling instead builds
a pool of agents, in which each of them accounts for theoperation of a single pair of nodes. Similar to the global
agent setting, every local agent uses a GRU to maintainthe memory hidden state based on the local state dynamics
of the corresponding node pair. All agents share the samelearnable parameters to generalize across different graphsizes. The memory hidden states are updated as follows:(i,j)(i,j) ht
= GRUL vti , vtj , ht−1 .

3.2. Execution Agents

Based on the hidden states, at every step t, the agent(i,j) (i,j)
proposes an action distribution q(zt |ht ). The sufficient statistic of the action distribution is generated from(i,j) MLP ht.

Starting from the initial states s0 , an effective information
aggregation process will execute the local message passing operations to iteratively edit, transform and control thegraph states, leading to a states sequence (s0 , s1 , ..., sT ).
The graph states st should be informative about target Y andinduce an easy mapping for the output networks (describedlater) fo : V → Y. We define the graph states at each time
step t to be the collection of node states and message states
st = {vti }. To derive such a sequence of states, we build
an agent that models the whole reasoning process between
all pairs of nodes across the T steps. The agent tries tolearn a general inductive bias of performing operations, selecting order and composing actions across reasoning steps.

Action head and differentiable sampling: The action distribution q(at |s≤t ) at each step is defined over all the local(i,j)
actions at = {zt }. We use a factorized distribution
Q
(i,j)
(i,j)
q(at |s≤t ) = i,j q(zt |s≤t ), where each of themis derived from either the global agent or the local agent.
(i,j)
Given the distributions q(zt |s≤t ), we can now use themto perform operations by sampling actions from these distributions. In our algorithm, we define the action z according to the choice of message function set: (1) for standard

Modeling Trajectories for Probabilistic Graph Inference

feature-based message functions, we use the Categorical
distribution to represent the probabilities of selecting eachtype of functions; (2) for universal messages, where theaction operation is the combination of function type, co(i,j)
efficient scale and shift, we factorize q(zt |s≤t ) as the
product of Categorical distribution over function types, and
two Gaussian distributions over scalar coefficients α and β.
For Gaussian distributions, we use reparameterization trick(Kingma & Welling, 2013) for backpropagating the gradients. For Categorical distribution, although the sampling
process of actions is straightforward, differentiating throughthe sampling requires reparameterization tricks on discretevariables. In our implementation, the Gumbel softmax trick(Jang et al., 2016) is used for differentiable sampling:

variational distribution π(τ |s, φ) doesn’t have access to theground truth variable y, which enforce it to learn the general inductive bias of reasoning over graph only based on
the input
feature and graph states. By re-writing the term
KL π(τ |s)||p(τ |x, y) , we derive the following equations:
log p(y|x) − KL π(τ |s)||p(τ |x, y)hi= Eτ ∼π log p(y|τ , s) − KL π(τ |s)||p(τ )
h
i
Y

Y
= Eτ ∼π log p(y|τ , s) − KLq(at |s≤t )||p(at ) .
t

t

(i,j)

exp((log(zt,k ) + k )/λ)(i,j),
z̃t,k = PK(i,j)
m=1 exp((log(zt,m ) + m )/λ)ij
where zt,kis the probability value for each category, k issample drawn from the Gumbel(0,1) distribution and λ is
the temperature parameter of softmax. As λ approaches 0,
the distribution over z̃tij converges to a Categorical distribution with one-hot encoding, leading to the discrete action
selections for agents at each timestep. The sampled actions
are used to perform set-based message passing and update
node state st according to Sec. 3.1.

Output networks Similar to standard GNNs, we use an
output network to map the graph states v to the target output
variable y through fo : V → Y. As proposed in (Kipf et al.,2018; Palm et al., 2018; Wei et al., 2016), mapping from thegraph states st to the target variable at every step insteadof final step helps to solve the gradient vanishing issue andimprove the performance. We have similar observations
in our experiments and follow the same setup for training.
Formally, our log-likelihood term to be optimized is:X
log p(y|x, τ ) =log p(yt |s≤t , τ ).

In the above equation, the prior distribution over reasoning
process τ can be derived from the graph structure or manually specified. The decomposed KL term between q(at |s≤t )and p(at ) are approximated by Monte Carlo sampling as
the reasoning process progresses. Since KL Divergence is
non-negative, the right hand side of the equation is known
as Evidence Loader Bound (ELBO). During training, forevery graph data point (x, y, E), we initialize the graphstates as s0 from x, use the agent to sample actions fromproduced the action distributions, execute the action throughthe set-based message passing, and map the graph states tothe target y through the output network to maximize thelikelihood. The proposed action distribution is regularized
by the prior distribution induced from graph structure E.

4. Discussions and Theoretical insight of PMP
In this section, we show that our PMP algorithm generalizes
the existing GNNs by modeling the information aggregation
process as a stochastic reasoning process. The resulting
PMP based GNNs can potentially lead to a more effective
decoding process and more efficient optimization.

t

The output network is realized by an MLP with learnable
parameters. For example, in the node classification setting,the MLP generates the logits of Multinomial distibutions.
3.3. Learning objective
Given the node features x, the target variable y, the graphstructure E, and the agents π(τ |s, φ), we formulate themodel under the Variational Inference framework (Kingma& Welling, 2013): we would like to infer and approximate
the distribution as close to the true posterior distributionp(τ |x, y) over the reasoning process τ as possible. We
use the parameterized agent π to represent a variational
distribution family and minimize the KL Divergence between π and the true posterior: p(τ |x, y). Note that the4.1. Coding theory perspective
It is well-known that, from an information-theoretic view,
there is a close connection between Variational Inference,
Bits-Back Coding (Hinton & Van Camp, 1993; Honkela &Valpola, 2004), and Minimum Description Length (MDL)(Rissanen, 1978). In the following, we show that the advantage of our algorithm can be interpreted from the perspectiveof coding theory.
Let x be the node features for a graph, y be node-level
labels, and τ be the inference/aggregation trajectories ofgraph neural networks. As τ is not observed, we treat it asa latent variable to be inferred using Variational Inference
framework. Assume that we have an encoder q(τ |x, φ), a
standard Variational Inference framework optimizes (mini-Modeling Trajectories for Probabilistic Graph Inferencemize) the following quantity:h−E(x,y)∼data Eτ ∼q(τ |x;φ) log p(y|x, τ ; θ)i+ log p(τ ) − log q(τ |x; φ)

(3)
(4)

From the coding perspective, if we view the above equation
as a coding process, the coding cost contains two parts: C =Crec + Ccoding to be optimzied, where Crec corresponds
to likelihood term log p(y|x, τ ; θ) in (3) and Ccoding is thecoding cost with “Bits-Back” part subtracted in (4). Notethat during testing, we only care about the first part of thecoding.

considering actions (i.e., indicator functions in GNNs), we
have a stronger learning model; (2) if one set of parameters is more sensitive, it’s computationally more efficient todirectly optimize on them.
For the optimization problem, we consider fixing the parameter space for θ ∈ Θ when maximizing the likelihood:hiL(θ, φ) = Eτ ∼q(τ |x,y,φ) p(y|x, τ , θ)On the other hand, when the parameter spaces for φ areinclusive, i.e. Ψ1 ⊂ Ψ2 , then we havemax L(θ, φ) ≤ max L(θ, φ).
We use q(τ ; φ) to represent a general encoder. The reconφ∈Ψ1 ,θ∈Θφ∈Ψ2 ,θ∈Θ
struction coding cost Crec could be re-written as:
Here Ψ1 corresponds to the existing GNNs and Ψ2 correhisponds to our PMP based GNNs. This guarantees that our
Crec = − E(x,y)∼data log p(y|x; θ)hi learning procedure can result in a lower risk due to the more− E(x,y)∼data Eτ ∼q(τ ;φ) (p(τ |x, y, θ) − log p(τ )) expressive,modeling capability.
where p(τ |x, y, θ) and p(y|x; θ) are posterior distribution
over τ and likelihood term under parameters θ. If we set
q(·) as the optimal distribution which is the true posterior
distribution under θ, then we have
h
i
1
Crec
= − E(x,y)∼data log p(y|x; θ)hi− E(x,y)∼data KL(p(τ |x, y, θ)||p(τ )) .
Instead, if we set q(·; φ) as the prior distribution over τ (i.e.,
we don’t provide any encoding information from τ side),then we have
h
i
2
Crec
= − E(x,y)∼data log p(y|x; θ)hi+ E(x,y)∼data KL(p(τ )|p(τ |x, y, θ)) .

From the perspective of parameter sensitivities, we can alsoachieve that a significantly smaller number of parameters
are required in optimization. This is because the parametersin q(τ |x, y, φ) and p(y|x, τ ; θ) have different sensitivities
to the final model performance. If we don’t have a flexible
q(·) to be learned, the burden of reconstruction all falls onto
the first term p(·), which is modeled by the GNN architecture. When we slightly increase the number of parametersin the sensitive parameter set (q(·) here), we may significantly reduce the number of parameters for the less sensitive
parameter set (p(·) here). This can significantly reduce thecomputation complexity for tuning parameters. We empirically show that by adding the q(·; φ) distribution, we canreduce the overall number of parameters by > 80% andmaintain, even increase, the performance.
4.3. Compare with GNN variants

12This implies Crec
≤ Crec, where the equality holds when
the posterior distribution p(τ |x, y) equals the prior distribution p(τ ), which means the best aggregation process τ is
to assume nothing (this can sometimes appear in too simple
tasks/applications). In harder cases where GNN is trying totackle a highly complex task, we can see that if the distribution is close to the true posterior, it can always help to get ashorter length of coding over the reconstruction term. Thisis exactly what variational inference is optimizing, which
minimizes the distance between a distribution q(τ |x; θ) andthe true posterior distribution p(τ |x, y; φ).

4.2. Optimization perspective
If we treat the parameters from q(τ |x, y, θ) andp(y|x, τ ; φ) as two distinct sets of parameters, then θ controls the generation of reasoning process, while φ controls
the likelihood model. We demonstrate that: (1) throughExisting GNN variants provide two remedies to the basic
GNN architecture: (1) the first class of models use gatesor attentions to control the information during propagation(Veličković et al., 2017; Li et al., 2015; Deng et al., 2016).
Note that these methods, although inspiring, are still constrained to non-probabilistic, single-step operations, which
don’t consider the full search space of the sequential process and lack systematic interpretations. (2) Several GNN
variants also take a Bayesian perspective (Kipf et al., 2018;
Zhang et al., 2019; Ng et al., 2018). In NRI (Kipf et al.,2018), an encoder is used to approximate the posterior distribution over graph structures, the sample of which still leads
to fixed inference processes. (Zhang et al., 2019) focuseson distribution over graph convolution weights, while (Nget al., 2018) uses Gaussian Process to model data points
correlations. Both are orthogonal to our approach and could
potentially benefit from our algorithm.

Modeling Trajectories for Probabilistic Graph Inference
Position: (2, 2)

model will be given patch features and the position for
one patch, and classify the rest. We use ResNet-18 (Heet al., 2016) as our base network architecture for extracting
features. The base network is jointly trained from scratch.

Reassembly

Models(Mena et al., 2018)
PMPL - Feature
Figure 2. Illustration of the image puzzle reassembly task. An
image is first split into patches which are then shuffled. A fullyconnected graph is built to connect all patches. Given the positionof one patch, the task requires model to infer positions for the rest.

5. Experiments
We test our proposed model on two tasks with different
modalities of data: the complex visual reasoning task andthe node classification on citation network benchmarks. Visual data resides in a high-dimensional space, leading torich and diverse interactions when organized in graphs. Citation network data has fewer dimensions per node, but larger
scales on the size of graphs.
5.1. Visual Reasoning — Image Puzzle Reassembly
Image puzzle reassembly is a classic problem that requires
understanding complex patterns or pattern combinationsin images, linking possible pairs of tiles, correcting errors
and iteratively solving the puzzle. The images are firstly
segmented into tiles with shuffled orders, see fig. 2. The
task is to reassemble the tiles into original full images withcorrect positions and orders.
5.1.1. DATASETS AND EXPERIMENTAL SETUP
Datasets We generate image puzzles from two large-scale
benchmark database: (1) Visual Genome (Krishna et al.,2017) - The Visual Genome (VG) database is known forimages with complex object layouts or interactions, and is
standard benchmark data for scene graph generation task. It
contains 64346 images for training and 43903 images fortesting. (2) COCO (Lin et al., 2014) - We use the COCO
dataset (2014 version) that contains 82783 images for training and 40504 images for testing. COCO is also standard
dataset used for “recognition in context” and contains diverse objects and semantic information.
Task setup To generate a d × d puzzle from an image withsize H ×W , we divide the image into d×d non-overlapping
patches. Each patch has size H/d × W/d. The d2 patches
are then shuffled into a set with size d2 . In our experiments,
we first resize the images into 256 × 256 before generating
puzzles. Three puzzle reassembly tasks are set up: 3 ×
3, 4 × 4, and 6 × 6, leading to graphs with size 9, 16 and
36 respectively. See figure 2 for illustration. The inference

Accuracy(%) ↑ / Kendall-tau ↑
2x2
3x3
4x4
5x5
1.0 / 1.0 .97 / .96 .9 / .88.79 / .78
1.0 / 1.0 1.0 / 1.0 .99 / .99 .99 / .99

Table 1. We compare our model’s performance with state-of-theart deep learning based puzzle reassembly model on Celeba
dataset.

5.1.2. BASELINES AND IMPLEMENTATION DETAILS
We compare our model with four powerful graph-structured
models as baselines: (1) Graph Neural Network (GNN)— we use the GNN described in Sec.2 as our basic baseline model to compare with. This GNN corresponds to ourmodel without the PMP components and uses a deterministic informagion aggregation process. Models with similararchitectures are widely used in other applications (Gilmer
et al., 2017; Xu et al., 2017). (2) Attention-based Graph
Neural Network (Veličković et al., 2017; Wang et al., 2018)— the attention-based GNN can be viewed as a special caseof our model, with single step soft message selections. (3)Recurrent Relational Network (RRN) (Palm et al., 2018) —RRN is a GNN model that performs long-term reasoning
via maintaining hidden memory on the node states. (4) Neural Relational Inference model (NRI) (Kipf et al., 2018) —the NRI model is the closest one to our model. It aims at
inferring the graph structure from the data and is formulated
under the Variational Inferece framework. However, the aggregation process is still fixed based on sampled structures.
Implementation details We randomly sample 30% images
from training set as validation set and search hyperparameters. Numbers reported in tables are for test set. For node dimension, we found 200 can achieve maximum performancefor baselines, while only 50 dims are needed for our PMP
models. Number of functions is set to 4 for GNN-Attention,
NRI and PMPL/G -Feature and 32 for PMPL/G -Universal.
We use d + 2 inference steps for all models, where d issquare root of graph size. Learning rate is set to 5e−4 . Models are trained for 60 epochs. More architecture details andnumber of parameters per model can be found in Appendix.
5.1.3. E XPERIMENTAL RESULTS
Results for image puzzle reassembly on two benchmark
datasets are summarized in Table 2. Performance are measured by per node (patch) classification accuracy and correlation score Kendall-tau. We find that our PMP-basedmodels can achieve > 90% on classification accuracy, andoutperform the GNN without PMP by up to 40%. Local
agent PMP overall performs better than global agent ones,

Modeling Trajectories for Probabilistic Graph Inference
ModelsrandomGNN
GNN-attention
RRN
NRI
PMPG - Feature
PMPG - Universal
PMPL - Feature
PMPL - Universal

Visual Genome
Accuracy(%) ↑Kendall-tau ↑
3x3
4x4
6x6
3x3
4x4
6x6
11.11
6.25
2.78
88.76 76.18 56.11 90.32 82.56 75.07
90.96 82.05 72.25 92.12 86.75 84.66
90.83 83.03 62.66 92.48 89.51 82.01
92.31 82.21 59.90 93.00 86.78 78.15
93.31 89.18 84.10 94.95 93.91 89.77
94.26 88.62 78.23 95.28 93.56 85.84
94.40 96.55 91.40 94.44 97.22 95.33
95.31 94.55 96.88 96.01 96.64 98.27

COCO
Accuracy(%) ↑Kendall-tau ↑
3x3
4x4
6x6
3x3
4x4
6x6
11.11
6.25
2.78
90.46 79.94 40.50 91.47 84.75 68.18
93.75 82.77 79.36 93.55 86.75 84.66
91.51 81.26 48.88 92.76 86.24 73.92
93.91 82.25 55.51 94.73 86.97 75.49
96.82 88.52 80.82 96.56 93.50 90.18
94.11 89.16 80.16 95.44 93.19 89.29
96.68 97.05 91.06 96.65 97.69 94.59
96.01 94.84 93.23 96.90 95.99 95.81

Table 2. Performance results on Visual Genome and COCO datasets, measured by per-patch classification accuracy and correlation score
Kendall-tau. PMPL and PMPG correspond to local agent version and global agent version respectively. “Feature” represents feature-based
message set passing and “Universal” represents universal message set passing. Our model achieve up to 40% improvement over baselines.

indicating that memorizing actions and history for the wholegraph is harder than factorizing the modeling into local ones.
Actually local agent PMP has higher parameter efficiency
too. We observe that universal messages can achieve similar (or higher) results compared to feature-based messages,
showing that the agent is capable of learning a set of featureagnostic and task-specific messages to perform reasoning.
Message function analysis Number of message functions(NMF) is highly correlated to the expressive power and flexibility of PMP models. With higher NMF, the PMP model.
We find that, with feature-based message functions, themodel achieves state-of-the-art through four different function types. For PMP with universal messages, the algorithm
needs more “prototypical” vectors (32 in the experiments) tocover possible component in the message space. We also observe that universal message based models have more stabletraining process. This is potentially due to that the message
set is disentangled from high-dimensional image featuresand is less sensitive to perturbations. We also find that message function composition appears across the consecutive
operations. For example, we measure the composition ratereq(kt =i,kt+1 =j)of two function types through (#f#freq(k=i)+#f req(k=j)) andfind the metric ranges from 0.0522 to 0.1804, indicating theusage of compound messages for reasoning.
ModelsGCNGATPMPL -F50%
26.55 ± 1.3
58.83 ± 1.35
65.37 ± 1.36

Noisy edges
100%
23.53 ± 1.44
45.82 ± 2.55
61.42 ± 1.31

200%
22.62 ± 1.84
34.53 ± 0.45
50.3 ± 1.13

Table 3. The performance of node classification for Cora dataset
under varied noisy edge settings. The experiment follows thestandard transductive setup in (Yang et al., 2016).

Parameter efficiency Through modeling the reasoning process τ , we find that PMP largely reduces the burden on
GNNs. Specifically, a basic GNN requires 321,0809 parameters, through modeling the reasoning process with an extra

6,6592 parameters using PMP, the number of parameters
for GNN can be reduced by 5 times smaller to 54,7709,
resulting in a model with 18.44% of original model size.
Compare with state-of-the-art method: We compared our
method with state-of-the-art deep learning based method forimage puzzle assembly task on CelebA dataset. The results
are shown in Table 1. Our model shows near perfect puzzle
reassembly results across four settings.
5.2. Node classification with noisy edges
We further inspect on models’ ability on handling noisyedges on large scale graphs. Current graph neural network
benchmark datasets on citation networks, such as Cora (Senet al., 2008), heavily rely on well caliberated edge matrix
as a strong prior information for inference. The burden
on aggregation process are lessened, compared to moregeneral cases where adjacency matrix information are often noisy or missing. In this experiment, we compare ourfeature-based PMP with local agents to two state-of-theart graph-structured models: (1) Graph Convolutional Networks (GCN) and (2) Graph Attention Networks (GAT)
on a standard citation network benchmark Cora. Insteadof using only the well designed adjacency matrix information, we generate perturbation to the matrix by adding{%50, %100, %200} ∗ |E| more noisy edges to test the robustness of algorithms w.r.t. edge noises, |E| is the numberof edges in the original graph. For implementation, we use
the official release of the codes on two baseline models anddefault settings. In our PMP model, we follow the same
hyper-parameters in baselines and uses recurrent network
agent with 32 hidden dimensions. Results summarized intable 3 show that our model is more robust to noisy edges in
performing inference over large scale graphs. Note that as
we are building local agents for each pair of nodes, generalizing to even larger scale graphs will lead to high memory
costs. Sampling-based methods could potentially help tosolve this problem. We leave it as future work.

Modeling Trajectories for Probabilistic Graph Inference

6. Conclusion
We propose PMP and formulate the reasonings over graphs
as learnable processes under Variational Inference framework. We show that our algorithm consistently outperforms
baseline GNN variants while maintaining high parameter
efficiency.

References
Bahdanau, D., Cho, K., and Bengio, Y. Neural machinetranslation by jointly learning to align and translate. arXiv
preprint arXiv:1409.0473, 2014.
Chen, S.-F., Chen, Y.-C., Yeh, C.-K., and Wang, Y.-C. F.
Order-free rnn with visual attention for multi-label classification. In Thirty-Second AAAI Conference on Artificial
Intelligence, 2018.
Chung, J., Gulcehre, C., Cho, K., and Bengio, Y. Empiricalevaluation of gated recurrent neural networks on sequencemodeling. arXiv preprint arXiv:1412.3555, 2014.
Defferrard, M., Bresson, X., and Vandergheynst, P. Convolutional neural networks on graphs with fast localized
spectral filtering. In Advances in Neural Information
Processing Systems, pp. 3844–3852, 2016.
Deng, Z., Vahdat, A., Hu, H., and Mori, G. Structure inference machines: Recurrent neural networks for analyzing
relations in group activity recognition. In Proceedingsof the IEEE Conference on Computer Vision and Pattern
Recognition, pp. 4772–4781, 2016.
Duvenaud, D. K., Maclaurin, D., Iparraguirre, J., Bombarell,R., Hirzel, T., Aspuru-Guzik, A., and Adams, R. P. Convolutional networks on graphs for learning molecular fingerprints. In Advances in neural information processing
systems, pp. 2224–2232, 2015.
Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., andDahl, G. E. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference
on Machine Learning-Volume 70, pp. 1263–1272. JMLR.
org, 2017.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE
conference on computer vision and pattern recognition,
pp. 770–778, 2016.
Henaff, M., Bruna, J., and LeCun, Y. Deep convolutional networks on graph-structured data. arXiv preprint
arXiv:1506.05163, 2015.
Hinton, G. E. and Van Camp, D. Keeping the neural networks simple by minimizing the description length of theweights. In Proceedings of the sixth annual conference
on Computational learning theory, pp. 5–13, 1993.

Honkela, A. and Valpola, H. Variational learning and bitsback coding: an information-theoretic view to bayesianlearning. IEEE transactions on Neural Networks, 15(4):800–810, 2004.
Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. arXiv preprint
arXiv:1611.01144, 2016.
Kingma, D. P. and Welling, M. Auto-encoding variational
bayes. arXiv preprint arXiv:1312.6114, 2013.
Kipf, T., Fetaya, E., Wang, K.-C., Welling, M., and Zemel,
R. Neural relational inference for interacting systems.
arXiv preprint arXiv:1802.04687, 2018.
Kipf, T. N. and Welling, M. Semi-supervised classification with graph convolutional networks. arXiv preprint
arXiv:1609.02907, 2016.
Krishna, R., Zhu, Y., Groth, O., Johnson, J., Hata, K.,
Kravitz, J., Chen, S., Kalantidis, Y., Li, L.-J., Shamma,
D. A., et al. Visual genome: Connecting language andvision using crowdsourced dense image annotations. International Journal of Computer Vision, 123(1):32–73,
2017.
Lafferty, J., McCallum, A., and Pereira, F. C. Conditionalrandom fields: Probabilistic models for segmenting andlabeling sequence data. 2001.
Li, Y., Tarlow, D., Brockschmidt, M., and Zemel, R.
Gated graph sequence neural networks. arXiv preprint
arXiv:1511.05493, 2015.
Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. Microsoft coco:
Common objects in context. In European conference on
computer vision, pp. 740–755. Springer, 2014.
Louizos, C., Ullrich, K., and Welling, M. Bayesian compression for deep learning. In Advances in Neural Information
Processing Systems, pp. 3288–3298, 2017.
Mena, G., Belanger, D., Linderman, S., and Snoek, J. Learning latent permutations with gumbel-sinkhorn networks.
arXiv preprint arXiv:1802.08665, 2018.
Ng, Y. C., Colombo, N., and Silva, R. Bayesian semisupervised learning with graph gaussian processes. In
Advances in Neural Information Processing Systems, pp.
1683–1694, 2018.
Palm, R., Paquet, U., and Winther, O. Recurrent relational
networks. In Advances in Neural Information Processing
Systems, pp. 3368–3378, 2018.
Rissanen, J. Modeling by shortest data description. Automatica, 14(5):465–471, 1978.

Modeling Trajectories for Probabilistic Graph Inference

Santoro, A., Raposo, D., Barrett, D. G., Malinowski, M.,
Pascanu, R., Battaglia, P., and Lillicrap, T. A simple neural network module for relational reasoning. In Advancesin neural information processing systems, pp. 4967–4976,
2017.
Sen, P., Namata, G., Bilgic, M., Getoor, L., Galligher, B.,
and Eliassi-Rad, T. Collective classification in network
data. AI magazine, 29(3):93–93, 2008.
Sha, F. and Pereira, F. Shallow parsing with conditionalrandom fields. In Proceedings of the 2003 Conference of
the North American Chapter of the Association for Computational Linguistics on Human Language TechnologyVolume 1, pp. 134–141. Association for Computational
Linguistics, 2003.
Taskar, B., Guestrin, C., and Koller, D. Max-margin markov
networks. In Advances in neural information processing
systems, pp. 25–32, 2004.
Teney, D., Liu, L., and van Den Hengel, A. Graph-structured
representations for visual question answering. In Proceedings of the IEEE Conference on Computer Vision
and Pattern Recognition, pp. 1–9, 2017.
van den Oord, A., Vinyals, O., et al. Neural discrete representation learning. In Advances in Neural Information
Processing Systems, pp. 6306–6315, 2017.
Veličković, P., Cucurull, G., Casanova, A., Romero, A.,
Lio, P., and Bengio, Y. Graph attention networks. arXiv
preprint arXiv:1710.10903, 2017.
Wang, X., Girshick, R., Gupta, A., and He, K. Non-local
neural networks. In Proceedings of the IEEE Conference
on Computer Vision and Pattern Recognition, pp. 7794–
7803, 2018.
Wei, S.-E., Ramakrishna, V., Kanade, T., and Sheikh, Y.
Convolutional pose machines. In Proceedings of the IEEE
conference on Computer Vision and Pattern Recognition,
pp. 4724–4732, 2016.
Xu, D., Zhu, Y., Choy, C. B., and Fei-Fei, L. Scene graphgeneration by iterative message passing. In Proceedingsof the IEEE Conference on Computer Vision and Pattern
Recognition, pp. 5410–5419, 2017.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How
powerful are graph neural networks? arXiv preprint
arXiv:1810.00826, 2018.
Yang, Z., Cohen, W. W., and Salakhutdinov, R. Revisiting
semi-supervised learning with graph embeddings. arXiv
preprint arXiv:1603.08861, 2016.
You, J., Ying, R., and Leskovec, J. Position-aware graphneural networks. arXiv preprint arXiv:1906.04817, 2019.

Zhang, Y., Pal, S., Coates, M., and Ustebay, D. Bayesian
graph convolutional neural networks for semi-supervisedclassification. In Proceedings of the AAAI Conference on
Artificial Intelligence, volume 33, pp. 5829–5836, 2019.