A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining

TL;DR

MoleculeSDE addresses the challenge of rich, multi-modal molecule representation by directly maximizing the mutual information between 2D topology and 3D conformation using dual SE(3)-aware diffusion processes. It introduces an $SE(3)$-equivariant topology-to-conformation diffusion and an $SE(3)$-invariant conformation-to-topology diffusion to generate the other modality in input space, resulting in a tighter MI bound than proxy methods like VRR. The work provides theoretical connections between diffusion-based training and MI estimation, and demonstrates state-of-the-art performance across a broad set of downstream tasks (26 of 32) on PCQM4Mv2, including both 2D-only and 3D-augmented predictions. Overall, MoleculeSDE advances multi-modal molecule pretraining by leveraging group symmetry and score-based diffusion to yield more faithful representations and transferable conformational knowledge for drug discovery tasks.

Abstract

Molecule pretraining has quickly become the go-to schema to boost the performance of AI-based drug discovery. Naturally, molecules can be represented as 2D topological graphs or 3D geometric point clouds. Although most existing pertaining methods focus on merely the single modality, recent research has shown that maximizing the mutual information (MI) between such two modalities enhances the molecule representation ability. Meanwhile, existing molecule multi-modal pretraining approaches approximate MI based on the representation space encoded from the topology and geometry, thus resulting in the loss of critical structural information of molecules. To address this issue, we propose MoleculeSDE. MoleculeSDE leverages group symmetric (e.g., SE(3)-equivariant and reflection-antisymmetric) stochastic differential equation models to generate the 3D geometries from 2D topologies, and vice versa, directly in the input space. It not only obtains tighter MI bound but also enables prosperous downstream tasks than the previous work. By comparing with 17 pretraining baselines, we empirically verify that MoleculeSDE can learn an expressive representation with state-of-the-art performance on 26 out of 32 downstream tasks.

Figures (4)

• Figure 1: Illustration of the MoleculeSDE pretraining. It is composed of one contrastive learning and two generative learning objectives. Contrastive learning aims to align the 2D topological and 3D conformational representations for the same molecule. The two generative learning objectives are molecule conditional generation from 2D topology to 3D conformation and from 3D conformation to 2D topology, respectively. The generative modeling from topology to conformation is an SE(3)-equivariant diffusion process that satisfies the physical symmetry in molecule geometries. The other direction from conformation to topology is an SE(3)-invariant diffusion process since only the invariant type-0 features (nodes and edges) are considered. We further include https://chao1224.github.io/MoleculeSDE, showing the trajectory of this process.
• Figure 2: Illustration on three downstream tasks. The first two cover single-modal information only, and we fine-tune the pretrained 2D and 3D GNN from MoleculeSDE, respectively. The last downstream tasks contain topology only, then we use the pretrained 2D GNN and SE(3)-equivariant SDE model to generate conformations, followed by a 3D GNN for property prediction.
• Figure 3: Pipeline for the SE(3)-equivariant SDE model from topology to conformation.
• Figure 4: Pipeline for the SE(3)-invariant SDE model from conformation to topology.