On Representation Knowledge Distillation for
Graph Neural Networks

Graph Neural Networks (GNNs) take as input an unordered set of nodes and the graph connectivity among them, and learn latent node representations for them via iterative feature aggregation or message passing across local neighborhoods [5]. Consider a graph $\mathcal{G}=\left(\mathcal{V},\mathcal{E}\right)$, where $\mathcal{V}$ is a set of $n$ nodes, and $\mathcal{E}$ is a set of edges associated with the nodes. Each node $i\in\mathcal{V}$ is associated with a $d$-dimensional initial feature vector $f_{i}$, which is iteratively updated by the backbone GNN via a generic message passing or graph convolution operation:

where $\textsc{Msg},\textsc{Upd}$ are learnable transformations such as multi-layer perceptrons, $e_{ij}$ are optional edge features, and Agg is a permutation-invariant local neighborhood aggregation function such as summation, maximization, averaging, or weighted averaging via attention.

We are usually interested in making predictions for each node in the graph, e.g. node classification on social networks or semantic segmentation of 3D point clouds. Thus, the final feature vectors after $L$ layers of message passing $f_{i}^{L}$ are passed to a linear classifier to obtain logits $z_{i}$, and the neural network is trained end-to-end via a cross-entropy loss $\mathcal{H}$ with the groundtruth class labels $y_{i}$: $\mathcal{L_{\textsc{Sup}}}=\mathcal{H}\left(y_{i},\text{softmax}\left(z_{i}\right)\right)$. For graph-level tasks, such as molecular property prediction, the feature vectors of all nodes are pooled via a permutation-invariant function Pool to obtain a global feature vector $f_{G}=\textsc{Pool}\left(f_{i}^{L}|i\in\mathcal{V}\right)$, and then passed to a linear classifier. Notably, representation learning still occurs at the node-level for graph-level prediction tasks [38], unlike 2D image classification where pooling layers learn feature vectors for entire images.

GNNs and 3D point clouds. The message passing framework in (1) can be used to present a unified ‘geometric’ view of deep learning on non-Euclidean data structures such as 3D point clouds, irregular voxel grids, meshes, and sets [39, 40]. In this work, we consider 3D point cloud networks [41, 42] which process sets of points as nodes. The edges originating from each point are heuristically determined by the $k$-nearest neighbors or radius ball queries in the 3D coordinate space. We additionally consider sparse 3D voxel networks [43, 44, 45] and use trilinear interpolation to convert voxel-level features to point-level features before distillation.

Knowledge Distillation (KD) transfers dark knowledge from high capacity teacher models (which are cumbersome to deploy) to efficient students via matching their output logits in addition to supervised learning. At each node $i\in\mathcal{V}$, logit-based KD [30] uses the cross-entropy loss or KL-divergence to match the output logits of the student $z_{i}^{S}$ and teacher $z_{i}^{T}$, scaled by temperature $\tau_{1}$:

In this work, we study auxiliary representation distillation techniques that augment or replace logit-based distillation with representational knowledge using the teacher and student features, $f^{T}\in\mathbb{R}^{d^{T}}$ and $f^{S}\in\mathbb{R}^{d^{S}}$, respectively.¹¹1 Following best practices [36, 46], we consider feature vectors from the penultimate layer before the prediction head for representation distillation. Our overall objective function for training the student model via the logits as well as latent feature vectors of the teacher is as follows:

where $\alpha,\beta$ are balancing weights for the logit-based KD loss $\mathcal{L}_{\text{KD}}$ and auxiliary representation distillation loss $\mathcal{L}_{\text{Aux}}$, respectively. Next, we will describe how $\mathcal{L}_{\text{Aux}}$ is instantiated as different representation distillation techniques such as $\mathcal{L}_{\text{LSP}}$, $\mathcal{L}_{\text{GSP}}$, and $\mathcal{L}_{\text{G-CRD}}$.

We briefly summarize LSP [32], a GNN representation distillation objective which trains the student model to preserve the local structure of graph data from the teacher’s node embedding space. The local structure around each node is defined as the set of parameterized pairwise distances to its neighboring nodes in the latent feature space. The similarity between a pair of linked nodes can be computed by kernel functions (which we tune for):

Thus, the local structure around each node $i\in\mathcal{V}$ is the softmax probability distribution of the similarities among $f_{i}$ and its neighbors’ features $f_{j}\ \forall(i,j)\in\mathcal{E}$. LSP trains the student model to mimic the local structure from the teacher’s embedding space via KL-divergence:

