Lune

KDD2025Top-tier venue

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

2025Year

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.

Questions to start from

Your agent calls

Lunesearch_papers

Ask in Lune

Free to start. No credit card required.

Related papers

Dusk over the sea between two cliffs drawn in fine vertical lines