KDD2025 GNN-SKAN:基于SwallowKAN的分子表征学习新范式

B站影视 电影资讯 2025-06-05 03:33 2

摘要:在药物研发和材料科学等领域,分子的理化性质、药理活性和材料特性直接决定其应用价值。因此,如何精准预测分子的各类属性,成为推动新药发现、分子筛选与材料设计等任务的关键步骤。近年来,随着人工智能辅助药物发现(AIDD)技术的快速发展,分子性质预测作为其中最基础也是

在药物研发和材料科学等领域,分子的理化性质、药理活性和材料特性直接决定其应用价值。因此,如何精准预测分子的各类属性,成为推动新药发现、分子筛选与材料设计等任务的关键步骤。近年来,随着人工智能辅助药物发现(AIDD)技术的快速发展,分子性质预测作为其中最基础也是最关键的任务之一,受到了广泛关注。而实现这一目标的核心在于分子表示学习:即将复杂多样的分子结构转换为计算机能够理解和处理的高维表示,为下游任务如性质回归、活性分类和反应预测提供支撑。

尽管已有众多方法被提出用于分子建模,目前最主流的路线仍是基于图神经网络(GNN)的表示学习方法。这类方法通常将分子看作图结构,其中原子表示为节点,化学键作为边,通过消息传递机制(Message Passing)在节点间传递信息、更新特征向量。典型代表如图卷积网络(GCN)、图注意力网络(GAT)和图同构网络(GIN)等,在建模原子局部邻域结构方面表现出色。然而,这些方法在实际应用中仍面临多重挑战:一方面,训练所需的高质量标注数据常常难以获得;另一方面,分子结构本身具有高度多样性,使得模型难以学习到稳定有效的表示。更为严重的是,GNN 架构普遍存在所谓的“过度压缩(over-squashing)”问题,即在多层消息传递过程中,远距离节点之间的信息被过度压缩,导致关键的结构交互信息丢失,从而降低了分子表示的精度和下游预测性能。

近年来,Kolmogorov–Arnold 网络(KAN)因其在小规模 AI 与科学计算任务中表现出的出色数据拟合能力和精确性而受到广泛关注。KAN 的设计理念突破了传统神经网络在激活函数层面的限制,通过对激活机制的重新定义,展现出强大的泛化能力与表达力。受到这些特性的启发,我们认为 KAN 尤其适合用于解决分子表示学习中面临的一个关键难题——高质量标注数据的严重不足。相比于传统方法,KAN 具备更强的少样本学习能力,因此在数据稀缺场景中具有天然优势。

此外,KAN 对细粒度结构信息的捕捉能力也为其在分子建模中的应用提供了坚实基础。在分子图中,不同原子及官能团之间可能存在复杂的非局部相互作用,而这类结构信息在标准 GNN 中往往因“过度压缩(over-squashing)”而被严重削弱。KAN 强大的拟合能力使其能够有效保留这些关键结构特征,从而在理论上有潜力缓解 GNN 在深层消息传递中的表达瓶颈。

尽管已有部分研究将 KAN 引入图结构任务,初步验证了其在图学习场景中的可行性与潜力,但它在效率上的短板始终未能被充分解决。KAN 本身计算开销较大、参数量较多,导致其难以直接应用于规模较大或结构复杂的分子图建模任务。因此,如何将 KAN 高效、可扩展地整合进图神经网络框架中,用于实际的分子表示学习任务,依然是一个具有挑战性但前景广阔的问题。

为此,我们提出了一个全新设计的 GNN 系列:GNN-SKAN,以及其增强版本GNN-SKAN+。这两个模型架构旨在继承 KAN 的表达能力优势,同时克服其在效率和可扩展性方面的瓶颈。如下图所示,我们方法的核心在于提出了一种专门为图结构任务设计的 KAN 变体,命名为SwallowKAN(SKAN)。与原始 KAN 使用固定 B-spline 激活函数不同,SKAN 引入了可学习的自适应径向基函数(RBF),不仅大幅降低了参数数量和计算复杂度,同时保留了原有模型对复杂结构分布的强拟合能力

图1

通过将 SKAN 深度嵌入到 GNN 的更新机制与分类器模块中,我们构建出一个既高效又具表现力的分子表示学习架构。实验结果表明,该模型在多个分类、回归与小样本任务上均取得显著性能提升,证明了这一方向的有效性和可行性

图2

GNN-SKAN 架构设计

GNN-SKAN 架构由图神经网络骨架与 SKAN 模块组成。整体计算流程包括四个阶段:

(1)Aggregation(邻居信息聚合)
首先,对于每一个图中的节点,模型根据其邻接原子节点的信息,通过传统 GNN 的聚合机制(如加权平均或注意力机制)生成初始的邻居消息向量。该过程用于捕捉局部结构信息。

(2)Update(节点特征更新)
随后,模型将聚合得到的信息输入到嵌入了 SKAN 模块 的更新函数中。与传统 MP-GNN 使用 MLP 或简单激活函数不同,SKAN 能够通过其残差激活结构,对节点信息进行更复杂的非线性变换。其更新函数为:

其中

是一个参数。

SKAN 在第d层的残差激活函数定义如下:

SKAN 的自适应参数调整能力使其能够有效应对多样化的分子结构特性,从而使模型更加稳健并具有更强的泛化能力,适用于各种分子数据集。

(3)Readout(图级表示提取)
在节点表示更新完成后,模型使用 MEAN pooling(或其他聚合方法)对所有节点向量进行汇总,获得分子的整体表征,用于图级任务如性质预测。

(4)Prediction(预测)
在 GNN-SKAN 中,最终预测由标准 MLP 分类器完成;在 GNN-SKAN+ 中,进一步使用 SKAN 替代 MLP 作为分类器,从而提升对复杂分子结构的拟合能力。

参数效率与计算复杂度分析

相较于原始 KAN,SKAN 显著优化了参数与计算效率。在相同层数 L、节点数 N 下,KAN 的参数复杂度约为 O(N^2LG),其中 G 是 B-spline 的节点数。而 SKAN 的复杂度为 O(N^2LM),其中 RBF 数量 M 通常远小于 G,因而在保持模型表达能力的前提下,显著减少了计算资源消耗。此外,实验表明仅使用两层的 GNN-SKAN 模型即可达到甚至超过五层传统 GNN 模型的表现,进一步验证了其高效性

表1

表2

在分子分类任务中,GNN-SKAN 和 GNN-SKAN+ 展现出显著的性能优势。实验结果表明,在几乎所有评估数据集上,它们都明显优于对应的传统图神经网络基线模型(如 GCN、GAT、GINE)。特别是增强版本 GNN-SKAN+,在多个关键任务中实现了最优性能,表明其具备更强的分子结构表达能力和分类判别能力。

具体来说,在 BBBP 数据集上,GCN-SKAN 相较于原始 GCN 模型的 ROC-AUC 提高了约 4.82%,而 GCN-SKAN+ 则进一步提升至 8.68%。在 HIV 数据集上,GNN-SKAN+ 获得了 0.786 的 ROC-AUC 分数,优于主流的自监督学习方法 GraphMAE(0.772)和 MGSSL(0.774)。此外,在 BACE 数据集中,GINE-SKAN 相较于 GINE 的提升达到 25.49%,凸显了其在应对分子结构高度多样化数据时的强大泛化能力

这些结果充分证明,GNN-SKAN 系列模型不仅在整体性能上超越传统 GNN,而且在实际应用中具备更强的稳健性与适应性

来源:科学兄弟连

相关推荐