Nuclear-Charge-Guided Mamba with KAN Dynamic Mixture for Molecular Property Prediction

Molecular property prediction (MPP) is a critical yet challenging task in drug discovery, where existing methods face fundamental limitations: Graph Neural Networks (GNNs) suffer from over-smoothing and limited receptive fields, while Graph Transformers incur prohibitive quadratic complexity. To address these challenges, we propose KAN-NC-Mamba, a novel molecular representation learning framework that integrates three key innovations. (1) We introduce a Nuclear-Charge-guided node ordering, which serializes molecular graphs by ascending atomic number through a chemistry-driven approach to prioritize chemically salient atoms without additional parameters. (2) We design a Nuclear-Charge-guided Mamba module that employs selective state-space models to establish global long-range dependencies with linear complexity, effectively overcoming the limited receptive fields of traditional graph neural networks. (3) We develop a KAN Dynamic Mixture module based on Kolmogorov-Arnold Networks to achieve nonlinear adaptive fusion of local and global features, breaking through the limitations of traditional linear weighting or simple concatenation fusion paradigms. Experimental results on ten benchmark datasets demonstrate that KAN-NC-Mamba achieves state-of-the-art performance on both classification and regression tasks. Ablation studies further validate the effectiveness and complementary nature of each innovative module, showing that the nuclear-charge ordering strategy provides significant performance improvement across multiple tasks compared to popular degree-based ordering. The proposed KAN-NC-Mamba not only advances the accuracy-efficiency frontier in molecular machine learning but also provides chemically interpretable representations, as validated by our visualization studies.