The purely local LSP objective over pre-defined edges does not preserve the global topology of how the teacher embeds the graph, as it does not account for latent interactions among disconnected nodes. We would like to explicitly train the student to preserve the global topology in order to better distill representational knowledge from the teacher. To achieve this, we first introduce a simple extension of LSP, the Global Structure Preserving loss (GSP), which matches all pairwise similarities among node features via mean squared error:

While theoretically more powerful than LSP, GSP may be computationally inefficient and involves two key design choices: (See App.F for ablation studies.)

MSE over KL-divergence We omit normalizing the similarities via softmax followed by KL-divergence. Instead, we use MSE to match the raw teacher and student pairwise similarity matrices as it worked better empirically.

Scalability and sub-sampling. For a set of $n$ nodes, GSP needs to computer $n^{2}$ pairwise similarities for both the teacher and student features. We often need to sub-sample large graphs or 3D point clouds when computing $\mathcal{L}_{\text{GSP}}$ due to GPU memory constraints instead of storing all possible pairwise similarities. For each experiment, random sub-sampling was done to retain as many nodes as possible subject to GPU memory limitations.

We benchmark distillation techniques across a range of tasks on single large-scale networks, batches of small graphs, as well as batches of 3D point clouds, summarized in Tab.1. As we are motivated by real-world and real-time applications involving noisy and shifting data distributions, we make the following considerations for our benchmarks:

Challenging for student models. Past work on distillation for GNNs [32] used small-scale PPI [2] (node classification) and ModelNet [49] (3D point cloud classification), both of which are considered saturated by the community [11, 50]. The performance gap between cumbersome teachers and lightweight student models is negligible for these datasets, which makes them unsuitable benchmarks for comparing techniques. (See App.D for an investigation on PPI.)

We evaluate node classification on ARXIV and the Microsoft Academic Graph (MAG) [51, 11], which are 70$\times$ and 800$\times$ larger than PPI, respectively, and consider the more challenging semantic segmentation task on 3D point cloud scenes from S3DIS [37] that are over 10$\times$ denser than CAD object scans from ModelNet. We additionally evaluate graph classification on MOLHIV from OGB/MoleculeNet [52].

Out-of-distribution evaluation. Unlike PPI and ModelNet, all datasets involve realistic and carefully curated train-test splitting procedures to evaluate out-of-distribution generalization. ARXIV and MAG follow a temporal split, MOLHIV trains on common molecular scaffolds while testing on rare ones, and S3DIS is split according to 6 distinct areas.

Our choice of teachers and student architectures are made with the following considerations about deploying GNNs: (1) depth – more message passing rounds allow models to access larger sub-graphs around each node, at the cost of training and inference time; (2) hidden channels – especially for giant graphs, increasing hidden channels boosts performance as well as memory consumption, and may require specialized sampling procedures; and (3) architectural complexity – principled geometric priors [42] or attention mechanisms over edges [4] improve model expressivity but tend to have higher inference latency and memory usage.

Unlike in [32], we benchmark distillation techniques across heterogeneous architecture families with significant mismatch in terms of parameter count, expressive power, and inference time; e.g. for social networks, we distil from Graph Attention Networks [4] to simple GCNs [1]. For 3D point clouds, we distil from Kernel Point ConvNet [42] teachers to lightweight PointNet++ [41] and voxel-based MinkowskiNet [43]. We give further rationale for our choice of teacher-student pairs for each experiment in Sec.5. Details on the configurations, inference time and latency of teacher and student models are available in App.B

Our implementation is built upon PyTorch Geometric [54], PyTorch Points 3D [53], and Open Graph Benchmark [11]. We follow the best practices and guidelines for each dataset.

Training. Teacher models are trained via the conventional supervised learning paradigm and the final weights for distillation are selected by early stopping. Student models are trained via the knowledge distillation pipeline described in Sec.2. For the OGB datasets, we follow their example implementations and training setups for both teacher and student models. For S3DIS, the teacher models are trained for 600 epochs following their original learning rate strategies, while the student models are trained for 300 epochs and use an exponential learning rate strategy. We use the conventional data preparation and augmentation procedures for S3DIS via PyTorch Points 3D.

Evaluation. For the OGB datasets, we use the official evaluators and report the average test performance across 8/10 random seeds (each teacher-student pair is re-trained for each seed). For S3DIS, we follow best practices in the literature: all models are trained once, and we report the average mIoU and mAcc over 10 voting runs on the held-out Area-5 test set using the evaluation protocol from PyTorch Points 3D.

