# Unified Molecule Pre-training with Flexible 2D and 3D Modalities: Single and Paired Modality Integration

arXiv: https://arxiv.org/abs/2510.07035 · HTML: https://mlsanity.com/paper/2510.07035/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 41 references that are in the map.
- First posted on arXiv 8 Oct 2025.

## Builds on (most cited first, 23 of 23)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Align before Fuse: Vision and Language Representation Learning with Momentum Distillation](https://mlsanity.com/paper/2107.07651/) (2021)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [ChemRL-GEM: Geometry Enhanced Molecular Representation Learning for Property Prediction](https://mlsanity.com/paper/2106.06130/) (2021)
- [Learning Gradient Fields for Molecular Conformation Generation](https://mlsanity.com/paper/2105.03902/) (2021)
- [GeoMol: Torsional Geometric Generation of Molecular 3D Conformer Ensembles](https://mlsanity.com/paper/2106.07802/) (2021)
- [Molecular geometry prediction using a deep generative graph neural network](https://mlsanity.com/paper/1904.00314/) (2019)
- [Learning Neural Generative Dynamics for Molecular Conformation Generation](https://mlsanity.com/paper/2102.10240/) (2021)
- [One Transformer Can Understand Both 2D & 3D Molecular Data](https://mlsanity.com/paper/2210.01765/) (2022)
- [A Generative Model for Molecular Distance Geometry](https://mlsanity.com/paper/1909.11459/) (2019)
- [Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching](https://mlsanity.com/paper/2206.13602/) (2022)
- [An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming](https://mlsanity.com/paper/2105.07246/) (2021)
- [Unified 2D and 3D Pre-Training of Molecular Representations](https://mlsanity.com/paper/2207.08806/) (2022)
- [A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining](https://mlsanity.com/paper/2305.18407/) (2023)
- [DrugCLIP: Contrastive Protein-Molecule Representation Learning for Virtual Screening](https://mlsanity.com/paper/2310.06367/) (2023)
- [Multimodal Molecular Pretraining via Modality Blending](https://mlsanity.com/paper/2307.06235/) (2023)
- [UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning](https://mlsanity.com/paper/2405.10343/) (2024)
- [Pre-training with Fractional Denoising to Enhance Molecular Property Prediction](https://mlsanity.com/paper/2407.11086/) (2024)
- [Artificial Intelligence for Drug Discovery: Are We There Yet?](https://mlsanity.com/paper/2307.06521/) (2023)
- [Molecule Joint Auto-Encoding: Trajectory Pretraining with 2D and 3D Diffusion](https://mlsanity.com/paper/2312.03475/) (2023)
- [Advancing Molecular Graph-Text Pre-training via Fine-grained Alignment](https://mlsanity.com/paper/2409.14106/) (2024)

## Abstract (arXiv)

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.