Paper: VQGraph: Rethinking Graph Representation Space for Bridging GNNs and MLPs (ICLR-2024)

Reviewer comment (Experiments):
It would be better if more details about the production scenario were included in the paper. If the authors are not experts in the graph knowledge distillation domain, they are unlikely to know the details of the experiments. So, more details about the tasks need to be described.

Author response:
A1: Thanks for your constructive suggestion, we have added the following details into **the appendix of the updated manuscript denoted in red**. Regarding to the production scenario, we recognize that real-world deployment of models necessitates the ability to generate predictions for new data points as well as reliably maintain performance on existing ones. Therefore, we apply our method in a realistic production setting to better understand the effectiveness of GNN-to-MLP models, encompassing both transductive and inductive predictions.

In a real-life production environment, it is common for a model to be retrained periodically. The holdout nodes in the inductive set represent new nodes that have entered the graph between two training periods. To mitigate potential randomness and to evaluate generalizability more effectively, we employ a test dataset $V^U_{ind}$, which contains 20% of the test data, and another dataset $V^U_{obs}$, containing the remaining 80% of the test data. This setup allows us to evaluate the model's performance on both observed unlabeled nodes (transductive prediction) and newly introduced nodes (inductive prediction), reflecting real-world inference scenarios.

We further provide three sets of results: "tran" refers to results on $V^U_{obs}$, "ind" refers to $V^U_{ind}$, and "prod" refers to an interpolated value between tran and ind, according to the split rate, indicating the performance in real production settings.

Reviewer comment (Experiments):
Compared with previous method, extra memory space is required due to the introduced learnable codebook. As the authors show, the codebook size has influence on the performance of the model. How to effective tune this hyperparameter is worth exploring.

Author response:
The size of the codebook inherently possesses an optimal point, a fact substantiated by our experimental findings. Notably, this optimum tends to scale **proportionally with the growth in the number of nodes** and features within a graph. To address practical concerns, we propose an efficient approach involving **a few key steps**. Specifically, employing a binary search process, we can determine the optimal codebook size within 3-4 steps, where it can be determined **solely by convergence** without running the entire training process, especially in large-scale datasets. Meanwhile, minor adjustments to the codebook size, especially within a range not comparable with the scale of graph complexity, have negligible impact on the overall performance. This streamlined process ensures that adapting the codebook size to new datasets is both manageable and expeditious.

Reviewer comment (Experiments):
The authors propose a generative approach learn codebook embedding. Is there any other way to learn the codebooks (such as self-supervised learning paradigm) or use other generative models (such as GAN) ?

Author response:
Any self-supervised learning or generative modeling methods that use graph reconstruction objective can learn a structure-aware codebook for graph data. However, self-supervised learning methods with mask-prediction can **only recover partial structures** which are not expressive enough for a structure-aware codebook. Other generative models like GAN pay more attention to how to learn **a good decoder not a good encoder (codebook)**. In contrast to them, our graph VQ-VAE effectively combine  representation learning with generative modeling for learning a expressive codebook that maximally preserves the entire graph structures.

Reviewer comment (Experiments):
Could the proposed method be applied to heterophilic graphs [1,2]?  

[1]  Lim et al., Large scale learning on non-homophilous graphs. NeurIPS 2021.  
[2] Platonov et al, A critical look at the evaluation of GNNs under heterophily: Are we really making progress? ICLR 2023.

Author response:
Capturing local structures is also crucial for heterophilic graphs, thus we apply our method to heterophilic graphs. Specifically, we employ two state-of-the-art Graph Neural Network (GNN) models known for their efficacy on heterophilic graphs, namely ACMGCN and GCNII, along with MLP. Incorporating our structure-aware code embeddings, we evaluate these models on two challenge datasets—Texas and Cornell. 
The results of these experiments, displayed below, demonstrate that our graph VQ-VAE **can generalize to different heterophilic GNN/MLP architectures and consistently improve their performances**. These remarkable results underscore the effectiveness of our method in effectively capturing local structural information, and demonstrate the versatility and broad applicability of our proposed approach across different graph domains. **We have added this part into the updated manuscript.**

