Unleash Graph Neural Networks from Heavy Tuning

We denote an undirected graph with $N$ nodes as $\mathcal{G}\{\mathcal{V},\mathcal{E},\mathbf{A}\}$, where ${\mathcal{V}}$ and ${\mathcal{E}}$ are sets of nodes and edges and $\mathbf{A}\in{\mathbb{R}}^{N\times N}$ is the adjacency matrix containing information of relationships. $\mathbf{A}$ can be either weighted or unweighted, normalized or unnormalized. The graph signals are stored in a matrix ${\mathbf{X}}\in{\mathbb{R}}^{N\times D_{f}}$, where $D_{f}$ is the number of features. Traditional Multi-Layer Perceptrons (MLPs) integrate data information solely from the feature direction for representation learning, overlooking potential connections between nodes. This limitation is addressed by Graph Neural Networks (GNNs), a deep-learning architecture that incorporates additional relationship information through graph convolutions. Typically, there are two types of GNNs: spatial GNNs that aggregate neighboring nodes with message passing hamilton2017inductive ; kipf2017semisupervised ; gasteiger2019predict and spectral GNNs developed from the graph spectral theory defferrard2016convolutional ; lin2023magnetic ; zou2023simple . With both types, the most important component of their architecture is the convolutional layer, which usually consists of graph convolution, linear transformation, and non-linear activation functions. For example, the very classic convolutional layer in GCN kipf2017semisupervised is formulated as $\sigma(\mathbf{A}{\mathbf{X}}{\mathbf{W}})$, where graph convolution involves node aggregation with graph adjacency $\mathbf{A}$, linear transformation is conducted with linear operator ${\mathbf{W}}$, and $\sigma$ is the activation function. The GNN training process often involves learning parameters in the linear transformation and for some architectures such as ChebNetdefferrard2016convolutional and the graph convolution.

Diffusion models generate a step-by-step denoising process that recovers data in the target distribution from random white noises. A typical diffusion model adopts a pair of forward-backward Markov chains, where the forward chain perturbs the observed samples eventually to white noises, and then the backward chain learns how to remove the noises and recover the original data. Assume that ${\mathbf{w}}_{0}\sim q_{{\mathbf{w}}}({\mathbf{w}}_{0})$ is the original data (e.g., vectorized network parameters) from the target distribution $q_{{\mathbf{w}}}({\mathbf{w}}_{0})$. Then, the forward-backward chains are formulated as follows.

To produce high-performing GNN parameters, we need a set of high-quality parameter samples such that the diffusion module can generate from their underlying population distribution. It is quite intuitive that good model performance is associated with appropriate selection of hyperparameters. So, we start with a coarse search with a relevantly small search space to determine a suitable configuration (details on the search space are discussed in Section 5). This includes the selection of hyperparameters in the optimizer such as learning rate, model architecture such as hidden size, and other model-specific factors such as the teleport probability $\alpha$ in APPNP gasteiger2019predict . Since parameter initialization may also influence the training outcome, we provide 10 random initializations for training with each configuration. A good configuration from coarse search is defined as the one leading to the best validation accuracy. After the coarse search, we save model checkpoints from the training process with the selected configuration as parameter samples. We discard the parameters from some start-up epochs where convergence has not been achieved. Parameters are reorganized and vectorized before they are further processed by PAE.

GAE is designed to encode graph information from graph signals ${\mathbf{X}}$ and structure $\mathbf{A}$, which will be employed as a condition for GNN generation. In this paper, we mainly focus on the node classification task to evaluate our method, so the GAE structure is oriented by the specific task, which means GAE encoder aims to produce a graph representation that works well on node classification. Straightforwardly, a GNN architecture should be considered. The general principle of GNN propagation is to update the node feature by aggregating its neighboring information. Thus, GNN graph representation has similar features for connected nodes, which will eventually lead to similar label prediction. This is well-suited for homophilic graphs, where connected nodes are normally from the same classes. However, for heterophilic graphs, where nodes with different labels are prone to be linked, such representation leads to an even worse classification outcome than MLP han2024from ; zheng2022graph . Hence, enlightened by some previous works shao2023unifying ; thorpe2022grand ; zhu2020beyond , we introduce the following GAE encoder, which is capable of handling both homophilic and heterophilic graphs:

which is composed of the concatenation of a two-layer GCN, a one-layer GCN, and an MLP, followed by a linear transformation ${\mathbf{W}}_{4}\in{\mathbb{R}}^{D_{\mathcal{G}}\times D_{p}}$, where $D_{\mathcal{G}}$ and $D_{p}$ are the dimensions of the concatenation and the latent parameter representation, respectively. $\text{G-Decoder}({\mathbf{h}})={\mathbf{h}}{\mathbf{W}}_{5}\in{\mathbb{R}}^{D_{c}}$ is a single layer MLP with $D_{c}$ being the number of node classes. Thus graph information is fully encapsulated in the output of the encoder. Finally, GAE is trained with cross entropy loss on node labels. The graph condition is then given by mean-pooling the node features

where ${\mathbf{h}}(i)$ is the i-th row of graph representation. The purpose of applying mean pooling on nodes is to construct an overall graph condition regardless of the graph size, which will later be combined with latent representation ${\mathbf{z}}_{0}\in{\mathbb{R}}^{D_{p}}$ as input of diffusion denoising network. We highlight that our method can be naturally extended to other graph tasks by altering the design and loss function of GAE.

G-LDM is developed from LDM with additional graph conditions from GAE to provide consolidated guidance on GNN generation. The generative target of G-LDM is latent GNN parameters ${\mathbf{z}}_{0}$. The forward pass of G-LDM is the same as traditional DDPM, while the backward pass uses a graph conditional backward kernel for data recovering from white noises. The corresponding G-LDM loss function is given by

For both PAE and denoising network ${\mathbf{\epsilon}}_{\mathbf{\theta}}$, we adopt the Conv1d-based architecture in wang2024neural . We include the graph condition by using the sum of ${\mathbf{z}}_{t}$, ${\mathbf{c}}_{\mathcal{G}}$, and $\text{Embedding}(t)$ as input to ${\mathbf{\epsilon}}_{\mathbf{\theta}}$. Our method is flexible with PAE and ${\mathbf{\epsilon}}_{\mathbf{\theta}}$ architectures. So, we choose an architecture that has been empirically validated by previous research. We observe in experiments that the most important factor in high-quality generation is to reduce parameter dimension sufficiently while ensuring reconstruction tightness. This can be done by setting the appropriate kernel size and stride for Conv1d layers in PAE. We discuss more details about PAE and ${\mathbf{\epsilon}}_{\mathbf{\theta}}$ architecture and how to ensure high-quality generation in Appendix B and Appendix F, respectively.

The last step before the final prediction is to obtain parameters and reconstruct the target GNN. G-LDM generates samples of latent parameters from white noises with the learned denoising network. Next, the learned PAE decoder is applied to reconstruct parameters back to the original dimension. These parameters are returned to the target GNN thus we have a group of GNNs with generated parameters. Lastly, we determine the optimal GNN depending on validation performance, which will be used for future predictions.

In Table 1, we present the results of GNN-Diff on generating GNNs for node classification. We compare the test accuracy of three models selected respectively from grid search, coarse search, and GNN-Diff generation based on validation accuracy. For simplicity, we will refer them as grid model, coarse model, and GNN-Diff model. Overall, GNN-Diff models outperform grid and coarse models on most GNN architectures and datasets. On Cora and Actor, GNN-Diff models successfully achieve the best test accuracy for all target GNNs. GNN-Diff models improve the test accuracy from 0.20% to 1.30% on Cora and from 0.06% to 1.77% on Actor. On Citeseer and Chameleon, GNN-Diff models also produce good test results for most GNN architectures. Notably, GNN-Diff boosts the test performance for H2GCN on Citeseer by 1.60% and for APPNP on Chameleon by 1.32%. However, GNN-Diff fails to outperform grid results for GIN and GAT on Citeseer and for GIN and SAGE on Chameleon. We have tracked the validation and test accuracy of GIN during the training for grid search, and noticed that good accuracy in grid search could be considered as outliers in overall GIN performance. They appeared abruptly in the training history and had much higher value than neighboring epochs. Thus, it is hard to capture the corresponding parameters based on coarse search samples. We suppose it is reasonable to treat this as special situation. For GAT on Citeseer and SAGE on Chameleon, heavy parameterization of their architectures lays the burden on PAE reconstruction. We suggest that a more sophisticated architecture of PAE may be helpful, though the current architecture still works well on these two GNNs for other datasets.

