PaperlayerView exact revision on arXiv

Source-linked AI summary

Train Your Own GNN Teacher: Graph-Aware Distillation on Textual Graphs

Costas Mavromatis, Vassilis N. Ioannidis, Shen Wang, Da Zheng, Soji Adeshina, Jun Ma, Han Zhao, Christos Faloutsos, George Karypis

arXiv:2304.10668v1cs.LG

TL;DR

Textual-graph GNNs can perform well but incur costly inference when combined with language models. GraD jointly trains a GNN teacher and graph-free student through a shared LM; across eight node-classification benchmarks, it outperforms conventional distillation and reports an accuracy–inference-time advantage for GraDBERT.

  • Problem

    Combining GNNs with LMs achieves strong node-classification performance, but its inference cost creates scalability challenges for practical deployment.

    Source: §1 · p. 2

  • Method

    GraD jointly optimizes a GNN teacher and graph-free student via a shared LM, which lets them exchange graph and textual information.

    Source: §1 · p. 2

  • Results

    Across eight node-classification benchmarks in transductive and inductive settings, GraD outperforms conventional knowledge distillation; GraDBERT is 3.24% more accurate than fine-tuned BERT and as effective as BERT with a GNN at 2.4x–13x smaller inference time.

    Source: §1 · p. 3, §8 · p. 14

  • Takeaways & Limitations

    GraD balances inference efficiency and classification effectiveness for textual graphs.

    Source: §8 · p. 14

    GraD relies on informative input node features, and its generalization to featureless graphs is uncertain.

    Source: §9 · p. 15

Abstract

from arXiv · show

How can we learn effective node representations on textual graphs? Graph Neural Networks (GNNs) that use Language Models (LMs) to encode textual information of graphs achieve state-of-the-art performance in many node classification tasks. Yet, combining GNNs with LMs has not been widely explored for practical deployments due to its scalability issues. In this work, we tackle this challenge by developing a Graph-Aware Distillation framework (GRAD) to encode graph structures into an LM for graph-free, fast inference. Different from conventional knowledge distillation, GRAD jointly optimizes a GNN teacher and a graph-free student over the graph's nodes via a shared LM. This encourages the graph-free student to exploit graph information encoded by the GNN teacher while at the same time, enables the GNN teacher to better leverage textual information from unlabeled nodes. As a result, the teacher and the student models learn from each other to improve their overall performance. Experiments in eight node classification benchmarks in both transductive and inductive settings showcase GRAD's superiority over existing distillation approaches for textual graphs.

1 Introduction

Textual-graph GNNs can achieve strong node-classification performance, but combining GNNs with LMs makes inference costly. GraD jointly trains a GNN teacher and graph-free student through a shared LM, with reported accuracy and inference-time advantages.

1 Introduction

  • LM-GNN inference can be expensive because each K-layer mini-batch fetches exponentially many nodes whose raw text must be encoded on the fly.The resulting cost can be prohibitive when fast inference is vital.Source: §1 · p. 2
  • GraD jointly optimizes a GNN teacher and graph-free student via a shared LM, exchanging graph-aware updates, soft labels, student imitation, and unlabeled-node text.The teacher transfers graph knowledge through soft labels, while the student’s use of unlabeled-node text helps fine-tune the shared LM.Source: §1 · p. 2
  • GraD is formulated as multi-task learning with three optimization strategies that vary teacher–student coupling and student flexibility.The framework uses a powerful LM student for large graphs and can substitute an MLP for the LM on small graphs.Source: §1 · p. 3
  • GraDBERT is 3.24% more accurate than fine-tuned BERT and as effective as BERT with a GNN at 2.4x–13x smaller inference time.The comparison is for node classification on ogbn-arxiv.Source: Figure 1 · p. 2, §1 · p. 3
  • The introduction frames GraD as addressing limitations of conventional knowledge distillation and reports improved classification for seen and unseen nodes across large- and small-scale graphs.The authors also describe the framework as achieving state-of-the-art node-classification performance and scaling to large graphs.Source: §1 · p. 3

2 Background

The task is to learn a graph-free predictor from node text that approximates a graph-based model. In textual-graph GNNs, text is converted to node features and then combined through neighborhood aggregation, with inference cost growing with neighborhood size and model depth.

2.1 Problem Formulation

  • The graph-free model τ′ takes node text Xv at inference and is intended to approximate graph-based predictions while avoiding use of graph structure then.The task defines labeled and unlabeled nodes and distinguishes transductive from inductive graph splits.Source: §2.1 · p. 3

