Graph Neural Networks with a Distribution of Parametrized Graphs
Traditionally, graph neural networks have been trained using a single
observed graph. However, the observed graph represents only one possible
realization. In many applications, the graph may encounter uncertainties, such
as having erroneous or missing edges, as well as edge weights that provide
little informative value. To address these challenges and capture additional
information previously absent in the observed graph, we introduce latent
variables to parameterize and generate multiple graphs. We obtain the maximum
likelihood estimate of the network parameters in an Expectation-Maximization
(EM) framework based on the multiple graphs. Specifically, we iteratively
determine the distribution of the graphs using a Markov Chain Monte Carlo
(MCMC) method, incorporating the principles of PAC-Bayesian theory. Numerical
experiments demonstrate improvements in performance against baseline models on
node classification for heterogeneous graphs and graph regression on chemistry
datasets.
See Hian Lee, Feng Ji, Kelin Xia, Wee Peng Tay
# Graph Neural Networks with a Distribution of Parametrized Graphs

###### Abstract

Traditionally, graph neural networks have been trained using a single observed graph. However, the observed graph represents only one possible realization. In many applications, the graph may encounter uncertainties, such as having erroneous or missing edges, as well as edge weights that provide little informative value. To address these challenges and capture additional information previously absent in the observed graph, we introduce latent variables to parameterize and generate multiple graphs. We obtain the maximum likelihood estimate of the network parameters in an Expectation-Maximization (EM) framework based on the multiple graphs. Specifically, we iteratively determine the distribution of the graphs using a Markov Chain Monte Carlo (MCMC) method, incorporating the principles of PAC-Bayesian theory. Numerical experiments demonstrate improvements in performance against baseline models on node classification for heterogeneous graphs and graph regression on chemistry datasets.

## 1 Introduction

Graph Neural Networks (GNNs) have facilitated graph representational learning by building upon Graph Signal Processing (GSP) principles and expanding their application in the domains of machine learning. Moreover, GNNs have demonstrated their effectiveness across a wide range of tasks in domains such as chemistry (Gilmer et al., 2017), recommendation systems (Ying et al., 2018; Chen et al., 2022), financial systems (Sawhney et al., 2021) and e-commerce settings (Liu et al., 2022), among others. However, GSP and GNNs conventionally rely on a fixed graph shift operator, such as the adjacency or Laplacian matrix, to analyze and learn from graph data, assuming that the given graph is accurate and noise-free. This approach has inherent limitations, considering that graph data is often uncertain.

The uncertainty is due to the existence of multiple potential variations in graph constructions as a universal optimal method does not exist. Furthermore, structural noise, which includes missing or spurious edges, and the absence of informative edge weights, can also contribute to the uncertainty in graph data (Zhang et al., 2019; Dong and Kluger, 2023). It is important to handle this uncertainty as the graph directly influences the results of both GSP and GNNs (Li et al., 2021).

Several GNN works have recognized that the provided graph in benchmark datasets is suboptimal. For example, in Topping et al. (2022), a method was introduced to enhance the provided graph by rewiring it at graph bottlenecks. Similarly, in Li et al. (2021) and Ye et al. (2020), approaches were developed to reweigh edges, to reduce information flow at cluster boundaries. Another perspective involves considering the given or observed graph as a particular realization of a graph model, as discussed in Zhang et al. (2019). In their work, a Bayesian framework was adopted to learn a more robust model that can withstand perturbations in graph topology. These collective efforts underscore the common observation that the observed graph is often imperfect, and determining the optimal graph is a non-trivial task, as it depends on both the physical connections and the edge weights, which regulate the rates of information transmission (Ji et al., 2023).

Our work aligns with the viewpoint presented in Zhang et al. (2019). We conceptualize the observed graph as an individual instance originating from a distribution of graphs, which is influenced by one or more latent parameters. Nevertheless, in contrast to Zhang et al. (2019) which proposed a Bayesian framework, we propose an EM framework for graph learning and name our model EMGNN. Even though both are probabilistic frameworks, the focus is distinctly different. In the case of the Bayesian framework of Zhang et al. (2019), the focus is on estimating the posterior distribution
of model parameters given the data. As such, model parameters are deemed as random variables trained by a series of characteristically similar graphs. Meanwhile, in our EM framework, we seek to maximize the log-likelihood of the observed data in conjunction with the latent variables. Additionally, we permit the generated graphs to demonstrate more pronounced variations. Our main contributions are as follows:

* We introduce a general framework for modeling the distribution of graphs to handle the uncertainty in graph data. The learned distribution then serves as a valuable tool for comprehending our model and offering insights into its behavior.
* We formulate the graph learning problem as a maximum likelihood estimation (MLE) so that tools from statistical learning can be applied. The new objective subsumes the classical objective of minimizing empirical loss if the graph is deterministic.
* We conduct evaluations of our model using real datasets, carry out experiments in two distinct applications, and observe promising performance compared to the respective baseline methods.
* We inspect the learned graph distribution, confirming that it effectively captures the intricacies of heterogeneous graph datasets, thus validating the utility of our model and framework.

## 2 Preliminaries

### Graph Neural Networks

Graph neural networks (Chen et al., 2020; Kang et al., 2023; Brody et al., 2022; Lee et al., 2021; Zhao et al., 2023), which are neural networks designed to operate on graphs, typically employ the message-passing framework. Within this framework, the features of each node are integrated with those of its neighboring nodes to update its feature representation.

More specifically, suppose that we have a graph \(G=(V,E)\), where \(V\) is the set of vertices and \(E\) is the set of edges. Moreover, each node \(v\in V\) is associated with (initial) node features represented by \(x_{v}^{0}\). The node features can then be updated in the \(k\)-th layer as follows:

\[x_{v}^{k}=\sigma(W^{k}\text{AGGR}(\{x_{v}^{k}\mid v\in\mathcal{N}(v)\})) \tag{1}\]

