Preference-driven Knowledge Distillation for Few-shot Node Classification

AI summary

Overview

Research area: Graph machine learning, specifically few-shot node classification on text-attributed graphs (TAGs), combining large language models (LLMs), graph neural networks (GNNs), and knowledge distillation (KD).

Technical level: Advanced. The paper assumes familiarity with GNN message passing, KL divergence, knowledge distillation, prompt-based LLM fine-tuning, and reinforcement learning (PPO).

One-sentence scope: The paper proposes PKD (preference-driven knowledge distillation), a framework that uses an LLM to annotate the most informative nodes for a set of GNN teachers and then selects, per node, the single most appropriate teacher GNN from which a student GNN should learn.

What This Paper Is About

Node classification on text-attributed graphs normally depends on large amounts of human-annotated labels for GNNs, while LLMs handle zero-/few-shot settings well but are too expensive to run at graph scale. The paper's goal is to blend the two: use an LLM's few-shot judgment to supply labels only where they matter most, and then distill from multiple, structurally different GNN teachers into one student GNN so that each node is taught by the message-passing mechanism that suits its own local topology.

Key Contributions

1. The PKD framework — a preference-driven knowledge distillation framework that unites an LLM with multiple GNN architectures for few-shot node classification on TAGs, built around two modules: the GNN-preference-driven Node Selector (GNS) and the Node-preference-driven GNN Selector (NGS).

2. GNN-preference-driven Node Selector (GNS) — a selector based on a defined K-uncertainty (δ_K), computed as the pairwise KL divergence among the teacher GNNs' SoftMax outputs, which determines which nodes the LLM should annotate. The paper proves (Proposition 3.1) that nodes with higher K-uncertainty are beneficial for GNN enhancement. GNS also includes a Distance-based Neighbor Selector (DNS) module that performs KNN search around each selected node in the pretrained teacher GNNs' embedding spaces and removes repeated neighbors.

3. Node-preference-driven GNN Selector (NGS) — a reinforcement learning formulation in which the fine-tuned LLM acts as the agent, node-specific attributes (semantic, structural, and prediction attributes) form the state, the choice of teacher is the action, and the student GNN's performance is the reward, enabling tailored per-node knowledge distillation.

4. Graph Topology Aware (GTA) prompts — four fine-tuning instruction types (Connectivity, Degree, Cycle Detection, Text Generation) used to fine-tune the LLM so it can comprehend graph topology, plus validation across nine TAGs showing PKD even defeats some state-of-the-art methods that use more node labels.

Main Findings