2.2 GNNs on Textual Graphs

  • A text encoder τ transforms each node’s raw text Xv into numerical features xv; a language model such as BERT may represent the sequence using its [CLS] token.The resulting features provide the inputs to the GNN.Source: §2.2 · p. 4
  • A GNN maps the K-hop computation graph centered on node v to a d-dimensional representation by aggregating the node’s and its neighbors’ representations layer by layer.Each update uses the previous-layer representation and those of direct neighbors; the described aggregation includes GraphSAGE and RGCN examples.Source: §2.2 · p. 4
  • Computing a node representation can cost O(SKC), where S is neighborhood size, K is GNN depth, and C is the LM inference cost.At the first layer, node representations are the text features xv.Source: §2.2 · p. 4

3 Towards Graph-Aware Knowledge Distillation

This section frames graph-free node classification as a distillation problem: transfer a GNN teacher’s knowledge to a graph-free student for effective classification without graph-based inference. It analyzes why conventional distillation can leave graph structure underused, especially for unlabeled and unseen nodes.

3 Towards Graph-Aware Knowledge Distillation

  • Knowledge distillation pairs a graph-based GNN teacher with a graph-free student, such as a language model, to learn effective node classification without graph-based inference.The standard procedure trains the GNN teacher, generates soft-labels, then trains another language model to mimic those predictions.Source: §3 · p. 4, §3.1 · pp. 4–5

3.1 Knowledge Distillation

  • Conventional KD trains the student on precomputed teacher soft-labels, so it does not use the graph structure during the student’s learning stage.Soft-labels can cover unlabeled nodes, but the student treats nodes independently and infers their interactions only through those labels.Source: §3.2 · p. 5, §3.1 · p. 5

3.2 What Does Knowledge Distillation Learn? An Analysis

  • Node-classification information can come from a node’s text, neighboring nodes’ text, and links; without neighborhood or graph use, only the node-text term can be maximized.The analysis decomposes information about a node’s label into information from its own text, other nodes’ text, and links.Source: §3.2 · pp. 5–6
  • A student’s performance depends on how informative teacher soft-labels are, and the teacher may not encode graph information optimally for unlabeled or unseen nodes.The teacher and student objectives use different node sets, and soft-labels cannot be precomputed for unseen nodes; informative textual features may help the student handle them.Source: §3.2 · pp. 6–7

Framework

GraD couples a GNN teacher and graph-free student through a shared text-encoding function, allowing graph structure to inform the student during training. It offers three optimization strategies and supports language-model and MLP student variants.

Framework

  • GraD couples the GNN teacher and graph-free student through a shared function, encoding graph structure into text representations used by the graph-free model at inference.The student’s soft-label training on unlabeled-node text also lets the teacher leverage textual information from those nodes.Source: Framework · pp. 7–8

4.1 GraD-Joint

  • GraD-Joint jointly optimizes teacher and student predictions through shared τ, combining teacher–student distillation with ground-truth supervision.The student mimics the teacher’s soft-labels, while a cross-entropy term encourages text representations useful to the teacher.Source: §4.1 · pp. 8–9

4.2 GraD-Alt

  • GraD-Alt alternates teacher and student updates to mitigate possible negative transfer and give the student more flexibility to fit teacher predictions.Alternating optimization can be computationally costly with a large language model and many graph nodes.Source: §4.2 · p. 9, §4.3 · p. 9

4.3 GraD-JKD

  • GraD-JKD first jointly trains teacher and student, then decouples and retrains the student to mimic the teacher’s predictions.The jointly trained teacher can better leverage textual information from unlabeled nodes.Source: §4.3 · p. 9

4.4 Student Models

  • GraDBERT uses a language model such as BERT for τ, while GraDMLP uses bag-of-words or TF-IDF vectors with trainable MLPs when graphs are small or lack raw text.GraDBERT requires raw node texts and assumes adequate graph data for training.Source: §4.4 · p. 10

5 Experimental Setup

The experiments use node-classification benchmarks with text or bag-of-words and TF-IDF features, and compare GraD with distillation and graph-aware LM methods. The setup also specifies different teacher, language-model, and optimization choices across settings.

5.1 Datasets

  • Experiments evaluate GraDBERT on Arxiv, Products, and Papers1.5M, and GraDMLP on Cora, Citeseer, Pubmed, A-Computer, and A-Photo, using accuracy.Product-Reviews and perturbed datasets are also used for GraDBERT in appendices.Source: §5.1 · p. 10

