Hierarchical Vector Quantized Graph Autoencoder with Annealing-Based Code Selection

TL;DR

Graph self-supervised learning often relies on perturbations that can distort structure. This paper investigates applying Vector Quantized VAE (VQ-VAE) to graphs and identifies two challenges: codebook underutilization and codebook space sparsity. It introduces HQA-GAE, which combines an annealing-based code selection to encourage broad code usage and a hierarchical two-layer codebook to capture relationships between embeddings. Across eight diverse datasets and multiple tasks, HQA-GAE achieves state-of-the-art results on link prediction and competitive performance on node classification, while maintaining scalable training. These findings demonstrate that discrete latent representations can effectively capture graph topology in SSL without perturbation-based view generation.

Abstract

Graph self-supervised learning has gained significant attention recently. However, many existing approaches heavily depend on perturbations, and inappropriate perturbations may corrupt the graph's inherent information. The Vector Quantized Variational Autoencoder (VQ-VAE) is a powerful autoencoder extensively used in fields such as computer vision; however, its application to graph data remains underexplored. In this paper, we provide an empirical analysis of vector quantization in the context of graph autoencoders, demonstrating its significant enhancement of the model's capacity to capture graph topology. Furthermore, we identify two key challenges associated with vector quantization when applying in graph data: codebook underutilization and codebook space sparsity. For the first challenge, we propose an annealing-based encoding strategy that promotes broad code utilization in the early stages of training, gradually shifting focus toward the most effective codes as training progresses. For the second challenge, we introduce a hierarchical two-layer codebook that captures relationships between embeddings through clustering. The second layer codebook links similar codes, encouraging the model to learn closer embeddings for nodes with similar features and structural topology in the graph. Our proposed model outperforms 16 representative baseline methods in self-supervised link prediction and node classification tasks across multiple datasets.

Figures (10)

• Figure 1: The schematic diagram of HQA-GAE, including graph encoding and decoding, annealing-based code selection, hierarchical codebook learning.
• Figure 2: The performance of node classification tasks for both normal GAE and VQ-GAE, where an MLP is used as the encoder and the decoder remains unchanged as a GNN.
• Figure 3: The effect of annealing-based code selection.
• Figure 4: Visualization of node representations and codebook embeddings on CiteSeer by t-SNE.
• Figure 5: Comparison of performance, model size, and pre-training time for link prediction on the Physics dataset. Each point's size represents the model size.