Joint task and distribution generalization via graph substructure prompting

Abstract

Driven by the remarkable task-level generalization ability of large language models, an emerging trend of graph learning is to enable fast adaptation to new tasks with limited annotations, and has found applications across a spectrum of domains. Graph meta-learning and graph prompting techniques have demonstrated potential in task generalization by transferring knowledge acquired from prior experiences to new tasks. However, these methods often overlook distribution shifts between training and testing data in real-world scenarios. To fill this gap, we delve into a novel and practical challenge, namely joint task and distribution generalization. Motivated by recent studies that explicitly identifying key substructures related to task prediction can help generalization, we introduce a refiner module to highlight key substructures robust to distribution shifts. To efficiently adapt the refiner to new tasks, we introduce a few extra parameters as prompt vectors to instruct its behavior. Specifically, we employ a global prompt to acquire universal knowledge and task-specific prompts to capture task-relevant information. We pretrain model parameters on known tasks, and efficiently adapt to a target task by merely learning a corresponding classifier and task-specific prompt. Extensive experiments in task generalization show that, the proposed Graph Substructure Prompting (GSP) significantly outperforms recent state-of-the-art (SOTA) methods on both in-distribution (ID) and out-of-distribution (OOD) data, instead of a trade-off between them. GSP also enjoys comparable or even less computational cost as compared to baselines.

Graphical abstract

Keywords

graph neural networks / graph prompting / task generalization / out-of-distribution generalization / few-shot learning

References

[1] Wang Y, Abuduweili A, Yao Q, Dou D. Property-aware relation networks for few-shot molecular property prediction. In: Proceedings of the 35th International Conference on Neural Information Processing System. 2021, 1334

[2] Guo Z, Zhang C, Yu W, Herr J, Wiest O, Jiang M, Chawla N V. Few-shot graph learning for molecular property prediction. In: Proceedings of the Web Conference 2021. 2021, 2559−2567

[3] Huang K, Zitnik M. Graph meta learning via local subgraphs. In: Proceedings of the 34th International Conference on Neural Information Processing Systems. 2020, 492

[4] Hamilton W L, Ying Z, Leskovec J. Inductive representation learning on large graphs. In: Proceedings of the 31st International Conference on Neural Information Processing Systems. 2017, 1025–1035

[5] Zitnik M, Leskovec J . Predicting multicellular function through multi-layer tissue networks. Bioinformatics, 2017, 33( 14): i190–i198

[6] Wang H, Zhang F, Wang J, Zhao M, Li W, Xie X, Guo M. RippleNet: propagating user preferences on the knowledge graph for recommender systems. In: Proceedings of the 27th ACM International Conference on Information and Knowledge Management. 2018, 417−426

[7] Wang H, Zhang F, Zhang M, Leskovec J, Zhao M, Li W, Wang Z. Knowledge-aware graph neural networks with label smoothness regularization for recommender systems. In: Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2019, 968−977

[8] Bose A J, Jain A, Molino P, Hamilton W L. meta-graph: few shot link prediction via meta learning. 2020, arXiv preprint arXiv: 1912.09867

[9] Santoro A, Bartunov S, Botvinick M, Wierstra D, Lillicrap T. Meta-learning with memory-augmented neural networks. In: Proceedings of the 33rd International Conference on Machine Learning. 2016, 1842−1850

[10] Sun M, Zhou K, He X, Wang Y, Wang X. GPPT: graph pre-training and prompt tuning to generalize graph neural networks. In: Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2022, 1717−1727

[11] Liu Z, Yu X, Fang Y, Zhang X. GraphPrompt: unifying pre-training and downstream tasks for graph neural networks. In: Proceedings of the ACM Web Conference 2023. 2023, 417−428

[12] Sun X, Cheng H, Li J, Liu B, Guan J. All in one: multi-task prompting for graph neural networks. In: Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2023, 2120−2131