5.2 Implementation Details

  • GraphSAGE is the teacher; GraDBERT uses one layer and samples 8 or 12 neighbors, with SciBERT for Arxiv and Papers1.5M and DistilBERT for Products.GraD-Alt alone is used for small-scale graphs.Source: §5.2 · p. 10

5.3 Compared Methods

  • Comparisons include conventional GNN-to-BERT KD, GNN-to-MLP distillation, graph-aware BERT methods, E2EG, and GLEM, with GraphSAGE matched among GNN-based methods.Table 1 compares graph-free methods’ performance and reports trainable LM parameters.Source: §5.3 · p. 10, Table 1 · p. 11

6 Experimental Results

The experiments compare graph-free GraDBERT and its GNN teacher with other methods, then assess GraDMLP strategies on small graphs and inductive nodes. They also examine inference time and qualitative examples in appendices.

6.1 GraDBERT Results

  • GraDBERT exceeds BERT+KD by 0.66, 0.98, and 1.75 percentage points on Arxiv, Products, and Papers, respectively.These are large-scale graph results for methods that do not use graph structure during inference.Source: §6.1 · p. 11
  • GraD’s GNN teacher performs best among the compared LM-GNN methods, and the comparison reports GraphSAGE with backbone LM parameter counts.The authors interpret the teacher result as evidence that GraD’s student improves teacher effectiveness.Source: §6.1 · p. 12, Table 2 · p. 12
  • In the inductive setting, GraDBERT beats BERT+KD by more than 0.7 percentage points on average; conventional KD methods show a larger performance drop on inductive nodes.Table 3 holds out 50% of validation and test nodes; ∆acc reports relative degradation between full and inductive settings.Source: §6.1 · p. 12, Table 3 · p. 13

6.2 GraDMLP Results

  • On small graphs, GraD-Alt performs best on Citeseer and A-Comp, while GraD-Joint performs best on Cora and Pubmed.The authors suggest tighter coupling may help when node texts, such as paper keywords, are more informative.Source: §6.2 · p. 12, Table 4 · p. 13
  • GraDMLP consistently outperforms MLP+KD at inductive rates from 10% to 75%; near 90%, small graphs leave few training-node connections, so GraDMLP has no useful graph information to leverage.Source: §6.2 · p. 13

7 Related Work

Related work combines language models with GNNs through structural or token-level pretraining and direct node-classification fine-tuning. GraD is positioned among graph-to-graphless distillation methods and work using GNNs in other tasks.

Learning on Textual Graphs

  • Textual-graph methods pair LMs with GNNs, using graph-data pretraining, token-level pretraining, or direct node-classification fine-tuning.Source: §7 · p. 13
  • GNNs for textual graphs have also been applied to question answering, sponsored search, and document classification.Source: §7 · p. 14

Distillation Approaches

  • GraD-Joint self-distills deeper GNN knowledge into shallower LM or MLP layers, but targets node interactions rather than model compression; GLNN and ColdBrew also distill graph knowledge to graph-free models.Source: §7 · p. 14
  • GraD learns better graph information than GLNN in transductive and inductive settings, while ColdBrew applies only to transductive settings; graph-regularized MLPs underperform KD approaches.Source: §7 · p. 14

8 Conclusion

GraD jointly trains a teacher GNN and graph-free student through a shared LM. Across eight node-classification benchmarks in transductive and inductive settings, it outperforms conventional knowledge distillation and balances efficiency with effectiveness.

8 Conclusion

  • GraD jointly trains a teacher GNN and graph-free student via a shared LM, allowing the models to learn from each other.Source: §8 · p. 14
  • Across eight node-classification benchmarks in both transductive and inductive settings, GraD outperforms conventional knowledge distillation.Source: §8 · p. 14
  • The paper concludes that GraD balances efficiency and effectiveness for textual graphs.Source: §8 · p. 14

9 Limitations & Ethical Statement

GraD relies on informative node features, trades accuracy for computational efficiency, and does not adapt to edge changes as a GNN could. The supplied implementation and inference analysis describe its graph-dataset scope, training setups, and efficiency–effectiveness comparison.

9 Limitations & Ethical Statement

  • GraD’s generalization to featureless graphs is uncertain; it also trades accuracy for computation efficiency, cannot adapt to edge changes like a GNN, and requires retraining to address biases in the training graph.The authors identify informative input node features as central to generalizing to unseen nodes.Source: §9 · p. 15

