Preference-driven Knowledge Distillation for Few-shot Node Classification

Thanks for reviewing this manuscript and providing valuable comments. We appreciate the positive feedback from all reviewers and your constructive suggestions for further improvement. W and Q represent weakness and question.

W1: While creating the GNS to avoid the high cost of labeling all nodes with an LLM is a reasonable approach, the subsequent need to use the LLM again to select a GNN for each node's prediction seems to reintroduce significant costs. This appears to offset the initial advantage.

Reply: Thanks for your comments. We use the fine-tuned LLM as the RL agent to select for each node the most appropriate teacher GNN. The prompt is constructed by combining the node's semantic, structural, and prediction attributes derived from the enhanced teacher GNNs. Note that the RL agent should be able to comprehend the node's structural attributes. If replacing the fine-tuned LLM with a small model such as BERT, we should fine-tune BERT again for comprehending the node's structural attributes. In addition, introducing an extra model that may hinder the synergy of the complementary strengths of LLMs and various GNNs is inelegant. Since the parameters of the fine-tuned LLM (the RL agent) are frozen, the time bottleneck is the inference time. Therefore, we can reduce the inference time by designing efficient attention mechanisms, using speculative decoding, etc. Additionally, considering that PKD outperforms the comparison method by a large margin, we can reduce the label ratio from 48% to, say, 10%, which depends on the node category induction capability of LLM, thereby decreasing the time cost caused by the teacher selection.

W2: Mentioning the differences between this paper's task and traditional meta-learning-based graph few-shot learning tasks (e.g., N-way, K-shot settings) in the 'Related Work' section could enhance the clarity of the task definition.

Reply: Thanks for your suggestions. As defined in [1], the nodes in the graph meta-learning task for node classification can be divided into the training set and the disjoint test set. Nodes in the test set are associated with totally different new classes. Unlike this, the labeled nodes in this paper's task already cover the entire range of categories, which is the most different setting. In the final version of this paper, we will add this point in the Related Work part.

[1] Zhou, Fan, Chengtai Cao, Kunpeng Zhang, Goce Trajcevski, Ting Zhong, and Ji Geng. ``Meta-gnn: On few-shot node classification in graph meta-learning.'' In Proceedings of the 28th ACM international conference on information and knowledge management, pp. 2357-2360. 2019.

Thanks for reviewing this manuscript and providing valuable comments. We appreciate the positive feedback from all reviewers and your constructive suggestions for further improvement. W and Q represent weakness and question.

W1: The computational complexity analysis, while provided in the appendix, reveals significant overhead that may limit practical applicability.

W4: The approach requires multiple training phases (LLM fine-tuning, teacher pretraining, RL training, student training), making it computationally expensive and potentially impractical for large-scale applications.

Reply: Thanks for your concerns. We agree with you. In the future, we can first distill knowledge from LLM with 72B parameters to LLM with 1.5B parameters without much performance degradation. Then, fine-tuning the 1.5B LLM will cost less time compared with the 8B LLM used in the paper. The time cost for training the teacher GNNs (see the below table) and the student GNN is acceptable. In the paper, we use the fine-tuned LLM as node label annotators for up to 48% of the total nodes. This step can be parallelized in the future, i.e., we can use multiple threads to query LLM for multiple nodes each time. The fine-tuned LLM is also used as the RL agent to select for each node the most appropriate teacher GNN. Since the parameters of the fine-tuned LLM are frozen, the time bottleneck of LLM annotation is the inference time. For this, we can reduce the inference time by designing efficient attention mechanisms, using speculative decoding, etc. Additionally, the size of the expanded training data increases the PKD’s per-epoch time. Considering that PKD outperforms the comparison method by a large margin, we can reduce the label ratio from 48% to, say, 10%, which depends on the node category induction capability of LLM, thereby decreasing the time cost caused by the RL training. Through the above modifications and innovations, PKD can be scalable to large graphs.

The time cost for training the teacher GNNs on the Cora, Ogbn-Arxiv and Amazon Ratings.

W2: There is a lack of discussion on the selection of student models. Can the ordinary convolution mechanism acquire the knowledge of multiple teacher models with multiple architectures?

Reply: Thanks for your insightful considerations. For homophily and heterophily graphs, we just use the basic and simple GNNs, i.e., GCN and H$_2$GCN, for the student GNN model. Following the settings in the manuscript, we choose GAT, APPNP, H$_2$GCN as the teacher GNNs for homophily graphs and select DirGNN, GPRGNN, HoloNets as the teacher GNNs for heterphily graphs.

