PaperlayerView exact revision on arXiv

Source-linked AI summary

VQGraph: Rethinking Graph Representation Space for Bridging GNNs and MLPs

Ling Yang, Ye Tian, Minkai Xu, Zhongyi Liu, Shenda Hong, Wei Qu, Wentao Zhang, Bin Cui, Muhan Zhang, Jure Leskovec

arXiv:2308.02117v3cs.LGcs.AIcs.CV

TL;DR

GNN-to-MLP distillation uses efficient MLP inference, but class-label targets may not represent the diverse local graph structures that GNNs encode. VQGraph learns a structure-aware codebook and distills soft code assignments from GNNs to MLPs. Across seven datasets, it reports higher average accuracy than GNNs, MLPs, and NOSMOG, with inference 828× faster than GNNs.

  • Problem

    Existing GNN-to-MLP methods distill predictions over a few class labels, which may not express the diverse local graph structures needed for structural knowledge transfer.

    Source: §1 · p. 2

  • Method

    VQGraph learns a codebook of discrete node-structure codes with a VQ-VAE variant and trains an MLP to match the teacher GNN’s soft code assignments.

    Source: §1 · p. 2, §4.2 · p. 6

  • Results

    VQGraph improves average accuracy over GNNs by 3.90%, MLPs by 28.05%, and NOSMOG by 1.39%, while inferring 828× faster than GNNs.

    Source: §6 · p. 11

  • Takeaways & Limitations

    The reported results show that VQGraph combines higher average accuracy than GNNs with substantially faster inference than GNNs.

    Source: §6 · p. 11

Abstract

from arXiv · show

GNN-to-MLP distillation aims to utilize knowledge distillation (KD) to learn computationally-efficient multi-layer perceptron (student MLP) on graph data by mimicking the output representations of teacher GNN. Existing methods mainly make the MLP to mimic the GNN predictions over a few class labels. However, the class space may not be expressive enough for covering numerous diverse local graph structures, thus limiting the performance of knowledge transfer from GNN to MLP. To address this issue, we propose to learn a new powerful graph representation space by directly labeling nodes' diverse local structures for GNN-to-MLP distillation. Specifically, we propose a variant of VQ-VAE to learn a structure-aware tokenizer on graph data that can encode each node's local substructure as a discrete code. The discrete codes constitute a codebook as a new graph representation space that is able to identify different local graph structures of nodes with the corresponding code indices. Then, based on the learned codebook, we propose a new distillation target, namely soft code assignments, to directly transfer the structural knowledge of each node from GNN to MLP. The resulting framework VQGraph achieves new state-of-the-art performance on GNN-to-MLP distillation in both transductive and inductive settings across seven graph datasets. We show that VQGraph with better performance infers faster than GNNs by 828x, and also achieves accuracy improvement over GNNs and stand-alone MLPs by 3.90% and 28.05% on average, respectively. Code: https://github.com/YangLing0818/VQGraph.

1 INTRODUCTION

GNN-to-MLP distillation aims to retain GNN structural knowledge while using efficient MLP inference, but class-label targets may not capture diverse local structures. VQGraph labels local structures with discrete codes and distills soft code assignments; the paper reports average accuracy gains over GNNs, MLPs, and NOSMOG alongside faster inference than GNNs.

1 INTRODUCTION

  • MLPs offer efficient graph inference because they use node features without explicit message passing, which otherwise requires fetching neighboring nodes’ topology and features.The passage describes neighbor fetching for target-node inference as time-consuming and computation-intensive.Source: §1 · pp. 1–2
  • Class-label targets may not express the diverse local structures of nodes, limiting transfer of structural knowledge from GNNs to MLPs.The earlier representation space may miss fine-grained structural differences between nodes of the same class.Source: §1 · p. 2
  • VQGraph uses a VQ-VAE variant to encode node substructures as discrete codes and distills GNN-to-MLP soft code assignments over the resulting codebook.The codebook is proposed as a more expressive space for identifying different local neighborhood structures.Source: §1 · p. 2
  • VQGraph improves average accuracy over GNNs by 3.90%, over MLPs by 28.05%, and over NOSMOG by 1.39%, while inferring 828× faster than GNNs.The reported averages span seven datasets.Source: §1 · p. 2

