# Molecule Joint Auto-Encoding: Trajectory Pretraining with 2D and 3D Diffusion

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 59 of its 74 references that are in the map.
- First posted on arXiv 6 Dec 2023.

## Builds on (most cited first, 30 of 59)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations](https://mlsanity.com/paper/2108.01073/) (2021)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [What Makes for Good Views for Contrastive Learning?](https://mlsanity.com/paper/2005.10243/) (2020)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [Self-supervised Learning: Generative or Contrastive](https://mlsanity.com/paper/2006.08218/) (2020)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [Score-based Generative Modeling in Latent Space](https://mlsanity.com/paper/2106.05931/) (2021)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [Diffusion Autoencoders: Toward a Meaningful and Decodable Representation](https://mlsanity.com/paper/2111.15640/) (2021)
- [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization](https://mlsanity.com/paper/1908.01000/) (2019)
- [Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One](https://mlsanity.com/paper/1912.03263/) (2019)
- [GPT-GNN: Generative Pre-Training of Graph Neural Networks](https://mlsanity.com/paper/2006.15437/) (2020)
- [Graph Self-Supervised Learning: A Survey](https://mlsanity.com/paper/2103.00111/) (2021)
- [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations](https://mlsanity.com/paper/2202.02514/) (2022)
- [Effective Data Augmentation With Diffusion Models](https://mlsanity.com/paper/2302.07944/) (2023)
- [Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise](https://mlsanity.com/paper/2208.09392/) (2022)

## Built on by (most cited first, 10 of 10)

- [UniGEM: A Unified Approach to Generation and Property Prediction for Molecules](https://mlsanity.com/paper/2410.10516/) (2024)
- [UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning](https://mlsanity.com/paper/2405.10343/) (2024)
- [MolSpectra: Pre-training 3D Molecular Representation with Multi-modal Energy Spectra](https://mlsanity.com/paper/2502.16284/) (2025)
- [A Multi-Grained Symmetric Differential Equation Model for Learning Protein-Ligand Binding Dynamics](https://mlsanity.com/paper/2401.15122/) (2024)
- [GeomCLIP: Contrastive Geometry-Text Pre-training for Molecules](https://mlsanity.com/paper/2411.10821/) (2024)
- [Manifold-Constrained Nucleus-Level Denoising Diffusion Model for Structure-Based Drug Design](https://mlsanity.com/paper/2409.10584/) (2024)
- [Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation](https://mlsanity.com/paper/2603.23398/) (2026)
- [Sculpting Molecules in Text-3D Space: A Flexible Substructure Aware Framework for Text-Oriented Molecular Optimization](https://mlsanity.com/paper/2403.03425/) (2024)
- [Unified Molecule Pre-training with Flexible 2D and 3D Modalities: Single and Paired Modality Integration](https://mlsanity.com/paper/2510.07035/) (2025)
- [MolGA: Molecular Graph Adaptation with Pre-trained 2D Graph Encoder](https://mlsanity.com/paper/2510.07289/) (2025)

## Abstract (arXiv)

Recently, artificial intelligence for drug discovery has raised increasing interest in both machine learning and chemistry domains. The fundamental building block for drug discovery is molecule geometry and thus, the molecule's geometrical representation is the main bottleneck to better utilize machine learning techniques for drug discovery. In this work, we propose a pretraining method for molecule joint auto-encoding (MoleculeJAE). MoleculeJAE can learn both the 2D bond (topology) and 3D conformation (geometry) information, and a diffusion process model is applied to mimic the augmented trajectories of such two modalities, based on which, MoleculeJAE will learn the inherent chemical structure in a self-supervised manner. Thus, the pretrained geometrical representation in MoleculeJAE is expected to benefit downstream geometry-related tasks. Empirically, MoleculeJAE proves its effectiveness by reaching state-of-the-art performance on 15 out of 20 tasks by comparing it with 12 competitive baselines.