Learning Invariant Representations of Graph Neural Networks via Cluster Generalization

Donglin Xia,Xiao Wang,Nian Liu,Chuan Shi
2024-03-06
Abstract:Graph neural networks (GNNs) have become increasingly popular in modeling graph-structured data due to their ability to learn node representations by aggregating local structure information. However, it is widely acknowledged that the test graph structure may differ from the training graph structure, resulting in a structure shift. In this paper, we experimentally find that the performance of GNNs drops significantly when the structure shift happens, suggesting that the learned models may be biased towards specific structure patterns. To address this challenge, we propose the Cluster Information Transfer (CIT) mechanism (Code available at
Machine Learning
What problem does this paper attempt to address?
### What problem does this paper attempt to solve? This paper aims to solve the problem that the performance of Graph Neural Networks (GNNs) drops significantly when facing structure shift. Specifically, when the structure of the test graph is different from that of the training graph, the performance of GNNs will be severely affected, indicating that existing GNNs may be biased towards specific structural patterns. To address this challenge, the authors propose the **Cluster Information Transfer (CIT)** mechanism to learn invariant representations, thereby improving the generalization ability of GNNs on different and unknown test graphs. #### Main problem description 1. **Impact of structure shift**: - GNNs learn node representations by aggregating local structural information, so their performance highly depends on the local structure of the graph. - In the real world, graph structures often change, such as user relationships in social networks, citation relationships in citation networks, etc. - When the structure of the test graph changes, the performance of GNNs will drop significantly, resulting in unstable performance of the model on new data. 2. **Limitations of existing methods**: - Existing methods usually assume that the data distributions of training and test graphs are the same, but in practical applications, this assumption often does not hold. - Some methods attempt to alleviate the structure shift problem by generating new graph structures or sampling unbiased test data, but these methods are difficult to implement in complex real - world scenarios. 3. **Proposed new method**: - The authors propose the CIT mechanism, which generates diverse node representations by transferring cluster information of nodes in the embedding space, thereby learning invariant representations that are robust to structural changes. - The CIT mechanism can simulate different structural environments without changing the graph structure and enhance the generalization ability of GNNs. #### Formula summary - **Node representation update formula**: \[ Z^{(l)}=\sigma(\tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}Z^{(l - 1)}W^{(l - 1)}_{\text{GNN}}) \] where \( Z^{(l)} \) is the node representation of the \( l \) - th layer, \(\tilde{A}=A + I\), \(\tilde{D}\) is the degree matrix of \(\tilde{A}\), \(\sigma\) is a nonlinear activation function, and \( W^{(l - 1)}_{\text{GNN}} \) is the trainable parameter of GNN. - **Cluster information transfer formula**: \[ Z'^{(l)}_i=\frac{\sigma(H^c_j)(Z^{(l)}_i - H^c_k)}{\sigma(H^c_k)}+H^c_j \] where \( H^c_k \) is the representation of the \( k \) - th cluster, \( Z^{(l)}_i \) is the representation of node \( i \), and \( j \) is a randomly selected target cluster. Through these methods, the CIT mechanism can effectively alleviate the impact of structure shift on the performance of GNNs and improve the stability of the model on different test graphs.