Learning Graph Quantized Tokenizers

Thank you very much for the constructive feedback. Please find our responses below.

Other Graph Tasks

Please refer to our new experiments discussed in the general response to this question. In summary, we have conducted new experiments on a diverse set of graph tasks including graph classification and regression, inductive node classification, and inductive link prediction on long-range graph benchmarks demonstrating the strong performance of our model.

Training Losses

We have updated the paper to address this point. The reason we used three SSL objectives in a multi-task manner is to ensure that we capture different aspects of information by using different high-performing families of SSL objectives: student-teacher distillation, masked autoencoding, and Infomax. We also cite two papers that show this is indeed the case (Lines 203-204). As we show empirically, these objectives each contribute to the overall performance.

WL Test

Thank you for your comments. In our initial version, we included these statements to highlight the significant expressive power (theoretically) of Graph Transformers (GTs), which is one of our motivations for focusing on GTs in this paper. We agree that the transformer used in Kim et al. differs from the one in our GQT, and the statements made were not entirely appropriate. As you suggested, we have updated the paper and uploaded the new version. In the new introduction, we analyze the pros and cons of GNNs and GTs, highlighting the unique advantage of combining them, i.e., GNNs can provide locality information into tokens such that GTs do not need to perform dense pair-wise attention but only rely on a few tokens to learn the long-range information of graphs, which is the motivation for our proposed GQT.

Transferable Components of GQT

The contribution of semantic edges can be inferred from Tables 4 and 5. Specifically, by keeping everything intact and just dropping semantic edges, we observe a sharper downstream performance degradation in the heterophily setting (Table 5) compared to the homophily setting (Table 4). This indicates that introducing semantic edges helps in both scenarios but is more important for the heterophily setting. We speculate that this can be considered as a general augmentation even for GNNs improving their performance. For PPR, previous works have already shown that using PPR by replacing the adjacency matrix [1] improves GNNs, and also when used as positional encoding [2] or when used to generate sequences [3] can improve the performance of Transformers. Also Tables 4 and 5 suggest that using PPR scores (structural gating) can contribute to the performance of the model. This is also transferable to baselines. We will add more experiments to the camera-ready version comparing PPR sequences with BFS, DFS, and other random-walk based sequences.

Missing Standard Deviations in Tables 4 and 5

Thanks for catching this; we have updated them in the paper.

Structural Gating

Yes, that is correct. We introduce structural gating in Lines 343-349. We have updated the paper to make it clearer.

Once again, we would like to thank you for the helpful suggestions, and we look forward to addressing any other questions you may have.

References

[1] Gasteiger et al. Diffusion Improves Graph Learning. NeurIPS 2019

[2] Khasahmadi et al. Memory-Based Graph Networks. ICLR 2019.

[3] Fu et al. VCR-graphormer: A mini-batch graph transformer via virtual connections. ICLR 2024.

Thanks for your prompt response! We tuned on GAT since we used GAT as the graph encoder in our GQT. We selected Computer, wikiCS, Squirrel, Minesweeper, on which the tuned ClassicGNN (GAT) achieves largest improvement as reported in the paper [1] (we exclude Chameleon as the dataset size is very small).

We also noticed that GCN in [1] achieves quite good results on Squirrel, even more than 7.8% improvement over the tuned SAGE and GAT, which is quite large on this long-used dataset.

As you suggested, we will try to reproduce the results of GCN on heterophilous graphs and will update the results soon.

[1] Luo, Yuankai, Lei Shi, and Xiao-Ming Wu. "Classic GNNs are Strong Baselines: Reassessing GNNs for Node Classification." arXiv preprint arXiv:2406.08993 (2024).

We would like to thank the reviewer for their constructive feedback. Please find our responses below:

Performance Comparison with Classic GNN Paper

