TAG2M- A Task-Agnostic Knowledge Distillation Framework for Distilling GNN to MLP
Ram Ganesh V, Ayush Singh, Aditi Rai, Harsh Pal, Deepanshu Bagotia, Akshay Sethi, Aakarsh Malhotra, Sayan Ranu
Abstract
Graph Neural Networks (Gnns) have achieved remarkable success in various downstream tasks, such as node classification and link prediction. Yet, efficiently deploying Gnns remains a challenge due to their computational complexity. Graph knowledge distillation aims to address this by transferring task-specific structural knowledge from teacher Gnns to lightweight student Gnns or Multi-Layer Perceptrons (MLPs). Despite its promise, existing distillation approaches suffer from several limitations: (i) they require extensive task-specific supervision(ii) they must be retrained separately for each downstream task, and (iii) they often struggle in heterophilous settings. To overcome these challenges, we propose TAG2M, a Task-Agnostic Gnn-to-MLP distillation framework designed for efficient and accurate few-shot inference. TAG2M introduces several novel strategies, including a self-supervised contrastive loss that captures topological information solely from node attributes. Additionally, it leverages Lipschitz embeddings to encode positional information with provable distortion bounds, ensuring robust representation learning. To further enhance adaptability for few-shot inference, TAG2M incorporates a learnable prompt head, which facilitates rapid task adaptation even in label-scarce settings. Unlike prior methods, TAG2M generalizes well across both homophilous and heterophilous datasets while delivering a significant computational advantage, achieving up to a 20X -200X speed-up. Extensive evaluations on 11 public datasets demonstrate its superior accuracy across diverse tasks, including node classification, link prediction, and node regression, outperforming state-of-the-art approaches.
Ask about this paper
Ask your agent about it.
Lune has read the top-tier papers around this one, so every answer names the papers it rests on.
Related papers
- Extracting Low-/High- Frequency Knowledge from Graph Neural Networks and Injecting It into MLPs: An Effective GNN-to-MLP Distillation FrameworkLirong Wu, Haitao Lin, Yufei Huang, Tianyu Fan et al.AAAI 2023 · 52 citations
- ProGMLP: A Progressive Framework for GNN-to-MLP Knowledge Distillation with Efficient Trade-offsWeigang Lu, Ziyu Guan, Wei Zhao, Yaming Yang et al.AAAI 2026 · 1 citation
- AdaGMLP: AdaBoosting GNN-to-MLP Knowledge DistillationWeigang Lu, Ziyu Guan, Wei Zhao, Yaming YangKDD 2024 · 10 citations
- Leap of FAITH from GNN-to-MLP: Fairness Aware Inference via DisTillation of GrapH KnowledgeVipul Kumar Singh, Jyotismita Barman, Sandeep Kumar, Tapan K. Gandhi et al.AAAI 2026
- Preference-driven Knowledge Distillation for Few-shot Node ClassificationXing Wei, Chunchun Chen, Rui Fan, Xiaofeng Cao et al.NeurIPS 2025 · 2 citations
