Mean-field theory of graph neural networks in graph partitioning
Tatsuro Kawamoto, Masashi Tsubaki, Tomoyuki Obuchi
Introduction
Deep neural networks have been subject to significant attention concerning many tasks in machine learning, and a plethora of models and algorithms have been proposed in recent years. The application of the neural network approach to problems on graphs is no exception and is being actively studied, with applications including social networks and chemical compounds . A neural network model on graphs is termed a graph neural network (GNN) . While excellent performances of GNNs have been reported in the literature, many of these results rely on experimental studies, and seem to be based on the blind belief that the nonlinear nature of GNNs leads to such strong performances. However, when a deep neural network outperforms other methods, the factors that are really essential should be clarified: Is this thanks to the learning of model parameters, e.g., through the backpropagation , or rather the architecture of the model itself? Is the choice of the architecture predominantly crucial, or would even a simple choice perform sufficiently well? Moreover, does the GNN generically outperform other methods?
To obtain a better understanding of these questions, not only is empirical knowledge based on benchmark tests required, but also theoretical insights. To this end, we develop a mean-field theory of GNN, focusing on a problem of graph partitioning. The problem concerns a GNN with random model parameters, i.e., an untrained GNN. If the architecture of the GNN itself is essential, then the performance of the untrained GNN should already be effective. On the other hand, if the fine-tuning of the model parameters via learning is crucial, then the result for the untrained GNN is again useful to observe the extent to which the performance is improved.
For a given graph , where is the set of vertices and is the set of (undirected) edges, the graph partitioning problem involves assigning one out of group labels to each vertex. Throughout this paper, we restrict ourselves to the case of two groups (). The problem setting for graph partitioning is relatively simple compared with other GNN applications. Thus, it is suitable as a baseline for more complicated problems. There are two types of graph partitioning problem: One is to find the best partition for a given graph under a certain objective function. The other is to assume that a graph is generated by a statistical model, and infer the planted (i.e., preassigned) group labels of the generative model. Herein, we consider the latter problem.
Before moving on to the mean-field theory, we first clarify the algorithmic relationship between GNN and other methods of graph partitioning.
Graph neural network and its relationship to other methods
The goal of this paper is to examine the graph partitioning performance using a minimal GNN architecture. To this end, we consider a GNN with the following feedforward dynamics. Each vertex is characterized by a -dimensional feature vector whose elements are (, ), and the state matrix \mbox{\boldmathX}=[x_{i\mu}] obeys
The GNN architecture described above can be thought of as a special case of the following more general form:
The belief propagation (BP) algorithm (also called the message passing algorithm) in Bayesian inference also falls under the framework of Eq. (2). While the domain of the state consists of the vertices () for GNNs, this algorithm deals with the directed edges , where is obtained by putting directions to every undirected edge. In this case, the state represents the logarithm of the marginal probability that vertex belongs to the group with the missing information of vertex at the th iteration. With the choice of matrix and activation functions shown in Table 1 (EM+BP), Eq. (2) becomes exactly the update equation of the BP algorithmPrecisely speaking, this is the BP algorithm in which the stochastic block model (SBM) is assumed as the generative model. The SBM is explained below. . The matrix \mbox{\boldmathM}=\mbox{\boldmathB}=[B_{j\to k,i\to j}] is the so-called non-backtracking matrix , and the softmax function represents the normalization of the state .
The BP algorithm requires the model parameters \mbox{\boldmathW}^{t} and \mbox{\boldmathb}^{t} as inputs. For example, when the expectation-maximization (EM) algorithm is considered, the BP algorithm comprises half (the E-step) of the algorithm. The parameter learning of the model is conducted in the other half (the M-step), which can be performed analytically using the current result of the BP algorithm. Here, \mbox{\boldmathW}^{t} and \mbox{\boldmathb}^{t} are the estimates of the so-called density matrix (or affinity matrix) and the external field resulting from messages from non-edges , respectively, and are common for every . Therefore, the differences between the EM algorithm and GNN are summarized as follows. While there is an analogy between the inference procedures, in the EM algorithm the parameter learning of the model is conducted analytically, at the expense of the restrictions of the assumed statistical model. On the other hand, in GNNs the learning is conducted numerically in a data-driven manner , for example by backpropagation. While we will shed light on the detailed correspondence in the case of graph partitioning here, the relationship between GNN and BP is also mentioned in .
Mean-field theory of the detectability limit
We analyze the performance of an untrained GNN on the stochastic block model (SBM). This is a random graph model with a planted group structure, and is commonly employed as the generative model of an inference-based graph clustering algorithm . The SBM is defined as follows. We let , and each of the vertices has a preassigned group label , i.e., . We define as the set of vertices in a group , , and represents the planted group assignment of vertex . For each pair of vertices and , an edge is generated with probability , which is an element of the density matrix. Throughout this paper, we assume that , so that the resulting graph has a constant average degree, or in other words the graph is sparse. We denote the average degree by . Therefore, the adjacency matrix \mbox{\boldmathA}=[A_{ij}] of the SBM is generated with probability
2 Dynamical mean-field theory
In an untrained GNN, each element of the matrix \mbox{\boldmathW}^{t} is randomly determined according to the Gaussian distribution at each , i.e., . We assume that the feature dimension is sufficiently large, but . Let us consider a state that represents the average state within a group, i.e., . The probability distribution that \mbox{\boldmath\mathsf{x}}^{t}=[\mathsf{x}_{\sigma\mu}^{t}] is expressed as
where denotes the average over the graph , the random linear transform \mbox{\boldmathW}^{t}, and the state \mbox{\boldmathX}^{t} of the previous layer. Using the Fourier representation, the normalization condition of Eq. (5) is expressed as
where is an auxiliary variable that is conjugate to , and \mathcal{D}\hat{\mbox{\boldmath\mathsf{x}}}^{t+1}\mathcal{D}\mbox{\boldmath\mathsf{x}}^{t+1}\equiv\prod_{\sigma\mu}(\gamma_{\sigma}d\hat{\mathsf{x}}^{t+1}_{\sigma\mu}d\mathsf{x}^{t+1}_{\sigma\mu}/2\pi i).
After taking the average of the symmetric SBM over as well as the average over \mbox{\boldmathW}^{t} in the stationary limit with respect to , the following self-consistent equation is obtained with respect to the covariance matrix \mbox{\boldmathC}=[C_{\sigma\sigma^{\prime}}] of \mbox{\boldmath\mathsf{x}}=\mbox{\boldmath\mathsf{x}}^{t}:
where . The detailed derivation can be found in the supplemental material. The reader may notice that the above expression resembles the recursive equations in . However, it should be noted that Eq. (7) is not obtained as an exact closed equation. The derivation relies mainly on the assumption that the macroscopic random variable \mbox{\boldmath\mathsf{x}}^{t} dominates the behavior of the state \mbox{\boldmathX}^{t}. It is numerically confirmed that this assumption appears plausible. This type of analysis is called dynamical mean-field theory (or the Martin-Siggia-Rose formalism) .
When the correlation within a group is equal to the correlation between groups (), the GNN is deemed to have reached the detectability limit. Beyond the detectability limit, Eq. (7) is no longer a two-component equation, but is reduced to an equation with respect to the variance of one indistinguishable group.
Normalized mutual information error function
Here, we instead propose the use of the normalized mutual information (NMI) as an error function for the readout classifier. The NMI is a comparison measure of two group assignments, which naturally eliminates the permutation degrees of freedom. Let \mbox{\boldmath\sigma}=\{\sigma_{i}=\sigma|i\in V_{\sigma}\} be the labels of the planted group assignments, and \hat{\mbox{\boldmath\sigma}}=\{\sigma_{i}=\hat{\sigma}|i\in V_{\hat{\sigma}}\} be the labels of the estimated group assignments. First, the (unnormalized) mutual information is defined as
where the joint probability is the fraction of vertices that belong to the group in the planted assignment and the group in the estimated assignment. Furthermore, and are the marginals of , and we let H(\mbox{\boldmath\sigma}) and H(\hat{\mbox{\boldmath\sigma}}) be the corresponding entropies. The NMI is defined by
This measure is permutation invariant, because the NMI counts the label co-occurrence patterns for each vertex in and \hat{\mbox{\boldmath\sigma}}.
Experiments
First, the consistency between our mean-field theory and a specific implementation of an untrained GNN is examined. The performance of the untrained GNN is evaluated by drawing phase diagrams. For the SBMs with various values for the average degree and the strength of group structure , the overlap, i.e., the fraction of vertices that coincide with their planted labels, is calculated. Afterward, it is investigated whether a significant improvement is achieved through the parameter learning of the model. Note that because even a completely random clustering can correctly infer half of the labels on average, the minimum of the overlap is .For this reason, the overlap is sometimes standardized such that the minimum equals zero. As mentioned above, we adopt as the specific choice of activation function.
We evaluate the performance of the untrained GNN in which the resulting state is read out using the k-means (more precisely k-means++ ) classifier. In this case, no parameter learning takes place. We set the dimension of the feature space to and the number of layers to , and each result represents the average over samples.
Figure 3a presents the corresponding phase diagram. The overlap is indicated by colors, and the solid line represents the detectability limit estimated by Eq. (7). The dashed line represents the mean-field estimate of the detectability limit of the spectral methodAgain, there are several choices for the matrix to be adopted in the spectral method. However, the Laplacians and modularity matrix, for example, have the same detectability limit when the graph is regular or the average degree is sufficiently large. , and the shaded area represents the region above which the inference is information-theoretically impossible. It is known that the detectability limit of the BP algorithm coincides with this information-theoretic limit so long as the model parameters are correctly learned. For the present model, it is also known that the EM algorithm can indeed learn these parameters . Note that it is natural that a Bayesian method will outperform others as long as a consistent model is used, whereas it may perform poorly if the assumed model is not consistent.
It can be observed that our mean-field estimate exhibits a good agreement with the numerical experiment. For a closer view, the overlaps of multiple graph sizes with are presented in Fig. 3b. For , the estimate is , and this appears to coincide with the point at which the overlap is almost . It should be noted that the performance can vary depending on the implementation details. For example, while the k-means method is performed to \mbox{\boldmathX}^{T} in the present experiment, it can instead be performed to \phi\left(\mbox{\boldmathX}^{T}\right). An experiment concerning such a case is presented in the supplemental material.
2 GNN with backpropagation and a trained classifier
Now, we consider a trained GNN, and compare its performance with the untrained one. A set of SBM instances is provided as the training set. This consists of SBM instances with , where an average degree and strength of the group structure are adopted. For the validation (development) set, graph instances of the same SBMs are provided. Finally, the SBMs with various values of and the average degree are provided as the test set.
We evaluated the performance of a GNN trained by backpropagation. We implemented the GNN using Chainer (version 3.2.0) . As in the previous section, the dimension of the feature space is set to , and various numbers of layers are examined. For the error function of the readout classifier, we adopted the NMI error function described in Section 4. The model parameters are optimized using the default setting of the Adam optimizer in Chainer. Although we examined various optimization procedures for fine-tuning, the improvement was hardly observable.
We also employ residual networks (ResNets) and batch normalization (BN) . These are also adopted in . The ResNet imposes skip (or shortcut) connections on a deep network, i.e., , where is the number of layers skipped, and is set as . The BN layer, which standardizes the distribution of the state \mbox{\boldmathX}^{t}, is placed at each intermediate layer . Finally, we note that the parameters of deep GNNs (e.g., ) cannot be learned correctly without using the ResNet and BN techniques.
The results using the GNN trained as above are illustrated in Fig. 4. First, it can be observed from Fig. 4a that a deep structure is important for a better accuracy. For sufficiently deep networks, the overlaps obtained by the trained GNN are clearly better than those of the untrained counterpart (see Fig. 3b). On the other hand, the region of where the overlap suddenly deteriorates still coincides with our mean-field estimate for the untrained GNN. This implies that in the limit , the detectability limit is not significantly improved by training. To demonstrate the finite-size effect in the result of Fig. 4a, the overlaps of various graph sizes are plotted in Fig. 4b. The variation of overlaps becomes steeper around as the graph size is increased, implying the presence of detectability phase transition around the value of predicted by our mean-field estimate.
The untrained and trained GNNs exhibit a clear difference in overlap when \mbox{\boldmathX}^{T} is employed as the readout classifier. However, it should be noted that the untrained GNN where \phi(\mbox{\boldmathX}^{T}) is adopted as the readout classifier exhibits a performance close to that of the trained GNN. The reader should also bear in mind that the computational cost required for training is not negligible.
Discussion
In a minimal GNN model, the adjacency matrix is employed for the connections between intermediate layers. In fact, there have been many attempts to adopt a more complex architecture rather than . Furthermore, other types of applications of deep neural networks to graph partitioning or related problems have been described . The number of GNN varieties can be arbitrarily extended by modifying the architecture and learning algorithm. Again, it is important to clarify which elements are essential for the performance.
The present study offers a baseline answer to this question. Our mean-field theory and numerical experiment using the k-means readout classifier clarify that an untrained GNN with a simple architecture already performs well. It is worth noting that our mean-field theory yields an accurate estimate of the detectability limit in a compact form. The learning of the model parameters by backpropagation does contribute to an improved accuracy, although this appears to be quantitatively insignificant. Importantly, the detectability limit appears to remain (almost) the same.
The minimal GNN that we considered in this paper is not the state of the art for the inference of the symmetric SBM. However, as described in Section 2, an advantage of the GNN is its flexibility, in that the model can be learned in a data-driven manner. For a more complicated example, such as the graphs of chemical compounds in which each vertex has attributes, the GNN is expected to generically outperforms other approaches. In such a case, the performance may be significantly improved thanks to backpropagation. This would constitute an interesting direction for future work. In addition, the adequacy of the NMI error function that we introduced for the readout classifier should be examined in detail.
The authors are grateful to Ryo Karakida for helpful comments. This work was supported by the New Energy and Industrial Technology Development Organization (NEDO) (T.K. and M. T.) and JSPS KAKENHI No. 18K11463 (T. O.).
References
Appendix A Derivation of the self-consistent equation
In this section, the detailed derivation of the self-consistent equation of the covariance matrix is derived. Here, we recast our starting-point equation:
We first take the average of over \mbox{\boldmathW}^{t}. The Gaussian integral with respect to \mbox{\boldmathW}^{t} yields
where we introduce the following quantities:
Note that as we will see below, \mbox{\boldmathu}^{t+1}, \mbox{\boldmathv}^{t}, \mbox{\boldmath\psi}^{t}, and their conjugates are related to \mbox{\boldmathX}^{t}, and thus the average over \mbox{\boldmathX}^{t} is taken outside of their integral.
We next take the average over a random graph. In Eq. (17), only the final term in the exponent is relevant to . We denote this term as . We also let . Because the graph is generated from the SBM, we have that
At the second line, we used the fact that . Then, at the third line we used . Finally, at the last line we used the symmetry of the undirected graph, .
Note here that the degrees of freedom with respect to the feature dimension are factored out, and thus the dependence on can be omitted. Hereafter, the same notation will be employed for the variables without the -dependence. We also introduce the notation . The factor inside of the average over \mbox{\boldmathX}^{t} in Eq. (17) can be written as follows:
As in the main text, we have defined .
where F_{\sigma\sigma^{\prime}}\left(\mbox{\boldmathX}^{t}\right) is defined as
A.2 Stochastic process with a correlated noise
Here, we compare Eq. (24) with a Markovian discrete-time stochastic process , in which each element is correlated via a random noise, i.e., , and for any . The corresponding normalization condition reads
Analogously to the case of the GNN, we have defined \mathcal{D}\hat{\mbox{\boldmathy}}^{t+1}\mathcal{D}\mbox{\boldmathy}^{t+1}\equiv\prod_{\sigma}\gamma_{\sigma}d\hat{y}^{t+1}_{\sigma}dy^{t+1}_{\sigma}/2\pi i.
A.3 Self-consistent equation
Finally, we compare Eqs. (24) and (26). However, note that these are not of exactly the same form, because the average over \mbox{\boldmathX}^{t} is taken outside of the exponential in Eq. (24). Two approximations are made in order to derive the self-consistent equation, and the assumptions that justify these approximations are discussed afterward.
holds in the stationary limit, then the group-wise state \mbox{\boldmathx}^{t} can be regarded as a Gaussian variable whose correlation matrix obeys
This equation is still not closed, because the right-hand side of Eq. (28) depends on the statistic of \mbox{\boldmathX}^{t}, rather than \mbox{\boldmath\mathsf{x}}^{t}. However, because the vertices within the group are statistically equivalent, are expected to obey the same distribution with mean , which itself is a random variable. If holds, then the right-hand side of Eq. (28) can be evaluated as the average with respect to the group-wise variable \mbox{\boldmath\mathsf{x}}^{t}. Then, within this regime we arrive at the following self-consistent equation with respect to the covariance matrix \mbox{\boldmathC}=[C_{\sigma\sigma^{\prime}}]:
Let us consider the first approximation that we adopted in Eq. (27). In the terminology of physics, this is the replacement of a free energy with an internal energy, or the neglect of the entropic contribution. It is difficult to evaluate this residual in general. However, note that this becomes closer to equality as every approaches the same value. Therefore, this implies that the self-consistent equation is more accurate as we approach the detectability limit, and yields an adequate estimate of the critical value.
Let us next consider the second approximation we adopted in Eq. (28). Although the law of large numbers with respect to (not ) ensures that has a certain value characterized by the group, this may be different from . In fact, the relation between these is in general an inequality (Jensen’s inequality) when the activation function is a convex function. The (exact) equality holds only when is constant or the function is linear within the considered domain.
The second approximation can be justified in the following cases. The first case is when the fluctuation of is negligible compared to the magnitude of . Note that this is the same assumption as we made in the first approximation. To see this precisely, let us express as for . We can formally write the probability distribution of in a hierarchical fashion as follows:
where P_{\mbox{\boldmath\sigma}}(\{\mathsf{x}_{\sigma}\}) is the probability distribution with respect to . Thus, the expectation \left\langle f(\mbox{\boldmathX})\right\rangle_{\mbox{\boldmathX}} can be expressed as
where , which can be a nontrivial function. However, whenever the contributions from the average with respect to are negligible, Eq. (31) implies that the expectation in Eq. (28) can be evaluated using only the group-wise variables . Another case is when the activation function is almost linear within the domain over which fluctuates. For example, in the case that , the present approximation does not deteriorate the accuracy even when . When either of these assumption holds, the equality of Jensen’s inequality is approximately satisfied, and our derivation of the self-consistent equation is justified.
Appendix B K-means classification using ϕ(𝑿)italic-ϕ𝑿\phi(\mbox{\boldmath$X$})
Instead of \mbox{\boldmathX}^{T}, \phi(\mbox{\boldmathX}^{T}) can be adopted to perform the k-means classification after the feedforward process. Again, we employ as the nonlinear activation function. The results of an untrained GNN and a trained GNN under the same experimental settings as in the main text are illustrated in Fig. 5 and Fig. 6, respectively. In Fig. 5a, the reader should note that the range of the color gradient is different from that in the phase diagram in the main text. For the untrained GNN, the obtained overlaps are clearly better than that using \mbox{\boldmathX}^{T}. It can be understood that the error is reduced because the nonlinear function drives each element of the state \mbox{\boldmathX}^{T} to either or , making the classification using the k-means method easier and more accurate. On the other hand, for the trained GNN, differences between the overlaps using \mbox{\boldmathX}^{T} and \phi(\mbox{\boldmathX}^{T}) are hardly observable.
Particularly for the case of an untrained GNN in which \phi(\mbox{\boldmathX}^{T}) is adopted for the readout classifier, the overlap gradually changes around the estimated detectability limit. This may be as result of the strong finite-size effect. Again, note that our estimate of the detectability limit is for the case that .