Modeling Graph Node Correlations with Neighbor Mixture Models

Linfeng Liu Michael C. Hughes Li-Ping Liu

Tufts University Tufts University Tufts University

###### Abstract

We propose a new model, the Neighbor Mixture Model (NMM), for modeling node
labels in a graph. This model aims to capture correlations between the labels
of nodes in a local neighborhood. We carefully design the model so it could be
an alternative to a Markov Random Field but with more affordable computations.
In particular, drawing samples and evaluating marginal probabilities of single
labels can be done in linear time. To scale computations to large graphs, we
devise a variational approximation without introducing extra parameters. We
further use graph neural networks (GNNs) to parameterize the NMM, which
reduces the number of learnable parameters while allowing expressive
representation learning. The proposed model can be either fit directly to
large observed graphs or used to enable scalable inference that preserves
correlations for other distributions such as deep generative graph models.
Across a diverse set of node classification, image denoising, and link
prediction tasks, we show our proposed NMM advances the state-of-the-art in
modeling real-world labeled graphs.

##  1 Introduction

Many real-world applications have data organized in a graph. One important
task in the analysis of graph-structured data is to model the discrete labels
of nodes, which are either observed or hidden. Probabilistic graphical models
(Wainwright et al.,, 2008) describe dependencies between nodes in an elegant
and extensible way, useful in practice for predicting node labels or edges in
ways that account for uncertainty.

Among graphical models, a Markov Random Field (MRF) (Kindermann,, 1980) is
often used to describe a distribution of node labels on an undirected graph,
capturing dependencies between nodes via unnormalized potential functions. One
well-known variant of the MRF is the Conditional Random Field (CRF) (Lafferty
et al.,, 2001), which explicitly includes observed features (Murphy,, 2012) in
potential functions. While an MRF has a convenient model definition, running
inference for an MRF is a hard computational problem. For example, it is hard
to evaluate the marginal probability for even a single node’s label; it is
also hard to sample a single node’s label from its marginal distribution.
These computations often require the inference over the entire graph, whose
runtime scales exponentially with the number of nodes for the kinds of graphs
commonly used to model rich dependencies in applications (Wainwright et al.,,
2008). To scale up, practitioners appeal to approximations such as mean-field
variational inference (Blei et al.,, 2017) that make strict simplifying
assumptions. There is a need for scalable methods that adequately account for
correlations between labels.

Recently, graph neural networks (GNN) (Wu et al.,, 2019) have made remarkable
progress in learning representations from graph data. These models take in
node-specific features and then propagate messages along graph edges through
several hidden layers (Kipf and Welling,, 2017; Hamilton et al.,, 2017;
Veličković et al.,, 2018; Xu et al.,, 2019). In the last layer, each node gets
a vector representation that encodes information about the local graph
structure and its observed features. While GNNs focus on learning informative
node representations useful for label or link prediction, they typically
assume labels are _independent_ given all observed node features.

A natural idea is to combine graph neural networks with a probabilistic model
to learn an expressive model for node labels given node features. The recent
Graph Markov Neural Network (GMNN) (Qu et al.,, 2019) combines a CRF model
with a graph neural network and shows the obvious benefit of modeling node
labels in a probabilistic manner. However, due to the difficulty of MRF
inference, GMNN requires substantial approximations such as pseudo-likelihood
maximization (Besag,, 1975) and mean-field variational inference, the latter
of which depends on an inference neural network that is not easy to train. In
this work, we aim to devise an alternative probabilistic model for graph-
structured data that can utilize GNNs for effective inference while overcoming
these difficulties in inference.

We propose the Neighbor Mixture Model (NMM), a new probabilistic model that
captures correlations in the distribution of node labels through an attention
mechanism (Veličković et al.,, 2018). This new model assumes every node
attends one of its neighbors in the graph (including itself) to get the latent
parameter that determines its label distribution. This parameter sharing
induces correlation between nodes that share an edge. Our NMM model itself, a
joint distribution over the labels of nodes, can be evaluated or sampled
without any specific node ordering. NMM parameters can be efficiently
parameterized by a GNN or another task-relevant neural network. The NMM thus
leverages the synergistic strengths of its two key components: neural networks
for extracting useful representations from the input features and
probabilistic graphical models for capturing correlations between nodes.

To scale to large graphs, we contribute an efficient variational inference
method for the NMM, with a tractable lower bound on the likelihood of node
labels in a training set which can be efficiently optimized to estimate
parameters. We empirically show our proposed variational lower bound on this
likelihood to be tight and thus capable of capturing useful correlations
between labels, while not requiring any additional parameters beyond the
original model.

Beside modeling observed node labels directly, the NMM can also serve as an
amortized inference method to efficiently approximate other models (e.g. MRFs
or deep generative models of graphs). Previous methods of amortized
variational inference (Kingma and Welling,, 2014; Shu et al.,, 2018; Zhang et
al.,, 2018) mostly reuse parameters across independent instances. GNNs are
recently used to do inference on small graphs (Yoon et al.,, 2018; Satorras
and Welling,, 2020), but training these GNNs requires “ground-truth” results
from another inference method. Unlike previous methods, an NMM can reuse
parameters for dependent variables on large graphs. We show how to train an
NMM to minimize its KL-divergence to the target distribution, without
requiring any ground-truth inference results or extra parameters that scale
with the size of the graph.

We evaluate our proposed NMM on node label classification, prediction of joint
label configurations for connected node pairs, image denoising, and link
prediction. When directly fit to observed labels in a graph, NMM improves
probabilistic predictions over state-of-the-art baselines. When used to
approximate complex distributions such as deep generative models for graphs
(Mehta et al.,, 2019) where exact inference is difficult, our flexible NMM
captures correlations that commonly-used mean-field independence assumptions
cannot, while maintaining scalability.

##  2 Background and Notation

Here we establish the learning problem. Let $G=(V,E)$ denote an undirected or
directed graph with $N$ vertices, where $V=\\{1,\ldots,N\\}$ is the vertex set
and $E$ is the edge set. Each node $i\in V$ is associated with a tuple
$(\mathbf{x}_{i},y_{i})$, where $\mathbf{x}_{i}\in\mathbb{R}^{F}$ is a real-
valued feature vector (observed for all nodes) and $y_{i}\in\\{1,2,\ldots
C\\}$ is a discrete label. We stack all feature vectors into matrix
$\mathbf{X}$ of size $N\times F$, and stack all labels in a column vector
$\mathbf{y}$ with length $N$. Our goal is to model the distribution
$p(\mathbf{y}|\mathbf{X})$, in a way that captures dependencies between labels
$y_{i},y_{j}$ that are “neighbors” in the graph.

Let $n(i)$ denote the _neighborhood_ of node $i$. The definition of the
neighborhood, which determines the range of correlation, is a model choice. In
this work, we focus on the neighborhood containing the first order neighbors
and the node itself: $n(i)=\\{j:(i,j)\in E\\}\cup\\{i\\}$. Throughout this
text, let $D$ denote the _maximum degree_ across all nodes in the graph, and
let $\Delta^{k}$ denote the space of probability vectors with length $k$.

##  3 Method

###  3.1 The Neighbor Mixture Model

The NMM is a generative model for labels $\mathbf{y}$ given node features
$\mathbf{X}$ that requires two key parameters:
$\bm{\alpha}=\\{\alpha_{i}\in\mathbb{R}_{+}^{C}:i\in V\\}$ and
$\mathbf{L}=\\{\mathbf{L}_{i}\in\Delta^{n(i)}:i\in V\\}$. Each
$\bm{\alpha}_{i}$ is a label concentration vector for node $i$. Each
$\mathbf{L}_{i}$ is a probability vector over neighbors of node $i$. These two
parameters are computed from a neural network,
$(\bm{\alpha},\mathbf{L})=f(\mathbf{X})$, which we will discuss later in Sec.
3.4. Given these parameters, the model formally defines
$p(\mathbf{y}|\mathbf{X})=p(\mathbf{y}|(\bm{\alpha},\mathbf{L})=f(\mathbf{X}))$
by

$\displaystyle\mathbf{z}_{j}$
$\displaystyle\sim\mathrm{Dirichlet}(\bm{\alpha}_{j}),$
$\displaystyle\mathbf{z}_{j}\in\Delta^{C},~{}~{}j=1,\ldots,N,$ $\displaystyle
c_{i}$ $\displaystyle\sim\mathrm{Categorical}(\mathbf{L}_{i}),$ $\displaystyle
c_{i}\in n(i),~{}~{}i=1,\ldots,N,$ $\displaystyle y_{i}$
$\displaystyle\sim\mathrm{Categorical}(\mathbf{z}_{c_{i}}),$ $\displaystyle
y_{i}\in\\{1,2,\ldots,C\\}.\quad\quad~{}$ (1)