2 RELATED WORK

Prior work accelerates GNN inference or distills GNNs into graph-independent students, but each approach leaves constraints that motivate structure-aware GNN-to-MLP distillation.

2 RELATED WORK

  • Pruning and parameter quantization can accelerate GNN inference but do not eliminate neighbor-fetching latency.Graph-MLP bypasses neighbor fetching with an MLP and neighbor contrastive loss, but is transductive rather than inductive.Source: §2 · pp. 2–3
  • GNN-to-MLP methods avoid message passing at inference, while recent GLNN and NOSMOG distill class predictions that may not sufficiently express graph structure.The passage contrasts these methods with GNN-to-GNN distillation and graph-dependent approaches that still require neighbor fetching.Source: §2 · p. 3
  • VQGraph proposes directly labeling nodes’ local neighborhood structures to facilitate structure-aware knowledge distillation.This addresses the stated limitation of class-prediction distillation targets.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries establish message-passing notation for GNNs and describe VQ-VAE’s discrete-latent reconstruction framework, which the paper explores for graph-structured data.

3 PRELIMINARIES

  • A graph is denoted G = (V, A), with A indicating whether each node pair is connected; GNNs aggregate neighbor messages and update node embeddings layer by layer.The notation also defines node features and one-hot classification targets.Source: §3 · p. 3
  • VQ-VAE maps inputs to encoder vectors, quantizes them to nearest codebook entries, and reconstructs inputs from the discrete codes.Its objective includes reconstruction, codebook, and commitment losses, with a straight-through gradient estimator.Source: §3 · p. 3
  • The paper explores VQ-VAE for representing discrete graph-structured data rather than its original continuous-data applications.The stop-gradient operation blocks gradients from flowing into its argument.Source: §3 · pp. 3–4

4 VQGRAPH

VQGraph learns a codebook that represents nodes’ local structures, then uses soft code assignments to distill structural knowledge from a GNN to an MLP. Tokenizer training reconstructs graph attributes and topology while preserving class information, and the distillation target is used during training rather than deployment.

4 VQGRAPH

  • VQGraph first learns a structure-aware graph tokenizer and codebook, then applies a tailored soft-code-assignment objective for GNN-to-MLP distillation.The codebook’s discrete codes label diverse local neighborhood structures.Source: §4 · p. 4

4.1 GRAPH TOKENIZER TRAINING

  • The tokenizer encodes teacher-GNN node embeddings by assigning each node its nearest codebook embedding, producing one of M discrete codes per node.The codebook embeddings are randomly initialized and optimized in pretraining.Source: §4.1 · p. 4
  • Tokenizer optimization reconstructs node attributes and graph topology, adds cross-entropy on labels and code-based GNN predictions, and includes VQ and commitment losses.The learned codebook is described as preserving class information while identifying local graph structures; η is set to 0.25 in the experiments.Source: §4.1 · pp. 4–5
  • The authors distinguish their graph-data VQ-VAE variant from continuous-data VQ-VAE and state that it avoids large variance and posterior collapse associated with graph autoencoders.They refer to experimental comparisons in Sec. 5.3.Source: §4.1 · p. 5
  • For large-scale tokenizer training, the paper uses subgraph-wise mini-batches and target-node connections as topology-reconstruction targets to approximate global structural information.The example samples fixed-size neighbor sets for feature aggregation.Source: §4.1 · p. 5

4.2 STRUCTURE-AWARE CODE-BASED GNN-TO-MLP DISTILLATION

  • The student MLP is trained to match the teacher GNN’s soft assignments over codebook entries using KL divergence, alongside classification and class-based distillation losses.The code assignment is used only during training and removed for deployment, so it does not increase inference time.Source: §4.2 · pp. 5–6

5 EXPERIMENTS

Experiments evaluate VQGRAPH on seven benchmarks, including transductive, inductive, and heterophilic settings, and analyze its efficiency, representations, and design choices.