• Strong overall accuracy: With Llama-3.1-8B-Instruct and 5 labeled nodes per class (# LN 5), PKD reaches Cornell 80.95 ± 1.1, Washington 83.74 ± 0.4, Texas 86.31 ± 0.5, Wisconsin 76.89 ± 0.9, Amazon Ratings 66.79 ± 0.3, Ogbn-Arxiv 61.03 ± 0.7, Wiki CS 81.39 ± 0.4, Pubmed 85.69 ± 0.3, and Cora 91.14 ± 0.3. The paper states it "almost achieves the best or second-best accuracy results."

• Performance under extreme label scarcity: At # LN 1 the same configuration still reaches Cornell 74.60 ± 2.1, Washington 76.64 ± 0.9, Texas 80.36 ± 1.3, Wisconsin 69.32 ± 2.8, Amazon Ratings 64.11 ± 1.7, Ogbn-Arxiv 53.67 ± 1.6, Wiki CS 79.31 ± 0.8, Pubmed 83.75 ± 1.1, and Cora 85.64 ± 2.1.

• Robust across LLMs: Replacing Llama-3.1-8B-Instruct with Qwen2.5-7B-Instruct or Mixtral-7B-Instruct-v0.3 (Table 2) still yields competitive results; for example, PKD Qwen at # LN 5 gives Cora 90.07 ± 0.4 and Texas 85.71 ± 0.2, and PKD Mixtral at # LN 5 gives Cora 88.56 ± 0.7 and Cornell 81.58 ± 2.1. The paper reports that PKD "consistently achieves superior node classification results across all datasets, irrespective of the specific type of LLM."

• One baseline wins on one dataset: GAugLLM achieves the best result on Pubmed (85.98 ± 0.6) but is outperformed by PKD on other datasets.

• Why competing KD methods struggle: KDGA and BGNN are described as excessively sensitive to GNN selection; MSKD uses a fixed message-passing mechanism and underperforms on all datasets compared to PKD; AGST goes Out-Of-Memory on Ogbn-Arxiv; IceBerg performs poorly on heterophily graphs; GCNII and EGNN are limited by extreme label insufficiency.

• Simple selection strategies lose: Random / # LN 5 and Voting / # LN 5 are reported to be defeated by PKD on all datasets.

• LLM pseudo-labels vs. real labels: Table 3 compares three baselines (GCNII, IceBerg, MSKD) and PKD under # LN 5, 48% LLM-generated labels, and 48% real labels on Cora, Wiki CS, Washington, and Wisconsin. For GCNII and IceBerg, LLM-annotated labels improve performance over # LN 5 on all datasets, but the same number of real labels performs better. PKD with Llama at # LN 5 shows Cora 90.27, Wiki CS 81.39, Washington 83.74, and Wisconsin 76.89 in that table.

• Fine-tuned LLM is strong zero-shot: The LLM fine-tuned with GTA prompts "demonstrates pretty zero-shot node classification performance, surpassing some semi-supervised GNNs," as shown in Figure 2 for homophily and heterophily graphs.

• Ablation confirms all three components matter: Under default settings (# LN = 3, K = 4), the GTA prompts, the DNS module, and the preference ranking V_PR each produce accuracy increments. Without DNS, neighbors are chosen directly from the adjacency matrix; without V_PR, training data is expanded by random selection. The table reports percentage increments for Cora and Amazon Ratings, though the presented values are partially garbled in the supplied text.

• All three reward parts help: Comparing R1 (student accuracy only), R2 (accuracy plus negative cross-entropy loss), and R3 (R2 plus negative distillation loss) in Figure 4, the full reward combining all three parts performs best.

• K matters and K = 4 is the sweet spot: Varying K over {1, 2, 3, 4, 5} for homophily and heterophily graphs, zero-shot node classification accuracy of the fine-tuned LLM is highest at K = 4 on most graphs. The text is truncated at the point where the accuracy trend is described further.

Methodology in Plain English

The researchers start by teaching an LLM about graphs. Since off-the-shelf LLMs cannot read raw graph structure, they fine-tune it with "Graph Topology Aware" prompts containing four kinds of small tasks: deciding whether two nodes are connected, reporting a node's degree from the adjacency matrix, detecting whether a cycle exists in a node sequence, and generating text for nodes based on a random walk.

Next they pick which nodes are worth the expense of an LLM annotation. Instead of trusting one GNN's uncertainty, they run four architecturally different teacher GNNs (GCN, GAT, APPNP, and H2GCN for homophily graphs; DirGNN, GPRGNN, HoloNets, and H2GCN for heterophily graphs) and measure how much the four disagree on each node using pairwise KL divergence. Nodes with the greatest disagreement have the highest K-uncertainty and are ranked highest for annotation. To help the LLM label a selected node well, a Distance-based Neighbor Selector searches each teacher's embedding space for nearby nodes and turns their text plus the node's own text into a prompt. The LLM's answers become extra training labels for the teachers.

Finally, the system decides which teacher should teach which node. It builds a per-node prompt combining that node's semantic, structural, and prediction attributes, and feeds it to the fine-tuned LLM. The LLM, wrapped with two small MLP projectors and trained with a simplified PPO with a CLIP-style constraint, outputs a probability distribution over the teachers and samples one. The reward mixes the student GNN's classification accuracy, its cross-entropy loss, and its distillation loss. The student GNN is then trained with a combined objective of distillation loss, cross-entropy on the few real labels, and an entropy term that pushes the student's logits toward one-hot vectors. Training alternates: while distilling, the selector's parameters are fixed; while the selector learns, the student's parameters are fixed.

Why This Matters

Impact on research: The paper argues that a single fixed message-passing mechanism cannot handle the diverse local topologies of real-world TAGs, and that prior multi-teacher distillation methods (BGNN, MTAAM, DMKD, etc.) ignore node-specific topology, which can cause no improvement or even performance degradation. PKD reframes distillation as two preference problems — which nodes to label, and which teacher to trust per node — and shows that both can be learned with an LLM in the loop, achieving competitive accuracy with as few as 1 to 5 labeled nodes per class.

Real-world applications (drawn from the TAG domains the paper names, such as citation, webpage, and product graphs):

• Academic citation networks, where labeling papers by field is expensive but text abstracts are abundant.
• Product graphs in e-commerce, where only a few items per category carry reliable labels.
• Web page classification and organization, where nodes carry page text and hyperlinks.
• Any text-attributed graph setting where manual annotation is "tedious, expensive, and time-consuming" and only a handful of labels per class exist.

Industry relevance: The work targets the practical trade-off between LLM quality and LLM cost. By querying the LLM only on high-uncertainty nodes rather than all unlabeled nodes, it reduces LLM inference costs, and by using a compact student GNN at test time it avoids the LLM's "large parameter scale" inference burden. The reported results also show the framework is not tied to one LLM vendor, since Qwen2.5-7B-Instruct and Mixtral-7B-Instruct-v0.3 produce comparable outcomes.

Future Directions

• Closing the gap with real labels: Table 3 shows that, for GCNII, IceBerg, and MSKD, 48% real labels still beat 48% LLM-generated labels, indicating room to improve LLM annotation quality.
• Understanding the truncated sensitivity results: The reported analysis of K is cut off after the observation that K = 4 is best on most graphs, leaving the behavior at other K values on heterophily versus homophily graphs as an open question in the supplied text.
• Scaling and cost: The paper does not report results on graphs larger than the nine studied (Cornell, Washington, Texas, Wisconsin, Amazon Ratings, Ogbn-Arxiv, Wiki CS, Pubmed, Cora), and AGST already fails with Out-Of-Memory on Ogbn-Arxiv, so behavior at much larger scale remains untested.
• Generalization of the selector: Whether the RL-based teacher selection and the K-uncertainty ranking remain effective when the number of teachers B grows, or when teacher GNNs are weak or highly correlated, is not reported in the paper content provided.

Target Audience

Researchers and graduate students working on graph representation learning, few-shot learning, graph knowledge distillation, or LLM-GNN integration will benefit most, along with practitioners building label-efficient node classification systems on text-attributed graphs who need to weigh LLM inference cost against GNN deployment cost. Readers should be comfortable with GNN architectures, KL divergence, prompt-based LLM fine-tuning, and reinforcement learning, since the paper's method section builds directly on all four.

Authors’ abstract

Graph neural networks (GNNs) can efficiently process text-attributed graphs (TAGs) due to their message-passing mechanisms, but their training heavily relies on the human-annotated labels. Moreover, the complex and diverse local topologies of nodes of real-world TAGs make it challenging for a single mechanism to handle. Large language models (LLMs) perform well in zero-/few-shot learning on TAGs but suffer from a scalability challenge. Therefore, we propose a preference-driven knowledge distillation (PKD) framework to synergize the complementary strengths of LLMs and various GNNs for few-shot node classification. Specifically, we develop a GNN-preference-driven node selector that effectively promotes prediction distillation from LLMs to teacher GNNs. To further tackle nodes' intricate local topologies, we develop a node-preference-driven GNN selector that identifies the most suitable teacher GNN for each node, thereby facilitating tailored knowledge distillation from teacher GNNs to the student GNN. Extensive experiments validate the efficacy of our proposed framework in few-shot node classification on real-world TAGs. Our code is be available.