The NMM captures correlations between labels $y_{i},y_{j}$ by sharing a
distribution parameter, vector $\mathbf{z}_{j}$, between neighboring nodes.
Vector $\mathbf{z}_{j}$ provides the probabilities for a categorical
distribution over $C$ possible node labels. Each node has its own
$\mathbf{z}_{j}$, generated a priori from a Dirichlet distribution. To
generate its label, each node $i$ then selects one node $c_{i}$ from its
neighborhood $n(i)$ and then “borrows” the vector $\mathbf{z}_{c_{i}}$ from
this chosen neighbor to sample its node label $y_{i}$ given this probability
vector. This formulation with explicit $c_{i}$-s is easier for inference
later, while the marginalization of all $c_{i}$-s below gives a clearer
understanding of how this model induces correlations between neighbors
$y_{i},y_{j}$.

Marginalizing neighbor indicators: After integrating $c_{i}$-s away, our model
for $y_{i}$ given $\mathbf{z}$ is:

$\displaystyle
y_{i}\sim\mathrm{Categorical}(\mathbf{u}_{i}),\quad\mathbf{u}_{i}=\textstyle\sum_{j\in
n(i)}L_{ij}\mathbf{z}_{j}.$ (2)

In this equivalent formulation, we view vector $\mathbf{L}_{i}$ as node $i$’s
attention weights over its neighbors. Each node computes label probability
vector $\mathbf{u}_{i}$ as a weighted sum of probability vectors
$\mathbf{z}_{j}$ in its neighborhood. Each $\mathbf{z}_{j}$ has a mixing
weight $L_{ij}$. Two nodes with overlapping neighborhoods are tied together by
sharing some mixture components $\mathbf{z}_{j}$. The joint density of
$\mathbf{y}$ given $\mathbf{Z}=(\mathbf{z}_{j})_{j=1:N}$ is decomposable:
$p(\mathbf{y}|\mathbf{Z})=\prod_{i}p(y_{i}|\mathbf{z}_{n(i)})$. If we further
marginalize out $\mathbf{Z}$, labels $y_{i}$ and $y_{j}$ are correlated if
$n(i)\cap n(j)\neq\emptyset$. While we focus on positive correlations, our
model can be modified to model negative correlations (e.g. using a transform
of $\mathbf{z}_{c_{i}}$).

Sampling node labels: We can draw samples of $\mathbf{y}$ by following the
model definition in (1): first draw samples of $\mathbf{z}$ and samples of
$\mathbf{c}$ from their respective priors, then use these values to draw
samples of $\mathbf{y}$. We can also leverage (2) to sample $\mathbf{y}$
without sampling $\mathbf{c}$.

By (2), drawing samples for a single node $i$ from its marginal $p(y_{i})$
only requires a few vectors $\mathbf{z}_{n(i)}$ in its neighborhood, so it is
_efficient_. In contrast, drawing an exact marginal sample for a single node
from an MRF usually requires inference over the entire graph.

###  3.2 Exact computation of marginal probability for small node sets

Consider observing labels for a _subset_ of nodes $\tau\subseteq V$. Given
known parameters $\bm{\alpha},\mathbf{L}$, we wish to compute the marginal
probability of the subset’s labels $\mathbf{y}_{\tau}$. We can do this using
the sum rule:

$\displaystyle
p(\mathbf{y}_{\tau}|\bm{\alpha},\mathbf{L})=\textstyle\sum_{\mathbf{c}_{\tau}}p(\mathbf{y}_{\tau},\mathbf{c}_{\tau}|\bm{\alpha},\mathbf{L})$
(3)

We can further write each term in that sum as marginalizing away all
$\mathbf{z}_{j}$ variables related to any neighbor $j$ of any node in the set
$\tau$. Let $n(\tau)=\cup_{i\in\tau}n(i)$ denote all such unique neighbor
nodes. Then the probability of interest becomes:

$p(\mathbf{y}_{\tau},\mathbf{c}_{\tau}|\bm{\alpha},\mathbf{L})=\int_{\mathbf{z}_{n(\tau)}}p(\mathbf{z}_{n(\tau)}|\bm{\alpha})p(\mathbf{c}_{\tau}|\mathbf{L})\prod_{i\in\tau}p(y_{i}|c_{i},\mathbf{z}_{n(i)})\\\
=\prod_{i\in\tau}L_{i,c_{i}}\cdot\prod_{j\in\mathbf{c}_{\tau}}\frac{\mathrm{B}(\bm{\alpha}_{j}+\mathbf{s}_{j}(\mathbf{y}_{\tau},\mathbf{c}_{\tau}))}{\mathrm{B}(\bm{\alpha}_{j})},$
(4)

where the integral simplifies due to Dirichlet-Categorical conjugacy (see the
appendix for derivation). We count the number of nodes using neighbor $j$ of
each class with vector
$\mathbf{s}_{j}(\mathbf{y}_{\tau},\mathbf{c}_{\tau})=\sum_{i\in\tau:c_{i}=j}\mathrm{onehot}(y_{i})$.
$\mathrm{B}(\bm{\alpha})=\frac{\Gamma(\alpha_{1})\Gamma(\alpha_{2})\ldots\Gamma(\alpha_{C})}{\Gamma(\sum_{c=1}^{C}\alpha_{c})}$
denotes the multivariate Beta function.

The computation of
$p(\mathbf{y}_{\tau},\mathbf{c}_{\tau}|\bm{\alpha},\mathbf{L})$ in (4) takes
time $O(C|\tau|)$. Thus, we can compute
$p(\mathbf{y}_{\tau}|\bm{\alpha},\mathbf{L})$ in (3) by summing over all the
$D^{|\tau|}$ possible configurations of $\mathbf{c}_{\tau}$. When $\tau$ is a
small set, then the overall computation is manageable (e.g. for one node it is
$O(CD)$). In contrast, for an MRF the marginal of even a single node requires
the inference over the entire graph, which is often much more expensive.

###  3.3 Scalable approximation of marginal for large node sets

For large node sets $\tau$, we need to appeal to approximate inference, as the
exact computation of the marginal likelihood becomes infeasible. We use
variational inference (Wainwright et al.,, 2008; Blei et al.,, 2017) and
derive a variational lower bound $\mathcal{L}(\bm{\alpha},\mathbf{L})\leq\log
p(\mathbf{y}_{\tau}|\bm{\alpha},\mathbf{L})$:

$\displaystyle\mathcal{L}(\bm{\alpha},\mathbf{L}):={\mathbb{E}_{q}}\left[\log\frac{p(\mathbf{y}_{\tau},\mathbf{c}_{\tau}|\bm{\alpha},\mathbf{L})}{q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})}\right].$
(5)

Note that the bound becomes an equality when
$q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$ is the true
posterior $p(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$, but
this is intractable. The key technical challenge is choosing an approximate
posterior distribution
$q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$ close to
$p(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$ that makes the
bound tight yet has more affordable computation.

To address this challenge, we define the distribution
$q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$ as a directed
graphical model with conditional probabilities derived from the true joint
$p(\mathbf{c}_{\tau},\mathbf{y}_{\tau}|\bm{\alpha},\mathbf{L})$. Given an
ordering $\pi$ of nodes in $\tau$, we define $q$ as:

$\displaystyle
q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})=\textstyle\prod_{i\in\tau}p(c_{i}|y_{i},\mathbf{c}_{<i},\mathbf{y}_{<i},\bm{\alpha},\mathbf{L}),$
(6)

where we denote the parents of node $i$ within $\pi$ as
$(<i)=\\{i^{\prime}\in\tau:\pi(i^{\prime})<\pi(i)\\}$.

The conditional
$p(c_{i}|y_{i},\mathbf{c}_{<i},\mathbf{y}_{<i},\bm{\alpha},\mathbf{L})$ is
calculated exactly by computing the joint for each of the possible $c_{i}$
values in $n(i)$ and then normalizing. Recall that neighborhoods have bounded
size ($|n(i)|\leq D$), so this is affordable. Each needed joint
$p(c_{i},y_{i},\mathbf{c}_{<i},\mathbf{y}_{<i}|\bm{\alpha},\mathbf{L})$ can be
efficiently computed from (4) using the subset $\tau=\\{i\\}\cup(<i)$. Thus,
our chosen $q$ is easy to evaluate and sample from.