To verify whether the student GNN can acquire the knowledge of multiple teacher GNNs with multiple architectures, we separately distill the knowledge from different pretrained teacher GNNs to the student GNN only by the KD loss and compare their performance. The results are shown in the below table. It is evident that the teacher GNNs and the student GNN perform similarly on the same graph. Therefore, the ordinary convolution mechanism of the student GNN can acquire the knowledge of multiple teacher GNNs with multiple architectures. We will add the experimental results in the final version of this paper.

The comparison between multiple teacher GNNs and the student GNN on Cora and Amazon Ratings.

W3: The paper lacks discussion of teacher GNN pretraining procedures.

Reply: For the teacher GNN pretraining procedures, we just utilize the few given node labels. After pretraining, we collect their logit outputs to conduct the GNN-preference-driven node selection based on the uncertainty. Thanks for your reminder. We will add this detailed description in the final version of this paper.

Q1: In Figure 2, what exactly does ''LLM with structural attributes'' refer to? Additionally, why are results for OGBN-ARXIV missing from this analysis?

Reply: We utilize the original LLM for node label annotation using the structural attribute, which refers to the prompt consisting of the node's neighbors chosen by the Distance-based Neighbor Selector and their semantic attributes. In Figure 2, the ''LLM with structural attributes'' refers to it.

For the results on Ogbn-Arixv, we present them in the table below and will add them in the final version of this paper. It is evident that our method is also effective on Ogbn-Arxiv.

The node classification accuracies of different methods on Ogbn-Arxiv.

Q2: Given that different nodes require different message-passing mechanisms (the core motivation for NGS), why are fixed architectures (GCN for homophily, H$_2$GCN for heterophily) chosen as student models? Wouldn't the student's message-passing mechanism also impact the effectiveness of knowledge distillation?

Reply: Thanks for your insight again. Since nodes have intricate local topologies, such as the various node degree, a single GNN cannot capture the essence of nodes completely. To this end, we use multiple teacher GNNs and select for each node the most appropriate teacher GNN. As our response to W2, the ordinary convolution mechanism of the student GNN can acquire the knowledge of multiple teacher GNNs with multiple architectures. Thus, using the fixed architectures for the student GNN is simple and effective.

The versatility of PKD. Indeed, our method has the versatility across various student GNN architectures. To further verify this, we compare the performance of the teacher GNN with the student GNNs distilled from it on two datasets. Specifically, for homophily graphs, we select GCN, GAT, and APPNP as the student GNNs, with H$_2$GCN chosen as the teacher GNN. For heterophily graphs, we choose DirGNN, GPRGNN, and HoloNets as the student GNNs, with H$_2$GCN selected as the teacher GNN. The results are presented in the following table. Although the message-passing mechanisms of student GNNs are different, they can achieve similar performance on the same graph after distillation from the teacher. Notably, we do not provide any node labels for the student GNNs during knowledge distillation. Therefore, the student's message-passing mechanism does not impact the effectiveness of knowledge distillation. These experimental results will be added to the final version of this paper.

The comparison between the teacher GNN and multiple student GNNs on Cora and Amazon Ratings.

Thanks for reviewing this manuscript and providing valuable comments. We appreciate the positive feedback from all reviewers and your constructive suggestions for further improvement. W and Q represent weakness and question.

W1: Although the paper provides a time-complexity analysis and a runtime table, PKD’s iterative LLM queries and RL loops make each training epoch slower than most baselines, thereby limiting its scalability to million-node graphs.

W2: Due to its iterative LLM queries and RL loops, PKD’s per-epoch time is higher than that of most baselines.

Reply: Thanks for your concerns. We agree with you. In the future, we can first distill knowledge from LLM with 72B parameters to LLM with 1.5B parameters without much performance degradation. Then, querying LLM will cost less time compared with that when querying an LLM with 8B parameters as used in the paper. The GNS iteratively queries LLM for node label annotations to expand the training data to 48% of the total nodes. This step can be parallelized, i.e., we can use multiple threads to query LLM for multiple nodes each time. The fine-tuned LLM is also used as the RL agent to select for each node the most appropriate teacher GNN. Since the parameters of the fine-tuned LLM are frozen, the time bottleneck of LLM query is the inference time. For this, we can reduce the inference time by designing efficient attention mechanisms, using speculative decoding, etc. Additionally, the size of the expanded training data also increases the PKD's per-epoch time. Considering the fact that PKD outperforms the comparison methods by a large margin, we can reduce the label ratio, e.g., from 48% to 10%, which depends on the node category induction capability of LLM, thereby decreasing the time cost caused by the PPO loop. Through the above modifications and innovations, PKD can be scalable to large graphs.

W3: Some hyperparameters are missing from the analysis, e.g. $\alpha,\beta,\gamma,\eta$.