| Dataset | ACMGCN | ACMGCN with Graph VQ-VAE | MLP   | MLP with Graph VQ-VAE | GCNII | GCNII with Graph VQ-VAE |
| :-----| :----: | :-----:| :----: | :-----:|  :-----:| :----: | 
| texas   | 86.49  | **87.03**      | 75.68 | **77.93** | 76.73 | **79.04**    |
| cornell | 84.05  | **84.59**          | 76.38 | **78.29** | 76.49 | **78.29**    |

Reviewer comment (Experiments):
The size of codebook is sensitive to different datasets, this may introduce some practical difficulty.

Author response:
The size of the codebook inherently possesses an optimal point, a fact substantiated by our experimental findings. Notably, this optimum tends to scale **proportionally with the growth in the number of nodes** and features within a graph. To address practical concerns, we propose an efficient approach involving **a few key steps**. Specifically, employing a binary search process, we can determine the optimal codebook size within 3-4 steps, where it can be determined **solely by convergence** without running the entire training process, especially in large-scale datasets. Meanwhile, minor adjustments to the codebook size, especially within a range not comparable with the scale of graph complexity, have negligible impact on the overall performance. This streamlined process ensures that adapting the codebook size to new datasets is both manageable and expeditious.

Reviewer comment (Theory):
It would be better to provide a theoretical understanding of this approach.

Author response:
We provide a theoretical understanding for VQGraph, assuming the data adimts the following hierarchical generation mechanism:
$$\begin{array}{rl}
  &c_i \sim q(c_i)\\\\
  &r_i \sim q(r_i|c_i)\\\\
  &G \sim q(G|(r_1,r_2,...r_N)), \\\\
\end{array}$$
where $c_i,r_i$ are the label and latent variable (code) for the i-th node in the graph $G=(V,A)$. In other words, we assume there exist an explicit latent representation for each node that can be used to generate the observed graph. Importantly, these code representations encode not only feature information but also structure information, so that the adjacency matrix can be generated. VQGraph can be viewed as using generative pre-training to recover the latent representation and facilitate the label prediction, as learning
$$q(c_i|r_i)$$ is easier than learning $$q(c_1,c_2,...c_N|G) = \int \prod_i q(c_i|r_i)q(r_1,r_2,...r_N|G)dr$$

Reviewer comment (Evaluation):
It seems that for APPNP, all three KD approaches can't outperform the teacher model. Is there any discussion or explanation for this phenomenon?

Author response:
This is due to the fact that the APPNP directly uses node features for prediction prior to the message passing on the graph, which is very similar to what the student MLP does in KD approaches. Therefore using APPNP as teacher model can not provide student MLP with much additional information.

Reviewer comment (Evaluation):
It would be better to provide an analysis what codebook entry is most informative to a class label, which can make this approach more convincing and intuitive.

Author response:
In the training stage, we tokenize each node with different neighborhood structures as discrete codes using Graph VQ-VAE. Now we take a closer look at codebook entries for different classes on Pubmed for in-depth analysis. There is little difference in the number of codes with each class and there is a uniform distribution of the nodes for each class, demonstrating the learning capacity of our model. Meanwhile, the results  demonstrate that there is small code overlap among different classes, which indicates the code distributions learned for different classes are mutually independent. This finding validates that our graph VQ-VAE enables distinct separation of the nodes with different class labels in the representation space, facilitating a better knowledge transfer in GNN-to-MLP distillation. We have added these analysis into the updated manuscript denoted in red.

Code Entries for each class of Pubmed:

|              | Class 1 | Class 2 | Class 3 |
| ------------ | ------- | ------- | ------- |
| Code Entries | 2490    | 2623    | 2591    |

Number of overlapping codes for each pair of classes:

|         | Class 1 | Class 2 | Class 3 |
| ------- | ------- | ------- | ------- |
| Class 1 | *       | 24      | 21      |
| Class 2 |         | *       | 19      |
| Class 3 |         |         | *       |