Our ultimate lower bound objective (5) requires an expectation with respect to
$q$. We can estimate this expectation using $T$ Monte Carlo samples
$\\{\mathbf{c}_{\tau}^{(t)}\\}_{t=1}^{T}$ from
$q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$. Each sample
can use a different ordering of nodes $\pi^{(t)}$, sampled from a uniform
distribution over permutations. Then the objective can be viewed as an
estimate of the average of lower bounds derived from all possible node orders.

Explanation of the chosen $q$ distribution: The parameterization of
$q(\mathbf{c}|\mathbf{y},\bm{\alpha},\mathbf{L})$ in (6) reuses the same
$\bm{\alpha},\mathbf{L}$ parameters as the NMM model, introducing _no_ extra
parameters. Though this choice may not be the most flexible distribution
possible, it has several advantages. First, it can capture correlations
between neighboring nodes (we do not assume each $c_{i}$ is independent, as
typical mean field approximations would). Second, it reduces the number of
variables that need to be estimated in the later optimization of the
variational lower bound. Otherwise, we may need another neural network to
parameterize $q$. Training two neural networks jointly can lead to severe
solution quality issues such as posterior collapse (Lucas et al.,, 2019).
Therefore, our chosen $q$ makes training our method far easier and more
reliable than alternatives.

###  3.4 Parameterization using Neural Networks

We wish to achieve a parameterization of our NMM that is both _scalable_ and
_informed_ by observed node features $\mathbf{X}$. As mentioned earlier, we
use a neural network $f$ to determine the parameters:
$(\bm{\alpha},\mathbf{L})=f(\mathbf{X};\theta)$. Here symbol $\theta$ denotes
all network parameters; we assume $\theta$ has fixed size that does not depend
on graph size (unlike $\bm{\alpha},\mathbf{L}$). By leveraging the strong
representational power of neural networks, this construction can compute
parameters informed by node features $\mathbf{X}$.

When the model is defined on a graph, we use a GNN as the backbone of
$f(\cdot;\theta)$. A GNN takes node features and propagates messages between
graph nodes to produce a vector representation for each node. For our NMM, let
the GNN produce vectors of size $C+H$, where $C$ is the number of possible
labels and $H$ is a free hyperparameter. Denote the first $C$ entries of this
output vector as $\mathbf{u}_{i}\in\mathbb{R}^{C}$, the remainder as
$\mathbf{v}_{i}\in\mathbb{R}^{H}$.

We deterministically transform these GNN outputs $\mathbf{u},\mathbf{v}$ to
produce our parameters $\bm{\alpha},\mathbf{L}$:

$\displaystyle\\{(\mathbf{u}_{i},\mathbf{v}_{i}):i\in
V\\}\leftarrow\mathrm{GNN}(\mathbf{X};\theta),\quad$ (7)
$\displaystyle\bm{\alpha}_{i}\leftarrow\sigma(\mathbf{u}_{i}),L_{ij}\leftarrow\textstyle\mathrm{softmax}\left(\omega^{2}\frac{\mathbf{v}_{i}^{\top}\mathbf{v}_{j}}{\|\mathbf{v}_{i}\|\|\mathbf{v}_{j}\|}+\gamma\delta(i=j)\right).$

To construct $\bm{\alpha}_{i}\in\mathbb{R}^{C}_{+}$, we cast the embedding
vector $\mathbf{u}_{i}$ of size $C$ to positive values using activation
function $\sigma(\cdot)$. Recall from (1) that positive vector
$\bm{\alpha}_{i}$ determines the label marginal associated with node $i$.

To construct $\mathbf{L}_{i}\in\Delta^{n(i)}$, the probability vector which
determines correlations with neighbors, we use a self-attention transformation
of the embedding $v_{i}$, where parameter $\gamma>0$ increases the probability
that $i$ attends to itself. Scalar $\omega^{2}$ controls the sharpness of the
distribution. The softmax function is taken over all indices in the current
neighborhood $n(i)$.

For our GNN parameterization, the learnable model parameters $\theta$ include
all GNN weights, $\omega^{2}$, and $\gamma$. For special graphs such as grids
for modeling image pixels, we can also use CNNs as the backbone.

###  3.5 Parameter learning

Given a graph with known labels $\mathbf{y}_{\tau}$ for a large subset $\tau$,
we could estimate parameters $\bm{\alpha},\mathbf{L}$ by maximizing
$\mathcal{L}(\bm{\alpha},\mathbf{L})$ in (5), which is a lower bound of the
NMM’s label marginal likelihood $\log
p(\mathbf{y}_{\tau}|\bm{\alpha},\mathbf{L})$. With the new scalable
parameterization, we maximize over $\theta$ directly, and our objective is a
lower bound of $\log p(\mathbf{y}_{\tau}|\mathbf{X})$:

$\displaystyle\hat{\theta}\leftarrow\arg\max_{\theta}~{}\mathcal{L}(\bm{\alpha},\mathbf{L}),\quad(\bm{\alpha},\mathbf{L})=f(\mathbf{X};\theta).$
(8)

The gradient of $\mathcal{L}(\bm{\alpha},\mathbf{L})$ with respect to
$\bm{\alpha},\mathbf{L}$ requires Monte Carlo estimation of gradients of
expectations over discrete random variables (since our indicators
$\mathbf{c}_{\tau}$ will be sampled from $q$ given parameters
$\bm{\alpha},\mathbf{L}$). We estimate these gradients using the well-known
REINFORCE estimator (Williams,, 1992), also known as the score function trick.
The optimization procedure converges well, but we could also explore improved
estimators such as Rebar and RELAX (Tucker et al.,, 2017; Grathwohl et al.,,
2018) in future work.

###  3.6 Node classification with the NMM

In graph node classification tasks, a subset of nodes $\tau$ have observed
labels, and we need to predict the labels of the remaining nodes
$\kappa\subset V$. In training, we learn an NMM defined by $\theta$ by
maximizing the marginal likelihood $p(\mathbf{y}_{\tau}|\mathbf{X})$ of
observed labels in (8). In prediction, we predict unknown labels with the
conditional $p(\mathbf{y}_{\kappa}|\mathbf{y}_{\tau},\mathbf{X})$. We stress
that we want to account for correlations within $\mathbf{y}_{\kappa}$.

We first suggest that for this per-node classification task, our NMM model
should be as good or better than standard GNN classifiers. Consider the
special case of our NMM, if we enforce that each $\mathbf{L}_{i}$ is set as
$\mathbf{L}_{ii}=1$ and $\mathbf{L}_{ij}=0$ for $j\neq i$ by setting a large
value to $\gamma$ in (7), then all node labels $y_{i}$ are conditionally
independent given $\mathbf{X}$, and the NMM would be equivalent to a standard
GNN model. With this analysis, we are certain that our NMM, when $L$ is
flexible (not constrained to enforce independence), will always match or beat
a GNN classifier in terms of training error.

Next, we describe our prediction of unknown labels given known labels for node
classification with the conditional
$p(\mathbf{y}_{\kappa}|\mathbf{y}_{\tau},\mathbf{X})$, where $\kappa$ contains
nodes with unknown labels. The conditional can be computed as follows (we omit
$\mathbf{X}$ in notation below for readability):

$p(\mathbf{y}_{\kappa}|\mathbf{y}_{\tau})=\sum_{\mathbf{c}_{\tau}}p(\mathbf{c}_{\tau}|\mathbf{y}_{\tau})p(\mathbf{y}_{\kappa}|\mathbf{c}_{\tau},\mathbf{y}_{\tau})\\\
\quad=\sum_{\mathbf{c}_{\tau}}p(\mathbf{c}_{\tau}|\mathbf{y}_{\tau})\int
p(\mathbf{y}_{\kappa}|\mathbf{z}_{n(\kappa)})p(\mathbf{z}_{n(\kappa)}|\mathbf{c}_{\tau},\mathbf{y}_{\tau})d\mathbf{z}_{n(\kappa)}.$
(9)

The posterior $p(\mathbf{z}_{n(\kappa)}|\mathbf{c}_{\tau},\mathbf{y}_{\tau})$
is computed using conjugacy and conditional independence such that for each
node $j\in n(\kappa)$ we have
$p(\mathbf{z}_{j}|\mathbf{c}_{\tau},\mathbf{y}_{\tau})=\text{Dir}(\bm{\alpha}^{\prime}_{j})$,
with parameter vector
$\bm{\alpha}^{\prime}_{j}=\bm{\alpha}_{j}+\mathbf{s}_{j}(\mathbf{y}_{\tau},\mathbf{c}_{\tau})$,
with $\mathbf{s}_{j}$ defined in (4). The integral above defines
$p(\mathbf{y}_{\kappa}|\mathbf{c}_{\tau},\mathbf{y}_{\tau})$ is then computed
using the conjugacy again. If $\tau$ is a small set of only a few nodes, the
summation can be computed in closed form as in the discussion of exact
marginals. Otherwise, the exact calculation of
$p(\mathbf{c}_{\tau}|\mathbf{y}_{\tau})$ is not feasible, and neither is the
sum over all $\mathbf{c}_{\tau}$. We instead approximate the sum using $T$
Monte Carlo samples from $q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau})$ as derived
in (6):

