Zero-Shot Generalization of GNNs over Distinct Attribute Domains

Authors: Yangyi Shen, Jincheng Zhou, Beatrice Bevilacqua, Joshua Robinson, Charilaos Kanatsoulis, Jure Leskovec, Bruno Ribeiro
Primary area: learning on graphs and other geometries & topologies
Keywords: GNN, zero-shot, graph foundation models

Abstract:
Inductive GNNs are able to generalize across graphs with the same set of node attributes. However, zero-shot generalization across attributed graphs with disparate node attribute domains remains a fundamental challenge in graph machine learning. Existing methods are unable to effectively make use of node attributes when transferring to unseen attribute domains, frequently performing no better than models that ignore attributes entirely. This limitation stems from the fact that models trained on one set of attributes (e.g., biographical data in social networks) fail to capture relational dependencies that extend to new attributes in unseen test graphs (e.g., TV and movies preferences). Here, we introduce STAGE, a method that learns representations of _statistical dependencies_ between attributes rather than the attribute values themselves, which can then be applied to completely unseen test-time attributes, generalizing by identifying analogous dependencies between features in test. STAGE leverages the theoretical link between maximal invariants and measures of statistical dependencies, enabling it to provably generalize to unseen feature domains for a family of domain shifts. Our empirical results show that when STAGE is pretrained on multiple graph datasets with unrelated feature spaces (distinct feature types and dimensions) and evaluated zero-shot on graphs with yet new feature types and dimensions, it achieves a relative improvement in Hits@1 between 40% to 103% for link prediction, and an 10% improvement in node classification against state-of-the-art baselines.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730534999854

thread_type:
official_review

<thread_object>
creator: reviewer_iea1

<content>
summary:
The paper introduces STAGE, a novel approach designed to enable zero-shot generalization for Graph Neural Networks (GNNs) across graphs with varying node attribute domains. STAGE aims to learn representation of statistical dependencies between attributes rather than their absolute values. This allows the model to transfer knowledge to unseen domains by leveraging analogous dependencies. Through experiments on multiple datasets, the paper demonstrates STAGE's superior performance in link prediction and node classification tasks, especially in terms of zero-shot cross-domain generalization.

strengths:
- Generalization from a node attributes view and the two stages for processing graph representation are interesting.
- The paper provides a theoretical analysis, linking STAGE with maximal invariants and statistical dependency measures, which provides theoretical support for the model's generalization capabilities.
- The paper shows STAGE's robustness when facing different attribute domains, which is a very important characteristic in the varied real-world data

weaknesses:
- Although the paper presents some quantitative results and shows good performance in different link prediction and node classification domains, it lacks some qualitative analysis. For example, it could demonstrate how the model learns that "income level is positively correlated to phone price" from the training set and then discovers that "height is positively correlated with clothing size" in a new domain, thus generalizing to the new domain.
- STAGE is capable of capturing and leveraging feature dependencies in graph data, rather than relying on specific attribute values. The article can illustrate which feature dependencies are effective on the test set after pre-training the model.

questions:
How well does this model handle larger graph datasets.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730552608654

thread_type:
official_review

<thread_object>
creator: reviewer_u75f

<content>
summary:
The paper presents STAGE, a method that enables zero-shot generalization of graph neural networks (GNNs) across graphs with different attribute domains. STAGE constructs STAGE-edge-graphs to capture statistical dependencies between attributes instead of absolute values, facilitating transferability to unseen domains. The method shows substantial improvement in zero-shot tasks like link prediction and node classification on graphs with entirely new feature spaces.

strengths:
1. STAGE's strategy to use statistical dependencies rather than raw attribute values to enhance zero-shot generalization in GNNs is novel for graph machine learning.
2. The theoretical basis for STAGE, connecting maximal invariants and statistical dependencies, is well-articulated and provides a sound foundation for the empirical results.
3. STAGE is shown to be adaptable across domains of varied attribute types and dimensions, a crucial quality for real-world applicability.