Reviewer comment (Writing):
The connection between GNNs-to-MLPs distillation and the proposal of a local-structure-aware codebook is less developed, since they are actually separate goals; more explanation is needed to make the main claim of the paper clearer.

Author response:
We here clarify the connections between GNNs-to-MLPs distillation and the proposition of a local-structure-aware codebook. The primary goal of GNNs is to establish an embedding space encompassing all nodes, wherein structural information is consolidated across multiple layers of graph convolutional layers. Kindly recall that the aim of GNN-to-MLP distillation is to facilitate the transfer of sufficient structure-based information to the MLP, enabling the latter to contain graph-related information. Unlike GNNs, MLPs are inherently limited in their capacity to capture structural information independently. Therefore, the knowledge distillation process serves for transferring local structural information to the MLP. The integration (a codebook) of these local-structure-aware codes work as an explicit embedding space for storing and conveying this local structural information, thereby enhancing the MLP's capacity to acquire such knowledge.

Reviewer comment (Experiments):
In practice, one of the main concerns is that the hyper-parameter for the code distillation in equation 7 is quite small compared to the class distillation one. It makes the whole proposal more intended to be a node embedding booster for GNNs, the graph tokenizer, while the codebook is not significantly effective in the distillation phase.

Author response:
Regarding the disparity in hyperparameters between the code distillation and class distillation in Equation 7, it is critical to consider the calculation processes for both the class-based distillation loss and code-based distillation loss. The class-based distillation loss entails computing the KL divergence between the student predictions and soft labels. On the other hand, the code-based loss involves comparing the target node representations with all M codes of the codebook embeddings. 
It is deserved to notice that **the class number and code number are extremely different regarding the order of magnitude**, which results in a considerable scale gap of approximately 10^7-fold in large-scale datasets. Consequently, despite the smaller hyperparameter, the corresponding code distillation loss is proportionally larger, effectively ensuring that **the gradients** propagated during backpropagation are maintained **at a comparable level with class distillation loss**.

Reviewer comment (Reproducibility):
The volume of the codebook is very likely to be the same scale as the original node features, which makes the proposal still in the trade-off of time and space cost for sure; therefore, more information about the parameter volume should be provided.

Author response:
Indeed, the trade-off between time and space costs is a long-term challenge in designing efficient graph neural network architectures. From this perspective, our proposed VQGraph is able to effectively capture the structural diversity of graph data with a compact codebook. Our method focuses on achieving a balance where the size of the codebook is sufficiently large to encapsulate the essential substructure patterns yet remains much smaller than the original substructures in graph data. The codebook’s volume is not directly proportional to the original node features but is mainly determined by the complexity of the graph data, considering both nodes and edges which produce different local substructures. Taking Cora dataset for example, while the **codebook size is 2048**, this allows us to efficiently represent a manifold larger number of possible 1-hop substructures. Considering Cora’s 2485 nodes with an average degree of about 4, there is **a theoretical limit of ( O(2485^4) ) distinct 1-hop substructure patterns**. Thus, our codebook provides a highly compressed representation that encapsulates these patterns with a significantly lower memory footprint.
Moreover, the empirical performance of our method demonstrates that despite the compact size of the codebook, there is no compromise on the expressiveness and the ability to facilitate meaningful representations. We want to highlight that the codebook size represents a new effective embedding space that strikes a balance between memory efficiency and the capacity to represent diverse structural information.

Reviewer comment (Theory):
It is better to explain more about why it is appropriate to use VQ-VAE to encode graph structure, since it was originally designed for continuous data. For graph-structured data, it might be easier to convert to frequency information, i.e., each position in the "codebook" stores the volume of a specific frequency. I am just curious why VQ-VAE is suitable for encoding discrete and non-Euclidean data.

Author response:
Whether it is for continuous data or for discrete data, the core of the encoder is to learn a optimal embedding space for maximally preserving the intrinsic information of the raw data, which is also the primary goal of VQ-VAE. Thus, the utilization of VQ-VAE is essentially not related to the data format. Besides, the discrete nature of graph-structured data aligns with the VQ-VAE's ability to handle categorical and non-continuous information with indexed codes. Each position in the codebook can encapsulate essential structural information, thereby facilitating the encoding of complex graph structures.