$\displaystyle\tilde{p}(\mathbf{y}_{\kappa}|\mathbf{y}_{\tau})=\textstyle\frac{1}{T}\sum_{t=1}^{T}p(\mathbf{y}_{\kappa}|\mathbf{c}_{\tau}^{(t)},\mathbf{y}_{\tau}),~{}~{}\mathbf{c}_{\tau}^{(t)}\sim
q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau}).$

When $\kappa$ is small, we check all configurations of $\mathbf{y}_{\kappa}$
and use the mode of the approximate conditional as the prediction. When
$\kappa$ is large, we predict labels in $\mathbf{y}_{\kappa}$ one by one using
each node’s marginal mode. Each subsequent predicted label is then merged into
$\mathbf{y}_{\tau}$ to decide the next label.

###  3.7 Time complexity

In _training_ , the extra cost of NMM compared to a GNN is the computation of
$q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$ in (6) and
extra optimization iterations. Getting a sample from (6) takes time $O(|E|T)$
with $T$ being the number of samples and $|E|$ the number of edges. The
computation of $q(\mathbf{c}_{\tau}|\mathbf{y}_{\tau},\bm{\alpha},\mathbf{L})$
should be faster than a GNN when $T$ is small because it does not deal with
features and hidden units as the GNN does. In practice, however, an optimal
implementation would need substantial effort for parallelization. In our
experiment, we have prioritized a simple implementation without
parallelization.

In _prediction_ for a single node’s label, the runtime cost of our NMM is
$O(CD)$ given the GNN output, with number of classes $C$ and maximum degree
$D$. We assume $D$ is small ($D\ll N$), so prediction time is not burdensome
compared to the cost of producing GNN’s prediction output. Runtime that is
linear in neighborhood size is the most affordable complexity one can have
when allowing correlations from all neighbors.

##  4 Approximating Other Distributions with the Neighbor Mixture Model

As a distribution with easy marginal distributions over nodes, NMM can also
approximate another distribution $p^{*}(\mathbf{y})$ over node labels
$\mathbf{y}$, such as an MRF. We only assume $p^{*}(\mathbf{y})$ to be an
unnormalized distribution and to have the ability to evaluate its unnormalized
log likelihood $\log p^{*}(\mathbf{y})$.

We seek to minimize the KL-divergence from our model’s
$p(\mathbf{y}|\mathbf{X};\theta)$ to the target $p^{*}(\mathbf{y})$ by
optimizing our parameters $\theta$. The input features $\mathbf{X}$ can
contain information about the target graph and the inference problem, e.g. we
can put MRF potential parameters into $\mathbf{X}$. If $p^{*}(\mathbf{y})$
uses features in the graph (e.g. CRF), we can also put these features to
$\mathbf{X}$. We propose the approximation as

$\min_{\theta}\textrm{KL}\left(p(\mathbf{y}|\mathbf{X};\theta)\;\|\;p^{*}(\mathbf{y})\right)\\\
=\min_{\theta}{\mathbb{E}_{p}}\left[\log p(\mathbf{y}|\mathbf{X};\theta)-\log
p^{*}(\mathbf{y})\right].$ (10)

The expectation of the first term, ${\mathbb{E}_{p}}\left[\log
p(\mathbf{y}|\mathbf{X};\theta)\right]$, is intractable as discussed in
previous sections. Here we have an upper bound for this term in the spirit of
hierarchical variational inference (HVI) (Ranganath et al.,, 2016; Louizos and
Welling,, 2017). The actual optimization problem becomes:

$\min_{\theta}{\mathbb{E}_{p}}\left[\log
p(\mathbf{y},\mathbf{c}|\mathbf{X};\theta)-\log
q(\mathbf{c}|\mathbf{y},\mathbf{X};\theta)-\log p^{*}(\mathbf{y})\right]\\\
\geq\textrm{KL}\left(p(\mathbf{y}|\mathbf{X};\theta)\;\|\;p^{*}(\mathbf{y})\right).$
(11)

Here we use
$q(\mathbf{c}|\mathbf{y},\mathbf{X};\theta)=q(\mathbf{c}|\mathbf{y},(\bm{\alpha},\mathbf{L})\leftarrow
f(\mathbf{X};\theta))$ as defined in (6) by setting $\tau=V$. Then we can use
Monte Carlo samples to estimate this bound. The gradients of the objective
with respect to parameters $\theta$ are again estimated by the REINFORCE
estimator.

We want to emphasize that using the NMM as the approximating family will have
better KL divergence than mean-field distributions in general. If we fix every
$\mathbf{L}_{i}$ to a one-hot vector such that $p(c_{i}=i)=1$, then both
$p(\mathbf{c}|\mathbf{y},\mathbf{X})$ and $q(\mathbf{c}|\mathbf{y})$ also
become deterministic distributions. In this special case, the NMM is a mean-
field distribution and the upper bound in (11) is exactly the KL-divergence of
the mean-field distribution. By optimizing $\mathbf{L}_{i}$-s, we can easily
improve the upper bound to be smaller than the mean-field KL-divergence, and
thus the underlying KL-divergence will be even smaller.

For tasks where the target contains both labels and latent variables (e.g.
Mehta et al., (2019)), which we write as $p^{*}(\mathbf{z},\mathbf{y})$, we
can use the NMM distribution $p(\mathbf{z},\mathbf{y}|\mathbf{X})$ to
approximate the target distribution:

$\displaystyle\min_{\theta}{\mathbb{E}_{p(\mathbf{z},\mathbf{y}|\mathbf{X})}}\left[\log
p(\mathbf{z},\mathbf{y}|\mathbf{X})-\log p^{*}(\mathbf{z},\mathbf{y})\right].$
(12)

Here we can compute $\log p(\mathbf{z},\mathbf{y}|\mathbf{X})$ efficiently and
need no further approximation. We can also use different priors for
$\mathbf{z}$ without requiring conjugacy. Using our NMM distribution is better
than assuming mean-field independence because our NMM includes mean-field as a
special case, while allowing more flexible correlations if needed.

Our method contributes a way to perform _amortization_ across labels that are
_dependent_ under a target model. It greatly reduces the number of
optimization parameters and speeds up inference. Even in cases when the target
model is being dynamically updated (e.g. during model learning), our NMM can
be updated alongside in an integrated way.

##  5 Experiment

Dataset | Task | Nodes | Edges | Features | Classes | Training n/e | Validation n/e | Test n/e | Label rate n/e
---|---|---|---|---|---|---|---|---|---
Cora | NC/LP | 2708 | 5429 | 1433 | 7 | 140/4616 | 500/271 | 1000/ | 0.05/0.85
Citeseer | NC/LP | 3327 | 4732 | 3703 | 6 | 120/4023 | 500/236 | 1000/473 | 0.04/0.85
Pubmed | NC/LP | 19717 | 44338 | 500 | 3 | 60/37689 | 500/2216 | 1000/4433 |  $<$0.01/0.85
NIPS12 | LP | 2037 | 3134 | - | - | -/2665 | -/156 | -/313 | -/0.85
Yeast | LP | 2361 | 6646 | - | - | -/5650 | -/332 | -/664 | -/0.85

Table 1: Dataset statistics. NC is node classification; LP is link prediction;
n/e is node/edge counts.

We evaluate our NMM on two types of tasks: modeling observed node labels (node
classification and denoising) and inference for other models (an MRF and a
model for graph generation). Due to limited space, full details of protocols
and hyperparameters are in the appendix.

###  5.1 Graph Node Classification

Algorithm | Cora | Citeseer | Pubmed
---|---|---|---
GCN | 81.3$\pm$0.9 | 71.0$\pm$0.7 | 79.0$\pm$0.4
GAT | 83.0$\pm$0.7 | 72.5$\pm$0.7 | 79.0$\pm$0.3
APPNP | 84.1$\pm$0.9 | 71.7$\pm$0.8 | 79.4$\pm$0.3
GMNN | 82.1$\pm$1.4 | 71.4$\pm$1.1 | 80.6$\pm$0.7
NMM-GCN | 84.3$\pm$0.5 ✓ | 72.0$\pm$0.3 ✓ | 79.2$\pm$0.2
NMM-GAT | 84.4$\pm$0.2 ✓ | 73.0$\pm$0.5 | 79.4$\pm$0.2 ✓
NMM-APPNP |  85.9$\pm$0.5 ✓ | 72.6$\pm$0.3 ✓ |  80.4$\pm$0.2 ✓

