Technical Summary: Graph is a Natural Regularization: Revisiting Vector Quantization for Graph Representation Learning

1. Problem Statement

Vector Quantization (VQ) has emerged as a powerful technique for learning compressed, discrete representations in vision and language domains, offering a pathway toward scalable Graph Foundation Models (GFMs) by tokenizing graph structures. However, this paper identifies a fundamental and underexplored challenge specific to the graph domain: codebook collapse.

Codebook collapse occurs when the majority of input data points are mapped to only a few codewords in the learnable codebook, leaving the vast majority of the codebook underutilized. The authors observe that this phenomenon is systematic in Graph VQ, occurring even when applying state-of-the-art (SOTA) mitigation strategies developed for vision and language (e.g., Exponential Moving Average, codebook reset, pretrained encoders). This collapse severely limits the expressiveness and generalization of graph tokens, leading to coarse representations and degraded downstream task performance.

The paper diagnoses the root causes of this collapse from two perspectives:

1. Data Perspective: Graph data possesses inherent properties such as feature redundancy and high connectivity density. The study finds a positive correlation between feature redundancy (measured by PCA@95%) and the severity of collapse, and a negative correlation between connectivity density (average node degree) and codebook utilization. Dense connectivity and redundant features exacerbate the tendency for nodes to map to identical tokens.
2. Optimization Perspective: Deterministic VQ relies on hard assignment (nearest neighbor selection). This creates a self-reinforcing feedback loop: frequently selected codewords receive more gradient updates and become increasingly dominant, while rarely selected codewords receive no updates and remain inactive. This "cocoon effect" locks the encoder into a narrow set of tokens, preventing the exploration of the full codebook capacity.

2. Methodology: RGVQ

To address these issues, the authors propose RGVQ (Regularized Graph Vector Quantization), a framework that integrates graph topology and feature similarity as explicit regularization signals. The method consists of two core components:

A. Gumbel-Softmax Reparameterization

To break the self-reinforcing loop of deterministic hard assignment, RGVQ replaces the non-differentiable argmax operation with a Gumbel-Softmax reparameterization.

• Instead of selecting a single nearest codeword, the model computes assignment logits based on the distance between node embeddings and all codewords.
• These logits are perturbed with Gumbel noise and passed through a temperature-controlled softmax to produce a differentiable probability distribution over the codebook.
• The quantized embedding is computed as a weighted average of all codewords based on these probabilities.
• Effect: This ensures that gradients flow to all codewords proportionally to their assignment probabilities, not just the selected one, allowing inactive codewords to receive updates and be reactivated.

B. Structure-Aware Contrastive Regularization

To mitigate the effects of graph data redundancy and dense connectivity, RGVQ introduces a regularization term that explicitly penalizes overly concentrated token usage among dissimilar nodes.

• Positive Set ($N_P$): For a given anchor node, the positive set includes nodes that are either structurally connected (neighbors) or semantically similar (top-$K$ feature-similar nodes).
• Negative Set ($N_N$): Includes nodes that are neither connected nor feature-similar.
• Loss Function: The model employs an InfoNCE loss to encourage nodes in the positive set to share similar token assignment distributions while penalizing nodes in the negative set for having overlapping distributions.
• Effect: This regularization discourages the model from mapping structurally or semantically distinct nodes to the same token, thereby promoting token diversity and preventing collapse induced by feature redundancy.

The total loss function combines the standard reconstruction loss (feature and link reconstruction), codebook loss, commitment loss, and the new structure-aware regularization loss.

3. Key Contributions

The paper makes the following contributions:

1. Systematic Empirical Study: It provides the first systematic benchmarking of codebook collapse in Graph VQ, demonstrating that standard mitigation strategies from other domains fail to resolve the issue in graph data.
2. Diagnosis of Collapse: It identifies that codebook collapse in graphs is driven by intrinsic data properties (redundancy and connectivity) and the optimization dynamics of deterministic hard assignment.
3. RGVQ Framework: It proposes a novel, differentiable framework that combines Gumbel-Softmax reparameterization with structure-aware contrastive regularization to enhance codebook utilization.
4. Comprehensive Evaluation: It validates RGVQ across multiple SOTA Graph VQ backbones and downstream tasks, showing consistent improvements in both codebook utilization and task performance.

4. Experimental Results

The authors conducted extensive experiments on various homophilous and heterophilous graph datasets (e.g., Cora, PubMed, WikiCS, Amazon-Ratings, Roman-Empire).

• Codebook Utilization: RGVQ significantly outperforms vanilla Graph VQ and its variants (including EMA, affine parameters, codebook reset, and pretrained encoders) in terms of codebook perplexity. For instance, on the PubMed dataset, while vanilla VQ achieved a perplexity of ~4.14 (indicating severe collapse), RGVQ achieved ~319.09, approaching the theoretical maximum. RGVQ maintained high utilization even as codebook sizes increased.
• Transferability: When integrated into the Graph Foundation Model (GFT), RGVQ achieved state-of-the-art performance across cross-task and cross-domain benchmarks (node classification, link prediction, graph classification), outperforming both supervised GNNs and other self-supervised methods.
• Serialization: In the context of Graph Quantized Transformers (GQT), RGVQ-enhanced models consistently improved node classification accuracy compared to GQT with other anti-collapse methods.
• Ablation Studies:
  • Removing either the Gumbel-Softmax reparameterization or the structure-aware regularization led to a significant drop in performance, confirming that both components are mutually dependent for preventing collapse.
  • The method was robust to variations in temperature ($\tau$) and the number of contrastive samples, with low, fixed temperatures proving sufficient.
  • RGVQ demonstrated robustness across varying levels of graph heterophily.

5. Significance and Claims

The paper claims that codebook collapse is not an incidental issue but a fundamental bottleneck in discrete graph token learning. By treating the graph structure itself as a natural regularizer, RGVQ offers a flexible "plug-in" solution that enhances the expressiveness and transferability of graph tokens.

The authors emphasize that their work does not propose a new generative model architecture per se, but rather a methodological improvement to the quantization process that enables more effective learning of graph representations. The significance lies in enabling the construction of more robust Graph Foundation Models by ensuring that the learned token vocabulary is fully utilized, thereby capturing finer-grained structural and semantic patterns essential for downstream tasks. The paper concludes that RGVQ provides a flexible framework that maintains token diversity and scales effectively to larger codebook sizes, addressing a critical limitation in current graph representation learning.