Baselines and hyperparameters. Following Yang et al. [32], we compare the GNN representation distillation techniques G-CRD, GSP and LSP to logit-based KD [30] (2) as well as FitNet [47] and Attention Transfer (AT) [48], two feature mimicking baselines adapted from 2D computer vision which are formulated as: $\mathcal{L}_{\text{Aux}}=\sum_{i\in\mathcal{V}}\lVert P^{T}(f^{T}_{i})-P^{S}(f^{S}_{i})\rVert^{2}_{2}$ (FitNet uses L2 normalized projection head, AT uses attention mapping). We tune the loss balancing weights $\alpha,\beta$ in (3) for all techniques on the validation set. For KD, we tune $\alpha\in\{0.8,0.9\},\tau_{1}\in\{4,5\}$. For FitNet, AT, LSP, and GSP, we tune $\beta\in\{100,1000,10000\}$ and the kernel in (4) (only LSP, GSP). For G-CRD, we tune $\beta\in\{0.01,0.05\},\tau_{2}\in\{0.05,0.075,0.1\}$ and the projection head in (9). When comparing representation distillation methods, we set $\alpha$ to 0 in order to ablate performance, as in [36], and reduce $\beta$ by one order of magnitude.

We consider the graph-level property prediction task over batches of molecular graphs. Reducing the inference latency of GNNs for molecules speeds up high throughput virtual screening [52]. Additionally, virtual screening on proprietary data will require homomorphically encrypted models, which further demand low layer count and hidden size [55]. As teachers, we consider 5-layer deep GIN-E [18] and PNA [16] augmented with virtual nodes, two strong OGB baselines. Notably, PNA is the most expressive message passing GNN but explicitly materializes messages over graph edges, leading to higher memory requirement and inference latency. Our student architectures are 2-layer GCN [1] and GIN [14], which are comparatively less expressive and do not use virtual nodes (and edge features in GIN), while having low inference latency.

In Tab.2, we compare logit-based KD and the representation distillation techniques across a range of teacher-student pairs. We find that the implicit G-CRD technique consistently improves over the supervised student’s performance and outperforms the explicit global structure preserving approach, GSP. In turn, GSP outperforms the purely local approach, LSP. While all other distillation techniques do offer minor performance boosts over the supervised student, G-CRD is the only one which improves over the KD baseline for most teacher-student combinations.

Semantic segmentation of 3D scene graphs is a safety critical real-time task with applications in autonomous driving and robotics. Models process 3D scenes as sets of points or voxel grids, and make a dense prediction by assigning semantic categories to each point/unit. Recent state-of-the-art architectures involve strong geometric priors and complex architectures, leading to increased inference latency and GPU memory requirement. Here, we consider distilling two powerful models, Kernel Point Convolution (KP-FCNN, rigid kernel) [42] and Sparse Point-Voxel CNN (SPVCNN) [44, 45], into simpler models with low inference latency and memory usage: PointNet++ with single-scale grouping [41] and voxel-based MinkowskiNet [43] at low 20% channel ratio.

In Tab.3, we see trends that are consistent with the previous section: G-CRD consistently improves the performance of students compared to GSP, and is particularly effective for boosting PointNet++. Notably, due to a large mismatch in teacher-student representation capacity, feature mimicking losses, FitNet and AT, may worsen the student’s performance over purely supervised learning.

Reducing model depth and feature dimensions boosts inference and reduces memory usage for real-time applications on large-scale graphs. For ARXIV, a homogeneous citation network where models classify the subject of each node/paper, we use a strong 3-layer GAT [4] teacher. GAT has high memory requirement during training and inference due to its attention mechanism. We distil into 2-layer GCN [1] and GraphSage [2], which are comparatively more scalable but lack the expressivity of attention. We also distil to structure-agnostic SIGN [22] designed for parallelized large-graph processing. LSP cannot be used for SIGN as it does not make use of graph structure.

Experiments on the giant heterogeneous Microsoft Academic Graph (MAG) of authors, papers, and institutions further tests the scalability of distillation. We distil among Relational-GCNs [56] at different depth and feature sizes. We use the GraphSAINT [29] sampler to fit GPU memory.