A Dataset Statistics

  • The evaluated datasets include citation graphs, product co-purchase graphs, and the heterogeneous Product-Reviews graph.The passage names Arxiv, Papers, Cora, Citeseer, and Pubmed as citation graphs; Products, A-Computer, and A-Photo as co-purchase graphs.Source: A Dataset Statistics · p. 18

B Further Implementation Details

  • The teacher GNN is GraphSAGE for both GraDBERT and GraDMLP, and the authors provide code in a public repository.The passage supplies the repository URL: https://github.com/cmavro/GRAD.Source: B Further Implementation Details · p. 18

B.1 GraDBERT

  • GraDBERT uses distributed training on a p4d.24xlarge, with one-layer GraphSAGE and sampled neighbors; inference uses sequence length 128 for ablations or 512 for SOTA performance.Arxiv and Papers use 8 sampled neighbors and Products uses 12; initialization varies by dataset, including SciBERT, DistilBERT, and pretrained BERT.Source: §B.1 · p. 18

B.2 GraDMLP

  • GraD-Joint uses limited hyperparameter search and validation-based model selection, while GraDMLP uses full-graph training and tunes epochs with early stopping.GraDMLP uses 128 hidden dimensions; GraD-Alt performs T-S steps at each full-graph optimization step.Source: §B.1 · p. 18, §B.2 · p. 19

C Inference Timing Analysis

  • On ogbn-arxiv, GraDBERT is reported to best balance inference efficiency and accuracy, with timing compared across two GPU machines.The p4d.24xlarge is estimated to be 2.5× faster than the p3.16xlarge and has 2.5× more GPU memory; Fig. 5 plots accuracy against inference time.Source: Figure 5 · p. 19, C Inference Timing Analysis · p. 19

D.1 Ablation Study on GraD’s Loss

The ablations and examples examine how distillation, hyperparameters, graph-aware representations, and training conditions affect GraD. Reported outcomes include gains on Product-Reviews and downstream GNNs, alongside limitations of copying teacher predictions for inductive nodes.

D.1 Ablation Study on GraD’s Loss

  • KD loss benefits both GraD’s teacher and student, with the largest benefit for the student, indicating that the GNN distills useful graph-aware information.The ablation is reported for Arxiv in Table 6.Source: §D.1 · p. 20, Table 6 · p. 20

D.2 Hyper-Parameter Sensitivity Analysis

  • On Products, higher λ can overfit poor GNN predictions when unlabeled nodes greatly outnumber labeled ones; α = 0.2 favors the student, while α = 0.8 favors the teacher.The sensitivity study tests α ∈ {0.2, 0.8} and λ ∈ {20, 50}.Source: §D.2 · p. 20, Table 7 · p. 21

G Failing KD Case

  • With perfect GNN predictions, a graph-free student can achieve perfect transductive predictions, but this construction fails inductively when teacher predictions are unavailable.In the stated failing case, the student has no learned parameters and can only make random guesses for inductive nodes.Source: G Failing KD Case · pp. 21–22

E Qualitative Examples

  • In two Arxiv examples, GraDBERT predicts the correct label by using the majority-neighbor label pattern where BERT+KD is misled by target-text phrases.The examples concern Discrete Math versus neighboring Algorithm papers, and Databases versus neighboring Security papers.Source: E Qualitative Examples · p. 20, Figure 6 · p. 22

F SOTA GNN Performance

  • GraDBERT improves RevGAT and SAGN+ by 0.20% on average, beats static LMs by 3.7 percentage points, and outperforms existing approaches by 1%–1.4% on simpler GNNs.For Products, the authors combine GraDBERT with GIANT because the dataset has a very small label rate.Source: F SOTA GNN Performance · p. 21, Table 8 · p. 23

H Product-Reviews Dataset

  • GraDBERT classifies 1,968–7,926 more Product-Reviews products than BERT across data splits, despite an average product degree of only 0.7.The inductive comparison uses the same product nodes for BERT and GraDBERT during training; GraDBERT uses review nodes as side information.Source: H Product-Reviews Dataset · p. 23, Table 9 · p. 24

I Perturbed Node Texts

  • When half of each node’s text is replaced with random-node text, GraDBERT performs better as α increases, whereas BERT+KD performs poorly.The authors suggest that BERT+KD cannot learn an accurate mapping from text to teacher labels in this setting; graph neighbors retain similar semantics to the original text.Source: I Perturbed Node Texts · p. 24, Table 10 · p. 24
Loading 2304.10668v1…