weaknesses:
1. STAGE's two-stage process involving STAGE-edge-graphs, conditional probability matrices, and subsequent embeddings may be challenging to implement or optimize in practice. Details on computational overhead compared to baselines would enhance clarity.
2. While STAGE is effective for pairwise dependencies, it is unclear how it handles more complex dependencies in highly interconnected graphs.
3. The success of STAGE appears dependent on the architecture and expressivity of the underlying GNNs (M1 and M2). Sensitivity analysis on different GNN backbones might clarify robustness across architectures.
4. The evaluation focuses on e-commerce and social network datasets; examining STAGE’s generalizability on domains like biomedical or geospatial networks would strengthen claims of universality.

questions:
1. How does STAGE handle attribute domains with highly heterogeneous data types, such as unstructured or mixed media data?
2. Could the authors elaborate on the computational cost associated with STAGE compared to baselines, especially in large-scale graphs?
3. Does STAGE's reliance on GNN backbones like NBFNet affect its generalizability? Could alternative GNN architectures be equally effective?
4. Has STAGE been evaluated in terms of the interpretability of the learned dependencies? If so, what methods were used?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730656928449

thread_type:
official_review

<thread_object>
creator: reviewer_jvnb

<content>
summary:
This paper introduces STAGE, a method designed for zero-shot generalization across attributed graphs with distinct attribute domains. STAGE constructs what it calls STAGE-edge-graphs for each edge in a graph, embedding statistical dependencies between attributes at each node pair. The model achieves significant performance gains in zero-shot settings for tasks like link prediction and node classification on various datasets.

strengths:
1. This paper achieves SOTA results by embedding statistical dependencies rather than raw features.
2. The STAGE is a domain-agnostic framework, which can generalize across disparate attribute spaces.

weaknesses:
1. The STAGE-edge-graph is a fully connected weighted graph, so I am concerned about the complexity.
2. Edge-based embeddings may limit its ability to capture high-order interactions in graphs.
3. The motivation in the introduction is not presented well. The authors didn't analyze why their proposed method can address the limitations they mentioned before, so it's hard to understand the intrinsic research thinking.
4. The experiments are a little weak. For example, I believe the 4.2 and 4.3 belong to the same type of experiment, they didn't analyze the complexity and the ablation study, and they didn't include the limited research papers they mentioned in the introduction into baselines,  which weakens the convincing.
5. I noticed this paper was submitted to ICML workshop so there is authors' information leakage. Both papers present STAGE for zero-shot generalization of GNNs across different attribute domains, and this paper just extends some real-world testing datasets. 
6. The authors didn't release their code although this is not compulsory, which may limit their reproducibility.

questions:
Please see the weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730767878987

thread_type:
official_review

<thread_object>
creator: reviewer_nstq

<content>
summary:
This paper studies how to use a pre-trained graph model in any new domain with unseen attributes, enhancing the zero-shot generalization. The authors propose a new model STAGE by learning the representations of statistical dependency between attributes, instead of the attribute values themselves. They also conduct experiments to validate the performance of STAGE across several benchmark datasets.

strengths:
1.	The paper is well written and easy-to-follow.
2.	Extensive results validate the effectiveness of the proposed model.

weaknesses:
1.	More discussions on the variants on LLM should be made.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1732215936543

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response and New Experiment Results (2/2)

comment:
> **STAGE’s scalability to large graphs**

As per JvnB and U75f’s request, we offered a complexity analysis and a comparison of training time between STAGE and the baselines. We therefore clarify that,

(1) Complexity analysis:

Let  $p$ be the number of features, $d$ the dimension of internal node and edge embeddings, $|E|$ the number of edges, and $|V|$ the number of nodes in the input graph. For the link prediction task, STAGE consists of three steps:

1. **Fully Connected STAGE-Edge-Graph Construction:** This step requires  $O(|E| p^2)$ operations because each fully connected STAGE-edge-graph has 2p nodes.
2. **Inference (GINE Layers):** We use 2 shared layers of GINE [1] for all STAGE-edge-graphs. A single layer on one fully-connected STAGE-edge-graph has complexity $O(p d + p^2 d) = O(p^2 d)$ since we have 2p d-dimensional nodes and $(2p)^2$ d-dimensional edges in each graph. Aggregating edge embeddings across all graphs takes $O(|E| p^2 d)$.
3. **Inference (NBFNet):** We use NBFNet to perform message passing on the original graph, which requires $O(|E|d + |V|d^2)$ for one forward pass [2].

