GNN-SKAN: Harnessing the Power of SwallowKAN to Advance Molecular Representation Learning with GNNs

TL;DR

This work introduces a new class of GNNs that integrates the Kolmogorov-Arnold Networks (KANs), known for their robust data-fitting capabilities and high accuracy in small-scale AI + Science tasks, and introduces the first work to integrate KANs into GNN architectures tailored for molecular representation learning.

Abstract

Effective molecular representation learning is crucial for advancing molecular property prediction and drug design. Mainstream molecular representation learning approaches are based on Graph Neural Networks (GNNs). However, these approaches struggle with three significant challenges: insufficient annotations, molecular diversity, and architectural limitations such as over-squashing, which leads to the loss of critical structural details. To address these challenges, we introduce a new class of GNNs that integrates the Kolmogorov-Arnold Networks (KANs), known for their robust data-fitting capabilities and high accuracy in small-scale AI + Science tasks. By incorporating KANs into GNNs, our model enhances the representation of molecular structures. We further advance this approach with a variant called SwallowKAN (SKAN), which employs adaptive Radial Basis Functions (RBFs) as the core of the non-linear neurons. This innovation improves both computational efficiency and adaptability to diverse molecular structures. Building on the strengths of SKAN, we propose a new class of GNNs, GNN-SKAN, and its augmented variant, GNN-SKAN+, which incorporates a SKAN-based classifier to further boost performance. To our knowledge, this is the first work to integrate KANs into GNN architectures tailored for molecular representation learning. Experiments across 6 classification datasets, 6 regression datasets, and 4 few-shot learning datasets demonstrate that our approach achieves new state-of-the-art performance in terms of accuracy and computational cost.

Figures (6)

• Figure 1: The visualization of our models' performance compared to the SOTAs on the HIV dataset. The circle size represents per-epoch computational cost. Purple circles represent self-supervised methods, green and yellow circles represent supervised methods, and blue circles represent our proposed methods (GNN-SKAN, GNN-SKAN+).
• Figure 2: Comparison of three Kolmogorov-Arnold Network (KAN) variants, including (a) KAN liu2024kan designed with B-splines, (b) FastKAN li2024kolmogorov designed with RBFs, and (c) SwallowKAN (Our SKAN) designed with adaptable RBFs. The differences in base functions with each KAN variant are reflected in the shaded areas of each subfigure.
• Figure 3: Comparison of the network architecture between MP-GNNs and our GNN-SKAN and GNN-SKAN+.
• Figure 4: t-SNE visualization of molecular representations on the BACE dataset, extracted from (a) GINE and (b) GINE-SKAN. The green and orange dots represent molecules with labels -1 and 1, respectively.
• Figure 5: Predicted scores for TPPO (CM) and Baicalein (CI) using GCN and GCN-SKAN. The CM class is labeled as 1, and the CI class as -1.