Graph Structure Learning with Bi-level Optimization

This task aims to learn a classifier $f(\bm{A},\bm{X};\bm{\theta})$ from a set of labeled nodes to forecast the remaining nodes labels, where $\bm{\theta}$ denotes model parameters. Assuming there are $M$ labels, we index them from 1 to $M$ without loss of generality. Formally, $\bm{Y}=[y_{1},y_{2},\cdots,y_{N}]^{\top}\in\mathbb{R}^{N}$ are the labels of nodes, where $y_{i}\in\mathbb{R}$ is the label of node $i$. The target is achieved by optimizing the model parameter $\bm{\theta}$ over the labeled nodes, which is formulated as:

where $l(\cdot)$ is a classification loss, and $\lambda$ is a hyperparameter to adjust the strength of parameter regularization.

Towards the first purpose, the core idea is to decompose the graph structure information into connectedness (the edges in the adjacency matrix) and the strength of connection (the latent variable). In general, there are two ways to model the connection strength regarding whether relying on the inductive bias of translation invariant or not. On one hand, attentive mechanisms or noise detection models are translation invariant, which decode the connection strength of each edge from its local features. However, with the consideration of the local structure heterogeneity issue in most real-world graphs Xu et al. (2018), it is risky to rely on the translation invariant bias. On the other hand, probabilistic mechanisms separately model the connection strength for each edge, where each edge corresponds to a specific distribution. However, it is non-trivial to set a proper prior in practice. According to these pros and cons, we summarize two considerations for extending the graph convolution: (1) edge-specific modeling; and (2) optimization without prior.

Towards this end, we model the connection strength as a parameter matrix $\bm{Z}$ with the same size as $\bm{A}$. Formally, the generic structure extractor (GSE) is abstracted as:

where $\sigma(\cdot)$ is a non-negative activation function since the value of strength is always positive¹¹1In this work, we use the min[max[0,x],1] function to restrict the value within [0, 1].. COM and AGG are the combination and aggregation functions respectively.

Noteworthy, different from GNNs, GSE decouples the graph structure from GNNs and treats it as a learnable objective. Besides, GSE is a generic extractor, which can be instantiated over most existing graph convolutions (cf. Appendix A). In this way, as long as learning the connection strength is set properly, GSEBO can downweight the neighbors with inter-class connection during representation propagation, reducing the impact of inter-class with bi-level optimization.

To achieve the second purpose of learning graph structure from global information, it is essential to carefully design a proper training algorithm to optimize the connection strength matrix $\bm{Z}$. Assume that we construct GSEBO with $K$ layers, which is denoted as $f(\bm{A},\bm{X};\bm{\theta})$ with parameters $\bm{\theta}=\{\bm{Z};\bm{W}^{(k)}|k\in[1,K]\}$. We have three main considerations for designing the training algorithm: (1) Connection strength is a relative value, which changes across different views. As shown in Figure 1(b), the connection between node 3 and node 5 is weak from the local view, i.e., the edge is inter-class and should be assigned a low weight. However, this edge is essential for the classification of node 9, 10, 11, which deserves a high weight. Therefore, the optimization objective of $\bm{Z}$ should be the overall performance of node classification. (2) $\bm{Z}$ and $\bm{\mathcal{W}}=\{\bm{W}^{(k)}|k\in[1,K]\}$ play different roles, but are closely related. The role of $\bm{Z}$ is close to $\bm{A}$, which restricts the model space for the mapping from node feature to label, and the role of $\bm{\mathcal{W}}$ includes the global mapping information for classification, which would relieve the cons of local optimization. That is, an update on $\bm{Z}$ will adjust $\bm{\mathcal{W}}$ and its optimization procedure. Therefore, the optimization of $\bm{Z}$ and $\bm{\mathcal{W}}$ are at two different but dependent levels. (3) Directly minimizing the objective function of Eq.  (1) to obtain the parameters $\bm{Z}$ and $\bm{\mathcal{W}}$ is not able to achieve the desired purpose of learning the structure for the reason of over-fitting, which is shown on Figure 3.

