Zero-Shot Generalization of GNNs over Distinct Attribute Domains

We appreciate the reviewer for recognizing that our paper “should hold great interest to many in the field of Graph ML” and that STAGE can help alleviate the significant “impediment to the creation of graph foundation models.” We now address their remarks.

Q1. “There is still overlap in the semantic meaning of many node attributes.”

A1: Our Ecommerce dataset comprises five product categories with semantically different attributes. Even when attributes appear semantically similar, they represent fundamentally different attribute domains with different values and distributions. For instance,

• Bed material can be Wood, Metal, Composite
• Shoe materials can be Leather, Synthetic, Canvas

This highlights STAGE’s ability to transfer across attribute domains that do not overlap in the semantic meaning, making it well-suited for diverse real-world applications.

Q2. “Have you tested the ability of STAGE to generalize across graph domains (e.g., from E-commerce and social networks)? If not, would it be possible to test this using additional datasets? To be clear, this experiment isn't necessary.”

A2: Our E-commerce datasets focus on temporal link prediction, while the social networks involve node classification. The social network data lacks edge creation times, so we cannot perform the same link prediction task. Additionally, the social network classification task is binary, while the E-commerce data does not have a directly comparable binary task. That said, we share the reviewer’s curiosity, and if they have suggestions for a suitable shared task on a different dataset, we would be happy to explore them.

Q3. “The authors mention that NBFNet-Gaussian seems to plateau after 3 graphs. However, the same can be said for STAGE. In reality, the mean performance of STAGE when going from 3 to 4 graphs is barely different.”

A3: Thank you, we will revise our statement to be more precise. While not the only method showing improvement, STAGE demonstrates particularly favorable scaling properties with additional graph domains. Specifically, STAGE exhibits notably tighter interquartile ranges compared to NBFNet-Gaussian at higher domain counts, suggesting more reliable performance across different domain combinations. Additionally, STAGE's lower whiskers consistently rise with more domains, showing that even its worst-case scenarios improve with more training data - a pattern less pronounced in NBFNet-Gaussian.

Q4. “Would you be able to include the results when holding out each of the 5 E-Commerce dataset?”

A4: We will include them in the revision.

Q5. “STAGE is quadratic with regards to the number of node attributes” “This severely limits the real-world potential of STAGE, as it can only handle graphs with few attributes.”

A5: While the quadratic complexity may be perceived as a limitation, we argue that it does not substantially diminish the usefulness of STAGE. In fact, numerous popular machine learning algorithms exhibit similar complexity constraints, including the Transformer Attention mechanism, which has a quadratic constraint. In relational deep learning, methods designed for small data thrive and make huge impacts on the real world. A recent example is TabPFN [1], a foundation model for tabular data that has found extensive scientific and business applications in real-world settings despite being designed for small datasets. Certain eigen-decomposition-based GNNs also have quadratic or cubic complexity.

[1] Hollmann et al., 2025. Accurate predictions on small data with a tabular foundation model

Q6. “STAGE requires separating attributes into those that are ordered and unordered. As such, some manual pre-processing must be done beforehand.” “This can further be quite tedious when many node attributes exist. However, this is a smaller weakness as it only needs to be done once.”

A6: We agree that this step may require upfront effort. However, it's worth noting that for categorical attributes, this processing can be done in lieu of mapping them into one-hot encodings, which is a common practice in ML pipelines. Additionally, for unordered attributes that are not categorical, it is likely that a pipeline already exists and can be repurposed.

Q7. “Some more intuitive explanation can be given for why STAGE considers the conditional probabilities of different node attribute types between nodes.”

A7: Thank you, we will include it in the revision. The key intuition is that zero-shot generalization requires focusing on statistical relationships between attributes rather than absolute values. By modeling conditional probabilities, STAGE recognizes patterns across different attribute spaces. For instance, instead of learning specific rules like 'phones with 8GB RAM tend to be expensive,' STAGE learns abstract relationships such as 'high values in X_1 correlate with high values in X_2,' enabling knowledge transfer across domains with different attributes.

We appreciate the reviewer for recognizing our theory is “rigorous and well-founded.” We now address their comments.

Q1. “The approach is quadratic in complexity concerning the number of attributes, which may limit its applicability to very large graphs or datasets with high-dimensional attribute spaces.”

A1: While the quadratic complexity may be perceived as a limitation, we argue that it does not substantially diminish the usefulness of STAGE. In fact, numerous popular machine learning algorithms exhibit similar complexity constraints, including the Transformer Attention mechanism, which has a quadratic constraint. In the domain of relational deep learning, methods designed for small data thrive and make huge impacts on the real world. A recent example is TabPFN [1], a foundation model for tabular data that has found extensive scientific and business applications in real-world settings despite being designed for small datasets. Certain eigen-decomposition-based GNNs also have quadratic or cubic complexity.

[1] Hollmann et al., 2025. Accurate predictions on small data with a tabular foundation model