Table 2: Node classification results on three datasets. Bold numbers indicates
the best performance, and the $\checkmark$ mark indicates the combination with
NMM outperforms the corresponding baseline.

Algorithm |  | Cora |  |  |  | Citeseer |
---|---|---|---|---|---|---|---
0.2 | 0.4 | 0.6 |  | 0.2 | 0.4 | 0.6
GMNN | -1.34 | -1.30 | -1.14 |  | -1.62 | -1.53 | -1.41
NMM-GCN | -1.04 | -0.97 | -0.81 |  | -1.37 | -1.34 | -1.26
NMM-GAT | -1.01 | -0.84 | -0.71 |  | -1.35 | -1.34 | -1.25
NMM-APPNP | -0.93 | -0.83 | -0.70 |  | -1.33 | -1.30 | -1.24

Table 3: Marginal log-likelihood of label pairs $y_{i},y_{j}$ in the test set
across fractions of observed data included in training set (0.2, 0.4, and
0.6). NMM assigns higher heldout probability than the GMNN. Statistically
indistinguishable best values are bold. Full table with standard deviations
over multiple runs in appendix.

We follow previous literature (Kipf and Welling,, 2017; Veličković et al.,,
2018; Qu et al.,, 2019) and evaluate our methods on three node classification
benchmarks: Cora, Citeseer, and Pubmed. Summary statistics are listed in Table
1. Our data split follows Yang et al., (2016): we select 20 nodes from each
class from training. We assess each model’s _accuracy_ on the test set.

We use three popular graph neural networks as baselines: Graph Convolutional
Network (GCN) (Kipf and Welling,, 2017), Graph Attention Network (GAT)
(Veličković et al.,, 2018), and Approximate Personalized Propagation of Neural
Predictions (APPNP) (Klicpera et al.,, 2019). These methods all assume node
labels are independent given node features. We then build our proposed NMM
model using each one of these GNNs as the backbone neural network in (7). We
refer to each variant of our method as NMM-GCN, NMM-GAT, and NMM-APPNP. We
also compare against the recent Graph Markov Neural Network (GMNN) (Qu et
al.,, 2019), which combines MRFs with GNNs to model node label dependencies.
The temperature parameter of GMNN is set to 1.0, leading to a formal
probabilistic model as argued in the original paper.

Our implementation is based on the PyTorch Deep Graph Library (DGL) (Wang et
al.,, 2019). We reuse the GCN results from Wang et al., (2019) to consider the
standard deviation. To fairly compare with GAT, we take the best results from
either the original paper (Veličković et al.,, 2018) or our rerun using DGL
and a larger GAT model.

Table 2 gives the mean node classification accuracy as well as standard
deviation over five random repeats. Results of different algorithms are
compared via a $t$-test with $p<0.05$. The best results and any statistically
indistinguishable from the best are bold. We see that our model combined with
APPNP achieves the best performance in general. Other combinations also
outperform the GMNN on two out of three datasets. For all three backbone GNNs,
integrating it with NMM leads to significantly improved performances (marked
as $\checkmark$) on at least two of the three datasets. This result confirms
that incorporating neighborhood label dependencies into GNN models is
beneficial.

To assess robustness, we test all algorithms at different training set sizes
(varying the _fraction_ of nodes included in the training set). We vary this
fraction from 0.1 to 0.5 while fixing validation and test fractions to 0.2 and
0.3. Figure 1 shows test accuracy of all algorithms on Cora and Citeseer
datasets. Every NMM-GNN combination (solid lines) outperforms its
corresponding baseline GNN (same color, dashed). All NMM-GNN models also
deliver better performance than GMNN.

Accurate estimation of probabilistic dependencies is an important aspect of
probabilistic models. Node classification accuracy only reflects marginal
probabilities of single nodes. To better compare estimated correlations
between nodes, we examine the _pairwise marginal log likelihood_ (PLL)
reported from both GMNN and NMM models. Formally, we compute PLL as the
average log-likelihood per edge: $PLL=\frac{1}{|E_{test}|}\sum_{(i,j)\in
E_{test}}\log p(y_{i},y_{j})$. We randomly choose an edge set $E_{test}$ such
that each edge has two incident nodes from the test node set. We vary the
training ratio from 0.2 to 0.6 and fix the validation ratio (0.1) and test
ratio (0.3). Table 3 shows that our NMM achieves better predictions of
pairwise probabilities than GMNN, suggesting that our NMM can better capture
correlations between connected nodes.

Finally, we also inspect the quality of the approximation by the variational
distribution in (5). We do so on the pairwise marginals
$p(y_{i},y_{j}),(i,j)\in E_{test}$. Figure 2 compares the variational lower
bound against the exact log marginal. The bound is quite tight, which is a
strong evidence that the variational distribution
$q(\mathbf{c}|\mathbf{y},\bm{\alpha},\mathbf{L})$ is an accurate approximation
of the true posterior $p(\mathbf{c}|\mathbf{y},\bm{\alpha},\mathbf{L})$.

###  5.2 Image Denoising

Figure 1: Classification accuracy on heldout test nodes as training set size
increases. Training ratio at 0.5 results in the highest test accuracy. This
result is higher than the results in Table 2 where we used a smaller training
ratio.

Figure 2: Marginal log-likelihood of label pairs $y_{i},y_{j}$ in the test set
over many training iterations, using NMM-GCN with 0.2 data fraction. Our
scalable bound is a tight approximation of the expensive exact.

Algorithm | NIPS12 | Yeast | Cora | Citeseer | Pubmed
---|---|---|---|---|---
DGLFRM | 88.66 $\pm$ 0.56 | 84.03 $\pm$ 0.69 | 93.44 $\pm$ 0.40 | 94.31 $\pm$ 0.36 | 96.47 $\pm$ 0.18
NMM-DGLFRM | 89.82 $\pm$ 0.64 | 84.92 $\pm$ 0.51 | 94.16 $\pm$ 0.47 | 94.75 $\pm$ 0.38 | 96.45 $\pm$ 0.14

Table 4: Link prediction results, measured by Average Precision (AP, higher is better). |  |  |
---|---|---|---
Original Image | Clean Patch | Baseline | NMM
Figure 3: Image denoising result. NMM recovers the sky and sharp propeller
edge better than the baseline.

Method | BSD300 | Kodak
---|---|---
Baseline | 30.28 | 31.07
NMM | 30.33 | 31.55

Table 5: PSNR (the larger the better) performance for image denoising on
BSD300 testing set and Kodak.

| MF |  NMM-$\bm{\alpha},\mathbf{L}$ | NMM-CNN
---|---|---|---
KL/UB | 3179 | 3080 | 3105

Table 6: KL divergence or upper bounds on KL for methods that approximate an
MRF (lower is better).

We also evaluate NMM model on an image denoising task. Recent methods use
U-Nets (Ronneberger et al.,, 2015) for denoising (Mao et al.,, 2016; Lehtinen
et al.,, 2018), which do not model pixel correlations. In this experiment, we
use NMM to model pixel correlations and put a U-Net as the backbone of the
NMM. We treat pixel values as discrete labels, and each pixel’s neighborhood
is the $3\times 3$ area around it.

The baseline U-Net and the NMM are trained with noisy images as inputs and
clean images as targets. The training set is the BSD300 (Martin et al.,, 2001)
training set. Due to resource limitations, both models are trained on
$32\times 32$ patches, not full images. The trained models are evaluated on
the BSD300 test set and the Kodak dataset111http://r0k.us/graphics/kodak/.
Following Lehtinen et al., (2018), we inject Gaussian noise to obtain training
(std. dev. $\sigma\in[0,50]$) and test ($\sigma=25$) images.

Table 5 shows the peak Signal-to-Noise Ratio (PSNR) of denoised images by
different methods (the larger the better), and Figure 3 shows a typical
example on which the NMM outperforms the baseline. Note that the PSNR is
directly computed from the per-pixel Mean Squared Error (MSE) and the U-Net
directly minimizes MSE during training, so the evaluation measure favors the
U-Net. Nevertheless, our NMM still achieves better PSNR. The main reason is
that images of the natural world have strong correlations between neighboring
pixels, and NMM can better capture such correlations. In Figure 3, the NMM
achieves qualitatively better denoising results in smooth areas while the
attention mechanism avoids blurring to keep edges sharp.