5.1 MAIN RESULTS

  • In the standard setting, VQGRAPH improves average accuracy over its teacher GNN by 3.90% and over NOSMOG by 1.39%.The comparison covers small- and large-scale graph datasets.Source: Table 1 · p. 7, §5.1 · p. 7
  • Across production settings, VQGRAPH improves average accuracy over GNNs by 2.93%, MLPs by 25.81%, and NOSMOG by 1.6%.The production results combine inductive and transductive evaluations.Source: Table 2 · p. 7, §5.1 · p. 7
  • On heterophilic graphs, graph VQ-VAE consistently improves the performance of the evaluated GNN and MLP architectures.The evaluation uses Texas and Cornell.Source: Table 4 · p. 8, §5.1 · p. 8

5.2 MODEL ANALYSIS

  • On Citeseer, VQGRAPH reaches 76% accuracy at 1.45 ms inference time and is 828× faster than SAGE-L3.It is also 105× faster than SAGE-L2.Source: Figure 3 · p. 8, §5.2 · p. 8
  • VQGRAPH’s distilled representations are more compact, and its nearest-neighbor matches track subgraph similarity while distinguishing same-class structural differences.The analysis uses t-SNE visualizations and node-centered subgraph retrieval.Source: Figure 4 · p. 8, §5.2 · pp. 9–10, Figure 5 · p. 9
  • VQGRAPH has the highest average cut value, 0.9493, versus 0.9276 for SAGE, 0.8725 for GLNN, and 0.9348 for NOSMOG.The paper interprets cut value as prediction consistency with graph topology.Source: Table 5 · p. 9, §5.2 · p. 9

5.3 ABLATION STUDY

  • Codebook size affects distilled performance: too few codes limit structural expressiveness, while too many introduce redundancy; optimal sizes vary by dataset.Graphs with more nodes or edges tend to require larger codebooks for optimal results.Source: Figure 6 · p. 10, §5.3 · p. 10
  • Both the graph tokenizer and soft code assignments improve average accuracy over class-based distillation, and VQGRAPH exceeds Only-VQ across datasets.The paper reports that graph VQ-VAE improves distillation more than classic graph autoencoders.Source: Table 6 · p. 10, §5.3 · p. 10

6 CONCLUSION

VQGraph enriches graph representations by labeling diverse local structures with a codebook and using it for structure-aware GNN-to-MLP distillation. Across seven datasets, it improves average accuracy over GNNs, MLPs, and the state-of-the-art distillation method while retaining faster inference than GNNs.

6 CONCLUSION

  • VQGraph improves average accuracy over GNNs by 3.90%, MLPs by 28.05%, and the state-of-the-art GNN-to-MLP distillation method by 1.39%, with 828× faster inference than GNNs.The reported results average across seven datasets.Source: §6 · p. 11
  • VQGraph labels nodes’ diverse local structures with a codebook to improve graph-representation expressiveness and facilitate structure-aware GNN-to-MLP distillation.Source: §6 · p. 11
  • Visualization, statistical analyses, and ablation studies further support the proposed model’s superiority.Source: §6 · p. 11

A.1 ANALYSIS ON CODEBOOK ENTRIES

On Pubmed, codebook entries are distributed similarly across classes, with little overlap between classes’ codes. The authors interpret this pattern as evidence that the learned representation separates nodes by class.

A.1 ANALYSIS ON CODEBOOK ENTRIES

  • Pubmed shows little overlap between the codes assigned to different paired classes.Table 8 reports the overlapping codes for each class pair.Source: §A.1 · p. 16, Table 8 · p. 16
  • The number of codes and the node distribution are reported as similar across Pubmed classes.Table 7 lists codebook entries for each class.Source: §A.1 · p. 16, Table 7 · p. 16
  • The authors interpret the limited code overlap as class-specific code distributions and as evidence of class separation that facilitates knowledge transfer.Source: §A.1 · p. 16

B.1 DATASETS

Experiments use seven public graph benchmarks with splits following prior work or official dataset protocols. The benchmarks cover citation networks, Amazon product co-purchases, and arXiv papers.