Towards this end, we propose a bilevel optimization Domke (2012); Maclaurin et al. (2015); Franceschi et al. (2017) algorithm with inner and outer optimizations to learn $\bm{\mathcal{W}}$ and $\bm{Z}$, respectively. Figure 2 shows the overall procedure of the algorithm, where the inner and outer optimization steps are iteratively executed until stop condition.

Inner optimization for common parameters. This step is similar to the normal training of GNN for $\bm{\mathcal{W}}$ optimization. In particular, we update $\bm{\mathcal{W}}$ with a gradient descent based optimizer (e.g., Adam Kingma and Ba (2015)) over the training nodes by minimizing Eq. (3) (i.e., set $l_{inner}=l(f(\bm{A},\bm{X}),y;\bm{\theta})+\lambda\|\bm{\theta}\|$). When calculating the gradient of $\bm{\mathcal{W}}$, we treat the connection strength matrix $\bm{Z}$ as constant.

Assuming that $\tau$ times of gradient descent, an approximate solution $\bm{\mathcal{W}}_{1},\cdots,\bm{\mathcal{W}}_{\tau}$ of the inner optimization problem are obtained. For $t=1$ to $\tau$, the updated $\bm{\mathcal{W}}$ is calculated as $\bm{\mathcal{W}}_{t}=\bm{\mathcal{W}}_{t-1}-\eta_{i}\nabla{L}(\bm{\mathcal{W}}_{t-1},\bm{Z})$, where $\eta_{i}$ is the inner learning rate, and the process of updating $\bm{\mathcal{W}}$ is shown on line 3-5 in Algorithm 1.

Outer optimization for graph structure. Similarly, we treat the sequence of $\bm{\mathcal{W}}$ as constant to optimize the graph structure $\bm{Z}$. Ideally, the objective should be the overall classification loss of all nodes in graph. Formally,

Apparently, the calculation of the ideal objective is intractable. Similar as the global parameters optimization, we can approximate the ideal objective as the empirical risk over training nodes. However, it can easily suffer from the over-fitting issue, which is shown on Figure 3. We believe the empirical risk over validation nodes is a better approximation of the ideal objective since it reflects to what extent the parameter generalizes well. In this light, we set $l_{outer}$ as the classification loss and optimize $\bm{Z}$ over the validation set $\mathcal{V}_{val}$ by minimizing Eq. (2).

The details of updating $\bm{Z}$ is shown on line 8-11 in Algorithm 1 and the mathematical formulation is shown on Appendix B. With the parameter $\bm{Z}$ updated by the graph structure optimization, we reoptimize the $\bm{\mathcal{W}}$ with $\bm{Z}$ by global parameters optimization, and repeat this process until the early stopping is met. To summarize, Algorithm 1 shows the training procedure of GSEBO. Besides, we explain why the bi-level optimization could defeat the over fitting problem when training on train and validation dataset on Appendix C.

We select four widely used node real-world classification benchmark datasets with graph of citation networks (Cora and Citeseer Kipf and Welling (2017)), social networks (Terrorist) Zhao et al. (2006), and air traffic (Air-USA) Wu et al. (2019). We summarize the statistics of the datasets and describe in detail on Appendix F. We adopt the same data split of Cora and Citeseer as Kipf and Welling (2017), and a split of training, validation, testing with a ratio of 10:20:70 on other datasets Zhao et al. (2021b).

We apply GSEBO on 4 representative GNN architectures: GCN Kipf and Welling (2017), GAT Veličković et al. (2018), GraphSAGE Hamilton et al. (2017) and JK-Net Xu et al. (2018). For each GNN, we compare GSEBO with the vanilla version, and three variants with state-of-the-art connection modeling methods: AdaEdge Chen et al. (2020a), DropEdge Rong et al. (2020) and GAUG Zhao et al. (2021b). In addition, we compared the GSEBO with GSL methods: BGCN Zhang et al. (2019), VGCN Elinas et al. (2020), PTDNet Luo et al. (2021), and advanced attention mechanism: MAGNA Wang et al. (2021). Note that the base model of GSEBO, BGCN, and PTDNet are GCN.