###  5.3 Approximating other distributions

Approximating an MRF. We first use NMM to approximate inference for an MRF, as
described in Sec. 4. We create an Ising model on a $128\times 128$ grid
(Murphy,, 2012, section 21.3.2). As a baseline, we approximate the MRF with
the mean-field distribution (MF). We use two NMMs to approximate the MRF: one
that directly optimizes $\bm{\alpha}$ and $\mathbf{L}$
(NMM-$\bm{\alpha},\mathbf{L}$), and one that optimizes a CNN to compute
$\bm{\alpha}$ and $\mathbf{L}$ (NMM-CNN). The input to the CNN is a grid graph
with all MRF edge potential and node potential parameters collated at grid
nodes. For the two NMMs, we can only compute an upper bound in (11). Table 6
shows that the upper bounds of the two NMMs achieve noticeably better
approximations (lower KL) than the mean-field baseline. The relatively small
difference (<1%) between NMM-$\bm{\alpha},\mathbf{L}$ and NMM-CNN also
indicates that the CNN as an inference network is strong enough to get good
inference results despite its amortized parameterization.

Approximating deep generative graph models with latent variables. Next, we use
our NMM to do inference for the Deep Generative Latent Feature Relational
Model (DGLFRM) (Mehta et al.,, 2019). Formally, DGLFRM itself defines a deep
generative stochastic blockmodel with random variables $\mathbf{y}_{i}$ and
$\mathbf{z}_{i}$222The original paper uses notation
$p(\mathbf{v},\mathbf{b})$, which maps to our notation as
$\mathbf{z}=\mathbf{v}$, $\mathbf{y}=\mathbf{b}$. We omit notation of their
variable $\mathbf{r}$ for conciseness, as we keep the same parameterization of
$\mathbf{r}$ as the original work.. Binary vector
$\mathbf{y}_{i}\in\\{0,1\\}^{K}$ indicates node $i$’s memberships in $K$
communities, and latent variable $\mathbf{z}_{i}$ serves as stick-breaking
construction of the Indian Buffet Process (IBP) for $\mathbf{y}_{i}$ (Teh et
al.,, 2007). DGLFRM pursues variational inference with a structured mean-field
approximate posterior:

$\displaystyle
q(\mathbf{z},\mathbf{y})=\textstyle\prod_{k=1}^{K}\prod_{i=1}^{N}q(z_{ik})q\left(y_{ik}|\pi_{ik}(\mathbf{z}_{i})\right).$

Here, $z_{ik}\sim\mathrm{Beta}(a_{k},b_{k})$,
$y_{ik}\sim\mathrm{Bernoulli}\left(\pi_{ik}(\mathbf{z}_{i})\right)$, and
$\pi_{ik}(\mathbf{z}_{i})=\prod_{j=1}^{k}z_{ij}$. The probability vector
$\bm{\pi}_{i}(\mathbf{z}_{i})$ is a transformation of i.i.d. random variables
$\mathbf{z}_{i}$.

We adapt our NMM to approximate a target posterior distribution like
$q(\mathbf{z},\mathbf{y})$ with slight modifications: instead of sampling
$\mathbf{y}_{i}$ with probability $\mathbold{\pi}_{i}$, our NMM still
introduces a random variable $c_{i}$ for node $i$, and then we sample
$\mathbf{y}_{i}$ from $\mathbold{\pi}_{c_{i}}$. As in the original NMM
definition, the indicator variable $c_{i}$ introduces membership correlations
between neighboring nodes, which the original inference above cannot do.

We trained two models: DGLFRM with its original inference, and our NMM-DGLFRM.
We evaluate link prediction over five datasets: NIPS12, Yeast, Cora, Citeseer,
and Pubmed, following Mehta et al., (2019). For each dataset, 10% and 5% of
edges are held out as test and validation sets. We report Average Precision
(AP) as the evaluation metric.

Table 4 shows that NMM-DGLFRM outperforms the baseline model on four out of
five datasets. We also observe that our model achieves higher variational
lower bound than the baseline DGLFRM (see the appendix). Since neighboring
nodes often share similar memberships, information sharing between neighbors
clearly improves the approximation of the true posterior. Thus, our NMM
method’s better inference has improved the model’s overall probabilistic
representation of graph data (e.g. better link prediction).

##  6 Conclusion

In this paper, we have presented a new model, the Neighbor Mixture Model
(NMM), which captures probabilistic correlations among labels arranged in a
graph. Parameterized by a GNN, the model learns informative representations
from input features. It enables scalable computation via a tractable
variational lower bound that requires no additional free parameters. The model
can also serve as an approximate distribution to enable scalable inference for
other models. The NMM’s high-quality performance across multiple tasks
indicates promising ability to model node correlations at scale, while easily
integrating with many other models.