Q2. “Could the authors provide an empirical comparison or toy visual example that highlights how order statistics capture invariant relationships better than normalized raw values in Sec. 3.1? How does the use of order statistics handle outliers or tied values in attribute distributions?”

A2: Thanks for the suggestion. Next, we present a toy example adapted from the Ecommerce dataset to illustrate why order statistics are superior to normalization for domain transfer.

Domain 1 (Train - Computers have attribute power_supply_watts):

• Computer A: 800 W
• Computer B: 600 W
• Computer C: 450 W
• Computer D: 300 W

Domain 2 (Test - Refrigerators have attribute energy_rating):

• Refrigerator A: 4.0 [A]
• Refrigerator B: 3.0 [B]
• Refrigerator C: 2.0 [C]
• Refrigerator D: 1.0 [D]

Now assume there exists a correspondence between Computer A and Refrigerator A, Computer B and Refrigerator B, and so on. This correspondence arises because users who prefer high-powered computers tend to select refrigerators with better energy ratings. When encoding these attributes, we need a representation that preserves this user preference pattern across domains, despite the different attribute scales.

With z-score normalization, the values become different across domains:

• Product A: 1.42 vs 1.34
• Product B: 0.34 vs 0.45
• Product C: -0.47 vs -0.45
• Product D: -1.28 vs -1.34

With order statistics, namely STAGE, the values remain identical in both domains:

• Product A: 0.25 (1/4 products)
• Product B: 0.5 (2/4)
• Product C: 0.75 (3/4)
• Product D: 1.0 (4/4)

Therefore, while normalization produces different values for the corresponding products, order statistics preserve consistency across domains by capturing ranking information, which is invariant to monotonic transformations. Additionally, by focusing on ranks, order statistics are inherently robust to outliers, and ties are handled by assigning the average rank to all tied items.

Q3. “Could the authors provide more theoretical or experimental illustration on how this loss of identifiers affects model expressiveness, especially in cases where attribute order might carry implicit information?”

A3: Without attribute identifiers, the model may lose the ability to distinguish attributes. For instance, consider two edges e1 and e2 in the original graph, and two features f1 and f2. If f1 has CDF value v1 in e1 and v2 in e2, while f2 has CDF value v2 in e1 and v1 in e2, with all other features having identical CDFs in both edges, then the two STAGE edge-graphs become isomorphic. In this case, STAGE produces identical edge embeddings for e1 and e2, even though the actual attribute distributions differ, limiting the model's expressivity and ability to distinguish them. However, our empirical analysis suggests this theoretical limitation has minimal practical impact. As shown in Figure 6, STAGE effectively captures meaningful attribute relationships even without explicit identifiers.

Q4. “Could the authors discuss how sensitive STAGE is to the diversity of training domains theoretically or experimentally?”

A4: Figure 4 shows that STAGE benefits from diverse training domains. STAGE’s performance consistently improves as we increase the number of training domains, with higher median values and tighter confidence intervals for both Hits@1 and MRR metrics. This suggests that STAGE effectively leverages the diversity in multiple graph domains to learn more robust and transferable representations. In contrast, the other baselines show minimal or inconsistent improvements with additional domains.

Q5. “The limitation of STAGE seems not to be discussed in the manuscript.”

A5: We acknowledge the limitations of STAGE on larger graphs in Sections 1 and 4. We appreciate the reviewer's suggestion and will further clarify them in the revision.

We appreciate the reviewer’s recognition of STAGE as “marking a significant step towards foundational graph models” and their positive assessment of our theory and empirical results. We will incorporate their suggestions into the revised manuscript. Below, we address their remarks:

Q1. Building each edge a STAGE-edge-graph “increases computational overhead, limiting experiments to small-to-medium-sized graphs and preventing scalability to larger graphs.”

A1: The computational complexity of STAGE is linear with respect to the number of edges and quadratic with respect to the number of attributes, rendering it particularly suitable for small-to-medium datasets. While this characteristic may be perceived as a limitation, we argue that it does not substantially diminish the usefulness of STAGE. In fact, numerous popular machine learning algorithms exhibit similar complexity constraints, including the Transformer Attention mechanism, which has a quadratic constraint. In the domain of relational deep learning, methods designed for small data thrive and make huge impacts on the real world. A recent example is TabPFN [1], a foundation model for tabular data that has found extensive scientific and business applications in real-world settings despite being designed for small datasets. Certain eigen-decomposition-based GNNs also have quadratic or cubic complexity.

The effectiveness of STAGE is evident in our experiments, for instance on the Ecommerce Stores dataset, where zero-shot learning achieves a Hits@1 score of nearly 0.6 when predicting user behavior from desktop configurations described by just 12 attributes. This result underscores that the computational trade-off enables substantial performance gains in settings where computational complexity is less of a constraint than model quality and expressiveness.

[1] Hollmann et al., Nature 2025. Accurate predictions on small data with a tabular foundation model

Q2. STAGE’s “effectiveness on text attribute graphs is very limited, as shown in Appendix C3, where STAGE performs no better than other baseline methods on social network graphs.”