In Tab.4, when considering each technique independently, logit-based KD is the best while G-CRD is the second best. We believe this can be explained by the homophily phenomenon for node labels on social networks [57]. The soft labels from the teacher model thus providing a more informative signal for knowledge transfer than the latent representations. On combining KD with representation distillation, KD + G-CRD consistently outperforms all other pairs. GSP is the worst performing technique for MAG, demonstrating the inability of explicit global topology preservation to scale to large graphs. Note that performance gains from distillation may seem marginal, but are significant as metrics are averaged over several thousand nodes (48K for ARXIV, 42K for MAG).

Across Tab.2, 4, 3, we have observed that G-CRD, which implicitly preserves global topology from the teacher to the student embedding space, outperforms explicit structure preserving approaches LSP and GSP, as well as feature mimicking baselines FitNet and AT. In order to unpack the efficacy of G-CRD beyond performance metrics, we measure how well distillation techniques preserve both global and local topology from the teacher to the student node embedding space in Tab.5. To quantify global representational similarity between the teacher and student, we use the recently proposed Centered Kernel Alignment score (CKA) [58, 59] between two embedding sets, as well as the classical Mantel Test [60, 61]³³3 We use the implementation from scikit-bio (skbio.stats.distance.mantel). of Pearson correlation between all pairwise cosine distances from two embedding sets (as in the GSP loss). Additionally, we quantify local structural similarity via another Mantel Test which only considers distances over pre-defined edges (as in the LSP loss).

Our results largely follow the intuitions developed in Fig.1: In terms of global topology and representational similarity, students trained with FitNet, GSP, and G-CRD have high correlations to the teacher. On the other hand, LSP is the most correlated for local topology over existing graph edges, but relatively poorer at preserving global topology. Overall, we see that the implicit contrastive approach G-CRD strikes a balance between preserving both local and global relationships, while the explicit LSP/GSP are best at preserving only one or the other.

Quantization of model weights and activations to lower precision arithmetic is a complementary compression technique to distillation. Models at lower precision such as 8-bit integers have significantly faster inference latency and lower memory usage. However, quantization aware training (QAT) is known to degrade performance. In Tab.6, we explore whether distillation can improve the performance of a lightweight quantized 2-layer GIN model when trained with vanilla QAT as well as DegreeQuant [28], a GNN-specific QAT. We find that QAT with G-CRD enables quantized models to retain a large portion of their performance at 8-bit integer precision as compared to other distillation techniques. Significantly, GIN trained with G-CRD at INT8 (73.65% ROC-AUC) can still outperform its purely supervised counterpart at full 32-bit floating point precision (73.03%, see column 6 in Tab.2).

G-CRD is well suited for QAT because the projection heads ensure that contrastive distillation take place at full precision even when the student model is at low INT8 (the heads do not interfere with the inference phase, which always uses INT8). On the other hand, structure preserving approaches LSP and GSP are particularly ill suited as they match relationships from low INT8 feature vectors (student) to full precision (teacher).

In biomedical discovery, we are interested in models’ ability to generalize or extrapolate to unseen regions of chemical space. In addition to evaluating on out-of-distribution molecular scaffolds from MOLHIV’s test set, we further test the transferability of distilled models on smaller scale molecule datasets from OGB: MOLBACE (1,513 samples), MOLSIDER (1,427 samples), and MOLESOL (1,128 samples) in Tab.7. We perform linear probing on frozen node features from a 2-layer GCN model. The GCN feature extractor can either be initialized randomly, pre-trained on MOLHIV via supervised learning, or pre-trained via various distillation techniques with a 5-layer PNA teacher. Encouragingly, we find that distillation leads to more transferable representations than random or purely supervised initialization. Overall, G-CRD boosts the transferability of student models for 2 out of 3 datasets with structurally different molecules as well as task semantics.

Beyond clean test set performance, lightweight 3D segmentation models often have to deal with ‘dirty’ data such as sparse, occluded or noisy scans. In Fig.3, we evaluate the impact of distillation on the robustness of PointNet++ across three common challenging scenarios (we also show the KP-FCNN teacher for comparison): (1) Sparse scans – randomly dropout a percentage of points for each scan; (2) Partial and occluded scans – dropout all points sampled within a number of random spheres of fixed radius of 0.5m for each scan; and (3) Noisy scans – add independent random noise to the 3D coordinates of each point with a variance factor $\sigma$.

Training the lightweight student with G-CRD consistently improves its robustness compared to purely supervised training as well as logit-based KD [30] and LSP [32]. Promisingly, lightweight students can be more robust to sparse or noisy scans than cumbersome teachers.