where \(\sigma\) is an activation function, \(W^{k}\) are the learnable weights in the \(k\)-th layer and \(\mathcal{N}(v)\) is the set of neighbors of \(v\). AGGR is a message aggregation function. The choice of AGGR defines various variants of GNNs (Xu et al., 2019). For example, the mean operator yields Graph Convolutional Networks (GCN) (Kipf and Welling, 2017), while using the attention mechanism results in Graph Attention Networks (GAT) (Velickovic et al., 2018).

In a GNN with \(K\) layers, the last layer outputs features \(\{x_{1}^{K},\ldots,x_{|V|}^{K}\}\). For node classification, these features can be used directly. Meanwhile, for a graph-level task, a READOUT graph-level pooling function is needed to obtain the graph-level representation (Xu et al., 2019).

### Signal Processing over a Distribution of Graphs

GNN is closely related to the theory of GSP (Shuman et al., 2013). Briefly, given an undirected graph \(G\), we consider a fixed graph shift operator \(S\) such as its adjacency or Laplacian matrix. A graph signal is a vector \(\mathbf{x}=(x_{v})_{v\in V}\) that associates a number \(x_{v}\) to each node \(v\in V\). Intuitively, applying the linear transformation \(S\) to \(\mathbf{x}\) is considered as a "shift" of \(\mathbf{x}\). If \(S\) is the normalized adjacency matrix, then it amounts to the AGGR step of (1) for GCN. More generally, if \(P(\cdot)\) is a single variable polynomial, then plugging in \(S\) results in the matrix \(P(S)\), which is called a _convolution filter_ in GSP. This notion of convolution appears in Michael Defferrard (2016), and has become widely used since then.

On the signal processing side, Ji et al. (2023) has developed a theory that generalizes traditional GSP. The authors propose a signal processing framework assuming a family of graphs are considered simultaneously, to tackle uncertainties in graph constructions. Formally, it considers a distribution \(\mu\) of graph shift operators \(S_{\lambda}\) parametrized by \(\lambda\) in a sample space \(\Lambda\). The work develops corresponding signal processing notions such as Fourier transform, filtering, and sampling. In particular, a convolution takes the form \(\mathbb{E}_{\lambda\sim\mu}[P_{\lambda}(S_{\lambda})]\), where \(P_{\lambda}(\cdot)\) is a polynomial and \(P_{\lambda}(S_{\lambda})\) is an ordinary convolution with shift \(S_{\lambda}\). Our work is based on a similar idea but replaces \(P_{\lambda}(S_{\lambda})\) with a more general filter. Furthermore, we introduce an EM framework that is not present in Ji et al. (2023) to learn the distribution of graphs.

## 3 Proposed Method

### Problem Formulation: Maximum Likelihood Estimation

We consider a distribution \(\mu\) on a parameter (sample) space \(\Lambda\subset\mathbb{R}^{r}\) of graphs \(\{G_{\lambda},\lambda\in\Lambda\}\), with a fixed set of nodes \(V\). For each \(\lambda\), there is a corresponding shift operator \(S_{\lambda}\). We usually assume that \(\mu\) has a density function \(p(\cdot)\) w.r.t. a base measure on \(\Lambda\). For example, if \(\Lambda\) is a finite set in \(\mathbb{R}^{r}\), we can use the discrete counting
measure as the base measure. On the other hand, if \(\Lambda\) is a compact interval in \(\mathbb{R}\), then we can choose the Lebesgue measure as the base measure.

Assume that each node \(v\in V\) is associated with features \(x_{v}\). They are collectively denoted by \(\mathbf{x}\). Suppose a GNN model \(\Psi\) outputs the learned embeddings \(\mathbf{z}=\Psi(\lambda,\mathbf{x};\boldsymbol{\theta})\) given the node features \(\mathbf{x}\), the graph parameter \(\lambda\), and the GNN model parameter vector \(\boldsymbol{\theta}\). These in turn are used to determine a vector of labels \(\hat{\mathbf{y}}\). For a task-specific loss \(\ell(\cdot,\cdot)\) that compares predicted \(\hat{\mathbf{y}}\) and true label (vector) \(\mathbf{y}\), we may compute \(L_{\mathbf{X}}(\lambda,\boldsymbol{\theta})=\ell(\hat{\mathbf{y}},\mathbf{y})\). We use \(\mathbf{X}\) to denote the full information \(\{\mathbf{x},\mathbf{y}\}\). We interpret \(\mathbf{X}\) as a sample from a random variable, denoted by \(\mathfrak{X}\), of collective information of features and labels.

An example of \(\Psi\) is the model described by (1). The parameters \(\boldsymbol{\theta}=\{\theta_{1}^{k},\theta_{2}^{k}\mid 1\leq k\leq K\}\) and \(\lambda\) determine \(W^{k}\) in the form of a linear combination \(W^{k}=\theta_{1}^{k}+\lambda\theta_{2}^{k}\). Moreover, AGGR is determined by the shift \(S_{\lambda}\) associated with \(G_{\lambda}\).

In general, as \(\lambda\) follows an unknown distribution \(\mu\), it is hard to find \(\boldsymbol{\theta}\) by minimizing \(\mathbb{E}_{\lambda\sim\mu}[L_{\mathbf{X}}(\lambda,\boldsymbol{\theta})]\) directly. On the other hand, the EM algorithm (Bishop, 2006) allows us to jointly estimate \(\mu\) and \(\boldsymbol{\theta}\) provided we can reformulate the objective as an MLE.

To minimize the loss given \(\mathbf{X}\), the parameter \(\boldsymbol{\theta}\) is determined by \(\lambda\) and vice versa. Therefore, both of them can be viewed as random variables, and thus \(\Psi(\cdot,\mathbf{x};\cdot)\) becomes a random GNN model that depends on \(\lambda,\boldsymbol{\theta}\) and input \(\mathbf{x}\). We aim to _identify a realization of the random models that makes the observation \(\mathbf{X}\) likely_, i.e., there is less discrepancy between the estimator labels \(\hat{\mathbf{y}}\) and ground truth labels \(\mathbf{y}\) measured by the loss \(\ell(\hat{\mathbf{y}},\mathbf{y})\). Motivated by the discussions above, we consider the likelihood function \(p(\lambda,\mathbf{X}\mid\boldsymbol{\theta})\) on \(\boldsymbol{\theta}\) and formulate the following MLE as our objective:

\[\boldsymbol{\theta}^{*}=\arg\max_{\boldsymbol{\theta}}p(\mathbf{X} \mid\boldsymbol{\theta})=\arg\max_{\boldsymbol{\theta}}\mathbb{E}_{\lambda \sim\mu}[p(\lambda,\mathbf{X}\mid\boldsymbol{\theta})]. \tag{2}\]

Before discussing the main algorithm in subsequent subsections, we preview the roles of \(\mu\) and \(L_{\mathbf{X}}(\cdot,\cdot)\) in the algorithm. We shall see that the EM algorithm outputs a distribution \(\hat{\mu}\) of \(\lambda\), serving as an estimate of \(\mu\), by leveraging the PAC-Bayesian framework. In this framework, the density of \(\lambda\) is proportional to the Gibbs posterior, influenced by a risk function. Consequently, \(\hat{\mu}\) assigns higher probability density to \(\lambda\) when the loss \(L_{\mathbf{X}}(\lambda,\boldsymbol{\theta}^{*})\) is lower. Therefore, we still need to minimize the given loss as a main component of the algorithm.

For a simple example, assume that \(\mu\) is the delta distribution \(\delta_{\lambda_{0}}\) supported on \(\lambda_{0}\) so that the graph \(G_{\lambda_{0}}\) is deterministic. If we consider the Gibbs posterior, then \(p(\mathbf{X}\mid\boldsymbol{\theta})\propto\exp(-\eta\ell(\hat{\mathbf{y}}, \mathbf{y}))\), where \(\hat{\mathbf{y}}\) depends on both \(\mathbf{X}=\{\mathbf{x},\mathbf{y}\}\) and \(\boldsymbol{\theta}\). Thus, maximizing \(p(\mathbf{X}\mid\boldsymbol{\theta})\) is equivalent to the classical objective of minimizing \(\ell(\hat{\mathbf{y}},\mathbf{y})\).

### Expectation-Maximization for GNN

Optimizing (2) directly can be challenging, and we utilize the EM algorithm that employs an iterative approach alternating between the E-step and the M-step. Adapted to our setting, the process unfolds as follows:

1. E-step: Given parameters \(\boldsymbol{\theta}^{(t)}\) at the \(t\)-th iteration, we compute the expectation as the Q-function \[Q(\boldsymbol{\theta}\mid\boldsymbol{\theta}^{(t)})=\mathbb{E}_{\lambda\sim p (\cdot\mid\mathbf{X},\boldsymbol{\theta}^{(t)})}[\log p(\lambda,\mathbf{X} \mid\boldsymbol{\theta})].\]
2. M-step: \(\boldsymbol{\theta}^{(t+1)}\) is updated as \(\arg\max Q(\boldsymbol{\theta}\mid\boldsymbol{\theta}^{(t)})\).

_For the E-step_ in the \(t\)-th iteration, in the same spirit as the PAC-Bayesian framework (Guedj, 2019), we apply the Gibbs posterior and assume that

\[p(\lambda,\mathbf{X}\mid\boldsymbol{\theta})\propto\exp(-\eta^{(t)}L_{ \mathbf{X}}(\lambda,\boldsymbol{\theta}))\pi_{0}(\lambda,\mathbf{X}), \tag{3}\]

for a tunable hyperparameter \(\eta^{(t)}\), while \(\pi_{0}(\cdot)\) is a prior density of the joint \((\lambda,\mathbf{X})\) independent of \(\boldsymbol{\theta}\), representing our initial knowledge regarding \(\lambda\) and \(\mathbf{X}\). In this expression, \(L_{\mathbf{X}}(\lambda,\boldsymbol{\theta})\) implicitly depends on the observations \(\mathbf{X}\).

The normalization constant is given by

\[\begin{split} C(\boldsymbol{\theta})&=\int_{(\lambda, \mathbf{X}^{\prime})\in\Lambda\times\mathfrak{X}}\exp(-\eta^{(t)}L_{\mathbf{X}^ {\prime}}(\lambda,\boldsymbol{\theta}))\pi_{0}(\lambda,\mathbf{X}^{\prime}) \,\mathrm{d}(\lambda,\mathbf{X}^{\prime})\\ &=\mathbb{E}_{(\lambda,\mathbf{X}^{\prime})\sim\pi_{0}}\Big{[} \exp(-\eta^{(t)}L_{\mathbf{X}^{\prime}}(\lambda,\boldsymbol{\theta}))\Big{]}. \end{split} \tag{4}\]

As a prior belief, we treat the observed \(\mathbf{X}\) as a typical sample such that the above average is (approximately) the same as the average over graphs by fixing \(\mathbf{X}\). We assume that for each fixed \(\mathbf{X}\), there exists some prior distribution with density \(p_{0,\mathbf{X}}(\cdot)\) on \(\Lambda\) such that:

\[\begin{split}&\mathbb{E}_{(\lambda,\mathbf{X}^{\prime})\sim\pi_{0}} \Big{[}\exp(-\eta^{(t)}L_{\mathbf{X}^{\prime}}(\lambda,\boldsymbol{\theta})) \Big{]}\\ &\approx\int_{\lambda\in\Lambda}\exp(-\eta^{(t)}L_{\mathbf{X}}( \lambda,\boldsymbol{\theta}))p_{0,\mathbf{X}}(\lambda)\,\mathrm{d}(\lambda)\\ &=\mathbb{E}_{\lambda\sim p_{0,\mathbf{X}}}\Big{[}\exp(-\eta^{(t)} L_{\mathbf{X}}(\lambda,\boldsymbol{\theta}))\,\Big{|}\,\mathbf{X}\Big{]}.\end{split}\]

For simplification, we denote \(p_{0,\mathbf{X}}(\cdot)\) as \(p_{0}(\cdot)\) and \(\mathbb{E}_{\lambda\sim p_{0}}\big{[}\exp(-\eta^{(t)}L_{\mathbf{X}}(\lambda, \boldsymbol{\theta}))\,\big{|}\,\mathbf{X}\Big{]}\) as \(\mathbb{E}_{\lambda\sim p_{0}}\big{[}\exp(-\eta^{(t)}L_{\mathbf{X}}(\lambda, \boldsymbol{\theta}))\big{]}\). We write

\[C(\boldsymbol{\theta})=\mathbb{E}_{\lambda\sim p_{0}}\Big{[}\exp(-\eta^{(t)}L_{ \mathbf{X}}(\lambda,\boldsymbol{\theta}))\Big{]}. \tag{5}\]
On the other hand, given \(\mathbf{\theta}^{(t)}\), from (3), we have

\[p(\lambda\mid\mathbf{X},\mathbf{\theta}^{(t)}) =\frac{p(\lambda,\mathbf{X}\mid\mathbf{\theta}^{(t)})}{p(\mathbf{X}\mid \mathbf{\theta}^{(t)})}\] \[\propto\exp(-\eta^{(t)}L_{\mathbf{X}}(\lambda,\mathbf{\theta}^{(t)})) \frac{\pi_{0}(\lambda,\mathbf{X})}{p(\mathbf{X}\mid\mathbf{\theta}^{(t)})}.\]

We assume that there is a prior \(p^{\prime}_{0,t}(\cdot)\) such that \(p^{\prime}_{0,t}(\cdot)\propto\frac{\pi_{0}(\lambda,\mathbf{X})}{p(\mathbf{X }\mid\mathbf{\theta}^{(t)})}\), and which does not depend on \(\mathbf{\theta}\) but is a function of \(t\). By fixing \(\mathbf{X}\), the posterior is written as

\[p(\lambda\mid\mathbf{X},\mathbf{\theta}^{(t)})\propto\exp(-\eta^{(t)}L_{\mathbf{ X}}(\lambda,\mathbf{\theta}^{(t)}))p^{\prime}_{0,t}(\lambda). \tag{6}\]

In our framework, we do not need to estimate the normalization constant for (6).

**Remark 1**.: _From the above discussion, we see that priors \(p_{0}(\cdot)\) and \(p^{\prime}_{0,t}(\cdot)\) play important roles. We discuss their choices in Section 4 below. It is desirable to have a weaker prior assumption, under which the optimizer can still be readily estimated._

_For the M-step,_ we analyze the Q-function in more detail. For convenience, we use \(p_{t}(\lambda)\) to denote \(p(\lambda\mid\mathbf{X},\mathbf{\theta}^{(t)})\).

With (5) and (6), we can express the Q-function as:

\[Q(\mathbf{\theta}\mid\mathbf{\theta}^{(t)})\] \[= \mathbb{E}_{\lambda\sim p_{t}}\bigg{[}\log\frac{\exp(-\eta^{(t)}L _{\mathbf{X}}(\lambda,\mathbf{\theta}))\pi_{0}(\lambda,\mathbf{X})}{C(\mathbf{\theta})}\bigg{]}\] \[= -\eta^{(t)}\mathbb{E}_{\lambda\sim p_{t}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta})]+D-\log C(\mathbf{\theta}),\]

where \(D\) is a constant independent of \(\mathbf{\theta}\).

To estimate \(\log C(\mathbf{\theta})\), we make the following considerations. First of all, by Jensen's inequality, \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta}) ]\leq\log C(\mathbf{\theta})\). This means that if \(\log C(\mathbf{\theta})\) is small, then necessarily so is \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta})]\). On the other hand, Teh et al. (2006) proposes to use \(\mathbb{E}[\log Y]+\frac{\text{var}(Y)}{2\mathbb{E}(Y)^{2}}\) to approximate \(\log\mathbb{E}[Y]\) for a random variable \(Y\). This is derived from the second-order Taylor expansion of \(\log Y\) at \(\log\mathbb{E}[Y]\). In our case, we have

\[\log C(\mathbf{\theta}) \approx-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}( \lambda,\mathbf{\theta})] \tag{7}\] \[+\frac{\text{var}\big{(}\exp(-\eta^{(t)}L_{\mathbf{X}}(\lambda, \mathbf{\theta}))\big{)}}{2\big{(}\mathbb{E}_{\lambda\sim p_{0}}\big{[}\exp(-\eta ^{(t)}L_{\mathbf{X}}(\lambda,\mathbf{\theta}))\big{]}\big{)}^{2}}.\]

If \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta})]\) is the dominant component, then we may use \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta})]\) as a proxy for \(\log C(\mathbf{\theta})\), which is more manageable. In Section 4.3.2, we shall numerically verify that this is indeed the case for our applications.

Hence, \(Q(\mathbf{\theta}\mid\mathbf{\theta}^{(t)})\) is approximated by

\[-\eta^{(t)}\Big{(}\mathbb{E}_{\lambda\sim p_{t}}[L_{\mathbf{X}}(\lambda,\mathbf{ \theta})]-\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta})] \Big{)}+D.\]

_In summary,_ if we disregard \(\eta^{(t)}\) and \(D\), which are independent of \(\mathbf{\theta}\), we can minimize the following function in the M-step:

\[J(\mathbf{\theta}) =\mathbb{E}_{\lambda\sim p_{t}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta })]-\mathbb{E}_{\lambda\sim p_{0}}[L_{\mathbf{X}}(\lambda,\mathbf{\theta})] \tag{8}\] \[=\int_{\lambda\in\Lambda}\big{(}p_{t}(\lambda)-p_{0}(\lambda)\big{)} L_{\mathbf{X}}(\lambda,\mathbf{\theta})\,\mathrm{d}\lambda.\]

### The Proposed Algorithm: EMGNN

To minimize \(J(\mathbf{\theta})\) in (8), our strategy is to re-express it as an expectation. For this purpose, we introduce a proposal distribution. Let \(q(\cdot)\) be the density function of a probability distribution on the sample space \(\Lambda\) whose support includes that of \(p_{0}\). Then we have:

\[J(\mathbf{\theta}) =\int_{\lambda\in\Lambda}q(\lambda)\frac{p_{t}(\lambda)-p_{0}( \lambda)}{q(\lambda)}L_{\mathbf{X}}(\lambda,\mathbf{\theta})\,\mathrm{d}\lambda\] \[=\mathbb{E}_{\lambda\sim q}\bigg{[}\frac{p_{t}(\lambda)-p_{0}( \lambda)}{q(\lambda)}L_{\mathbf{X}}(\lambda,\mathbf{\theta})\bigg{]}.\]

We propose to minimize \(J(\mathbf{\theta})\) by first randomly drawing samples \(\Lambda_{T^{\prime}}=\{\lambda_{1},\ldots,\lambda_{T^{\prime}}\}\) according to the density \(q(\cdot)\). Following that, we successively apply gradient descent to \(\frac{p_{t}(\lambda_{t^{\prime}})-p_{0}(\lambda_{t^{\prime}})}{q(\lambda)}L_{ \mathbf{X}}(\lambda_{t^{\prime}},\mathbf{\theta})\) to update \(\mathbf{\theta}\). Finally, given (6), \(p_{t}(\lambda)\) can be approximated by an empirical distribution if we apply an MCMC method. The overall algorithm is summarized in Algorithm 1 and illustrated in Fig. 1.

**Remark 2**.: _In practice, the choices of the prior distributions \(p_{0}(\cdot),q(\cdot)\) and \(p^{\prime}_{0,t}(\cdot)\) are hyperparameters. Moreover, in our experiments, \(p^{\prime}_{0,t}(\cdot)\) is set to be the same for every \(t\). We also discretize the continuous sample space \(\Lambda\) for simplicity in analysis and computation._

**Remark 3**.: _If we algorithmically plug in the delta distribution supported on \(\lambda_{0}\) and \(p_{0}(\lambda_{0})=0\) for \(p^{\prime}_{0,t}(\cdot)\) and \(q(\cdot)\) respectively, then EMGNN reduces to the ordinary GNN model on the graph \(G_{\lambda_{0}}\)._

Figure 1: Illustration of EMGNN.
**Remark 4**.: _Note that for the coefficient \(\frac{p_{t}(\lambda)-p_{0}(\lambda)}{q(\lambda)}\), if \(p_{t}(\lambda)<p_{0}(\lambda)\), then the loss \(L_{\mathbf{X}}(\mathbf{\mathbf{\theta}})\) is to be made larger. Intuitively, in this case, a "bad" \(\lambda\) is chosen. For the choice of \(q(\cdot)\), in practice, we propose two options in Section 4: either the uniform distribution or \(q(\cdot)=p_{t}(\cdot)\). Nonetheless, \(q(\cdot)\) can also be other appropriate density functions._

As we do not minimize \(J(\mathbf{\mathbf{\theta}})\) directly, we justify the proposed approach under additional assumptions. We theoretically analyze the performance of the proposed (randomized) algorithm in lines 5-12 of Algorithm 1, denoted by \(\mathcal{A}\). With samples \(\Lambda_{T^{\prime}}\), the algorithm \(\mathcal{A}\) outputs \(\widehat{\mathbf{\mathbf{\theta}}}=\mathcal{A}(\Lambda_{T^{\prime}})\). The following expression is considered in algorithm \(\mathcal{A}\):

\[J_{\Lambda_{T^{\prime}}}(\widehat{\mathbf{\mathbf{\theta}}})=\frac{1}{T^{\prime}} \sum_{\lambda_{t^{\prime}}\in\Lambda_{T^{\prime}}}\frac{p_{t}(\lambda_{t^{ \prime}})-p_{0}(\lambda_{t^{\prime}})}{q(\lambda_{t^{\prime}})}L_{\mathbf{X}} (\lambda_{t^{\prime}},\widehat{\mathbf{\mathbf{\theta}}}).\]

We assume that after translation and scaling by positive constant of \(L_{\mathbf{X}}(\lambda,\cdot)\) if necessary, the expression \(\frac{p_{t}(\lambda)-p_{0}(\lambda)}{q(\lambda)}L_{\mathbf{X}}(\lambda,\mathbf{ \mathbf{\theta}})\) always belong to \([0,1]\). The following notions are well-known.

**Definition 1**.: _A differential function \(f\) is \(\alpha\)-Lipschitz if for all \(x\) in the domain of \(f\), we have \(\|\nabla f(x)\|\leq\alpha\). It is \(\beta\)-smooth if its gradient is \(\beta\)-Lipschitz._

Denote by \(1_{p_{t}(\cdot)\geq p_{0}(\cdot)}(\lambda)\) the indicator that is \(1\) if \(p_{t}(\lambda)\geq p_{0}(\lambda)\), and \(0\) otherwise. Let \(b_{1}=\mathbb{E}_{\lambda\sim q}1_{p_{t}(\cdot)\geq p_{0}(\cdot)}\). Intuitively, it computes the measure of \(\lambda\), for which \(p_{t}(\cdot)\) is larger. On the other hand, let \(b_{2}=\mathbb{E}_{\lambda\sim q}\frac{|p_{t}(\lambda)-p_{0}(\lambda)|}{q( \lambda)}\), and \(\gamma=\sup_{\lambda\in\Lambda}1/q(\lambda)\).

**Theorem 1**.: _Assume for any \(\lambda\), the loss \(L_{\mathbf{X}}(\lambda,\cdot)\) is convex, \(\alpha\)-Lipschitz and \(\beta\)-smooth. Let \(b_{1},b_{2},\gamma\) be defined as above. If for every \(t^{\prime}\leq T^{\prime}\), the non-increasing step-size in the algorithm \(\mathcal{A}\) satisfies \(a_{t^{\prime}}\leq\min\{2/(\beta\gamma),c/t^{\prime}\}\) for a constant \(c\), then there is a constant \(C\) independent of \(T^{\prime},\alpha\) such that_

\[\Big{|}\mathbb{E}_{\mathcal{A},\Lambda_{T^{\prime}}}\Big{[}J_{\Lambda_{T^{ \prime}}}(\widehat{\mathbf{\mathbf{\theta}}})-J(\widehat{\mathbf{\mathbf{\theta}}}) \Big{]}\Big{|}\leq\epsilon=C\bigg{(}\frac{b_{2}^{2}\alpha^{2}}{T^{\prime}} \bigg{)}^{\frac{1}{\beta\gamma(1-b_{1})+1}}.\]

Proof.: The proof is given in the Appendix. 

**Remark 5**.: _From the result, we see that if \(b_{1}\) is close to \(1\), i.e., the set \(\{\lambda\mid p_{t}(\lambda)\geq p_{0}(\lambda)\}\) has a large measure, then the expected error decays at a rate close to \(T^{\prime-1}\)._

### A Brief Discussion on Testing

As our framework deals with a distribution of graphs, during testing, we acquire our final learned representation as \(\mathbf{z}_{\mathrm{final}}=\mathbb{E}_{\lambda\sim p_{T}}\big{[}\Psi(\lambda, \mathbf{x};\mathbf{\mathbf{\theta}}_{T^{\prime}}^{(T)})\big{]}\). The learned model parameters are a particular realization of the possible random models that align with the observed data \(\mathbf{X}\) and the multiple graphs influence the final embeddings based on their respective likelihoods. The embedding \(\mathbf{z}_{\mathrm{final}}\) is subjected to a softmax operation to obtain \(\hat{\mathbf{y}}\) for node classification tasks, while a READOUT function is applied for graph-level tasks.

## 4 Experiments

To validate our algorithm, we explored two tasks: node classification in heterogeneous graphs and graph regression using chemical datasets. Code and hyperparameter details are provided in the Appendix.

### Heterogeneous Graphs

Heterogeneous graphs are graph structures characterized by the presence of multiple node types and multiple edge types, imparting a greater degree of complexity compared to homogeneous graphs, which consist of a single node and edge type. In these datasets, we leveraged latent parameter(s) to generate multiple graph instances with varying edge weights for each edge type, representing different information transmission rates.

In the case of a heterogeneous graph with \(\omega\) edge types, we introduce a vector of latent parameters \(\mathbf{\mathbf{\lambda}}=\{\lambda_{1},\ldots,\lambda_{\omega}\}\), where each \(\lambda_{i}\geq 0\) and \(\sum_{i=1}^{\omega}\lambda_{i}=1\)
Drawing a set of latent variables uniformly from the space that the parameters span, a graph adjacency matrix of \(G_{\mathbf{\lambda}}\) is generated as follows:

\[A_{\mathbf{\lambda}}=\sum_{i=1}^{\omega}\lambda_{i}A_{i}, \tag{9}\]

where \(A_{i}\) is the respective edge-type specific adjacency matrix. Note that when any \(\lambda_{i}=0\), the edges of the associated edge types are removed.

#### 4.1.1 Baselines and Datasets

The heterogeneous graph datasets used are the same as those employed in Yun et al. (2019) and Lee et al. (2022). These datasets include two citation networks named DBLP1 and ACM2, as well as a movie dataset named IMDB3. Within the DBLP dataset, there are three distinct node types (Paper (P), Author (A), Conference (C)) and four edge types (PA, AP, PC, CP). The ACM dataset also comprises three node types (Paper (P), Author (A), and Subject (S)) and four edge types (PA, AP, PS, SP). Similarly, the IMDB dataset has three node types (Movie (M), Actor (A), Director (D)) along with four edge types (MD, DM, MA, AM). The dataset statistics are given in the Appendix.

Footnote 1: [https://dblp.uni-trier.de/](https://dblp.uni-trier.de/)

Footnote 2: [http://dl.acm.org/](http://dl.acm.org/)

Footnote 3: [https://www.imdb.com/interfaces/](https://www.imdb.com/interfaces/)

We assessed our approach against five baseline models. Specifically, GAT and GCN are designed for homogeneous graphs, while GTN (Yun et al., 2019), SimpleHGN (Lv et al., 2021) and SeHGNN (Yang et al., 2023) are state-of-the-art models developed for heterogeneous graph settings. For our model, we treat edges of different directions as a single type, resulting in two symmetric edge-type specific adjacency matrices. We consider the sample space \(\Lambda\) to span the range \([0.0,1.0]\) and discretize it in increments of \(0.05\) to obtain \(\widehat{\Lambda}\).

We consider different variants of EMGNN with a GCN backbone, based on choices of \(p_{0}(\cdot),p^{\prime}_{0,t}(\cdot),q(\cdot)\), as summarized in Table 1. In particular, for EMGNN-PD and EMGNN-PH, we set \(\lambda_{0}\) for the delta function to be any \(\lambda\) such that \(\{\lambda\in\Lambda-\widehat{\Lambda}\}\). Consequently, \(p_{0}(\cdot)\) will be \(0\) with probability \(1\) w.r.t \(q(\cdot)\) on \(\widehat{\Lambda}\). Hence, for these variants, there is no "bad" \(\lambda\) such that the corresponding iteration increases \(L_{\mathbf{\mathrm{X}}}(\cdot,\cdot)\)(Refer to Remark 4).

#### 4.1.2 Results

Results are shown in Table 2. Similar to recent findings (Lv et al., 2021), GCN and GAT are observed to perform competitively against models designed for heterogeneous graphs such as GTN under appropriate settings. Meanwhile, EMGNN-PT consistently outperforms other variants in our framework, in both micro and macro F1 scores. In particular, the superior performance of EMGNN-PT, EMGNN-PO, EMGNN-PD and EMGNN-PH compared to GCN indicates the effectiveness of learning with multiple graphs. Moreover, EMGNN-PT outperforming EMGNN-PD, along with EMGNN-PO frequently outperforming EMGNN-PH, indicates that increasing the loss for a "bad" \(\lambda\) is beneficial as it penalizes deviations from desirable graphs.

EMGNN-PT also often surpasses baseline models with attention mechanisms, namely GAT, Simple-HGN, SeHGNN, and GTN, despite not incorporating any attention mechanisms. This could be attributed to the construction of multiple graphs, which may form instances whose information is similar to what is achieved with semantic attention. In addition, the model may be able to extract additional useful interactions in other instances of graphs which can enhance its performance.

### Chemical Datasets

Conventional molecular graph representations mirror a molecule's Lewis structure, with atoms as nodes and chemical bonds as edges. This representation falls short in capturing variations in molecular properties resulting from different three-dimensional (3D) spatial arrangements when molecules share the same topology, as seen in cases like cis-trans isomers. Moreover, molecules inherently possess uncertainty due to their quantum mechanical properties, particularly concerning electron orbitals. Hence, using a distribution of graphs for learning in such cases is a sensible choice.

The process of generating different molecular graphs begins with the acquisition of coarse 3D coordinates of the atoms in a molecule using RDKit4. Following that, the interatomic Euclidean distances between all atoms within the molecule are calculated. A parameter, \(\lambda\) is then introduced to define a threshold range, \([0,\lambda]\), for determining node connections and generating multiple graph instances. The notion of employing thresholding based on the interatomic distance between
nodes in molecular graphs has been previously documented in works such as Shui and Karypis (2020) and Shen et al. (2023). The former introduced a cut-off distance hyperparameter to construct heterogeneous molecular graphs. Meanwhile, our approach aligns more closely with the latter, where the Vietoris-Rips complex and thresholding are used to form a series of \(G_{\lambda}\) graphs. However, in Shen et al. (2023), they utilized five non-overlapping, manually adjusted intervals for thresholding and adopted a computationally intensive multi-channel configuration to learn from the five generated graphs. There are also works such as Thomas et al. (2018), where molecules are treated as 3D point clouds and a radius is set to specify interacting vertices.

The graph construction process may involve adding new edges, connecting distant nodes, or removing existing edges. Ideally, the model should prioritize "useful" graph realizations and assign a low probability to less beneficial ones, effectively discarding them.

#### 4.2.1 Baselines and Datasets

MoleculeNet5 is a popular benchmark for molecular machine learning, encompassing multiple datasets to be tested on a diverse of molecular properties. For our evaluation, we specifically selected the datasets FreeSolv, ESOL, and Lipophilicity, all of which are designed for graph regression tasks. Further elaboration on the chosen datasets can be found in the Appendix.

Footnote 5: [https://moleculeenet.org/datasets-1](https://moleculeenet.org/datasets-1)

We compared our approach against standard models for molecular properties prediction that do not incorporate transfer learning from a larger dataset such as Zinc156. The selected baseline models for this comparison included Weave (Kearnes et al., 2016), MPNN (Gilmer et al., 2017), AttentiveFP (Xiong et al., 2020), GIN (Xu et al., 2019), as well as the standard GCN and GAT models. For EMGNN, a GNN model that generalizes GCN with a degree-1 convolutional filter \(P_{\lambda}(S_{\lambda})\) (refer to Section 2.2) is utilized as the backbone of our model. The sample space \(\Lambda\) spans the range \([1,10]\)A and \(\widehat{\Lambda}\) is the discretized space with 0.05 increments.

Footnote 6: a database of purchasable drug-like compounds; [https://zinc.docking.org/tranches/home/](https://zinc.docking.org/tranches/home/)

#### 4.2.2 Results

In Table 3, the average test root mean square error (rmse) over ten runs with standard deviation is reported for the graph regression task, where the molecular properties of molecular graphs are to be predicted. The result shown is for the case of \(q(\cdot)=p_{t}(\cdot)\). We observe that EMGNN frequently performed better than the baselines. This may be due to EMGNN's training process, which exposes it to diverse graph realizations, allowing it to capture non-covalent interactions that are critical for characterizing the physical properties of molecules. In contrast, the baselines employ the conventional molecular graph representation. We note that our framework does not explicitly incorporate bond angles but it does expose the model to graphs with a broad range of connectivities. This exposure indirectly integrates geometric information, as the latent variable constructs graphs with bond lengths falling within specific ranges. This provides our model with additional 2D information regarding interatomic distances, which may offer insights into the underlying 3D structure.

### Model Analysis

#### 4.3.1 Distribution Learned

For the heterogeneous node classification task, we examined the learned empirical distributions, depicted in Fig. 2. Across all datasets, we notice that the empirical probability of \(\lambda\) falls within the range of approximately \([0.4,0.6]\) is relatively high. This observation suggests a possible explanation for the decent performance of GCN on a single graph with uniform edge weights.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline  & \multicolumn{2}{c}{IMDB} & \multicolumn{2}{c}{ACM} & \multicolumn{2}{c}{DBLP} \\  & Micro-F1 & Macro-F1 & Micro-F1 & Macro-F1 & Micro-F1 & Macro-F1 \\ \hline GCN & 61.91 \(\pm\) 0.67 & 60.91 \(\pm\) 0.57 & 91.92 \(\pm\) 0.40 & 92.00 \(\pm\) 0.41 & 94.60 \(\pm\) 0.31 & 93.88 \(\pm\) 0.36 \\ GAT & 63.54 \(\pm\) 1.10 & 61.87 \(\pm\) 0.95 & 92.61 \(\pm\) 0.36 & 92.68 \(\pm\) 0.36 & 94.48 \(\pm\) 0.22 & 93.74 \(\pm\) 0.27 \\ GTN & 60.58 \(\pm\) 2.10 & 59.12 \(\pm\) 1.58 & 92.12 \(\pm\) 0.62 & 92.23 \(\pm\) 0.60 & 94.17 \(\pm\) 0.26 & 93.59 \(\pm\) 0.40 \\ Simple-HGN & 58.91 \(\pm\) 1.06 & 58.30 \(\pm\) 0.34 & **92.73 \(\pm\) 0.21** & 92.56 \(\pm\) 0.42 & 94.48 \(\pm\) 0.38 & 93.69 \(\pm\) 0.32 \\ SeHGNN & 62.13 \(\pm\) 2.38 & 60.62 \(\pm\) 1.95 & 92.45 \(\pm\) 0.17 & 92.51 \(\pm\) 0.16 & 94.86 \(\pm\) 0.14 & 94.14 \(\pm\) 0.19 \\ \hline EMGNN-PT & **64.78 \(\pm\) 1.24** & **63.36 \(\pm\) 0.80** & 92.70 \(\pm\) 0.26 & **92.78 \(\pm\) 0.26** & **95.06 \(\pm\) 0.39** & **94.41 \(\pm\) 0.45** \\ EMGNN-PO & 63.35 \(\pm\) 0.79 & 62.25 \(\pm\) 0.59 & 92.35 \(\pm\) 0.38 & 92.45 \(\pm\) 0.38 & 94.95 \(\pm\) 0.24 & 94.28 \(\pm\) 0.28 \\ EMGNN-PD & 62.49 \(\pm\) 0.87 & 61.55 \(\pm\) 0.71 & 92.31 \(\pm\) 0.43 & 92.41 \(\pm\) 0.42 & 94.89 \(\pm\) 0.17 & 94.15 \(\pm\) 0.23 \\ EMGNN-PH & 62.01 \(\pm\) 0.55 & 61.15 \(\pm\) 0.46 & 92.18 \(\pm\) 0.52 & 92.29 \(\pm\) 0.52 & 95.02 \(\pm\) 0.19 & 94.34 \(\pm\) 0.20 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Heterogeneous node classification task. Results averaged over ten runs. The best performance is boldfaced and the second-best performance is underlined.
For IMDB, (9) is of the form \(\lambda A_{MD}+(1-\lambda)A_{MA}\). We observed that \(\lambda=1\) has a relatively lower probability compared to \(\lambda=0\). When \(\lambda=1\), it implies that edges in \(A_{MA}\) are all removed. This indicates that the MA relation is more crucial than the MD relation. This observation might be linked to the significant difference in edge density in the edge-type specific adjacency matrices where \(A_{MA}\) has three times more non-zero entries than \(A_{MA}\). A similar trend was observed in the ACM dataset, where the disparity in edge density is also substantial. In ACM's case, (9) takes the form \(\lambda A_{PA}+(1-\lambda)A_{PS}\). The high probability of \(\lambda=1\) indicates that the PS relation is more significant than PA for the task. Consequently, we can infer that edge types with relatively sparse connections have a limited impact on the task.

#### 4.3.2 The Dominant Component of (7)

In Section 3.2, we use \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}\big{(}L_{\mathbf{X}}(\lambda, \boldsymbol{\theta})\big{)}\) to approximate \(\log C(\boldsymbol{\theta})\). To justify this, we provide numerical evidence that the dominant component of (7) is \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}\big{(}L_{\mathbf{X}}(\lambda, \boldsymbol{\theta})\big{)}\). This assertion is supported by assessing the ratio

\[\varrho=\frac{\rho}{-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}\big{(}L_{ \mathbf{X}}(\lambda,\boldsymbol{\theta})\big{)}},\]

where \(\rho=\frac{\text{var}\left(\exp(-\eta^{(t)}L_{\mathbf{X}}(\lambda,\boldsymbol{ \theta}))\right)}{2\left(\mathbb{E}_{\lambda\sim p_{0}}\exp(-\eta^{(t)}L_{ \mathbf{X}}(\lambda,\boldsymbol{\theta}))\right)^{2}}\).

We plot the \(\varrho\) values for the heterogeneous graph datasets on EMGNN-PT over multiple \(t\) iterations in Fig. 3. We found that \(\varrho\) consistently exhibits small absolute values, supporting the postulation that \(-\eta^{(t)}\mathbb{E}_{\lambda\sim p_{0}}\big{(}L_{\mathbf{X}}(\lambda, \boldsymbol{\theta})\big{)}\) is the main component in (7).

## 5 Conclusion

In this paper, we explored employing a distribution of parametrized graphs for training a GNN in an EM framework. Through a probabilistic framework, we handle the uncertainty in graph structures stemming from various sources. Our approach enables the model to handle multiple graphs where the prediction loss is utilized to estimate the likelihood of the graphs. The model's performance is enhanced as we provide it with a wider array of graphs, which it can then sift through to acquire more valuable information or remove noise.

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline  & \multicolumn{3}{c}{Random split} & \multicolumn{3}{c}{Scaffold split} \\ Datasets & FreeSolv & ESOL & Lipophilicity & FreeSolv & ESOL & Lipophilicity \\ \hline GCN & \(1.157\pm 0.215\) & \(0.652\pm 0.073\) & \(0.707\pm 0.030\) & \(2.618\pm 0.298\) & \(0.876\pm 0.037\) & \(0.760\pm 0.009\) \\ GAT & \(1.873\pm 0.522\) & \(0.837\pm 0.101\) & \(0.704\pm 0.058\) & \(2.942\pm 0.591\) & \(0.907\pm 0.034\) & \(0.777\pm 0.037\) \\ Weave & \(1.497\pm 0.251\) & \(0.798\pm 0.088\) & \(0.789\pm 0.059\) & \(3.129\pm 0.203\) & \(1.104\pm 0.063\) & \(0.844\pm 0.031\) \\ MPNN & \(1.388\pm 0.404\) & \(0.703\pm 0.075\) & \(0.640\pm 0.025\) & \(2.975\pm 0.775\) & \(1.117\pm 0.058\) & \(\mathbf{0.735\pm 0.019}\) \\ AttentiveFP & \(1.275\pm 0.289\) & \(0.673\pm 0.085\) & \(0.719\pm 0.042\) & \(2.698\pm 0.297\) & \(0.855\pm 0.029\) & \(0.762\pm 0.022\) \\ GIN & \(1.678\pm 0.494\) & \(0.792\pm 0.097\) & \(0.716\pm 0.073\) & \(2.957\pm 0.696\) & \(0.990\pm 0.057\) & \(0.770\pm 0.021\) \\ \hline EMGNN & \(\mathbf{0.936\pm 0.162}\) & \(\mathbf{0.606\pm 0.041}\) & \(\mathbf{0.639\pm 0.028}\) & \(\mathbf{2.189\pm 0.128}\) & \(\mathbf{0.834\pm 0.027}\) & \(0.743\pm 0.013\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Graph regression task on molecular datasets. Average test rmse reported, the lower the better.

Figure 3: Plot of \(\varrho\) across \(t\) EM iterations

Figure 2: Empirical distribution of \(p_{t}(\cdot)\) that is obtained from the final E-step.