The grid search accuracy we report here for some GNNs may be lower than other relevant works. This is majorly caused by different experiment setting. The grid model evaluated by test set is the model with highest validation accuracy among all others from the entire comprehensive grid search process. Some relevant works run GNNs with each individual hyperparameter configuration, and report the best test accuracy among configurations. Validation data is only used for model selection during the training process with each configuration.

We further investigate the experiment results by visualization. With coarse search, we only visualize the saved checkpoints for GNN-Diff training to be directly compared with models generated by GNN-Diff. We use GCN2 on Cora as an example and try to answer three questions:

In Figure 2 (a), we present a scatter plot between validation and test accuracy for partial grid (green), coarse (orange), and GNN-Diff (red) models that achieve at least 70% prediction accuracy on both validation and test set. In our experiments, the models from three methods with the highest validation accuracy are selected for evaluation on test data. We use three black circles to indicate the position of the final models. Clearly, the grid model with the highest validation accuracy has the lowest test accuracy compared to the other two methods. This is a sign of overfitting on validation data, which eventually results in poor generalizability.

Although grid search generally guarantees to find a promising hyperparameter configuration, its reliance on validation may be counterproductive sometimes. Having a smaller search space or using other search algorithms, such as random search, may solve the problem. But how do we know the appropriate search space size? In addition, nearly 80% of our experiment results show that grid search can find better generalized model than coarse search, hence reducing search space or random search may not be good options in most cases. This shows the superiority of our method, which only requires samples from a good configuration and then further optimizes the model performance by exploring the underlying population distribution. This avoids finding a configuration that is optimal for validation but not for unseen data.

As we know, the parameters of coarse models are the training data of GNN-Diff. So, the parameters generated by GNN-Diff are expected to follow a similar distribution as coarse model parameters. In Figure 2 (b), we use a bubble isometric mapping plot to visualize coarse (orange) and GNN-Diff (red) parameter distributions via dimensionality reduction. Isometric mapping (Isomap) seeks to preserve the geodesic distances between all pairs of data points in the 2D embedding. It is particularly helpful for visualizing high-dimensional parameters by projecting them onto a low-dimensional space while maintaining the underlying geometry. Additionally, we plot each model point as bubbles with bubble size representing the corresponding test accuracy.

We observe that GNN-Diff parameters follow the pattern of the coarse distribution (a firework-like expanding spread with three directions). Besides, the generation further explores the potential population distribution by extending in two of the three directions. Notably, the generated models close to the center of coarse distribution are associated with lower accuracy, whereas the extended ones have better prediction outcomes. This may validate our conjecture that GNN-Diff has the ability to approximate population parameter distribution with a good hyperparameter configuration and generate better-performing samples from it.

In Figure 2 (c), we provide the Kernel Density Estimation (KDE) plot of test accuracy distributions for grid models (green), coarse models (orange), and GNN-Diff models (red). The test accuracy of grid models shows a broader distribution with two peaks, one around 20% and the other one around 80%. Grid search, while thorough, spends much unnecessary time and computation on training with sub-optimal or even unsatisfactory configurations. Coarse models present a test accuracy distribution with a peak around 80% but with some variability. Two factors have contributed to this distribution. Firstly, coarse models are collected from a good hyperparameter configuration of coarse search. Secondly, we discard some models from start-up epochs that are generally associated with low test accuracy. The test accuracy distribution of GNN-Diff models has a sharp peak at a comparably higher level, implying high-performing outcomes with less variability than the other two methods. This means GNN-Diff consistently finds models with comparably better generalizability, outperforming both grid search and coarse search in terms of finding high-quality parameters.