[13] Huang Q, Ren H, Chen P, Kržmanc G, Zeng D, Liang P, Leskovec J. PRODIGY: enabling in-context learning over graphs. In: Proceedings of the 37th Conference on Neural Information Processing Systems. 2023

[14] Fang T, Zhang Y, Yang Y, Wang C, Chen L. Universal prompt tuning for graph neural networks. In: Proceedings of the 37th International Conference on Neural Information Processing Systems. 2023, 2285

[15] Li H, Wang X, Zhang Z, Zhu W . OOD-GNN: out-of-distribution generalized graph neural network. IEEE Transactions on Knowledge and Data Engineering, 2023, 35( 7): 7328–7340

[16] Gui S, Li X, Wang L, Ji S. GOOD: a graph out-of-distribution benchmark. In: Proceedings of the 36th Conference on Neural Information Processing Systems. 2022, 2059−2073

[17] Liu J, Shen Z, He Y, Zhang X, Xu R, Yu H, Cui P. Towards out-of-distribution generalization: a survey. 2023, arXiv preprint arXiv: 2108.13624

[18] Sun Q, Li J, Peng H, Wu J, Fu X, Ji C, Yu P S. Graph structure learning with variational information bottleneck. In: Proceedings of the 36th AAAI Conference on Artificial Intelligence. 2022, 4165−4174

[19] Shang C, Chen J, Bi J. Discrete graph structure learning for forecasting multiple time series. In: Proceedings of the 9th International Conference on Learning Representations. 2021

[20] Jin W, Ma Y, Liu X, Tang X, Wang S, Tang J. Graph structure learning for robust graph neural networks. In: Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020, 66−74

[21] Miao S, Liu M, Li P. Interpretable and generalizable graph learning via stochastic attention mechanism. In: Proceedings of the 39th International Conference on Machine Learning. 2022, 15524−15543

[22] Wu Y, Wang X, Zhang A, He X, Chua T S. Discovering invariant rationales for graph neural networks. In: Proceedings of the 10th International Conference on Learning Representations. 2022

[23] Li H, Zhang Z, Wang X, Zhu W. Learning invariant graph representations for out-of-distribution generalization. In: Proceedings of the 36th International Conference on Neural Information Processing Systems. 2022, 859

[24] Chen Y, Zhang Y, Bian Y, Yang H, Ma K, Xie B, Liu T, Han B, Cheng J. Learning causally invariant representations for out-of-distribution generalization on graphs. In: Proceedings of the 36th Conference on Neural Information Processing Systems. 2022, 22131−22148

[25] Li H, Wang X, Zhang Z, Zhu W. Out-of-distribution generalization on graphs: a survey. 2022, arXiv preprint arXiv: 2202.07987

[26] Xu K, Hu W, Leskovec J, Jegelka S. How powerful are graph neural networks? In: Proceedings of the 7th International Conference on Learning Representations. 2019

[27] Kipf T N, Welling M. Semi-supervised classification with graph convolutional networks. In: Proceedings of the 5th International Conference on Learning Representations. 2017

[28] Thrun S, Pratt L. Learning to Learn. New York: Springer, 2012

[29] Liu G, Zhao T, Xu J, Luo T, Jiang M. Graph rationalization with environment-based augmentations. In: Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2022, 1069−1078

[30] Han X, Jiang Z, Liu N, Hu X. G-Mixup: graph data augmentation for graph classification. In: Proceedings of the International Conference on Machine Learning. 2022, 8230−8248

[31] Yu J, Liang J, He R. Mind the label shift of augmentation-based graph OOD generalization. In: Proceedings of 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023, 11620−11630

[32] Sui Y, Wang X, Wu J, Lin M, He X, Chua T S. Causal attention for interpretable and generalizable graph classification. In: Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2022, 1696−1705

[33] Fan S, Wang X, Mo Y, Shi C, Tang J. Debiasing graph neural networks via learning disentangled causal substructure. In: Proceedings of the 36th Conference on Neural Information Processing Systems. 2022, 24934−24946

