Mamba-driven multi-perspective structural understanding for molecular ground-state conformation prediction

TL;DR

The paper addresses predicting molecular ground-state conformation from 2D structures by introducing MPSU-Mamba, a Mamba-based framework that builds a comprehensive, multi-perspective understanding of molecular structure. It employs dedicated scanning strategies for atom types, connections, and positions, plus a bright-channel guided mechanism to focus on critical atoms, enabling accurate 3D predictions. Across QM9 and Molecule3D benchmarks, it achieves substantial improvements over state-of-the-art methods and shows strong generalization in few-shot scenarios. The approach has potential implications for faster molecule design and larger biomolecular applications by explicitly codifying structure-focused perception into the prediction model.

Abstract

A comprehensive understanding of molecular structures is important for the prediction of molecular ground-state conformation involving property information. Meanwhile, state space model (e.g., Mamba) has recently emerged as a promising mechanism for long sequence modeling and has achieved remarkable results in various language and vision tasks. However, towards molecular ground-state conformation prediction, exploiting Mamba to understand molecular structure is underexplored. To this end, we strive to design a generic and efficient framework with Mamba to capture critical components. In general, molecular structure could be considered to consist of three elements, i.e., atom types, atom positions, and connections between atoms. Thus, considering the three elements, an approach of Mamba-driven multi-perspective structural understanding (MPSU-Mamba) is proposed to localize molecular ground-state conformation. Particularly, for complex and diverse molecules, three different kinds of dedicated scanning strategies are explored to construct a comprehensive perception of corresponding molecular structures. And a bright-channel guided mechanism is defined to discriminate the critical conformation-related atom information. Experimental results on QM9 and Molecule3D datasets indicate that MPSU-Mamba significantly outperforms existing methods. Furthermore, we observe that for the case of few training samples, MPSU-Mamba still achieves superior performance, demonstrating that our method is indeed beneficial for understanding molecular structures.

Figures (9)

• Figure 1: The overview of the proposed MPSU-Mamba framework. a, Mamba-driven multi-perspective molecular structural understanding for 3D conformation prediction, mainly comprising bright-channel guided operations, MPSU-Mamba encoder, MPSU-Mamba decoder, and 3D conformation prediction head. Particularly, given a 2D input molecule, we flatten it into a sequence of atoms. After obtaining the projected features, a bright-channel guided operation is first performed to focus on critical conformation-aware information. Then, we define multiple Mamba layers to extract the encoding representation. Here, we utilize the scanning strategies based on atom types and connections between atoms. Next, a conformation prediction head is defined to estimate the initial 3D positions. Finally, we still design a multi-layer Mamba as the decoder to achieve the corresponding 3D molecular conformation. Here, the scanning strategies are selected based on atom types, connections between atoms, and atom positions. b, The details of Bright-Channel guided operation. We observe that by stacking multi-layer guided operations, the accuracy of conformation prediction could be improved effectively. c, View of a MPSU-Mamba operation. For type- and connection-based strategies, we only utilize the forward direction. Differently, for the position strategy, the forward and backward directions are used.
• Figure 2: Visualization examples of different scanning strategies. Here, the number in each atom depicts the corresponding scanning order. To obtain a comprehensive understanding of molecular structures, a series of dedicated scanning mechanisms are defined from two various perspectives, i.e., atom types and connections between atoms. Extensive experimental results on multiple benchmarks demonstrate that these strategies could effectively improve the performance of conformation prediction.
• Figure 3: Scanning based on 3D atom coordinates (Descending order). Since the coordinates are produced based on the encoding representation, using position-based scanning could not only enhance the understanding level but also improve the accuracy of the final prediction.
• Figure 4: Visual comparison between our MPSU-Mamba and existing state-of-the-art methods.
• Figure 5: More comparison results between our MPSU-Mamba and existing state-of-the-art methods.