**Total Complexity:** The overall forward pass has a complexity of $O(|E| p^2 d + |E|d + |V|d^2)$. 

(2) STAGE does not pose a heavy computation overhead:

To assess the computational overhead of STAGE, we measured the average wall time per training epoch on the E-Commerce Stores dataset. This dataset, with 17,463 nodes, 183,618 edges, and up to 16 node attributes, represents the largest in our experiments and thus presents the most demanding scenario for STAGE-edge-graph construction. Utilizing an 80GB A100 GPU, we obtain the following results:

| Models          	| Wall Time per Training Epoch on E-Commerce (seconds) | Zero-shot Hits@1 Performance on H&M |
|---------------------|----------------------------------------|----------------------------------------|
| NBFNet-raw      	| 318.65 | 0.0005 $\pm$ 0.0004 |
| NBFNet-gaussian 	| 322.13 | 0.0925 $\pm$ 0.0708 |
| NBFNet-structural   | 322.31 |  0.2231 $\pm$ 0.0060 |
| NBFNet-llm      	| 316.55 | 0.2302 $\pm$ 0.0015 |
| NBFNet-normalized   | 316.87 | 0.2286 $\pm$ 0.0010 |
| **NBFNet-STAGE (Ours)** | **341.36** | **0.4666 $\pm$ 0.0020** |

NBFNet-STAGE is only 7.83% slower than the fastest baseline (NBFNet-llm), a reasonable tradeoff for its performance gains. The additional time is due to computing STAGE-edge-graph embeddings during each forward pass, while building the STAGE-edge-graphs is a one-time preprocessing step. In practice, the additional factor in the complexity has never prevented us from running in the datasets we considered.

We believe we have carefully addressed all reviewers’ feedback, providing compelling evidence for the effectiveness and generalizability of our STAGE approach. We thank the reviewer for their thoughtful comments, and we believe these revisions significantly strengthen our paper. Thank you! We hope the reviewers will find these changes satisfactory to reconsider their score. We are more than happy to address any further questions.


[1] Huang et al. "Prodigy: Enabling in-context learning over graphs." NeurIPS 2024.

[2] Liu, et al. “One for all: Towards training one graph model for all classification tasks.” ICLR 2024.

[3] Chen, et al. “LLaGA: Large language and graph assistant.” ICML, 2024.

[4] Zhao et al. "Graphany: A foundation model for node classification on any graph." ArXiv 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732216075513

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response and New Experiment Results (1/2)

comment:
We thank all reviewers for appreciating our work and providing valuable feedback. We are glad to see that our work is recognized as addressing a “novel” and “interesting” problem with important “real-world applicability” (**U75f**, **iea1**), having a “well-articulated” and “sound theoretical support” (**U75f**, **iea1**), with “extensive” and “SOTA empirical results” (**nStQ**, **JvnB**).

In the general response, we address the reviewers' common questions and present new experimental results.

> **New baselines for node classification using GraphAny**

As per reviewer JvnB's request, we clarify that:

1. Our submission included PRODIGY [1] as a baseline, presented as NBFNet-llm (link prediction, Table 1) and GINE-llm (node classification, Table 2). 
2. The request for OneForAll [2] and LLaGA [3] were excluded due to incompatibility with our setup. OneForAll is limited to edge type classification (queries like (s, ?, t)) and does not support tail node prediction (queries like (s, r, ?)) or ranking metrics (e.g., Hits@1, MRR). LLaGA's reliance on LLMs is constrained by context window sizes, making it impractical for large-scale datasets like E-Commerce Store. 
3. Adding an extra baseline: GraphAny [4].
Since GraphAny is specifically designed for the node classification task, we report its results in the zero-shot node classification experiment shown in Table 2 of the updated paper.

The following is the updated Table 2 incorporating GraphAny’s results:

| Models        	| Accuracy ($\uparrow$) | 
|-------------------|---------------------------------|
| GINE-structural | 0.564 $\pm$ 0.0466 |
| GINE-gaussian | 0.588 $\pm$ 0.0250 |
| GINE-normalized | 0.541 $\pm$ 0.0148 |
| GINE-age | 0.582 $\pm$ 0.0657 |
| GINE-llm [3] | 0.550 $\pm$ 0.0368 |
| GraphAny [4]      	| 0.591 $\pm$ 0.0083          	|
| **GINE-STAGE (Ours)** | **0.652 $\pm$ 0.0042**          	|

