Fine-Tuning Graph Neural Networks via Graph Topology induced Optimal Transport

20 Mar 2022  ·  Jiying Zhang, Xi Xiao, Long-Kai Huang, Yu Rong, Yatao Bian ·

Recently, the pretrain-finetuning paradigm has attracted tons of attention in graph learning community due to its power of alleviating the lack of labels problem in many real-world applications. Current studies use existing techniques, such as weight constraint, representation constraint, which are derived from images or text data, to transfer the invariant knowledge from the pre-train stage to fine-tuning stage. However, these methods failed to preserve invariances from graph structure and Graph Neural Network (GNN) style models. In this paper, we present a novel optimal transport-based fine-tuning framework called GTOT-Tuning, namely, Graph Topology induced Optimal Transport fine-Tuning, for GNN style backbones. GTOT-Tuning is required to utilize the property of graph data to enhance the preservation of representation produced by fine-tuned networks. Toward this goal, we formulate graph local knowledge transfer as an Optimal Transport (OT) problem with a structural prior and construct the GTOT regularizer to constrain the fine-tuned model behaviors. By using the adjacency relationship amongst nodes, the GTOT regularizer achieves node-level optimal transport procedures and reduces redundant transport procedures, resulting in efficient knowledge transfer from the pre-trained models. We evaluate GTOT-Tuning on eight downstream tasks with various GNN backbones and demonstrate that it achieves state-of-the-art fine-tuning performance for GNNs.

PDF Abstract

Results from the Paper


Task Dataset Model Metric Name Metric Value Global Rank Result Benchmark
Graph Classification BACE GTOT-Tuning ROC-AUC 83.4 # 1
Graph Classification BBBP GTOT-Tuning ROC-AUC 70 # 1
Graph Classification ClinTox GTOT-Tuning ROC-AUC 72 # 1
Graph Classification HIV GTOT-Tuning ROC-AUC 78.2 # 1
Graph Classification MUV GTOT-Tuning ROC-AUC 80 # 1
Graph Classification SIDER GTOT-Tuning ROC-AUC 63.5 # 1
Graph Classification Tox21 GTOT-Tuning ROC-AUC 75.6 # 2
Graph Classification ToxCast GTOT-Tuning ROC-AUC 64 # 2

Methods