Causal-aware Graph Neural Architecture Search under Distribution Shifts

We are grateful to the reviewer for the insightful comments and suggestions. We have meticulously reviewed each point and provide the following responses:

W1. Computational Efficiency Concerns: Section 3.3 reveals a potential limitation.

W2-2. Evaluation datasets: Performance on DrugOOD or GOOD to address concerns about memory consumption and computational time.

Q2. How does your method behave in the large graph compared to the previous fixed-network methods since you may search in a large network space? Will it be out of memory or the time computation will exponentially grow？

A1. Thank you for your valuable feedback. Since DrugOOD is similar to OGBG-Mol* (both involving molecule datasets with comparable graph sizes), we conducted additional experiments on the GOODSST2 dataset from GOOD. This dataset, converted from text sequences, presents a different domain compared to molecule datasets. Here, nodes represent words, edges indicate relations between words, and labels reflect sentence sentiment.

We compared performance, time, and memory costs across methods. To ensure fairness, we recorded results after 100 epochs for all methods. Additionally, since methods like DIR, GIL, and CIGA did not achieve optimal performance within 100 epochs, we also recorded their results after 200 epochs.

CARNAS achieves the best performance while maintaining competitive time and memory costs compared to non-NAS-based (fixed-network) graph OOD methods. The efficiency of the architecture search process stems from CARNAS's ability to simultaneously search for the architecture and optimize its parameters, minimizing additional computational overhead. This demonstrates the practical viability and efficiency of CARNAS, even for larger-scale graph datasets.

Following your advice, we have included these new results on the GOODSST2 dataset in Appendix F in the revised pdf. Additionally, we have cited both DrugOOD and GOOD for their valuable contributions to graph benchmarks.

We would like to express our sincere appreciation to the reviewer for providing us with detailed comments and suggestions. We have carefully reviewed each point and respond to the reviewer’s comments point by point as follows.

W1. “I believe the paper does not clearly explain why NAS can help adjust GNNs to identify causal information, which I consider the main issue of the paper. In my view, NAS optimizes the structure of GNNs, enhancing their efficiency or expressiveness, but it does not inherently enable GNNs to determine what type of data to model. At the very least, the authors did not provide a clear explanation of this point in the paper.”

Q1. Could you explain how NAS can guide GNNs to model specific data causal relationships?

A1. Thank you for your question. We believe there is a misunderstanding regarding the target problem of this work. We would like to clarify that this work utilizes causality to enhance NAS rather than employing NAS to identify causal information. We apologize for the confusion and will further clarify this in the revised manuscript.

To be specific, we frame the entire predicting pipeline in graph task as graph $\rightarrow$ architecture $\rightarrow$ label. Our work aims to unveil and utilize the causal graph-architecture relationship to address distribution shifts during the graph neural architecture search (NAS) process. This is an under-explored yet important issue, as previous works [1-3] emphasize the complex relationship between graph data and the optimal architecture. However, prior methods fix the network structure and solely focus on the graph-label relationship, neglecting the influence of the architecture itself, which can lead to suboptimal results under distribution shifts.

As you mentioned, NAS optimizes the structure of GNNs (based on graph data), and this process provides an opportunity to delve into the graph-architecture relationship. Therefore, we incorporate the architecture search process to improve generalization by identifying and leveraging the causal graph-architecture relationship to construct the optimal architecture (denoted as $A_c$ in our paper) for graph data.

Specifically, in our methods:

1. The causal subgraph $G_c$, extracted from the Disentangled Causal Subgraph Identification module, serves as a carrier of information causally relevant to the optimal architecture, rather than being merely causally related to the label $\hat{Y}$ as in prior approaches.
2. In the Invariant Architecture Customization module, we customize the optimal architecture $A_c$ using $G_c$ and simulated intervention architectures ${A_v}_j$s using intervention graphs. The loss term $\mathcal{L}\_{arch}$ for ${A_v}_j$s is used to regulate the influence of spurious components on the construction of the optimal architecture. This allows us to identify and leverage causal graph-architecture relationships to build the optimal architecture.

In conclusion, we first propose to jointly optimize causal graph-architecture relationship and architecture search by offering an end-to-end training strategy for extracting and utilizing causal relationships between graph data and architecture, which is stable under distribution shifts, during the architecture search process, thereby enhancing the model’s capability of OOD generalization.

