Learning Graph Quantized Tokenizers

Authors: Limei Wang, Kaveh Hassani, Si Zhang, Dongqi Fu, Baichuan Yuan, Weilin Cong, Zhigang Hua, Hao Wu, Ning Yao, Bo Long
Primary area: learning on graphs and other geometries & topologies
Keywords: Graph Learning, Graph Tokenization, Graph Transformer

Abstract:
Transformers serve as the backbone architectures of Foundational Models, where domain-specific tokenizers allow them to adapt to various domains. Graph Transformers (GTs) have recently emerged as leading models in geometric deep learning, outperforming Graph Neural Networks (GNNs) in various graph learning tasks. However, the development of tokenizers for graphs has lagged behind other modalities, with existing approaches relying on heuristics or GNNs co-trained with Transformers. To address this, we introduce GQT (\textbf{G}raph \textbf{Q}uantized \textbf{T}okenizer), which decouples tokenizer training from Transformer training by leveraging multi-task graph self-supervised learning, yielding robust and generalizable graph tokens. Furthermore, the GQT utilizes Residual Vector Quantization (RVQ) to learn hierarchical discrete tokens, resulting in significantly reduced memory requirements and improved generalization capabilities. By combining the GQT with token modulation, a Transformer encoder achieves state-of-the-art performance on 20 out of 22 benchmarks, including large-scale homophilic and heterophilic datasets. The implementation is publicly available at \href{https://github.com/limei0307/GQT}{https://github.com/limei0307/GQT}.

Decision: accept (poster)
Reviewer scores: 8, 5, 6, 5, 8

Reviews and discussion:
<thread>
creation date:
1730647027148

thread_type:
official_review

<thread_object>
creator: reviewer_d7ql

<content>
summary:
In this work, the authors propose GQT, a transformer-based framework for graph prediction tasks. GQT consists of a variety of proposed components, including a custom training loss and a tokenizer trained separately with a self-supervised learning approach. The tokenizer leverages vector quantization, resulting in a highly efficient tokenization that can be fed into a standard transformer encoder and applied to very large graphs with up to 2.5 million nodes. Empirically, GQT outperforms a variety of GNN and transformer baselines on 16 out of 18 transductive node classification problems, including both homophilic and heterophilic graphs.

strengths:
- **S1**: The authors evaluate their model on a wide variety of transductive node classification datasets, both homophilic and heterophilic and up to 2.5 million nodes. The empirical results are impressive and indicate that GQT is a versatile approach to transductive node classification.
- **S2**: The authors present an ablation study in Table 4 and Table 5 on the various components of their approach. The study indicates that most components help with improving the overall model performance; see **Q1** for a follow-up question on the ablation study.
- **S3**: The approach for learning a tokenizer using vector quantization is highly interesting and should be further explored in future research. In particular, I expect that GQT can be readily adapted to other types of graph learning problems (such as inductive settings) while being highly efficient.

weaknesses:
- **W1**: While the results are impressive, the authors evaluate exclusively on transductive node classification problems. In my opinion, the authors should be more upfront about this limitation or include results for inductive settings, e.g., the inductive tasks in [https://arxiv.org/abs/2206.08164](https://arxiv.org/abs/2205.12454).
- **W2**: Section 4.2 introduces the training loss with three different loss terms. However, the authors do not provide any intuition or justification for these choices. As such, it is difficult for a reader to understand why these choices were made and how crucial they are to the overall framework. A similar issue appears in Sections 5.1 and 5.2.
- **W3**: In the introduction in L35-36, the authors state that “GTs possess an expressive power equivalent to the 2-Weisfeiler-Lehman (WL) isomorphism test”, citing Kim et al. (https://arxiv.org/abs/2207.02505).  The transformer used in Kim et al. is entirely different from the transformers used in transductive node classification and in particular, different from GQT. As such, this reference does not serve as a motivation to the approach present in this work. Further, note that the GQT uses a GNN to learn local interactions and would suffer from the same expressivity limitations as GNNs, leading to the next statement in L37-38: “This [the 2-WL expressivity of GTs] surpasses the expressive power of message-passing GNNs, which are limited to the 1-WL test”. This statement is incorrect. The 1-WL and 2-WL tests as defined in Kim et al. are equivalent in expressive power. Finally, in L38-39, the authors state that “[…] a Transformer with sufficient attention heads can match or exceed the expressive power of a second-order invariant graph network, outperforming message-passing GNNs”, again citing Kim et al. Here, the authors omit the fact that transformers can only achieve said result when receiving as input $O(n^3)$ number of tokens, resulting in a transformer with $O(n^6)$ runtime and memory complexity. Again, this approach is entirely different to GQT and does not serve as a motivation for GQT. I strongly recommend the authors to adjust the introduction, remove the incorrect claims and provide a proper motivation for their approach.

questions:
- **Q1**: Which components of GQT are also applicable to baselines? For example, while the tokenizer is certainly unique to GQT, adding semantic edges and PPR features to nodes may also improve baseline performance. I wonder whether the authors can give any indication as to whether the improved performance of GQT stems largely from components unique to GQT or whether their proposed components may be useful in improving the results of transformers for transductive node classification in general.
- **Q2**: Table 4 and 5: Are the standard deviations missing from these tables or was the ablation study performed only on a single seed? In case the study was performed on a single seed, I strongly recommend to repeat the ablation study for at least three seeds to verify that the results are statistically significant.
- **Q3**: Figure 1, the experimental section, the ablation studies as well as the conclusion mention “structural gating”. However, this term is never formally introduced. I suppose the authors refer to the gating mechanism in Equation 16. If this is the case, I encourage the authors to define properly define the term there.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712189985

thread_type:
official_review

<thread_object>
creator: reviewer_fmi4

<content>
summary:
This work attempt to design a graph tokenizer that utilizes graph self-supervised learning to train a encoder. And then leverage the Transformer to capture long-range dependencies. Specifically, the authors adapts Residual Vector Quantization to learn hierarchical discrete tokens for reducing memory requirements and improving generalization capabilities.

strengths:
1. The proposed Graph Quantized Tokenizer utilizes discrete tokens, which can significantly reduce memory requirements.

weaknesses:
1. The major concern of this manuscript is the experiment part. According to [1], classic GNNs (e.g. GCN, GraphSAGE and GAT) can also achieve excellent performance on node classification tasks. However, the proposed GQT obtains lower performance on almost all datasets. This makes me worry about the actual contribution of this work.

2. The novelty of the proposed method is limited. There are already some works try to utilize vector quantization to learn discrete tokens for graph learning [2,3]. And there are also existing work utilize graph augmentation and Personalized PageRank to generate a sequence per node [4].

3. There are also some concept errors in the manuscript. For example, the authors mention that "Recent studies on applying Large Language Models (LLMs) to graph-related tasks have found the representing ..." in the introduction section. But most recent methods about combining LLMs with graph learning can only work on textual-attributed graph, and they just leverage LLMs to capture better textual  understanding [5, 6].


[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] Yang, Ling, et al. "VQGraph: Rethinking graph representation space for bridging GNNs and MLPs." The Twelfth International Conference on Learning Representations. 2024.

[3] Xia, Jun, et al. "Mole-bert: Rethinking pre-training graph neural networks for molecules." (2023).

[4] Fu, Dongqi, et al. "VCR-graphormer: A mini-batch graph transformer via virtual connections." arXiv preprint arXiv:2403.16030 (2024).

[5] Li, Yuhan, et al. "A survey of graph meets large language model: Progress and future directions." arXiv preprint arXiv:2311.12399 (2023).

[6] Jin, Bowen, et al. "Large language models on graphs: A comprehensive survey." IEEE Transactions on Knowledge and Data Engineering (2024).

questions:
1. Could the authors provide the explanation about why GQT obtains lower performance than classic GNNs?

2. I think the authors need to re-emphasize the innovation of this work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712382263

thread_type:
official_review

<thread_object>
creator: reviewer_cy7l

<content>
summary:
This paper proposes a self-supervised method, named GQT, for graph token learning. GQT adopts the residual vector quantization to generate graph tokens and trains the model with mutual information maximization, graph reconstruction loss and distillation loss minimization. Given the learned codebook, GQT augments the input graph with edges between similar nodes and generates a node sequence for each node with topK-related nodes in personalized PageRank. Extensive results on node classification are provided to demonstrate the effectiveness of GQT.

strengths:
- This paper decouples the graph tokenizer from transformer training, which poses an interesting thread for developing graph transformers.
- The proposed method achieves competitive performance on node classification, including heterophilic and homophilic datasets.
- GQT scales to benchmarks with large graphs.

weaknesses:
- As the core module in GQT, the authors did not provide a sufficient description of the residual vector quantization (nor in the appendix). How does GQT learn the mapping from graph nodes to the codebook tokens? The mapping relation $X_Q$ is introduced in L194/199 and never used again in the following context.
- The writing of section 5.2 is rather arbitrary. I cannot fully comprehend the embedding matrix $X_T$, which is 'trained end-to-end with the Transformer'. How does $X_T$ be learned in GQT? In L333-334, the authors argue that 'adding this explicit aggregated token leads to better performance compared to initializing $X_T$ with $C$. Does this mean that $X_T$ is initialized with $C$ or not?
- The operator $[\cdot,\cdot]$ in Eq 14/15 lacks a definition.
- The GQT (or tokenizer) is claimed to (1) capture long-range dependencies; (2) less memory footprint; (3) improved generalization capabilities, which lack of supportive analysis in the experiments.

questions:
-  How does GQT perform with different sizes of codebooks ($K$, $c$ and $d_c$).
- The graph codebook learning resembles node clustering. Can authors provide a comparison between these two types of methods? Or how does node clustering perform in the graph codebook learning?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719420910

thread_type:
official_review

<thread_object>
creator: reviewer_rtiq

<content>
summary:
The paper proposes Graph Quantized Tokenizer (GQT) that considers three key properties such as local interactions, memory efficiency, and robustness & generalization for graph Transformer. Different existing works that train tokenizer and Transformer architecture in an end-to-end manner, GQA decouples the training tokenizer from the Transformer. In other words, it first trains the tokenizer and then learns graph Transformer architecture given the tokenized graph. From the experimental results, GQA shows the best performance compared to other baselines on 7 out of 8 datasets. It also achieves the best performance on large-scale datsets.

strengths:
- The design of tokenizers in graph Transformers is really an important research topic. The paper well describes the existing works of graph Transformers and tokenizers in the related works.
- The paper is easy to follow.
- The experiments of the paper demonstrate that the proposed method is effective with multiple experiments including homophilious, heterophilious and large-scale datasets.

weaknesses:
- I think that the novelty of the paper is limited.
   - One of the main contributions highlighted in this paper is quantized tokenization for graph Transformers. However, the paper simply combines graph neural networks for dealing with graph-structured data and existing residual vector quantization for the quantization. 
   - Another contribution is multi-task self-supervised learning. The paper uses two self-supervised learning objectives such as DGI and GMAE2, which are not proposed self-supervised training scheme in this paper. I think that using two self-supervised objectives to improve the effectiveness of the model is hard to become a good contribution in top-tier ML conferences.

- More experiments need to be included.
   - The author claims that the proposed quantization method improves the robustness and generalization of the model. But, there is no experimental result concerning the robustness and generalization.
   - In the similar manner, the validation of efficiency needs to be added to empirically prove that the proposed quantization is effective and efficient.
   - Also, the paper only deals with node classification tasks. It would be better if the paper included other graph-related tasks, such as node regression and graph classification, which have been actively dealt with in other graph Transformer papers.

questions:
Please refer to the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731102599734

thread_type:
official_review

<thread_object>
creator: reviewer_ctgj

<content>
summary:
the paper proposes a graph tokenizer that utilizes multi-task graph self-supervised objectives to train a graph encoder, which can capture local interaction and have memory efficiency

strengths:
1. paper is well written and easy to follow
2. experiment demonstrate strong results compared with baseline methods.
3. motivation and methodology is clear

weaknesses:
NO

questions:
NO
</content>
</thread_object>

</thread>
<thread>
creation date:
1732255569458

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to AC and Reviewers - Part 1

comment:
We extend our gratitude to the reviewers for their constructive feedback and insightful comments. We are pleased that the reviewers recognized the interest and relevance of our work. We acknowledge the common concerns raised by the reviewers and would like to address them in this general response. Additionally, we have updated the paper based on the feedback. All changes are highlighted in **red** in the updated manuscript.

**Performance on diverse graph tasks (reviewers rTiq, cy7L, D7qL)**

We would like to thank the reviewers for this great suggestion, which we believe will significantly improve the impact and quality of our paper. We initially focused on **transductive node classification** tasks, as many graph Transformer baselines [1-3] focus on these tasks. To further evaluate our model, we assessed its performance on graph classification, graph regression, link prediction, inductive node classification, and long-range modeling benchmarks from the **Long Range Graph Benchmark (LRGB)** [4], which is proposed to evaluate the capability of GNNs and GTs to capture long-range interactions. We used the Peptides-Func dataset for **graph classification** with Average Precision (AP) metric, the Peptides-Struct dataset for **graph regression** with Mean Absolute Error (MAE) metric (**lower is better**), the COCO-SP for **inductive node classification** with macro F1 metric, and the PCQM-Contact for **link prediction** with Mean Reciprocal Rank (MRR) metric. We compare our results to baselines reported in [5]. The results presented in the table below suggest that GQT coupled with the Transformer encoder can achieve promising results on all these tasks. In the camera-ready version, we will provide more comprehensive results by reporting our model’s performance on more datasets for each task. Notably, we did not modify the tokenizer and made only slight adjustments to the Transformer encoder, as discussed in Lines 149-152.

| Task | **Graph classification** | **Graph regression** | **Inductive Node classification** | **Inductive Link prediction** |
|---|---|---|---|---|
| Dataset | Peptides-Func | Peptides-Struct | COCO-SP | PCQM-Contact |
|Metric | AP ↑ | MAE ↓ | F1 ↑ | MRR ↑ |
| GCN | 0.5930±0.0023 | 0.3496±0.0013 |  0.0841±0.0010 | 0.3234±0.0006 |
| Exphormer | 0.6258±0.0092 | 0.2512±0.0025 | 0.3430±0.0108 | 0.3587± 0.0025 |
| GPS | 0.6535 ± 0.0041 | 0.2500 ± 0.0005 | 0.3412 ± 0.0044 | 0.3337 ± 0.0006 |
| Graph-Mamba | 0.6739±0.0087 | 0.2478±0.0016 | 0.3960±0.0175 | 0.3395± 0.0013 |
| Ours | 0.6903±0.0085 | 0.2452±0.0018 | 0.4007±0.0125 | 0.3427± 0.0012 |

**Robustness and Generalization (reviewers rTiq and cy7L)**

We would like to thank the reviewers for this great suggestion. To measure robustness, we use **PRBCDA and GRBCD inference time attacks** [6] to assess the performance drop during inference. A **lower drop in performance indicates greater robustness**. To this end, we compare the encoder of our tokenizer with the encoder of a tokenizer trained with a reconstruction objective (RQ-VAE) [7]. The results shown below suggest that our proposed tokenizer is more robust to attacks due to the multi-task graph-specific SSL objectives, as opposed to the standard RQ-VAE, which uses a reconstruction objective. In camera-ready version, we will provide additional results on other datasets.
 
| | GQT (ours) | RQ-VAE |
|---|---|---|
| GR-BCD Attack: Accuracy drop (PubMed) ↓     | 15.8% | 20.4% |
| GR-BCD Attack:Accuracy drop (ogbn-arxiv) ↓ | 10.4% | 14.8% |
| PR-BCD Attack: Accuracy drop (PubMed) ↓     | 18.1% | 23.3% |
| PR-BCD Attack: Accuracy drop (ogbn-arxiv) ↓ | 11.3% | 17.2% |

To measure improved generalization, we follow the common practice of treating **downstream predictive performance as a proxy for generalization**. As shown in Table 4 of the paper, every component of the tokenizer, including both SSL objectives and the quantization layer, contributes to the downstream predictive performance, thereby improving the model’s generalizability. Furthermore, to evaluate the contribution of multi-task SSL objectives to downstream performance, we compare our results with those of a tokenizer trained based on the RQ-VAE [7] design, which employs a reconstruction objective. The results presented below indicate that using multi-task SSL objectives significantly improves downstream predictive performance, which is strongly correlated with the method's generalization. In the camera-ready version, we will provide additional results on other datasets.

| Tokenizer | GQT (ours) | RQ-VAE |
|---|---|---|
| ogbn-arxiv | 73.14±0.16 | 66.05 ± 0.48 |
| Minesweeper | 95.28±0.44 | 89.69 ± 0.35 |
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:**

The paper proposes a self-supervised method for graph token learning, Graph Quantized Tokenizer (GQT). The proposed framework requires two-stage training; the tokenizer and the model are separately trained. Specifically, the authors adopt Residual Vector Quantization to learn hierarchical discrete tokens. Extensive experiments demonstrate that the proposed framework improves memory efficiency and generalizability, achieving superior performance across various datasets including large-scale datasets.

**Strengths:**

1. **Strong empirical results and versatility.** The authors demonstrate the superiority of their model across various transductive node classification datasets, including both homophilic and heterophilic graphs, and large-scale graphs with 2.5 million nodes. Also additional experimental results on graph-level, and link-level tasks support the versatility of the proposed method.
2. **Memory-efficient graph transformer.** The proposed Graph Quantized Tokenizer utilizes discrete tokens, which can significantly reduce memory requirements.
3. **Motivation.** Designing and learning tokenizers in graph Transformers are important research topics. The authors proposed a learnable tokenizer for graphs using vector quantization. Decoupling the training of the graph tokenizer and graph transformer is interesting to further explore. 
4. **Presentation.** Paper is well-written and provides comprehensive related works.

**Weaknesses:**

1. **Limited technical novelty.** This paper successfully combines existing modules, but the modules are not newly developed in this paper. The proposed pipeline can be viewed as a simple combination of existing graph Transformers and residual vector quantization. Also loss functions for self-supervised learning have already been studied in the literature. 

**Main reasons:**

This paper studies a graph transformer equipped with a learnable tokenizer to improve performance and enable large-scale graph analysis. Although the modules are mainly from existing works in the literature, the proposed pipeline is well-motivated and shows strong performance. Also, the authors addressed most concerns raised by reviewers via responses with empirical results.