## References

  * Besag, (1975) Besag, J. (1975).  Statistical analysis of non-lattice data.  Journal of the Royal Statistical Society: Series D (The Statistician), 24(3):179–195.
  * Blei et al., (2017) Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. (2017).  Variational inference: A review for statisticians.  Journal of the American statistical Association, 112(518):859–877.
  * Grathwohl et al., (2018) Grathwohl, W., Choi, D., Wu, Y., Roeder, G., and Duvenaud, D. (2018).  Backpropagation through the void: Optimizing control variates for black-box gradient estimation.  In International Conference on Learning Representations.
  * Hamilton et al., (2017) Hamilton, W., Ying, Z., and Leskovec, J. (2017).  Inductive representation learning on large graphs.  In Advances in neural information processing systems, pages 1024–1034.
  * Kindermann, (1980) Kindermann, R. (1980).  Markov random fields and their applications.  American mathematical society.
  * Kingma and Ba, (2014) Kingma, D. P. and Ba, J. (2014).  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980.
  * Kingma and Welling, (2014) Kingma, D. P. and Welling, M. (2014).  Auto-encoding variational bayes.  In International Conference on Learning Representations (ICLR).
  * Kipf and Welling, (2017) Kipf, T. N. and Welling, M. (2017).  Semi-supervised classification with graph convolutional networks.  In International Conference on Learning Representations (ICLR).
  * Klicpera et al., (2019) Klicpera, J., Bojchevski, A., and Günnemann, S. (2019).  Predict then propagate: Graph neural networks meet personalized pagerank.  In International Conference on Learning Representations (ICLR).
  * Lafferty et al., (2001) Lafferty, J., McCallum, A., and Pereira, F. C. (2001).  Conditional random fields: Probabilistic models for segmenting and labeling sequence data.  In International Conference on Machine Learning.
  * Lehtinen et al., (2018) Lehtinen, J., Munkberg, J., Hasselgren, J., Laine, S., Karras, T., Aittala, M., and Aila, T. (2018).  Noise2noise: Learning image restoration without clean data.  In International Conference on Machine Learning, pages 2965–2974.
  * Louizos and Welling, (2017) Louizos, C. and Welling, M. (2017).  Multiplicative normalizing flows for variational bayesian neural networks.  In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 2218–2227. JMLR. org.
  * Lucas et al., (2019) Lucas, J., Tucker, G., Grosse, R., and Norouzi, M. (2019).  Understanding posterior collapse in generative latent variable models.  In ICLR Workshop DeepGenStruct.
  * Mao et al., (2016) Mao, X.-J., Shen, C., and Yang, Y.-B. (2016).  Image restoration using convolutional auto-encoders with symmetric skip connections.  arXiv preprint arXiv:1606.08921.
  * Martin et al., (2001) Martin, D., Fowlkes, C., Tal, D., and Malik, J. (2001).  A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics.  In Proceedings Eighth IEEE International Conference on Computer Vision. ICCV 2001, volume 2, pages 416–423. IEEE.
  * Mehta et al., (2019) Mehta, N., Duke, L. C., and Rai, P. (2019).  Stochastic blockmodels meet graph neural networks.  In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4466–4474, Long Beach, California, USA. PMLR.
  * Murphy, (2012) Murphy, K. P. (2012).  Machine learning: a probabilistic perspective.  MIT press.
  * Qu et al., (2019) Qu, M., Bengio, Y., and Tang, J. (2019).  Gmnn: Graph markov neural networks.  In International Conference on Machine Learning, pages 5241–5250.
  * Ranganath et al., (2016) Ranganath, R., Tran, D., and Blei, D. (2016).  Hierarchical variational models.  In International Conference on Machine Learning, pages 324–333.
  * Ronneberger et al., (2015) Ronneberger, O., Fischer, P., and Brox, T. (2015).  U-net: Convolutional networks for biomedical image segmentation.  In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer.
  * Satorras and Welling, (2020) Satorras, V. G. and Welling, M. (2020).  Neural enhanced belief propagation on factor graphs.
  * Shu et al., (2018) Shu, R., Bui, H. H., Zhao, S., Kochenderfer, M. J., and Ermon, S. (2018).  Amortized inference regularization.  In Advances in Neural Information Processing Systems, pages 4393–4402.
  * Teh et al., (2007) Teh, Y. W., Grür, D., and Ghahramani, Z. (2007).  Stick-breaking construction for the indian buffet process.  In Artificial Intelligence and Statistics, pages 556–563.
  * Tucker et al., (2017) Tucker, G., Mnih, A., Maddison, C. J., Lawson, J., and Sohl-Dickstein, J. (2017).  Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models.  In Advances in Neural Information Processing Systems, pages 2627–2636.
  * Veličković et al., (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. (2018).  Graph Attention Networks.  International Conference on Learning Representations.
  * Wainwright et al., (2008) Wainwright, M. J., Jordan, M. I., et al. (2008).  Graphical models, exponential families, and variational inference.  Foundations and Trends® in Machine Learning, 1(1–2):1–305.
  * Wang et al., (2019) Wang, M., Yu, L., Zheng, D., Gan, Q., Gai, Y., Ye, Z., Li, M., Zhou, J., Huang, Q., Ma, C., Huang, Z., Guo, Q., Zhang, H., Lin, H., Zhao, J., Li, J., Smola, A. J., and Zhang, Z. (2019).  Deep graph library: Towards efficient and scalable deep learning on graphs.  ICLR Workshop on Representation Learning on Graphs and Manifolds.
  * Williams, (1992) Williams, R. J. (1992).  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  Machine learning, 8(3-4):229–256.
  * Wu et al., (2019) Wu, Z., Pan, S., Chen, F., Long, G., Zhang, C., and Yu, P. S. (2019).  A comprehensive survey on graph neural networks.  arXiv preprint arXiv:1901.00596.
  * Xu et al., (2019) Xu, D., Cheng, W., Luo, D., Gu, Y., Liu, X., Ni, J., Zong, B., Chen, H., and Zhang, X. (2019).  Adaptive neural network for node classification in dynamic networks.  In 2019 IEEE International Conference on Data Mining (ICDM), pages 1402–1407. IEEE.
  * Yang et al., (2016) Yang, Z., Cohen, W., and Salakhudinov, R. (2016).  Revisiting semi-supervised learning with graph embeddings.  In Proceedings of The 33rd International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR.
  * Yoon et al., (2018) Yoon, K., Liao, R., Xiong, Y., Zhang, L., Fetaya, E., Urtasun, R., Zemel, R., and Pitkow, X. (2018).  Inference in probabilistic graphical models by graph neural networks.
  * Zhang et al., (2018) Zhang, C., Bütepage, J., Kjellström, H., and Mandt, S. (2018).  Advances in variational inference.  IEEE transactions on pattern analysis and machine intelligence, 41(8):2008–2026.

##  Appendix A Calculating of the Joint Probability of $\mathbf{y}$ and
$\mathbf{c}$

The closed-form of
$p(\mathbf{y}_{\tau},\mathbf{c}_{\tau}|\bm{\alpha},\mathbf{L})$ is needed for
marginalization (Sec. 3.2 of main paper), variational lower bound (Sec. 3.3),
and conditional probability (Sec. 3.6). Here we give a detailed derivation.
The general idea is to leverage the Dirichlet-categorical conjugacy to do the
computation.

$\displaystyle p(\mathbf{y}_{\tau},\mathbf{c}_{\tau}|\bm{\alpha},\mathbf{L})$
$\displaystyle=\int_{\mathbf{z}_{n(\tau)}}p(\mathbf{y}_{\tau},\mathbf{c}_{\tau},\mathbf{z}_{n(\tau)}|\bm{\alpha},\mathbf{L})$
$\displaystyle=\int_{\mathbf{z}_{n(\tau)}}p(\mathbf{z}_{n(\tau)}|\bm{\alpha})p(\mathbf{c}_{\tau}|\mathbf{L})p(\mathbf{y}_{\tau}|\mathbf{c}_{\tau},\mathbf{z}_{n(\tau)})$
$\displaystyle=p(\mathbf{c}_{\tau}|\mathbf{L})\int_{\mathbf{z}_{n(\tau)}}p(\mathbf{z}_{n(\tau)}|\bm{\alpha})\prod_{i\in\tau}p(y_{i}|c_{i},\mathbf{z}_{n(i)})$
$\displaystyle=\prod_{i\in\tau}L_{i,c_{i}}\int_{\mathbf{z}_{n(\tau)}}\prod_{j\in
n(\tau)}\text{Dir}(\mathbf{z}_{j}|\bm{\alpha}_{j})\prod_{i\in\tau}z_{c_{i},y_{i}}$
$\displaystyle=\prod_{i\in\tau}L_{i,c_{i}}\prod_{j\in
n(\tau)}\frac{\mathrm{B}(\bm{\alpha}_{j}+\mathbf{s}_{j}(\mathbf{y}_{\tau},\mathbf{c}_{\tau}))}{\mathrm{B}(\bm{\alpha}_{j})}$
$\displaystyle=\prod_{i\in\tau}L_{i,c_{i}}\prod_{j\in\mathbf{c}_{\tau}}\frac{\mathrm{B}(\bm{\alpha}_{j}+\mathbf{s}_{j}(\mathbf{y}_{\tau},\mathbf{c}_{\tau}))}{\mathrm{B}(\bm{\alpha}_{j})}.$
(13)

Here the integral is computed from the Dirichelet-multinomial conjugacy. The
vector
$\mathbf{s}_{j}(\mathbf{y}_{\tau},\mathbf{c}_{\tau})=\sum_{i\in\tau:c_{i}=j}\mathrm{onehot}(y_{i})$,
and
$\mathrm{B}(\bm{\alpha})=\frac{\Gamma(\alpha_{1})\Gamma(\alpha_{2})\ldots\Gamma(\alpha_{C})}{\Gamma(\sum_{k}\alpha_{k})}$
denotes the multivariate Beta function.

##  Appendix B Graph Node Classification: Details and Additional Results

###  B.1 Experimental Details

In node classification, we have used several possible GNN architectures (GCN,
GAT, and APPNP) to construct $\bm{\alpha}$ and $\mathbf{L}$. In the paper, we
described a single function $f(\mathbf{X};\theta)$ that used a GNN to compute
hidden representation vectors $(\mathbf{u}_{i},\mathbf{v}_{i})_{i\in V}$ for
every node $i$, and then compute $\bm{\alpha}$ from $(\mathbf{u}_{i})_{i\in
V}$ and compute $\mathbf{L}$ from $(\mathbf{v}_{i})_{i\in V}$. In the actual
implementation, we have used two separate graph neural networks to compute
$(\mathbf{u}_{i})_{i\in V}$ and $(\mathbf{v}_{i})_{i\in V}$. Essentially, this
means our function $f$ in practice does not share structure between
computation of $\mathbf{u}_{i}$ and $\mathbf{v}_{i}$, but nothing prevents
trying other implementations in the future.

We now review our concrete architectures for computing $(\mathbf{u}_{i})_{i\in
V}$ with each possible GNN architecture. For GCN, we use a two layer GCN with
16 hidden nodes as in Kipf and Welling, (2017). For GAT, we use a two layer
GAT with 8 attention heads in both layers and 16 hidden nodes (slightly larger
than the 8 hidden nodes used in Veličković et al., (2018)). Finally, for APPNP
we use a two layer APPNP with 0.1 teleport probability, 10 propagation steps,
and 64 hidden nodes as in Klicpera et al., (2019).

Similarly, we can compute $(\mathbf{v}_{i})_{i\in V}$ with each possible GNN
architecture. Our GCN uses a one layer GCN (64 hidden nodes). Our GAT uses a
one layer GAT (32 hidden nodes and 8 heads). APPNP uses the same architectures
to compute $\mathbf{u}_{i}$-s and $\mathbf{v}_{i}$-s, except that the output
dimension for a $\mathbf{v}_{i}$ is 32. Experiments are conducted in a Tesla
V100 GPU cluster.

Hyperparameters. Hyperparameters are selected by checking performance on a
fixed validation set. Models are trained on a maximum of 200 epochs, stopping
early if validation performance has not improved over a contiguous window of
100 epochs. We use Adam with a step size pool $\\{0.01,0.03,0.05,0.07,0.1\\}$.
We apply $L_{2}$ regularization with $\gamma=0.005$ and dropout rate from
$\\{0.5,0.7\\}$. The output activation for $\bm{\alpha}$ inference network is
validated from
$\\{\operatorname*{softplus}(\cdot),\mathrm{square(\cdot)}\\}+1.0$, to ensure
the $\bm{\alpha}$ defines valid Dirichlet densities with a unique mode. In our
experiment, we observe that $\operatorname*{softplus}(\cdot)$ is better at
stabilizing the training process while $\mathrm{square(\cdot)}$ enables faster
training, especially when the output logits of the $\bm{\alpha}$ inference
network are negative. For all node classification experiments, to compute
$\mathbf{L}$ we set the “self-attention” scalar $\gamma$ to be 0.0, and cosine
similarity scale $\sigma^{2}$ is optimized with an initial value 1.0.

Baselines. The GMNN model has a “temperature” parameter to sharpen the
variational distribution. This ad-hoc parameter leads to slightly better
performance but very poor probability estimation. Throughout our experiments,
we have tested the model with the temperature parameter setting to 1.0, which
yields “standard” variational inference.

###  B.2 Additional Results: Runtime and Convergence Analysis

Algorithm | Cora | Citeseer | Pubmed
---|---|---|---
GCN | 0.01s | 0.02s | 0.04s
GMNN | 0.03s | 0.07s | 0.14s
NMM-GCN | 0.04s | 0.03s | 0.05s
Table 7: Training time per epoch. Each epoch has a forward pass, a loss
computation, and a backward pass.

Figure 4: Convergence curves of GCN and NMM-GCN.

Runtime. In addition to time complexity (Sec. 3.7 of main paper), we report
mean training time per epoch for 100 epochs to probe practical runtime (Kipf
and Welling,, 2017). Table 7 summarizes the results. NMM-GCN introduces minor
runtime overhead its backbone GCN. NMM-GCN is faster than GMNN on Citeseer and
Pubmed. Similar results are observed when using other backbones for NMM (GAT
and APPNP)

Convergence Analysis. Figure 4 shows the convergence curves of GCN and NMM-
GCN. We see NMM-GCN has similar convergence behavior as its backbone GCN. We
did not include convergence curve of GMNN, as GMNN has a pre-training stage
beforehand.

###  B.3 Additional Results: Pairwise Probability Assessment

Table 8 gives full results for pairwise log likelihood (PLL, Table 3 of main
paper). Mean PLL and standard deviation are reported over five random runs.
The results indicate that the NMM has better calibrated probabilities than
GMNN.

Algorithm |  | Cora |  |  |  | Citeseer |
---|---|---|---|---|---|---|---
0.2 | 0.4 | 0.6 |  | 0.2 | 0.4 | 0.6
GMNN | -1.34 $\pm$ 0.06 | -1.30 $\pm$ 0.09 | -1.14 $\pm$ 0.04 |  | -1.62 $\pm$ 0.08 | -1.53 $\pm$ 0.05 | -1.41 $\pm$ 0.04
NMM-GCN | -1.04 $\pm$ 0.01 | -0.97 $\pm$ 0.08 | -0.81 $\pm$ 0.03 |  | -1.37 $\pm$ 0.02 | -1.34 $\pm$ 0.01 | -1.26 $\pm$ 0.01
NMM-GAT | -1.01 $\pm$ 0.01 | -0.84 $\pm$ 0.06 | -0.71 $\pm$ 0.01 |  | -1.35 $\pm$ 0.02 | -1.34 $\pm$ 0.02 | -1.25 $\pm$ 0.02
NMM-APPNP | -0.93 $\pm$ 0.01 | -0.83 $\pm$ 0.04 | -0.70 $\pm$ 0.02 |  | -1.33 $\pm$ 0.02 | -1.30 $\pm$ 0.02 | -1.24 $\pm$ 0.01

Table 8: Marginal log-likelihood of label pairs $y_{i},y_{j}$ in the test set
across fractions of observed data included in training set (0.2, 0.4, and
0.6).

##  Appendix C Image Denoising: Details

The baseline U-Net (Ronneberger et al.,, 2015) is trained to fit pixel values.
We also use such a network as the backbone of the NMM. We convert a scalar
prediction $u_{i}$ to a positive vector $\bm{\alpha}_{i}$ as follows by
$\alpha_{ik}=\exp(-(k-u_{i})^{2}/\sigma^{2})*s,k=0,\ldots,255$. Then the
Dirchlet distribution with $\bm{\alpha}_{i}$ can gives a probability vector
with probabilities concentrated around $u_{i}$. Vectors
$(\mathbf{v}_{i})_{i\in V}$ are computed from a two-layer MLP branching out
from the second-to-the-last layer of the U-net. To compute $\mathbf{L}$, we
set $\gamma=5.0$. We treat each pixel with a surrounding $3\times 3$ window as
its neighbors. Except for the learning rate as 0.0001, all hyperparameters are
kept in line with Lehtinen et al., (2018).

##  Appendix D Approximating Other Distributions: Details and Additional
Results

###  D.1 Approximating an MRF

We give more details about the main paper’s Table 6, an experimental
comparison between a mean-field variational inference baseline for the MRF,
and using our NMM to approximate the MRF. In this experiment, the function
$f(\mathbf{X};\theta)$ has two separate CNNs to parameterize $\bm{\alpha}$ and
$\mathbf{L}$. The input to both CNNs is the graph with MRF potentials attached
to graph nodes, so $\mathbf{X}$ includes all MRF potentials. The first layer
of the CNN has kernel size $11\times 11$ and filter size 128, followed by a
ReLU nonlinearity. The second layer of the CNN has the kernel size $1\times
1$, followed by an identity activation. The number of filters in the second
layer is set to be 2 for the $\bm{\alpha}$ and the number of neighbors (49)
for the $\mathbf{L}$. We treat each pixel with a surrounding $7\times 7$
window as its neighbors. We train our model using Adam optimizer (Kingma and
Ba,, 2014) with 1000 maximum epochs and $0.001$ step size.

###  D.2 Approximating Deep Generative Graph Models with Latent Variables

####  D.2.1 Experimental Details

We follow the same experiment setting as in DGLFRM (Mehta et al.,, 2019). The
only difference is that DGLFRM samples $\mathbf{y}_{i}$ from the probability
vector $\bm{\pi}_{i}$ directly, while we sample $\mathbf{y}_{i}$ from
$\bm{\pi}_{c_{i}}$ with $c_{i}\sim\mathrm{Categorical}(\mathbf{L}_{i})$ to
describe neighbor correlations. To compute $\mathbf{L}_{i}$ from
$\mathbf{v}_{i}$, we use the same amortization structure as DGLFRM.
Specifically, the inference network for $(\mathbf{v}_{i})_{i\in V}$ is a two
layer GCN. The first layer has 32/64 hidden nodes for Cora, Citeseer and
Pubmed, or 128/256 hidden nodes for NIPS12 and Yeast. The second layer has
50/100/200 hidden nodes to represent the community size. All the models are
trained for 500 to 1000 iterations. Adam optimizer is used at a learning rate
0.01.

Figure 5: ELBO curve as a function of number of iteration. Tested on NIPS12.

####  D.2.2 Additional Results: Link Prediction

Algorithm | NIPS12 | Yeast | Cora | Citeseer | Pubmed
---|---|---|---|---|---
DGLFRM | 86.47 $\pm$ 0.83 | 77.97 $\pm$ 0.79 | 93.00 $\pm$ 0.34 | 93.79 $\pm$ 0.42 | 96.11 $\pm$ 0.19
NMM-DGLFRM | 86.88 $\pm$ 0.67 | 79.06 $\pm$ 0.72 | 93.48 $\pm$ 0.41 | 93.81 $\pm$ 0.50 | 96.05 $\pm$ 1.16

Table 9: Results of AUC ROC.

The main paper’s Table 4 showed link prediction results as measured by Average
Precision (AP). Additionally, here we share similar link prediction results
evaluated by the Area Under the ROC Curve (AUC). Table 9 compares results from
the original DGLFRM and results from our inference method. NMM inference
significantly improves the performance of DGLFRM on two out of five datasets
(Yeast and Cora).

####  D.2.3 Additional Results: Likelihood Bound Quality

As we have mentioned in the submission, using the NMM for inference can better
maximize the variational lower bound on the marginal likelihood of observed
labels under the DGLFRM model than the original inference method. The trace
plot in Figure 5 shows evidence of this improvement, showing a noticeably
higher value of the evidence lower bound (ELBO). From this, we conclude that
sharing neighborhood information with our proposed NMM provides improved
modeling capabilities for graph data.