Thanks for pointing this interesting paper which provides a thorough empirical analysis of classic GNNs. We did not report the results from [1] due to two main reasons: (1) we, along with several other groups, have been unable to reproduce the results reported in [1] using their exact setting. Therefore, we were cautious about reporting that work before validating it. (2) Moreover, we reported the performance of GNNs from their original papers, as our goal was not to perform extensive hyperparameter tuning, unlike [1]. This is also the case with other Graph Transformer works that are consistently outperformed by [1] (e.g., Table 2 in [1]). Nonetheless, to address your concerns, we performed two sets of experiments. (1) Reproducing the results from [1] using their exact setting, and (2) performing the same hyperparameter tuning for the GNN encoder of the tokenizer. We emphasize that we did not perform hyperparameter tuning for the Transformer encoder. The results are shown below. As shown, (1) the reproduction results are inferior to those reported in [1], and (2) with the same hyperparameter tuning only on the GNN encoder of the tokenizer, we consistently outperform the results we obtained from their method.

LLMs

We agree that current LLMs for graph tasks are mostly applied to textual-attributed graphs (TAGs). What we meant here was to focus on LLM-based methods that directly take the structure into account and perform aggregation in in-context or struct-tuning settings; [2-3] are examples of such approaches. We rewrote the paragraph and updated the citations to make it clearer.

Innovation

We already cite VQGraph [4] in the paper (Lines 121-123), which differs from our approach as they use a VQ codebook (we use RQ) for distilling GNNs to MLPs. Thank you for suggesting Mole-bert [5]; we missed this one as it is specific to molecules whereas our method is domain-agnostic. We added this work to the updated literature review. The differences between Mole-bert and our approach are as follows: (1) unlike Mole-bert, which is specifically designed for molecules, our work is general and domain-agnostic, (2) Mole-bert uses VQ along with a GNN to perform BERT-style pretraining, whereas we only use the GNN module for tokenization and use a Transformer to learn long-range dependencies. This makes our proposed method more flexible, as it follows the standard tokenizer-transformer paradigm and is compatible with rapid advances in Transformer improvements; (3) similar to VQGraph, Mole-bert uses VQ whereas we use RQ, which yields hierarchical codebooks. We ran an experiment to show the difference between using VQ and RQ on downstream tasks. The results shown below indicate that RQ significantly outperforms VQ due to its hierarchical nature.

We also discussed VCR-Graphormer [7] in our model design and compared the performance in the experiments. In addition to our performance being better than VCR-Graphormer, our proposed tokenization is more advanced and novel than VCR-Graphormer for the following aspects:

Before sampling. Our token obtaining is semantic, hierarchical, and learned by the residual vector quantization method, which means (1) our token is shorter than the original input feature and more friendly to the attention mechanism, as the relationship between attention computation cost and feature length is beyond linear; (2) our token is more friendly to the inductive setting, because VCR-Graphormer’s tokenization is heuristic and needs to establish all prerequisites, such as identifying structure-aware and content-aware virtual connections and neighbors from scratch, whereas ours employs a trained tokenizer encoder, which can obtain token representations for new nodes in a one-shot manner without needing any rerunning.

After sampling. VCR-Graphormer takes the tokens individually to feed them into the attention mechanism, which increases attention costs. Our attention mechanism makes the following novel efforts: (1) for attention efficiency, we take the summarization intra-codebooks to save attention complexity; (2) for attention effectiveness, we introduce trained embedding matrix to attach to each codebook for the output expressiveness, since our number of codebooks is only 3, which operation will not induce too many computational costs.

Thank you.

In addition to presenting our model's strong performance across various benchmarks, we would like to take a step back and re-emphasize the broader significance of this paper. Two key aspects set our model apart from existing approaches:

(1) Our model offers greater flexibility compared to specialized Graph Transformers. Unlike these models, our tokenizer can be seamlessly integrated into future advances in Transformer design, which is an area of research that is progressing at a much faster pace than Graph Transformers or GNNs. This versatility makes our design more adaptable to rapid progress in the field. When compared to GNNs, our approach eliminates the need to iterate through multiple architectures for each dataset to find the optimal GNN layer for a given task.