W2. The paper lacks theoretical justification for the regulatory capability of NAS.

A2. Thank you for your comment. While our work primarily focuses on the empirical effectiveness of Neural Architecture Search (NAS), it is worth noting that NAS has been widely studied[4-8] for its ability to regulate and optimize neural network performance across diverse tasks.

Moreover, we have included additional theoretical analysis about the problem to further illustrate our method in a more rigorous way. The added theoretical analysis is highlighted in blue in Appendix C in the updated pdf.

---

[1] You et al. Design space for graph neural networks. NeurIPS2020

[2] Corso et al. Principal neighbourhood aggregation for graph nets. NeurIPS2020

[3] Xu et al. How neural networks extrapolate: From feedforward to graph neural networks. arXiv2020

[4] DARTS: Differentiable Architecture Search. ICLR 19.

[5] Auto-GNN: Neural Architecture Search of Graph Neural Networks. 2019.

[6] Graph Neural Architecture Search. IJCAI 20.

[7] Automated Machine Learning on Graphs: A Survey. IJCAI 21.

[8] One-shot Graph Neural Architecture Search with Dynamic Search Space. AAAI 21.

We would like to express our sincere gratitude to the reviewer for providing us with detailed comments and insightful questions. We have carefully considered the reviewer's feedback and would like to address each point as follows:

W1. While the paper shows good performance on the tested datasets, it lacks a detailed analysis of computational complexity and memory requirements. Specifically, the time complexity could become prohibitive for very large graphs. The authors should discuss how their method performs on graphs with millions of nodes and edges, which are common in real-world applications like social networks.

A1. Thank you for your comment. We provide a detailed analysis of computational complexity and memory requirements in Appendix E.2, E.4, and F. Specifically, as shown in Table 6, CARNAS consistently requires less time while achieving superior performance compared to the best-performing NAS baseline, demonstrating its enhanced efficiency and effectiveness. Additionally, Table 8 & 9 shows that CARNAS's overall time and memory costs are competitive with non-NAS-based (fixed-network) graph OOD methods, achieved by simultaneously searching for architectures and learning parameters. The practical time and memory efficiency of CARNAS makes it a viable and effective choice. We would also like to clarify that our work, like many previous Graph OOD studies [1,2], focuses on the task of graph classification, where the size of each graph instance does not typically contain millions of nodes and edges. Graphs with millions of nodes and edges, as you mentioned, are usually used for node classification (or link prediction). The task of node classification can be viewed as graph classification based on the ego-graphs of a node, where the K-hop neighbor graph of each node serves as a graph instance that are significantly smaller than the original whole graph, allowing our model to generalize to node classification with comparable complexity to traditional methods.

W2. The method requires careful tuning of four critical hyperparameters, which may significantly impact performance. In particular, the edge importance threshold $t$ in Eq.(6) and the intervention intensity $\mu$ in Eq.(10) show high sensitivity in experiments. While the authors provide some sensitivity analysis on BACE dataset, they don't fully explain how to effectively tune these parameters for new datasets or application domains.

A2. Thank you for your comment. The edge importance threshold $t$ is used to predefine the ratio of extracting causal subgraphs from original graph instance and $\mu$ is used to predefine the intervention intensity. The values of these two hyper-parameters can be roughly determined based on domain knowledge relevant to the dataset. In practical, we tune the parameters by giving a recommended range for each of them, and conduct hyper-parameter optimization on validation set. The range for SPMotif is $t\in [0.5,0.85], \mu\in[0.2,0.5], \theta_1\in[0.3,0.5], \theta_2\in[0.005,0.015]$ and for OGBG-Mol* is $t\in [0.4,0.7], \mu\in[0.4,0.7], \theta_1\in[0.1,1.0], \theta_2\in[0.001,0.020]$.

W3. The paper lacks formal theoretical guarantees for the causal discovery process. While the empirical results are strong, the authors should clarify under what conditions their method is guaranteed to identify true causal relationships and provide bounds on the probability of discovering spurious correlations. Additionally, the relationship between the intervention loss and causal invariance could be more rigorously established.

A3. Thank you for your invaluable advice! Although the rebuttal time is limited, we have tried to provide a detailed theoretical analysis to further illustrate our method in a more rigorous way. The added theoretical analysis is highlighted in blue in Appendix C in the updated pdf. It step by step introduces how we propose to jointly optimize the causal graph-architecture relationship and architecture search by offering an end-to-end training strategy. This strategy enables the extraction and utilization of causal relationships between graph data and architecture, which remain stable under distribution shifts during the architecture search process, thereby enhancing the model’s capability for OOD generalization.

