Unified Molecule Pre-training with Flexible 2D and 3D Modalities: Single and Paired Modality Integration
This addresses a bottleneck in drug discovery and material design by enabling robust pre-training with incomplete or expensive-to-generate molecular data, though it is incremental as it builds on existing multimodal approaches.
The paper tackles the problem of molecular representation learning by proposing FlexMol, a flexible pre-training framework that supports single-modality input, overcoming limitations of existing methods that require paired 2D and 3D data. It achieves superior performance across molecular property prediction tasks, as demonstrated in extensive experiments.
Molecular representation learning plays a crucial role in advancing applications such as drug discovery and material design. Existing work leverages 2D and 3D modalities of molecular information for pre-training, aiming to capture comprehensive structural and geometric insights. However, these methods require paired 2D and 3D molecular data to train the model effectively and prevent it from collapsing into a single modality, posing limitations in scenarios where a certain modality is unavailable or computationally expensive to generate. To overcome this limitation, we propose FlexMol, a flexible molecule pre-training framework that learns unified molecular representations while supporting single-modality input. Specifically, inspired by the unified structure in vision-language models, our approach employs separate models for 2D and 3D molecular data, leverages parameter sharing to improve computational efficiency, and utilizes a decoder to generate features for the missing modality. This enables a multistage continuous learning process where both modalities contribute collaboratively during training, while ensuring robustness when only one modality is available during inference. Extensive experiments demonstrate that FlexMol achieves superior performance across a wide range of molecular property prediction tasks, and we also empirically demonstrate its effectiveness with incomplete data. Our code and data are available at https://github.com/tewiSong/FlexMol.