ABKD: Graph Neural Network Compression with Attention-Based Knowledge Distillation

1. The idea of distilling hidden layers is very novel and is very common in knowledge distillation. This method also does not consider the graph structure in distillation.
2. GNNs typically are not deep. So it is unclear how much a small attention can contribute to the performance. We can also use some simple heuristics to identify the pair, such as aligning later layers with later layers and beginning layers with beginning layers.
3. The distillation process could be costly as the method distills every teacher-student pair.
4. I am curious how the method performs on other OGBN datasets, such as OGBN-product.

• The accuracy improvements over prior methods are quite small, just 1-2% on relatively small datasets. The gains may not be significant enough to demonstrate a strong advantage.
• Ablation studies are limited to analyzing the impact of different components. More in-depth empirical analysis could be beneficial.
• Novelty is somewhat limited since attention to layer alignment exists in other domains. Adaptations for graphs are incremental.

• Presentation/typesetting is poor and can be improved

  Typesetting: I feel authors have not adhered to the formatting instructions (e.g. using different bibliography style, table captions under tables, smaller font size for figure and table captions). Matrices and number sets use the same formatting family (e.g. not using $\mathbb{R}$, but $R$ for reals, etc.). Table formatting can also be improved link1

  Presentation: Table 1 feels unnecessary. $\S$3.1.3 defines only GCN formulation, but experiments use RGCN/GAT/GraphSage (why was the formulation included then?) Proofs and other mathematical details are not easy to follow. As a simple example, you can express the dissimilarity score $D_{ij}$ by first defining graph feature vectors for layers $i$ and $j$ of teacher and student (by, e.g. using $\sum$ operator) and then simply defining $D_{ij}$ as the Euclidean distance between the difference of those two feature vectors. Model format in Tables 2/3 is not provided.

• Although the approach is interesting, it considers pooled graph embeddings, and not aligning individual node embeddings (as in G-CRP).

• Please check equations carefully, as I am not confident they always match the code. E.g. in eq. 4 you calculate the squared Euclidean distance between graph feature vectors and then normalise it in eq.5, but in the code, you seem to normalise it trice: L170+L171 and L181 of distillation.py.

• Theorem proving is not formal. If you want to formalise theorem 1, please define "will be affected" mathematically and prove it.

• There are some experimental details missing: Learning rate, number of epochs, etc. README in the repository uses quite cryptic flag names to infer hyperparameter setup easily.

• Hardware is reported in the paper, but I found no runtimes reported in the main text or supplementary.

• Some experiments do not seem very sensible: Although past literature has studied knowledge distillation from deeper into a shallower network (e.g. 5 layer to 2 layer in G-CRD), performing distillation from a 64L network to a 4L network seems counterintuitive -- you are forcing the student GNN to capture properties of a 64-hop neighbourhood with a 4-hop filter. Even more surprisingly, according to Figure 5 all of the 64 layers' embeddings (each of dimensionality 64) get compressed into 1 layer embedding of dimensionality 4 for Pubmed. I believe the results here have an alternative explanation (e.g. something related to dataset complexity/diameter/etc).

• Some experimental results do not seem statistically significant: In Table 3, the difference to G-CRD is well within 2 standard deviations for 3/4 datasets. Similarly, in Table 6, projection vs no projection is quite close (could this be only due to the additional parameters used during the KD phase?)

• All experiments only cover transductive learning.

This work proposes an Attention-Based Knowledge Distillation (ABKD) approach for the compression of Graph Neural Networks (GNNs). Despite the clear effort and potential of this work, there are several concerns to be addressed:

Missing Existing Literature:

The manuscript does not sufficiently contextualize the proposed ABKD within the landscape of existing research on knowledge distillation for GNNs. As the references provided [1-7] illustrate, distilling knowledge from intermediate layers of GNNs is not a novel concept. The differentiation of ABKD from these established methods is not clearly articulated. A detailed comparison with these specific works would be crucial to highlight the unique contributions of ABKD.

Technical Novelty and Justification:

The paper would benefit from a clearer technical exposition of ABKD, especially regarding how it improves upon the limitations of the methods cited. There should be a specific discussion on the advantages that ABKD offers over these prior works, supported by empirical evidence. The proposed attention-based KD introduces additional learnable parameters. It is not clear how the supervision prevents convergence to trivial solutions.

Empirical Comparison and Evaluation:

The evaluation lacks a comprehensive comparison with the methods listed in the references. Including such comparisons would provide a clearer picture of the performance and efficiency gains of ABKD. We recommend an expanded evaluation section that not only benchmarks against these methods but also discusses the reasons for any observed improvements.

Clarity of Contribution:

The current contribution of the paper is overshadowed by the lack of distinction from prior works. It is crucial to clarify what aspects of ABKD are novel and how they advance the state of the art.

1. Kim, J., Jung, J. and Kang, U., 2021. Compressing deep graph convolution network with multi-staged knowledge distillation. Plos one, 16(8), p.e0256187.

2. Jing, Y., Yang, Y., Wang, X., Song, M. and Tao, D., 2021. Amalgamating knowledge from heterogeneous graph neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 15709-15718).

3. Yang, C., Liu, J. and Shi, C., 2021, April. Extract the knowledge of graph neural networks and go beyond it: An effective knowledge distillation framework. In Proceedings of the web conference 2021 (pp. 1227-1237).

4. Zheng, W., Huang, E.W., Rao, N., Katariya, S., Wang, Z. and Subbian, K., 2021. Cold brew: Distilling graph node representations with incomplete or missing neighborhoods. arXiv preprint arXiv:2111.04840.

5. Wang, C., Wang, Z., Chen, D., Zhou, S., Feng, Y. and Chen, C., 2021. Online adversarial distillation for graph neural networks. arXiv preprint arXiv:2112.13966.

6. Joshi, C.K., Liu, F., Xun, X., Lin, J. and Foo, C.S., 2022. On representation knowledge distillation for graph neural networks. IEEE Transactions on Neural Networks and Learning Systems.

7. Liu, J., Zheng, T. and Hao, Q., 2022. HIRE: Distilling high-order relational knowledge from heterogeneous graph neural networks. Neurocomputing, 507, pp.67-83.