Reviewer comment (Experiments):
According to the inner product in the reconstruction loss of equation 3 and the benchmarks used, it is quite important to emphasize that the effective domain of the proposal is probably homophilic graphs, otherwise further discussion or experiments on heterophilic graphs are necessary.

Author response:
Capturing local structures is also crucial for heterophilic graphs, thus we apply our method to heterophilic graphs. Specifically, we employ two state-of-the-art Graph Neural Network (GNN) models known for their efficacy on heterophilic graphs, namely ACMGCN and GCNII, along with MLP. Incorporating our structure-aware code embeddings, we evaluate these models on two challenge datasets—Texas and Cornell. 
The results of these experiments, displayed below, demonstrate that our graph VQ-VAE **can generalize to different heterophilic GNN/MLP architectures and  consistently improve their performances**. These remarkable results underscore the effectiveness of our method in effectively capturing local structural information, and demonstrate the versatility and broad applicability of our proposed approach across different graph domains. Moreover, We want to emphasize that the inner product is **specifically for measuring the relationships between neighboring nodes, not specifically for homophilic graphs**. And it can be replaced by any non-parametric or parameterized relation measurements. **We have added this part into the updated manuscript.**

| Dataset | ACMGCN | ACMGCN with Graph VQ-VAE | MLP   | MLP with Graph VQ-VAE | GCNII | GCNII with Graph VQ-VAE |
| :-----| :----: | :-----:| :----: | :-----:|  :-----:| :----: | 
| texas   | 86.49  | **87.03**      | 75.68 | **77.93** | 76.73 | **79.04**    |
| cornell | 84.05  | **84.59**          | 76.38 | **78.29** | 76.49 | **78.29**    |

Reviewer comment (Experiments):
What is the data set used for Figure 1? Cora? All other figures/tables need further clarification on which dataset is used.

Author response:
We appreciate your attention to detail and apologize for any confusion. The dataset employed for Figure 1 and Figure 4 is Citeseer. In our paper, we have provided the datasets used for other figures and tables.

Reviewer comment (Theory):
In equation 3, the node embeddings are still included, so the claimed local structural embedding is not that promising, can you compare or explain the difference of the two parts of the reconstruction function? Also, what's the nonlinear function here?

Author response:
Certainly. In the early stages of our experiments, we conducted comprehensive ablation studies on the distinct components of our reconstruction function. The result is presented in the following table. From the result, it is obvious that the edge reconstruction segment is significantly better than the node embeddings in terms of its contribution to the overall performance. While the node embeddings does lead to a certain improvement in results, we include them in our reconstruction loss for better results. The nonlinear function here is the sigmoid function. 

| Dataset    | Only node reconstruction | Only edge reconstruction | node reconstruction+edge reconstruction |
| :----: | :----: | :----: | :----: | 
| Citeseer   | 74.84                    | 75.56                    | 76.09                                   |
| Pubmed     | 77.52                    | 78.10                    | 78.40                                   |
| Cora       | 83.15                    | 83.72                    | 83.93                                   |
| A-computer | 84.16                    | 85.03                    | 85.17                                   |
| A-photo    | 93.63                    | 94.15                    | 94.21                                   |
| Arxiv      | 71.17                    | 72.19                    | 72.43                                   |
| Products   | 78.40                    | 78.84                    | 79.17                                   |

Reviewer comment (Presentation):
Figure 2 needs improvement—In equation 7, the loss of classification and class distillation should also be clarified.

Author response:
We have updated Figure 2 in the updated manuscript according to your suggestions.
$L_{cls}$ is the cross-entropy loss between the MLP prediction $\hat y_{{v}
}$ and the class label $y_{v}$.
$L_{labeldistill}$ is the KL-divergence loss between the student prediction $\hat y_{v}$ and the soft label $y^{soft}_{v}$

Reviewer comment (Novelty):
The codebook does not seem to be so different from VGAE from the point of view of structure reconstruction, but the node feature reconstruction part makes it so different. Also, the graph structure part is not well designed, but an inner product.

