Policy Message Passing: A New Algorithm for Probabilistic Graph Inference

1 INTRODUCTION . Not every path is created equal . Powerful sequential inference algorithms have been a core research topic across many tasks that involve partial information , large search spaces , or dynamics over time . An early algorithm example is dynamic programming which memorizes the local information and uses a fixed inference trace to acquire a global optimum . In the modern deep learning era , sequential memory architectures such as recurrent neural networks are used for sequential inference ( e.g . language generation ( Sutskever et al . ( 2014 ) ) ) to narrow down the search space . Neural-based iterative correction is developed in amortized inference for posterior approximation ( Marino et al . ( 2018 ) ) . Also , a large bulk of works are developed in reinforcement learning for sequential inference of the environment with partial observation , non-future aware states , and world dynamics . Graph-structured models are another class of models that heavily rely on local observations and sequential inference . It contains nodes and message functions with partial information , the final output of model relies on a sequence of communication operations which transforms the nodes information dynamically over iterations . A general formulation of learning graph-structured models is maximizing the likelihood : max θ p ( y|s0 , x ; Fθ , P ) ( 1 ) The objective describes the process where it takes in the information x over a graph , starts from initial states s0 and maps to the target output y . This process involves two key components , message functions F and inference process P . Message function designs focus on the choices of functions defined over the interaction between nodes or variables . Simple log-linear potential functions are commonly used in traditional probabilistic graphical models ( Sutton et al . ( 2012 ) ) . Kernelized potential functions ( Lafferty et al . ( 2004 ) ) , gaussian energy functions ( Krähenbühl & Koltun ( 2011 ) ) , or semantic-driven potential functions ( Lan et al . ( 2011 ) ) are developed in varied cases to adapt to specific tasks . Deep graph neural networks generalizes the message functions to neural networks . There are a variety of function types developed such as graph convolution operations ( Kipf & Welling ( 2016 ) ) , attention-based functions ( Veličković et al . ( 2017 ) ) , gated functions ( Li et al . ( 2015 ) ; Deng et al . ( 2016 ) ) . Inference procedures over graphs are mainly developed under the context of probabilistic graphical models , with the classic inference techniques such as belief propagation ( Pearl ( 1988 ) ) , generalized belief propagation ( Yedidia et al . ( 2001 ) ) , and tree-reweighted message passing ( Wainwright et al . ( 2003 ) ) . However , there has not been many research in developing powerful inference process in modern deep graph-structured models ( Graph Neural Networks ) . Most graph-structured models still follow a fixed hand-crafted rule for performing inference over nodes . In this paper , we take a different tack and instead propose to represent the inference process modeled by a neural network Pφ : max φ , θ Eτ∼Pφ [ p ( y|s0 , x , τ ; Fθ ) ] ( 2 ) The neural network Pφ represents the distribution over all possible inference trajectories in the graph model . By taking the distribution perspective of message passing , the whole inference is reformulated as stochastic sequential processes in the deep graph-structured models . This formulation introduces several properties : Distribution nature - the possible inference trajectories in graph , besides synchronous or asynchronous , are in a vast space with varied possibilities , leading to different results ; Consistent inference - maintaining a powerful neural network memorizing over time could lead to more effective and consistent inference over time ; Uncertainty - inference itself will affect prediction , with a distribution , such uncertainty over the reasoning between nodes can be maintained . The primary contributions of this paper are three folds . Firstly , we introduce a new perspective on message passing that generalizes the whole process to stochastic inference . Secondly , we propose a variational inference based framework for learning the inference models . Thirdly , We empirically demonstrate that having a powerful inference machine can enhance the prediction in graph-structured models , especially on difficult reasoning tasks , or graphs with noisy edges . 2 BACKGROUND : GRAPH NEURAL NETWORKS . A large volume of Graph Neural Network ( GNN ) models have been developed recently ( Defferrard et al . ( 2016 ) ; Li et al . ( 2015 ) ; Kipf & Welling ( 2016 ) ; Henaff et al . ( 2015 ) ; Duvenaud et al . ( 2015 ) ) for varied applications and tasks . This line of methods typically share some common structures such as propagation networks and output functions . In our paper , without loss of generality , we follow a unified view of graph neural network proposed in a recent survey ( Battaglia et al . ( 2018 ) ) and introduce the graph neural network architecture here . In this formulation , there are two key components : propagation networks and output networks . In the propagation networks , the fundamental unit is a “ graph-to-graph ” module . This module takes a graph with states st as input and performs information propagation ( message passing ) operations to modify the values to st+1 . The graph states at each step t consists of sets of variables representing the attributes of nodes Vt = { v } t and edges Et = { e } t , along with global properties ut . The graph-to-graph module updates the attributes of these nodes , edges , and global properties via a set of message passing , implemented as neural networks . et = fe ( { v } t , Et , E ) ; vt = fv ( { e } t , Vt , E ) ; u = fu ( { v } t , { e } t , E ) ( 3 ) The message passing in the standard graph neural networks is a fixed aggregation process defined by the graph structure E. The output network is a general part of model that maps the information in the graph states to the target output y : p ( y|sT ) . In general , there is no restriction on the form of target y or the output mapping function , like the target can be a set of labels for nodes , edges , the whole graph or even sentences . 3 GRAPH MODELS WITH POLICY MESSAGE PASSING . In this section we describe our stochastic graph inference algorithm under the formulation of graph neural networks . As the inference process under the algorithms is more of a reasoning procedure , we refer to it as reasoning in the later content . Our approach takes the initial state s0 of a graph representing observations and latent variables in a domain , uses a set of parameterized reasoning agents to pass messages over this graph , and infers the output prediction y . Our model contains three main parts : 1 . Graph st , representing states of observed and latent variables . 2 . A set of message functions { fk ( · , · ; θk ) } , parameterized by θk , that determine the possible manner in which messages can be passed over edges in the graph . 3 . An automatic reasoning agent A ( · ; φ ) , parameterized by φ , that determines which message function would be appropriate to use along an edge in the graph . For every timestep t during reasoning , our model takes the current graph state st and the previous reasoning history , then generates a distribution over actions for choosing which messages to pass ( or pass no information ) . The new graph state st+1 is obtained through performing message passing actions sampled from the distribution . The final prediction is thus inferred sequentially through the operation of the reasoning agents over the graph . Fig . 1 illustrates the algorithm . 3.1 MODEL DEFINITION . Our inference procedure operates on graphs . We start with an initial graph state s0 = ( v0 , e0 , u0 ) , where v0 , e0 , and u0 contain information about individual nodes , edges , and global information about the graph respectively . Nodes in this graph can represent observations or latent variables from our domain ; we assume the set of nodes is known and fixed . Edges in this graph represent possible relationships between nodes , over which we will potentially pass messages during inference . The edges are assumed known and fixed , though our method will permit flexibility via its inference . Given a graph state s , we can output y using a network fo ( s ; θo ) . Note that in general this formulation could have a variety of possible output functions , corresponding to different tasks . In the current instantiation we focus on a single such output function that is used for a single task . Given the initial graph state s0 , we perform a set of reasoning steps that pass messages over this graph to obtain a refined graph representation . The goal of this message passing is to successively refine the graph state such that the graph state can be accurately decoded to match the correct output for the given input graph . The update functions we define will modify the node attributes v , though the formulation could be generalized to edge and global attributes as well . Message Formulation : The message passing is decomposed into operations on edges in the graph . For the ith node in the graph , we update its state by passing messages from all nodes in its neighbour set N ( i ) . vi = ∑ j∈N ( i ) fzij ( v i , vj ) ( 4 ) The form of message to pass is a crucial component of our method . We select the function fzij ( · , · ) to use in this message passing for each edge . We define a common set of functions { fk ( · , · ) } 0≤k≤K that can be used across any edge in the graph . To allow for no information to be passed along an edge , we define the function f0 ( · , · ) = 0 . Each of the other functions fk ( · , · ) is a non-linear function ( e.g . MLP ) with its own learnable parameters θk . Reasoning Agents : Alongside these functions , we define a reasoning agent that determines which function to use on a particular edge in the graph at the current time step of inference . We run one copy of this reasoning agent for each edge of the graph . The reasoning agent takes as input the current graph state and for its edge ( i , j ) outputs a distribution over the message to pass on that edge ( computed by a network A ) . The reasoning agent maintains state hijt ( propagated via a network B ) over iterations of inference so that it can potentially remember which types of message passing actions were taken at different time steps t. We also include a global state hut = ∑ ij h ij t which can be used to share information across the agents operating on different edges . The networks are updated and distributions computed using the following equations : hijt+1 = B ( st , i , j , h ij t ; θB ) ( 5 ) π ( zijt+1|st ) ∝ A ( st , i , j , h ij t , h u t ; θA ) + 1 ( 6 ) where π ( zijt+1|st ) is a distribution over z ij t+1 , the message function to use for edge ( i , j ) on step t + 1 given the current graph state . Note that a uniform prior is added to message distribution for regularization , though in general other priors could be used for this purpose . The parameters φ = ( θA , θB ) for the reasoning agent are shared ; however , each edge maintains its own copy of these agents , via different states hijt . Action sampling : Given these distributions π , we can now use them to perform inference by sampling actions from these distributions . In our implementation , the Gumbel softmax trick is used for differentiable sampling . By executing a set of actions { zt } sampled from each of these distributions π , our model takes selected nodes and corresponding message functions to pass , and the graph state is updated to st+1 after one reasoning/passing step . Over a sequence of time steps , this process corresponds to a reasoning trajectory τ = ( z1 , z2 , . . . , zT ) that produces a sequence of graph states ( s0 , s1 , . . . , sT ) . At test time , we evaluate a learned model by sampling multiple such reasoning trajectory sequences and aggregating the results .

Summary: The paper presents a policy message passing algorithm for graph neural network where both node & edge energy functions and the message passing procedure are treated in a unified manner as outputs of neural nets. In particular, the message passing procedures for individual edges are individualized (i.e. not homogeneous) while a common memory pool is also shared among these procedures. It also aims at doing message passing in the distribution level instead of singling out a particular message passing rule. The idea is new (to my understanding) and interesting. Empirical results demonstrate the usefulness of the proposed methods. On the other hand, the main concerns are: