解决分子图神经网络深度学习中的全局与局部信息平衡问题
HSA-Net: Hierarchical and Structure-Aware Framework for Efficient and Scalable Molecular Language Modeling
- 分层自适应投影模块动态切换注意力与Mamba处理浅层与深层特征
- 在多个分子属性预测数据集上超越现有最佳方法,提升显著
- 适合需要高精度分子表征的药物发现与材料设计研究者
分子表征学习是分子描述生成和分子性质预测等下游任务的基础,依赖图神经网络(GNN)。然而,GNN存在过平滑问题,导致深层中节点特征退化。现有基于交叉注意力的特征投影方法虽缓解此问题,但在深层仍表现不佳。为此,我们探索使用Mamba作为替代投影器,因其擅长处理复杂序列。但发现其虽能保留深层全局拓扑信息,却忽视浅层细粒度特征,表现出全局与局部信息的权衡。为解决这一关键矛盾,我们提出分层结构感知网络(HSA-Net),包含两个模块:首先,分层自适应投影器(HAP)根据图层深度动态选择交叉注意力或结构感知图-Mamba投影器,生成多层级高质量特征;其次,源感知融合(SAF)模块依据聚合特征特性自适应选择融合专家,实现精准有效的最终表征融合。大量实验表明,本框架在定量与定性上均优于当前最先进方法。
原文摘要 · Abstract (English)
Molecular representation learning, a cornerstone for downstream tasks like molecular captioning and molecular property prediction, heavily relies on Graph Neural Networks (GNN). However, GNN suffers from the over-smoothing problem, where node-level features collapse in deep GNN layers. While existing feature projection methods with cross-attention have been introduced to mitigate this issue, they still perform poorly in deep features. This motivated our exploration of using Mamba as an alternative projector for its ability to handle complex sequences. However, we observe that while Mamba excels at preserving global topological information from deep layers, it neglects fine-grained details in shallow layers. The capabilities of Mamba and cross-attention exhibit a global-local trade-off. To resolve this critical global-local trade-off, we propose Hierarchical and Structure-Aware Network (HSA-Net), a novel framework with two modules that enables a hierarchical feature projection and fusion. Firstly, a Hierarchical Adaptive Projector (HAP) module is introduced to process features from different graph layers. It learns to dynamically switch between a cross-attention projector for shallow layers and a structure-aware Graph-Mamba projector for deep layers, producing high-quality, multi-level features. Secondly, to adaptively merge these multi-level features, we design a Source-Aware Fusion (SAF) module, which flexibly selects fusion experts based on the characteristics of the aggregation features, ensuring a precise and effective final representation fusion. Extensive experiments demonstrate that our HSA-Net framework quantitatively and qualitatively outperforms current state-of-the-art (SOTA) methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。