MOL-Mamba: Enhancing Molecular Representation with Structural & Electronic Insights

TL;DR

MOL-Mamba tackles the limitation of omitting electronic information in molecular graph representations by proposing a Mamba-enhanced architecture that fuses structural graphs with electronic descriptors. The framework comprises an Atom & Fragment Mamba-Graph ($\mathcal{G}_A$ MG) for fine-grained structural reasoning and a Mamba-Transformer ($\text{MT}$) fuser to integrate molecular structure with electronic correlations, guided by two training schemes: Structural Distribution Collaborative Training and E-semantic Fusion Training. It introduces a formal problem setup with a dual-graph representation ($\mathcal{G}_A$ and $\mathcal{G}_F$) and descriptors $\mathcal{D}_E$, and optimizes a total loss $\mathcal{L}_{total}=\lambda_d\mathcal{L}_d+\lambda_s\mathcal{L}_s+\lambda_f\mathcal{L}_f+\lambda_{mask}\mathcal{L}_{mask}$. Empirical results across 11 datasets show MOL-Mamba outperforms state-of-the-art baselines on 8 of 11 tasks, with notable performance on BBBP (ROC-AUC $=75.0$) and ESOL (MAE $=0.63$), while maintaining efficiency comparable to GNNs and better parameter efficiency than graph transformers, demonstrating strong potential for drug design and materials discovery.

Abstract

Molecular representation learning plays a crucial role in various downstream tasks, such as molecular property prediction and drug design. To accurately represent molecules, Graph Neural Networks (GNNs) and Graph Transformers (GTs) have shown potential in the realm of self-supervised pretraining. However, existing approaches often overlook the relationship between molecular structure and electronic information, as well as the internal semantic reasoning within molecules. This omission of fundamental chemical knowledge in graph semantics leads to incomplete molecular representations, missing the integration of structural and electronic data. To address these issues, we introduce MOL-Mamba, a framework that enhances molecular representation by combining structural and electronic insights. MOL-Mamba consists of an Atom & Fragment Mamba-Graph (MG) for hierarchical structural reasoning and a Mamba-Transformer (MT) fuser for integrating molecular structure and electronic correlation learning. Additionally, we propose a Structural Distribution Collaborative Training and E-semantic Fusion Training framework to further enhance molecular representation learning. Extensive experiments demonstrate that MOL-Mamba outperforms state-of-the-art baselines across eleven chemical-biological molecular datasets.

Figures (5)

• Figure 1: The pretraining framework of our approach MOL-Mamba (Molecular Mamba). It consists of three modules: fragment-level graph GNN$_F$, atom-level structural reasoning Mamba-graph (MG), and a molecular structural & electronic Mamba-Transformer (MT) fuser. We implement two pretraining stages: molecular structure learning with the Distribution Collaborative Training, then E-semantic Fusion Training is conducted for molecular structural & electronic fusion learning.
• Figure 2: Illustration of the Mamba-Graph (MG) module. It consists of the GNN layer and Mamba block. A novel graph node sorting strategy to exploit Mamba's context-aware reasoning. A new GSSM mechanism is designed to adaptively select relevant information from the graph context. Two fragment-related loss constraints $\mathcal{L}_{structure}$ and $\mathcal{L}_{fragment}$ facilitate molecular structure learning of the MG module.
• Figure 3: MOL-Mamba with different module settings on classification benchmarks BBBP and Tox21, regression benchmarks ESOL and Lipo, (lower is better for regression).
• Figure 4: T-SNE visualization of feature separation of our pretraining (PT) methods on BBBP and ESOL. The "Structure PT" refers to pretraining the model with structural losses ($\mathcal{L}_d, \mathcal{L}_s, \mathcal{L}_f$), the "E-semantic PT" pretrains the model with e-masked loss $\mathcal{L}_{mask}$. "ER" is Error Rate.
• Figure 5: Exemplary explanations for using MG and MT modules in molecular prediction tasks. The radius scales linearly with the node-feature sum. It highlights key structural and electronic adjustments.