(2) It is worth noting that all GNNs and specialized Graph Transformers require access to original node features during inference, which can be a significant limitation. In contrast, our design overcomes this challenge by learning discrete tokens. This is particularly important in large-scale industrial settings, where storing and retrieving node features for billions of users or items is not feasible due to memory constraints. We reported results comparing our model with a GAT model showing that with our design we can use a Transformer with half the memory footprint compared to GAT. This efficiency advantage is crucial in large-scale industrial applications.

In summary, our model not only demonstrates strong performance across benchmarks but also offers two significant advantages: adaptability to rapid developments in the Transformer family of architectures and improved memory efficiency. We believe these contributions make our paper an important contribution to the community.

We would be more than happy to address any additional questions you may have or provide further clarification on any aspects of our work to increase your confidence in your evaluation and in our work.

Thanks again for pointing out this interesting paper, which provides a thorough empirical analysis of classic GNNs on the node classification task.

1. As you suggested, we tried to reproduce the results of GCN and further tuned our GQT on heterophilous graphs. The results are shown below. As shown, with additional hyperparameter tuning, we can outperform or be comparable to the results we obtained from their method.

   We also listed the selected hyperparameters (# GNN layers, # GNN hidden dim) in the table. It can be observed that ClassicGNN typically requires a larger number of layers on these heterophilous graphs, which results in higher memory costs. In contrast, our GQT incorporates tokens learned from GNN (local) encoders and Transformer to capture long-range interactions, allowing us to use smaller GNN layers.

   Note that the authors of [1] recently updated the code for these medium datasets, therefore, our new (reproduced) run is based on this code (previously, we followed the code in the “large_graph” folder).

2. In addition, we have included additional results on the Long Range Graph Benchmark (LRGB) [2] in general response (including graph-level and link-level tasks), which is proposed to evaluate the capability of GNNs and GTs to capture long-range interactions. Our GQT (along with several other GT baselines) outperform GNNs on this benchmark, demonstrating the effectiveness of GTs (and our GQT) in capturing long-range dependencies.

Additionally, we would like to emphasize that, our primary goal is to investigate an effective and transferable graph tokenization strategy that enables the learned tokens to encode local interactions, allowing a Transformer encoder to focus on long-range dependencies. Unlike [1], our objective was not to perform extensive hyperparameter tuning, which is also the case with other Graph Transformer works that are consistently outperformed by [1]. However, we acknowledge the contribution of [1] and its strong results. Nonetheless, as shown we either outperform or on par with their results. our To ensure a fair comparison, we performed additional tuning on our method.

We sincerely thank you for your time! Hope we have addressed your concerns through the new experiments. We look forward to your reply and further discussions, thanks!

[1] Luo, Yuankai, Lei Shi, and Xiao-Ming Wu. "Classic GNNs are Strong Baselines: Reassessing GNNs for Node Classification." arXiv preprint arXiv:2406.08993 (2024).
[2] Dwivedi, Vijay Prakash, et al. "Long range graph benchmark." Advances in Neural Information Processing Systems 35 (2022): 22326-22340.

We thank the reviewer for their positive response and constructive feedback. We address the raised questions as follows.

Vector Quantization

We have added an algorithm and detailed explanation of how graph nodes are mapped into tokens in the Appendix. A brief description is also provided in Section 3, Lines 153-161. In summary, the process involves passing the input graph through a GNN encoder to learn node representations. Each node representation is then compared against the embeddings of the first codebook, and the index of the closest embedding to each node’s representation is assigned to it as its first discrete token. The difference between the node representation and the assigned codebook embedding is then compared to the embeddings of the second codebook to find the second discrete token. This process repeats for all codebooks. $X_Q$ is a matrix of size $n \times c$, where n is the number of nodes and c is the number of codebooks. This matrix essentially holds the discrete tokens for nodes. For example, if we have 1000 nodes and 3 codebooks with 256 embeddings, each node will be represented by 3 discrete tokens where each token is an integer between 0-255. Thus, $X_Q \in \{0, 1, \cdots, 255\}^{1000 \times 3}$. $T_v$ in Line 296 is a row of this matrix that represents the 3 discrete tokens of node $v$. We have slightly changed the notation to make it clearer.

Embedding Matrix

The embedding matrix is not learned during training the tokenizer; instead, it is learned during training the Transformer. Following the example provided in the previous response, if we set the input dimension to the Transformer to d, then $X_T$ is a matrix of size $3 \times 256 \times d$, which is essentially the embedding table for discrete tokens. The embeddings are initialized randomly. As mentioned in Lines 333-334, we found that initializing this embedding table randomly and introducing a fourth token by summing up the embeddings of the codebooks for each node (Eq. 14) yields better results compared to initializing the embedding table with the embeddings of codebooks. We have updated the paper to clarify this point.

[.,.] Operator

Thank you for pointing out this. We have added a sentence to explain it in the updated paper.

For Eq 14 $C[i, t_i^v]$, $i$ and $t_i^v$ are integers, and this operator is essentially the indexing operator on tensors. For example, $C[i, j]$ returns the embedding of the $j$-th entry in the $i$-th codebook.

For Eq 15, we use [ || ] to show the concatenated sequence.

Supportive Analysis

We would like to clarify that we do not claim that the tokenizer captures long-range dependencies. Instead, we claim that the tokenizer uses a GNN to model local interactions, allowing the Transformer to focus on capturing long-range dependencies. This is analogous to Vision Transformers (ViTs) [1], where instead of overwhelming the Transformer with pixels, a linear or convolutional layer first models the local information on image patches. Nevertheless, we have conducted new experiments on long-range graph benchmarks to show that our model enables capturing such interactions. We have also addressed generalization and memory footprint in the general response.

Different sizes of Codebooks

We treat the number of codebooks, codebook size, and embedding as hyperparameters and perform lightweight tuning for each dataset. As mentioned in the Appendix, we choose the number of codebooks from {1, 2, 3, 6, 9}, codebook size from {128, 256, 512, 1024}, and hidden dimensions from {128, 256, 512, 1024}. We have also added a Table to the appendix that reports the chosen hyperparameters for each dataset. To address your question, we have conducted experiments on ogbn-arxiv and Minesweeper datasets to investigate the effect of varying these hyperparameters. The results shown below suggest that as we increase each of these hyperparameters, the performance tends to increase first and then decrease, respectively. We will add more analysis to the camera-ready version.

Node Clustering

Vector Quantization (VQ) is closely related to K-Means clustering, and the embeddings of the codebook can be interpreted as centroids of the clusters. A good heuristic for initializing the VQ codebook is, in fact, performing K-Means clustering on the first batch of data. However, unlike classic algorithms such as K-Means, VQ performs representation learning and clustering end-to-end. In this work, we used Residual Quantization (RQ) which is related to hierarchical clustering, where the embeddings of the first codebook can be interpreted as centroids of coarse-grained clusters, and the embeddings of the last codebook are centroids of fine-grained clusters. To investigate this further, we trained a variant of our tokenizer by removing the quantization layer. Once the model was trained, we clustered the codebooks using K-Means and used the cluster indices as tokens. The results are shown below. As the results suggest, RQ achieves better performance because clustering and representation learning are performed end-to-end. Meanwhile, the results confirm our findings in Table 4, indicating that multi-task SSL objectives have a more impact on downstream performance. Additionally, they indicate that the learned representations by the tokenizer are expressive, allowing us to achieve even good performance by clustering them using K-Means.

References

[1] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2020.

I thank the authors for their efforts during the rebuttal and for providing detailed responses. After reviewing the clarifications and additional explanations, my initial concerns have been addressed.

However, I noticed that the authors have only added a comparison of inference time in the appendix. Given that GQT relies on a two-stage model training process, how does its training time compare to the single-stage GTs? Based on the current results, the performance improvement of GQT appears to be limited. It remains to be demonstrated whether the additional time cost for codebook training is justified. Can the learned codebooks generalize to other datasets?

Additionally, -philous graphs in the manuscript should be -philic graphs.

Thank you for the constructive engagement. We will definitely update the paper based on the feedback we received in rebuttal period. Please let us know if you had any further questions and we are more than happy to answer them. Thanks again.

Thank you for the feedback and questions. Please find our responses as follows:

Novelty

We would like to clarify that our goal is not to introduce a new graph self-supervised objective or a novel quantization layer, as each of these topics warrants its own independent paper. Instead, our goal is to investigate an effective and transferable graph tokenization strategy that enables the learned tokens to encode local interactions, allowing a Transformer encoder to focus on long-range dependencies. To achieve this, we propose using state-of-the-art graph self-supervised objectives in a multitask fashion to capture various aspects of the interactions (as shown in Table 4, each objective contributes to the overall performance). Additionally, we have conducted new experiments showing that our method is robust and achieves strong performance across a diverse set of graph tasks. We believe our proposed approach will have a significant impact on the graph learning community by enabling researchers to leverage the accelerating advancements from the NLP and CV communities in designing more efficient Transformers.

Robustness and Generalization

Please refer to our new experiments discussed in the general response to this question. In summary, we have performed new inference time attacks to evaluate robustness and conducted comparisons with RQ-VAE to assess generalization.

Efficiency

Please refer to our new experiments discussed in the general response to this question.

Other Graph Tasks

Please refer to our new experiments discussed in the general response to this question. In summary, we have conducted new experiments on a diverse set of graph tasks including graph classification and regression, inductive node classification, and inductive link prediction on long-range graph benchmarks demonstrating the strong performance of our model.

Dear Reviewer rTiq,

Thank you again for your constructive comments. As the rebuttal period is coming to an end in one day, we would like to take this opportunity to kindly ask if you have any further questions based on our responses. We are more than happy to answer any additional questions during your re-evaluation.

Thank you very much for your positive feedback. We would love to address any further questions you may have or clarify any part that is unclear to increase your confidence in your evaluation and in our work. Thanks again and looking forward to having a constructive discussion with you during the rebuttal period.

Efficiency (reviewers rTiq and cy7L)

To address concerns about memory footprint, we have already highlighted two examples in Lines 483-485 and 426-429, showing significant memory reduction when using discrete tokens instead of node features. For instance, on the ogbn-products dataset with 2,449,029 nodes and 100-dimensional node features, GQT requires only 3 codebooks of size 4096 each to represent the tokens, resulting in a remarkable 30-fold reduction in memory usage. Please note that this memory reduction occurs after training the tokenizer. As the encoder of the tokenizer is a GNN that processes the graph with original node features, its memory footprint is on par with any arbitrary GNN. However, since the Transformer encoder only consumes discrete token IDs, which are significantly fewer than the total number of nodes, we achieve a substantial reduction in memory footprint. In the camera-ready version, we will include a table in the appendix to illustrate the improvement for all datasets. As an additional experiment, we compare the inference time and memory usage between our Transformer encoder and a Graph Attention Network (GAT) when performing inference on all graph nodes. The results below show that while our Transformer is on par with a sparse implementation of GAT in terms of inference time, it requires half the GPU memory.

References

[1] Fu et al. VCR-Graphormer: A Mini-batch Graph Transformer via Virtual Connections. ICLR 2024

[2] Chen et al. NAGphormer: A Tokenized Graph Transformer for Node Classification in Large Graphs. ICLR 2023

[3] Wu et al. NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification. NeurIPS 2022

[4] Dwivedi et al. Long Range Graph Benchmark. NeurIPS 2022, Datasets and Benchmarks Track

[5] Wang et al. Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces. Arxiv 2024

[6] Geisler et al. Robustness of Graph Neural Networks at Scale. NeurIPS 2021

[7] ​​Lee et al. Autoregressive Image Generation using Residual Quantization. CVPR 2022