Author response:
The critical difference between our graph VQ-VAE and VGAE is that VQ-VAE utilizes an explicit latent space (codebook) for enhancing the structure-aware graph representation learning while VGAE learns an implict latent space for graph generative modeling. The former focuses on learning a better graph encoder and the latter pays more attention to a better graph decoder (generator). From this perspective, it is reasonable to use a somple inner product as graph decoder in training our graph VQ-VAE, because a well-designed decoder would make the model to **depend heavily on the decoder** to reconstruct graph structures, thus limiting the sufficient learning of graph encoder.

Reviewer comment (Evaluation):
An interesting point, which is worth to study further, is what the Figure 5(ii) implies: denser graph needs more space, which is parallel to the discovery from graph signal processing domain, that the structural information cannot be summarized only by some low-frequency filters, the high-frequency ones are also informative in such graphs.

Author response:
We are delighted to discover that our findings align with those in the graph signal processing domain. Indeed, low-frequency filters excel at capturing smooth or slow-changing signal shifts, while high-frequency filters are informative for capturing fast-changing signal shifts within a graph.In real-life production scenarios, incorporating these high-frequency filters with graph signal learning can greatly enrich the representation of local structures, which are often diverse in large graphs. 

To further extend our framework, we believe it can be applied to other domains as well. For instance, we can employ a 2-way tokenizer to encode the local shifting structures into codebooks. This constraint serves as a regularization function, guiding the model to focus on the "most important" substructure information that best reconstructs the graph. In this approach, we could utilize separate codebooks for low-frequency and high-frequency filters, allowing for different codebook sizes. We assign a smaller codebook to low-frequency filters to facilitate overall and smooth learning, while allocating a larger codebook to high-frequency filters to capture more detailed and localized information. By employing this strategy, we aim to capture both the local and global structural information embedded within graph signals.

In summary, our framework not only aligns with insights from the graph signal processing domain but also presents opportunities for further exploration. By incorporating a 2-way tokenizer and utilizing separate codebooks for low- and high-frequency filters, we can enhance the representation of both global and local structures in various domains.

Reviewer comment (Experiments):
Some details of the experimental setup are not clear enough to me. As far as I understand from the paper, in the inductive setting,  $G^L \cup G^U_{obs}$ is used for training the teacher GNN and for distillation. True labels $Y^L$ are used to train the teacher but are not used in distillation. Is it right? The paper mentions that "the edges between $G^L \cup G^U_{obs}$ and $G^U_{ind}$ are removed in training, while they are leveraged during inference to transfer positional features via average operator". What does this mean? Do you augment the node feature with the features of its neighbors for inference?

Author response:
**Q1: Some details of the experimental setup are not clear. Are true labels $Y_L$ used to train the teacher but not used in distillation? Do you augment the node feature with the features of its neighbors for inference?** 

A1:  In our inductive experimental settings, we align with the framework established by NOSMOG [1] to ensure a fair comparison. 

* To clarify, we utilize true labels $Y_L$ for both training the teacher GNN and for GNN-to-MLP distillation process. In training the student, we employ soft labels and soft code assignments within the labeled and observed subsets to compute the class-based distillation loss and the code-based distillation loss. The classification loss is defined as the cross-entropy loss between the MLP prediction and the true labels $Y_L$. 

* Indeed, in the inference process, we adopt the practice of augmenting the node feature with its neighbors' features, in line with the approach outlined in NOSMOG and  DeepWalk. This augmentation strategy aims to provide nodes with additional information, thereby enriching the inferencing capabilities. With this methodology, we aim to circumvent potential limitations, such as the occurrence of nodes with identical features but dissimilar neighbors, a scenario that could pose challenges for the MLP in effectively processing the information. This instance directly addresses the concern raised in your second question, and our elucidation stands as follows.

[1] Tian Y, Zhang C, Guo Z, et al. Learning mlps on graphs: A unified view of effectiveness, robustness, and efficiency[C]//The Eleventh International Conference on Learning Representations. 2023.