[34] Wu T, Ren H, Li P, Leskovec J. Graph information bottleneck. In: Proceedings of the 34th Conference on Neural Information Processing Systems. 2020, 20437−20448

[35] Houlsby N, Giurgiu A, Jastrzebski S, Morrone B, De Laroussilhe Q, Gesmundo A, Attariyan M, Gelly S. Parameter-efficient transfer learning for NLP. In: Proceedings of the 36th International Conference on Machine Learning. 2019, 2790−2799

[36] Alemi A A, Fischer I, Dillon J V, Murphy K. Deep variational information bottleneck. In: Proceedings of the 5th International Conference on Learning Representations. 2017

[37] Hu W, Fey M, Zitnik M, Dong Y, Ren H, Liu B, Catasta M, Leskovec J. Open graph benchmark: datasets for machine learning on graphs. In: Proceedings of the 34th Conference on Neural Information Processing Systems. 2020, 22118−22133

[38] Wang S, Dong Y, Huang X, Chen C, Li J. Faith: few-shot graph classification with hierarchical task graphs. In: Proceedings of the 31st International Joint Conference on Artificial Intelligence. 2022, 2284–2290

[39] Torres L, Arrais J P, Ribeiro B . Few-shot learning via graph embeddings with convolutional networks for low-data molecular property prediction. Neural Computing and Applications, 2023, 35( 18): 13167–13185

[40] Altae-Tran H, Ramsundar B, Pappu A S, Pande V . Low data drug discovery with one-shot learning. ACS Central Science, 2017, 3( 4): 283–293

[41] Huang R, Xia M, Nguyen DT, Zhao T, Sakamuru S, Zhao J, Shahane SA, Rossoshek A, Simeonov A. Tox21Challenge to build predictive models of nuclear receptor and stress response pathways as mediated by exposure to environmental chemicals and drugs. Frontiers in Environmental Science. 2016 Jan 14;3:85.

[42] Kuhn M, Letunic I, Jensen L J, Bork P . The sider database of drugs and side effects. Nucleic Acids Research, 2016, 44( D1): D1075–D1079

[43] Rohrer S G, Baumann K . Maximum unbiased validation (MUV) data sets for virtual screening based on PubChem bioactivity data. Journal of Chemical Information and Modeling, 2009, 49( 2): 169–184

[44] Richard A M, Judson R S, Houck K A, Grulke C M, Volarath P, Thillainadarajah I, Yang C, Rathman J, Martin M T, Wambaugh J F, Knudsen T B, Kancherla J, Mansouri K, Patlewicz G, Williams A J, Little S B, Crofton K M, Thomas R S . ToxCast chemical landscape: paving the road to 21st century toxicology. Chemical Research in Toxicology, 2016, 29( 8): 1225–1251

[45] Paszke A, Gross S, Massa F, Lerer A, Bradbury J, , . PyTorch: an imperative style, high-performance deep learning library. In: Proceedings of the 33rd Conference on Neural Information Processing Systems. 2019

[46] Fey M, Lenssen J E. Fast graph representation learning with PyTorch geometric. 2019, arXiv preprint arXiv: 1903.02428

[47] Kingma D P, Ba J. Adam: a method for stochastic optimization. In: Proceedings of the 3rd International Conference on Learning Representations. 2015

[48] You Y, Chen T, Shen Y, Wang Z. Graph contrastive learning automated. In: Proceedings of the 38th International Conference on Machine Learning. 2021, 12121−12132

[49] Yin Y, Wang Q, Huang S, Xiong H, Zhang X. AutoGCL: automated graph contrastive learning via learnable view generators. In: Proceedings of the 36th AAAI Conference on Artificial Intelligence. 2022, 8892−8900

[50] Xie Y, Xu Z, Zhang J, Wang Z, Ji S . Self-supervised learning of graph neural networks: a unified review. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023, 45( 2): 2412–2429

RIGHTS & PERMISSIONS

Higher Education Press