Here we focus on whether incorporating graph condition is useful for GNN generation. Previous works have discussed using data and tasks as conditions, but only for the purpose of tailoring their methods for future unseen data or similar tasks nava2022meta ; soro2024diffusion ; zhang2024metadiff . The role of intrinsic data characteristics in network generation is still unexploited. Nevertheless, with graph neural networks, the graph structure is very essential in the network architecture and has a consequential impact on graph signal processing. So, only considering data (graph signals) as the generative condition is insufficient. Previous study nava2022meta includes predictive tasks as the condition by encoding their language descriptions, which requires additional effort to form appropriate descriptions and embedding. Alternatively, we propose a more straightforward way via specific task supervision on condition construction. Therefore, we design the GAE discussed in Subsection 4.2, which incorporates both graph data and structural information in the generative condition and is trained with the task loss function. We conduct two experiments: (1) a comparison experiment between GNN-Diff and p-diff wang2024neural to show how graph condition improves GNN generation quality; (2) an ablation study on GNN-Diff generative condition to validate the current GAE architecture on homophilic and heterophilic graphs. In the first experiment, we chose p-diff because it has the same PAE and ${\mathbf{\epsilon}}_{\mathbf{\theta}}$ as GNN-Diff and no generative condition. Hence, we do not need to consider different performances caused by different generative model architectures. Both methods are trained with the same hyperparameters and experiment settings. For both experiments, we conduct sampling and prediction 5 times and report the best results for all methods/conditions to offset the influence of randomness in diffusion generation and ensure fair comparison.

Comparison between GNN-Diff and p-diff

In Figure 3, we compare the generated parameter and test accuracy distributions with and without graph conditions. This experiment is conducted with GCN2 on Cora. Based on the implementation of PyTorch Geometric, GCN2 has 4 learnable components, including biases and weights in two graph convolutional layers. For visualization purposes, we pick one parameter from each component and plot its distributions with sample models from coarse search (orange), GNN-Diff generation (blue, top row), and p-diff generation (blue, bottom row). It is clear that GNN-Diff generation captures the sample distribution and extends the sample distribution range to explore potential population distribution. By contrast, p-diff generation produces much centered distribution but fails to cover the sample distribution. This indicates that graph condition provides a useful guidance on GNN generation, and helps the generative model to better learn the target distribution. We also provide histograms of test accuracy distribution for both methods. p-diff models obtain around 20% to 80% accuracy when being applied to test data. GNN-Diff models, on the other hand, have accuracy more centered around 80%. Also, the worst-performing model generated by GNN-Diff still outperforms many p-diff models. Although high accuracy can be achieved by p-diff, the generation quality is unstable without the guidance of graph conditions.

Ablation Study on Generative Condition

In Figure 4, we present the test accuracy of GNN generation with different conditions. GAE refers to the architecture discussed in Subsection 4.2, which is adopted by GNN-Diff. According to Equation (2), GAE encoder combines GCN and MLP to handle both homophilic and heterophilic graphs. In this experiment, we apply ablation on graph conditions. GCN1 and GCN2 represent graph conditions with only graph convolution-based representation. MLP2 and MLP1 represent conditions with only data information. Lastly, we include no condition (p-diff) as the baseline. Different target GNNs are considered to show the generalizability of the conclusion. On the homophilic graph, Cora, GAE leads to the best test result, while conditions without graph structural information (MLP2 and MLP1) and no condition fail to outperform grid search results. On heterophilic graphs, Actor, MLP2, MLP1, and no condition produced better results than graph convolution-based conditions (GCN2 and GCN1). However, in GAE, the combination of graph convolution and MLP still generates the best prediction outcome. Therefore, we conclude that the current GAE architecture consistently leads to high-performing parameters regardless of the graph type.

Last but not least, we would like to discuss the time efficiency of GNN-Diff. In Figure 5, we use a bar chart to show the time of grid search (green), coarse search (orange), and GNN-Diff training (red) for different target GNNs and datasets. We stack bars of coarse search and GNN-Diff training as the time cost of the entire GNN-Diff process. Inference time is not shown as it is negligible (usually less than 10 seconds). The time is estimated based on the average search iteration time and average epoch time for three GNN-Diff modules. Implementation time may vary due to factors such as data pre-processing and validation. According to the plots, it is very clear that the GNN-Diff method is more efficient than a comprehensive grid search. The advantage is less obvious for simple GNN architectures such as GCN1, for which only a small search space is needed for grid search. However, it is very valuable for architectures with large hyperparameter space such as GAT and GNNs with cumbersome graph convolution such as APPNP.