In the experiments, the latent dimension of all the methods is set to the 16. The parameters for all baseline methods are initialized as the corresponding papers, and are carefully tuned to achieve optimal performances. The learning rate of inner and outer optimization are set to 0.01. The hyperparameter $\lambda$ is set to $5\times{10^{-4}}$, and we search for the inner learning depth $\tau$ over the range [5,10,15,20,25]. Due to bilevel optimization can not guarantee the convergence, we set the patience of early stopping to 20 to terminate the optimization of GSEBO. To prevent overfitting, the dropout ratio is set to 0.5 for all the methods. The network architectures of all the mehtods are configured to be the same as described in the original papers. Our experiments are conducted with Tensorflow running on GPU machines (NVIDIA 2080Ti). For all the compared methods, we report the average accuracy on the test set over 10 runs.

GSEBO outperforms the baselines in most cases. Considering the average performance over four datasets, the improvement of GSEBO over the baselines is in the range of [3.0%, 12.5%], which validates the effectiveness of the proposed method. In particular, (1) Probabilistic mechanisms. The performance gain of BGCN and VGCN over the vanilla GCN are limited, which might because of the unsatisfied assumption of the prior distribution. This result shows the rationality of relaxing the prior assumption and modeling the connection strength with parameters directly. (2) Connection modeling. DropEdge, AdaEdge, GAUG, and PTDNet achieve better performance than vanilla GNN, which modify the structure of the graph from different perspectives such the smoothness Chen et al. (2020a) and robustness Rong et al. (2020). These results reflect the benefit of connection modeling. However, there is still a clear gap between these methods and GSEBO, which is attributed to the global objective for learning the edge strength. (3) Attention mechanism. The performance of MAGNA is inferior, which is consistent with the analyzing in Knyazev et al. (2019).

On the four GNN architectures, GSEBO achieves better performance than the vanilla version in all cases. In particular, GSEBO achieves an average improvement across datasets of 4.2%, 4.4%, 4.0%, and 5.74% over the vanilla GCN, GAT, GraphSAGE, and JK-Net, respectively. These results justify the effectiveness of structure learning of the proposed GSEBO. Across the four architectures, GSEBO achieves the largest improvement over JK-Net. We postulate the reason is that the jump connection in JK-Net makes it aggregates more hops of neighbors than the other GNNs. As the hops increase, the homophily ratio of neighbors will decrease, i.e., more neighbors are in classes different from the target node. Therefore, optimizing the connection strength (i.e., $\bm{Z}$) is more beneficial on JK-Net.

From the perspective of dataset, GSEBO consistently performs better than the vanilla version on the four datasets. Specifically, the average improvement over the four across classic GNN architectures achieved by GSEBO are 1.9%, 3.1%, 5.3% and 8.0% on the four datasets, respectively. Considering that the four datasets come from different scenarios, these results are evidence for the potential of GSEBO to be widely applied in practice. Note that the trend of performance improvement is similar to the density of graph where Air-USA is the most dense graph with the largest performance improvement. As the number of neighbors increases, the percentage of neighbors essential for the classification of the target node will decrease. This result can reflect the rationality of optimizing the connection strength according to the overall classification objective.

Moreover, the training loss and test accuracy of GNN methods and the proposed GSEBO are shown in Figure 4. We observe that the loss of GSEBO tend to be stable after 30 epochs and is smaller than GCN, GAT, and GraphSAGE, showing the empirical convergence of GSEBO. This is because simultaneously optimizing the parameters of $\bm{\mathcal{W}}$ and $\bm{Z}$ would make the loss function more easily approach to the labels. However, the loss of JK-Net is smaller than GSEBO, we attribute it to that JK-Net is easier to trap in over-fitting.