B.1 DATASETS

  • Five small datasets use CPF’s random-seed-specific splits, while Arxiv and Products use the official OGB time-based and popularity-based splits, respectively.Source: §B.1 · p. 16
  • Cora and Citeseer are scientific-publication citation benchmarks, while Pubmed comprises diabetes-related articles with labels for diabetes type.Cora uses sparse bag-of-words features, and Citeseer has the largest feature count among the paper’s datasets.Source: §B.1 · pp. 16–17
  • A-computer and A-photo are Amazon co-purchase graphs whose product nodes use review-derived bag-of-words features and predefined product-category labels.Source: §B.1 · p. 17
  • Arxiv represents citation-linked computer-science papers with title-and-abstract embeddings, whereas Products represents Amazon co-purchases with product-description bag-of-words features.Source: §B.1 · p. 17

B.2 MODEL HYPERPARAMETERS AND IMPLEMENTATION DETAILS

The implementation evaluates production-style transductive and inductive predictions, and tests VQGraph with alternative relation modules and teacher GNN architectures. Across teacher architectures, VQGraph remains competitive with the teachers and other distillation methods.

B.2 MODEL HYPERPARAMETERS AND IMPLEMENTATION DETAILS

  • The production protocol evaluates observed unlabeled nodes as transductive predictions and newly introduced holdout nodes as inductive predictions.The inductive subset is 20% of the test data; the remaining 80% forms the observed subset, and “prod” interpolates between their results according to the split rate.Source: §B.2 · p. 17
  • Tokenizer pretraining combines a teacher GNN and learnable codebook, with VQ between encoder and classifier and separate decoders for node attributes and graph topology.The code embeddings are initialized uniformly.Source: §B.2 · pp. 17–18
  • The distillation setup freezes the pretrained teacher GNN and codebook, and the student MLP generally matches the teacher’s layer count and hidden dimension.The stated matching keeps their parameter counts the same unless a -wi or -Li variant is specified.Source: §B.2 · p. 18
  • During distillation, the teacher provides both class soft labels and soft code assignments, while code assignment for the MLP is used only in training.Source: §B.2 · p. 18
  • The code-assignment procedure is removed at test time, so it does not add inference time.Source: §B.2 · p. 18
  • Loss weights for categorical and structural information are adjusted because their relative importance can vary across graph datasets.Source: §B.2 · p. 18

C.1 VQGRAPH WITH DIFFERENT RELATION MODULES IN EQUATION (4)

  • L2 distance and cosine similarity both support improved code distillation, with L2 showing an approximately 0.2% average-accuracy advantage.The relation-module choice has minimal influence on overall performance.Source: Table 13 · p. 19, §C.1 · p. 19

C.2 VQGRAPH WITH DIFFERENT TEACHER GNNS

  • VQGraph consistently outperforms teacher GNNs and other GNN-to-MLP distillation methods across five datasets when distilled from GraphSAGE, GCN, GAT, or APPNP teachers.The four teacher architectures achieve comparable performance.Source: Figure 7 · p. 19, §C.2 · p. 19

C.3 ROBUSTNESS EVALUATION WITH NOISY FEATURES

VQGraph is comparable to or improves on GNNs and NOSMOG across feature-noise levels, while GLNN and MLP performance deteriorates as noise increases. At fully perturbed input features, VQGraph remains similar to GNNs.

C.3 ROBUSTNESS EVALUATION WITH NOISY FEATURES

  • The evaluation adds independent Gaussian noise to node features using Ẋ = (1−α)·X+α·n, with α controlling noise level.Source: §C.3 · p. 19
  • At α = 1, where perturbed and original features are independent, VQGraph maintains performance similar to GNNs, while GLNN and MLP perform poorly.Source: §C.3 · p. 20
  • Across feature-noise levels, VQGraph achieves comparable or improved performance to GNNs and outperforms NOSMOG, while GLNN and MLP deteriorate as noise increases.The evaluation compares average accuracy across datasets and noise levels.Source: §C.3 · pp. 19–20, Figure 8 · p. 20
Loading 2308.02117v3…