Reply: For this, we supplement the hyperparameters sensitivity analysis for $\alpha,\beta,\gamma,\eta$ on Cora and Amazon Ratings. Specifically, when analyzing $\alpha$, we set $\beta=1$, $\gamma=1$, $\eta=0.5$. This strategy also applies to the sensitivity analysis of $\beta$ and $\gamma$. As for $\eta$, we fix $\alpha$, $\beta$, and $\gamma$ to their respective best values. The results of various values for them are presented in the following four tables. Intuitively, $\beta$ and $\eta$ are more sensitive than the other two hyperparameters, because $\beta$ is related to the supervision signal provided by the LLM-generated annotation labels and $\eta$ is related to the student GNN's performance.

The influence of $\alpha$.

The influence of $\beta$.

The influence of $\gamma$.

The influence of $\eta$.

Q1: The current reward function combines two loss functions and accuracy, which raises the complexity, has a simpler expression of the reward function been explored?

Reply: To address this, we conduct an ablation study on each part of the reward function. The visualization results are shown in Figure 4 in the manuscript. The single usage of accuracy ($R_1$) does not perform as well as using more parts (combination with $L_{CE}$). Figure 4 also shows that the $L_{DL}$ part is beneficial to quickly improve model performance and maintain stability.

To further verify the necessity of using all three parts, we train the Node-preference-driven GNN Selector on Cora and Amazon Ratings using four kinds of combinations of the three parts. The results are presented in the below table. Therefore, the best reward function contains all three parts.

Results of the combination of various loss functions.

Q2: Is there any way to further minimize large language model calls?

Reply: Specifically, the LLM is mainly used for node category induction and teacher GNNs selection. Considering that PKD performs much better than the comparison methods, we can probably reduce the large language model calls by reducing the number of annotated nodes from 48% to, say, 10%, which depends on the node category induction capability of LLM. Otherwise, more human annotations for node labels are needed in order to further minimize large language model calls.

Q3: What is the need to use reinforcement learning?

Reply: Our final goal is to select the only one best teacher GNN for each node. That is, we need to explore the discrete action space and find a series of assignment actions to get the highest global reward across the expanded training data. Reinforcement learning (RL) allows us to navigate through the large action space by learning an optimal policy that maximizes the long-term reward. Through interaction with the environment, the selector progressively refines its decisions on node-to-teacher assignments, leading to a more efficient and effective assignment strategy. Besides, we can handle the challenges of discrete decision-making and achieve better performance on the task by leveraging RL. It can also incorporate heterogeneous metrics, such as distillation loss, cross entropy loss, and classification accuracy, into the reward function, and has a more flexible training objective.

Q4: How does the method proposed in this paper compare to the results without using reinforcement learning?

Reply: The results on Cora and Amazon Ratings of our RL-based method and the other three teacher selection methods (entropy-based ranking, i.e., selecting the teacher GNN with the highest prediction confidence, random selection, and end-to-end learning) are shown in the below table. It is obvious that the RL-based method significantly outperforms the other three methods.

The compared results on Cora and Amazon Ratings of our RL-based method and the other three methods .

Thanks for reviewing this manuscript and providing valuable comments. We appreciate the positive feedback from all reviewers and your constructive suggestions for further improvement. Considering the correlation between Weaknesses and Questions, we reply after integrating them as below.

[W1, Q1]: Please quantify the LLM budget: report how many calls and tokens were needed to reach the 48% training split on each dataset, give a rough GPU-hour cost, and show how a strong baseline performs if it receives the same number of LLM-generated labels.

Reply: As shown in Figure 2 in the manuscript, the fine-tuned LLM already has significant performance for zero-shot node classification. When conducting the LLM annotation, we only call LLM once for each node. That is, the number of LLM calls ($t$) is equal to the number of nodes selected by the GNS (GNN-preference-driven Node Selector), i.e., $t = 0.48 \cdot N - Q$, where $N, Q$ are the numbers of all nodes and the initial labeled nodes, respectively. The needed number of tokens ($T$) can be calculated by $T = m \cdot (0.48 \cdot N - Q)$, where $m$ is the LLM's max_tokens value, which is set to 2048 in this paper. For instance, the values of $t$ and $T$ for Cora are 1265 and 2590720 under the # LN 5 setting.

To reach the 48% training split on each dataset, we first select the nodes preferred by GNN based on the preference ranking and then conduct the node-wise category-induction prompts to query LLM for their pseudo labels. In this process, we only utilize one GPU card (NVIDIA A800-SXM4-80GB), and the time costs using different LLMs are listed in below table. The results show that that inference time varies across different LLMs, even when evaluated under the same environmental settings, with larger models exhibiting slower inference time.

The time costs of different LLMs.

