Fine-Tuning Graph Neural Networks via Graph Topology induced Optimal Transport
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Graph Classification | PROTEINS | Accuracy62.89 | 742 | |
| Molecular property prediction | MoleculeNet BBBP (scaffold) | ROC AUC71.5 | 117 | |
| Molecular property prediction | MoleculeNet SIDER (scaffold) | ROC-AUC0.633 | 97 | |
| Molecular property prediction | MoleculeNet MUV (scaffold) | ROC-AUC0.85 | 68 | |
| Molecular property prediction | MoleculeNet HIV (scaffold) | ROC AUC81.1 | 66 | |
| molecule property prediction | MoleculeNet (scaffold split) | BBBP71.5 | 58 | |
| Molecular property prediction | MoleculeNet Tox21 (scaffold) | ROC-AUC78.6 | 48 | |
| Molecular property prediction | MoleculeNet ClinTox (scaffold) | ROC-AUC0.779 | 45 | |
| Molecular Property Classification | BACE (MoleculeNet) scaffold (test) | ROC-AUC0.853 | 30 | |
| Graph Classification | IMDB-M | Accuracy0.5113 | 22 |