We observe that GraphAny is the best-performing model among all the baseline models. Nevertheless, our model GINE-STAGE significantly outperforms GraphAny with a 10.3% relative performance improvement. This demonstrates the effectiveness and superior performance of STAGE compared to the baseline node classification foundation model. 

 > **STAGE is effective across different GNN backbones**

As per reviewer JvnB and U75f’s request, we did another ablation studying the sensitivity of STAGE using different GNN backbones. We clarify that:
1.  STAGE works well with NBFNet and GINE as the backbone model when performing on link prediction and node classification tasks.
2. Below we present an extra experiment using Graph Convolutional Network (GCN)  as the backbone model on node classification.

To further demonstrate STAGE's adaptability, we replaced the GINE backbone that was used for our results shown in Table 2 with a GCN adapted to handle multi-dimensional edge attributes. Specifically, an MLP was integrated into GCN’s message-passing to process edge attributes. 

The results, shown in updated Table 6, confirm that GCN-STAGE significantly outperforms all baselines, achieving a 7.33% improvement in average zero-shot test accuracy and a much smaller standard deviation, showcasing its robustness and stability.

| Model            | Accuracy ($\uparrow$) |
|------------------|----------------------------|
| GCN-structural   | 0.547 $\pm$ 0.0658              |
| GCN-gaussian     | 0.567 $\pm$ 0.0382              |
| GCN-normalized   | 0.570 $\pm$ 0.0315              |
| GCN-llm          | 0.526 $\pm$ 0.0300              |
| **GCN-STAGE (Ours)** | **0.593** $\pm$ **0.0046**          |
 
These experimental results show STAGE's superior performance compared to all baseline methods, irrespective of the chosen backbone GNN architecture (GINE or GCN). This consistent improvement across different architectures underscores the versatility and broad applicability of STAGE, reinforcing its position as a robust and effective framework for graph representation learning.
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:** This paper introduces STAGE, a method designed to enable zero-shot generalization for Graph Neural Networks (GNNs) across attributed graphs with distinct attribute domains. By focusing on statistical dependencies between attributes rather than their raw values, the method aims to generalize effectively to unseen domains. The proposed framework includes STAGE-edge-graphs to model pairwise attribute dependencies and is evaluated on link prediction and node classification tasks, showcasing significant improvements over state-of-the-art baselines.


**Decision:** The paper introduces a novel framework for zero-shot generalization of GNNs using attribute dependency modeling but fails to address critical aspects that would ensure its broader applicability. Specifically, several reviewer raised concerns that due to the algorithm's complexity of $\mathcal{O}(p^2)$, its applicability to a broader range of scenarios is significantly constrained. In fact, I noticed that the datasets used in the paper all have a relatively small number of features. However, in many real-world applications, the dimensionality of node attributes is often much higher [1, 2, 3]. Similarly, another related issue is that this method relies heavily on a clear and explicit definition or encoding of the features (e.g., I don't think it can deal with word embeddings of Abstract in a citation network), which further limits the applicability of the method. For example, in the social network datasets, 54 out of 58 features were removed as they are "difficult to encode either because they are random texts input by the user or because there is no straightforward way to turn the features into totally ordered ones". In fact, Reviewer U75f also raised a similar concern, highlighting that the method was evaluated on relatively small data sets, with no exploration of challenging or diverse datasets like biomedical or geospatial graphs.
Although the authors included a discussion of the relevant limitations in the revised version, I believe this is a more fundamental issue that needs to be addressed. Based on this, I think the paper may not yet be ready for acceptance at ICLR.

[1] Tang, Jie, et al. "ArnetMiner: extraction and mining of academic social networks." Proceedings of the 14th ACM SIGKDD international conference on Knowledge Discovery and Data Mining, 2008.

[2] Shen, Xiao, et al. "Adversarial deep network embedding for cross-network node classification." Proceedings of the AAAI conference on Artificial Intelligence, 2020.

[3] Rozemberczki, Benedek, et al. "Multi-scale attributed node embedding." Journal of Complex Networks, 2021.