A2: We believe the reviewer is referring to Table 5, discussed in Appendix E, which presents an additional experiment focusing on zero-shot prediction of age on social networks. As explained in Appendix E, this experiment was designed to highlight the inherent difficulty of predicting age in a zero-shot setting, given that "age" and "gender" are the only shared attributes available for node labels. Our goal was to illustrate that age prediction is fundamentally challenging across models, justifying our focus on gender prediction (Table 2). Consistent with this, neither our STAGE nor text embeddings outperformed other approaches in age prediction. All models struggled due to the significant distributional shift between age attributes in the Friendster and Pokec networks, as visualized in Figure 5.

Nonetheless, we agree with the reviewer that STAGE is designed to handle non-text attributes, theoretically guaranteeing generalization when considering continuous and discrete, as well as ordered and unordered attributes. We believe that extending STAGE to handle text-attribute, for instance by coupling it with an initial text encoder, represents an interesting avenue for future research, although one that requires a separate effort.

Q3. Minor suggestions

A3: We thank the reviewer’s suggestion on manuscript change and will include it in the updated version.

We appreciate the reviewer’s recognition that STAGE “can be useful in many real-world adaptation scenarios,” and thank them for acknowledging STAGE’s versatility, as well as for their appreciation of our theoretical motivation and experiments. We now address their comments.

Q1: “..It might still be interesting to compare to foundation models that claim to have zero-shot generalization ability.”

A1: We already include GraphAny in Table 2, a baseline claimed to have zero-shot generalization ability. Per your suggestion, we include another foundation model, namely GCOPE [1] in Table 2 (reporting zero-shot test accuracy on Pokec, trained on Friendster), and report the results below. Notably, STAGE outperforms the GCOPE foundation model by 21.9%.

Q2: “How do you view the effectiveness of this method under the structure shift and will STAGE become problematic under the presence of structure shift?” And “Is it possible to compare with graph OOD works in the same settings?”

A2: Our experiments account for substantial structural variation, as shown in Table 3 (Appendix C), where the number of nodes, edges, and average degrees often double from training to testing. Despite this, STAGE consistently outperforms other baselines in all zero-shot settings (Figure 3), demonstrating its effectiveness in scenarios where both the attribute domain and the structure shift.

Nonetheless, we acknowledge that STAGE does not make explicit assumptions about structural shifts, but it also does not impose constraints on them, and can be integrated with any GNN supporting input edge embeddings. We anticipate that integrating STAGE with existing graph-OOD GNNs, which are explicitly designed for structural shifts (e.g., varying graph sizes), could further enhance performance in these settings.

Finally, to the best of our knowledge, existing graph OOD methods (e.g., those in [2]), focus on structural shifts but do not address the attribute space shifts we consider, which include changes in the number of attributes between train and test. As a result, existing graph OOD methods cannot be evaluated in our setting. We will clarify this distinction in our revision and include related works on graph OOD.

Q3: “The complexity of creating edge graphs still remains concerning and limits the usage on large datasets.”

A3: The computational complexity of STAGE is linear in the number of edges and quadratic in the number of attributes, rendering it particularly suitable for small-to-medium datasets. While this characteristic may be perceived as a limitation, we argue that it does not substantially diminish the usefulness of STAGE. In fact, numerous popular machine learning algorithms exhibit similar complexity constraints, including the Transformer Attention mechanism, which has a quadratic constraint. In the domain of relational deep learning, methods designed for small data thrive and make huge impacts on the real world. A recent example is TabPFN [3], a foundation model for tabular data that has found extensive scientific and business applications in real-world settings despite being designed for small datasets. Certain eigen-decomposition-based GNNs also have quadratic or cubic complexity.

Q4: “The current design only considers the pairwise relations, which is a rather simplified version of the attribute hypergraph”

A4: Our STAGE-edge-graph captures marginal probabilities (e.g., P(A), P(B), P(C)) and pairwise conditional probabilities (e.g., P(A|B)) among attributes of node pairs. However, we agree that this representation requires additional assumptions about independence to recover joint probabilities of three or more attributes. While extending STAGE to incorporate these higher-order interactions is theoretically possible, it would result in increased complexity, whose exploration represents an important research direction on its own.

Q5: “What design corresponds to assigning unique attribute identifiers to label the nodes of our STAGE-edge-graphs and what design corresponds to dropping the attribute identifiers?”

A5: In a STAGE-edge-graph, labeling each node with the attribute ID of an endpoint node corresponds to assigning unique attribute identifiers. In contrast, our design of representing each node solely by its probability density function, without the attribute ID, corresponds to dropping the attribute identifiers. We will incorporate it in the next revision, thank you.

[1] Zhao et al., 2024. All in One and One for All: A Simple yet Effective Method towards Cross-domain Graph Pretraining

[2] Zhang et al., 2025. A Survey of Deep Graph Learning under Distribution Shifts: from Graph Out-of-Distribution Generalization to Adaptation

[3] Hollmann et al., 2025. Accurate predictions on small data with a tabular foundation model