We would like to express our sincere appreciation to the reviewer for providing us with detailed comments and suggestions. We have carefully reviewed each comment and offer the following responses:

W1. Clarity in Sec 3.3: Given I’m having limited familiarity with Graph NAS, the dynamic graph neural network architecture production and optimization process described in Sec 3.3 remains somewhat unclear for me. A visual representation and a more detailed explanation would significantly improve the paper's readability.

A1. Thank you for your valuable feedback. We aim to intuitively illustrate the architecture production and optimization process below, referring to Figure 1, which visualizes Section 3.3, along with the algorithm in Appendix B.

• Production: Given a graph embedding, such as $\mathbf{H_c}$, our goal is to construct an architecture $A$ (with $K$ layers, each containing $|\mathcal{O}|$ candidate operators). For example, in Figure 1, the architecture consists of 2 layers, each with 3 candidate operators. We represent the architecture $A$ as a matrix $\mathbf{A} \in \mathbb{R}^{K \times |\mathcal{O}|}$, where $\mathbf{A}_{k,u} = \alpha_u^k$ is the mixture coefficient of operator $o_u(\cdot)$ in layer $k$, denoting the importance of different operators in different layers. In the figure, this coefficient is visualized as arrows of varying shades in the architecture. The constructed architecture $A$ is the result of combining these coefficients with the specific operators, as shown in Eq. 11. Additionally, we use a prototype vector $\mathbf{op}_u^k$ to represent an operator in a specific layer, allowing us to learn the coefficient $\alpha_u^k$ based on both $\mathbf{op}_u^k$ and $\mathbf{H_c}$.
• Optimization: Using the production process described above, we customize the optimal architecture $A_c$ for $G_c$ and simulate intervention architectures ${A_v}j$s for intervention graphs. We use $A_c$ as the optimal architecture for the graph instance to output the final label prediction and apply $\mathcal{L}_{arch}$ to ${A_v}_j$s to regulate the influence of spurious components when constructing the optimal architecture. This approach identifies and leverages the causal relationships between graphs and architectures to construct the optimal architecture.

W2. Causal-Aware Solution's Justification: While the paper presents a causal-aware solution for handling distribution shifts, some aspects require stronger theoretical support to underscore the novelty and significance of the approach.

W2-1. More Theoretical Support.

A2-1. Thank you for your feedback. Although the rebuttal time is limited, we have attempted to provide a detailed theoretical analysis of the problem to further illustrate our method in a more rigorous manner. The added theoretical analysis is highlighted in blue in Appendix C in the updated pdf.

W2-2. Reliability of Causal Subgraph in Latent Space.

A2-2. Thank you for your comment. Representing causal subgraphs and performing interventions in the latent space allow us to construct architectures based on the intervened latent features. Following your advice, the added theoretical analysis also reflects how our method obtains the causal parts. Additionally, we visualize the edge importance for selecting causal subgraphs in each graph in Figure 8 and the searched architectures for them in Figure 7 (Appendix G), which better demonstrate the learned causal graph-architecture relationships.

W2-3. Difference with PGExplainer.

A2-3. Thank you for your comment. While PGExplainer [26] aims to identify explanatory subgraphs based on edge features—a common step in various graph tasks—we would like to clarify the key differences between our method and PGExplainer.

In the Disentangled Causal Subgraph Identification module, we introduce the concept of disentangled learning [1,2] to represent distinct latent factors as unique vector representations. This approach allows us to comprehensively disentangle causal latent graph structural features from spurious ones, enabling a more precise identification of causal subgraphs that influence architecture construction. Together with the subsequent modules, our method effectively disentangles complex graph-architecture relationships, learning distinct latent factors that guide the customization of optimal architectures from intricate graph structure features.

To summarize, CARNAS differs in both objective and methodology: 1. Objective and Motivation: PGExplainer focuses on extracting explanatory subgraphs, while our method aims to disentangle causal features from spurious ones and extract causal subgraphs for architecture optimization. 2. Mechanism: Instead of merely selecting important edges based on graph features, we leverage a disentangled layer to separate causal features from spurious ones, offering a fundamentally different mechanism for subgraph identification.