Multi-Modal Molecular Representation Learning via Structure Awareness

TL;DR

This work tackles molecular representation learning with multi-modal data (2D graphs, 2D images, and 3D graphs) under a self-supervised framework. It introduces MMSA, a two-module architecture that first aligns and fuses multi-modal signals through auto-encoders with contrastive and reconstruction objectives, and then captures higher-order, invariant structure via a hypergraph network and a memory bank. The total objective combines the multi-modal auto-encoder loss L_ae with a structure-awareness loss L_sa, achieving state-of-the-art results on MoleculeNet across classification, regression, and retrieval tasks, including strong cross-dataset generalization. This approach demonstrates substantial practical impact by producing robust, transferable molecular embeddings and offering a versatile plugin for existing graph-based models, with promising directions for broader molecular data modalities and tasks. $L_{ae} = \lambda L_{cl} + (1-\lambda) L_{rl}$, $L_{sa} = \alpha L_{me} + (1-\alpha) L_{pre}$, and $L_{overall} = L_{ae} + L_{sa}$ are central formulations guiding the learning process.$

Abstract

Accurate extraction of molecular representations is a critical step in the drug discovery process. In recent years, significant progress has been made in molecular representation learning methods, among which multi-modal molecular representation methods based on images, and 2D/3D topologies have become increasingly mainstream. However, existing these multi-modal approaches often directly fuse information from different modalities, overlooking the potential of intermodal interactions and failing to adequately capture the complex higher-order relationships and invariant features between molecules. To overcome these challenges, we propose a structure-awareness-based multi-modal self-supervised molecular representation pre-training framework (MMSA) designed to enhance molecular graph representations by leveraging invariant knowledge between molecules. The framework consists of two main modules: the multi-modal molecular representation learning module and the structure-awareness module. The multi-modal molecular representation learning module collaboratively processes information from different modalities of the same molecule to overcome intermodal differences and generate a unified molecular embedding. Subsequently, the structure-awareness module enhances the molecular representation by constructing a hypergraph structure to model higher-order correlations between molecules. This module also introduces a memory mechanism for storing typical molecular representations, aligning them with memory anchors in the memory bank to integrate invariant knowledge, thereby improving the model generalization ability. Extensive experiments have demonstrated the effectiveness of MMSA, which achieves state-of-the-art performance on the MoleculeNet benchmark, with average ROC-AUC improvements ranging from 1.8% to 9.6% over baseline methods.

Figures (7)

• Figure 1: The proposed MMSA pre-training framework consists of two main modules: the multi-modal molecular representation learning module and the structure-awareness module. $\mathcal{F}_i$ represents the method for learning the feature representation of modality $m_i$, where $m_i$ can be 2D graphs, 2D images, 3D graphs, etc. In the multi-modal molecular representation learning module, the molecular embedding $c$ is generated by aggregating the modality embeddings $\{ c_i \}_{i=1}^M$. The $\mathcal{L}_{cl}$ loss is designed to capture shared information across modalities, while the $\mathcal{L}_{rl}$ loss focuses on enhancing the generalization ability of each modality embedding. In the structure-awareness module, hypergraphs capture higher-order correlations between molecular graphs, enabling the learning of generalizable molecular structures. The memory mechanism is developed to further enhance the generalization ability of multi-modal molecular embeddings by retaining and reusing crucial structural and contextual information, thereby improving the robustness and adaptability of the learned representations.
• Figure 2: Comparison of Graph and Hypergraph.
• Figure 3: RMSE ($\Downarrow$) performance of various methods on four regression datasets for molecular property prediction. "$\Downarrow$" indicates that lower RMSE values correspond to better model performance. The results are reported as the mean (standard deviation) of RMSE across five random seeds. "$\Delta$" represents the absolute percentage improvement, calculated as $\Delta = \left( 1 - \text{RMSE}_{\text{baseline }}/\text{RMSE}_{\text{MMSA}} \right) \times 100$.
• Figure 4: Performance (ROC-AUC %, $\Uparrow$ and RMSE, $\Downarrow$) on MoleculeNet across different baseline methods equipped with MMSA.
• Figure 5: Clustering results on the BACE task using molecular representations obtained by different methods. The colours represent the labels of the downstream tasks (discrete binary labels for the BACE task). For each subgraph, the performance of the clustering results is evaluated using the Davies-Bouldin Index (DBI, where smaller values are better) and the Normalized Mutual Information (NMI, where larger values are better). The MMSA method (the last one) demonstrates a significant ability to improve clustering performance.

Theorems & Definitions (1)

• Remark 1