The Influence of LLM-Generated Pseudo-Labels. To demonstrate the effect of LLM-generated pseudo-labels, we compare the node classification performance of PKD and three baselines under different label settings (# LN 5, 48% training ratio expanded by the annotated labels and real labels, respectively). The experiments are conducted on four datasets (Cora, Wiki CS, Washington, and Wisconsin). The results are presented in below table. For GCNII and IceBerg, which are proposed to tackle the challenge of sparse labels, using the LLM-anotated node labels can improve their performance on all datasets. However, using the same number of real labels achieves better performance.

Classification accuracy comparisons.

[W2, Q2]: Show that the PPO teacher selector is really needed by providing learning curves and variance across seeds and by comparing its accuracy and runtime with two simpler choices.

Reply: Figure 4 in the manuscript shows the learning curves (accuracy vs. epoch) of the NGS (i.e., the RL component). It is evident that the NGS with the $R_3$ reward function converges to a stable state. We also plan to provide the learning curves (loss vs. epoch) of the NGS in the final version of this paper. To demonstrate that the PPO teacher selector is really needed, we replace it with two simple teacher selectors and repeat the knowledge distillation five times using distinct random seeds. The first method selects a teacher GNN for each node randomly, while the second method chooses the teacher GNN for each node based on the entropy value of its logit outputs. The lower the entropy value, the higher the confidence of the teacher. We select the most confident GNN as the teacher. We record the average test classification accuracy and standard deviation of each method with parameters that lead to the peak validation accuracy. The below table presents the results on Cora which demonstrate that our method is both more stable and more effective than the other two methods.

The comparison between PPO and other methods on Cora.

[W3, Q3]: Complete the ablation study with a variant that omits the RL selector, one that keeps the LLM frozen, and a short sweep over the loss-weight coefficients $\alpha,\beta,\gamma,\eta$ to demonstrate robustness.

Reply: We compare our method PKD with two other variants on Cora and Amazon Ratings. The first variant replaces the RL selector with a random selector, and the second variant omits the fine-tuning of GTA prompts and directly utilizes the annotated node labels generated by the original LLM to expand the training data. The results are presented in the below table and will be added to the final version of this paper. It is evident that PKD outperforms the other two variants by a large margin.

The comparison between PKD and the other two variants.

Hyperparameter Sensitivity Analysis. Additionally, we perform the hyperparameter sensitivity analysis over the loss-weight coefficients $\alpha,\beta,\gamma,\eta$ on two datasets, and report the results in the below tables. For the sensitivity analysis of $\alpha$, we set $\beta=1$, $\gamma=1$, $\eta=0.5$. This strategy also applies to the sensitivity analysis of $\beta$ and $\gamma$. As for $\eta$, we fix $\alpha$, $\beta$, and $\gamma$ to their respective best values. We can see that $\beta$ and $\eta$ are more sensitive than the other two hyperparameters, because $\beta$ is related to the supervision signal provided by the LLM-generated annotation labels and $\eta$ is related to the student GNN's performance.

The influence of $\alpha$.

The influence of $\beta$.

The influence of $\gamma$.

The influence of $\eta$.

[W4, Q4]: Add practical resource figures: GPU type, peak memory and total wall-clock time for GTA fine-tuning, node annotation, teacher re-training and the PPO loop, at least on CORA and OGBN-ARXIV, so the scalability of PKD is clear.

Reply: We implement our proposed PKD with PyTorch (2.5.1), PyTorch Geometric (2.6.1), Python (3.10.16), Transformers (4.50.3), and vllm (0.7.0). We conduct all experiments on the NVIDIA A800-SXM4-80GB GPU and Intel(R) Xeon(R) CPU Max 9468. The peak memories of performing PKD on the Cora and Ogbn-Arxiv are listed in the below label.

The peak memory of performing PKD on the Cora and Ogbn-Arxiv.

The time costs for LLM annotation on Cora and Ogbn-Arixv can refer to the Reply to the [W1, Q1]. The time costs for GTA fine-tuning, the Distance-based Neighbor Selector (DNS), teacher re-training and PPO loop on Cora and Ogbn-Arixv are presented in the following tables, respectively. Notably, for the teacher re-training stage, we show all the teacher GNNs that are trained togethe, and their total training times and epochs.

The GTA fine-tuning configurations on Llama-3.1-8B-Instruct.

The times of DNS (seconds) on the Cora and Ogbn-Arxiv.

The time costs of retraining teacher GNNs on Cora, Amazon Ratings, and Ogbn-Arxiv.

Running times of PPO on Cora and Obgn-Arixv. "m" and "s" denote minute and second.

Thank you for your response - I am satisfied with the revisions provided, and I maintain my positive score.