# UniGEM: A Unified Approach to Generation and Property Prediction for Molecules

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 59 references that are in the map.
- First posted on arXiv 14 Oct 2024.

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [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)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [GraphMAE: Self-Supervised Masked Graph Autoencoders](https://mlsanity.com/paper/2205.10803/) (2022)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Understanding Diffusion Models: A Unified Perspective](https://mlsanity.com/paper/2208.11970/) (2022)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Cormorant: Covariant Molecular Neural Networks](https://mlsanity.com/paper/1906.04015/) (2019)
- [Pre-training Molecular Graph Representation with 3D Geometry](https://mlsanity.com/paper/2110.07728/) (2021)
- [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)
- [Geometric Latent Diffusion Models for 3D Molecule Generation](https://mlsanity.com/paper/2305.01140/) (2023)
- [3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction](https://mlsanity.com/paper/2303.03543/) (2023)
- [Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules](https://mlsanity.com/paper/1906.00957/) (2019)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)
- [3D Infomax improves GNNs for Molecular Property Prediction](https://mlsanity.com/paper/2110.04126/) (2021)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [Diffusion-based Molecule Generation with Informative Prior Bridges](https://mlsanity.com/paper/2209.00865/) (2022)
- [MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation](https://mlsanity.com/paper/2302.09048/) (2023)
- [Deconstructing Denoising Diffusion Models for Self-Supervised Learning](https://mlsanity.com/paper/2401.14404/) (2024)
- [SODA: Bottleneck Diffusion Models for Representation Learning](https://mlsanity.com/paper/2311.17901/) (2023)
- [Non-Denoising Forward-Time Diffusions](https://mlsanity.com/paper/2312.14589/) (2023)
- [Equivariant Energy-Guided SDE for Inverse Molecular Design](https://mlsanity.com/paper/2209.15408/) (2022)
- [Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics](https://mlsanity.com/paper/2302.00600/) (2023)
- [Bayesian Flow Networks](https://mlsanity.com/paper/2308.07037/) (2023)
- [GeomGCL: Geometric Graph Contrastive Learning for Molecular Property Prediction](https://mlsanity.com/paper/2109.11730/) (2021)
- [Diffusion-Based Representation Learning](https://mlsanity.com/paper/2105.14257/) (2021)

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

- [Straight-Line Diffusion Model for Efficient 3D Molecular Generation](https://mlsanity.com/paper/2503.02918/) (2025)
- [Unified Cross-Scale 3D Generation and Understanding via Autoregressive Modeling](https://mlsanity.com/paper/2503.16278/) (2025)
- [TFG-Flow: Training-free Guidance in Multimodal Generative Flow](https://mlsanity.com/paper/2501.14216/) (2025)
- [Synergistic Benefits of Joint Molecule Generation and Property Prediction](https://mlsanity.com/paper/2504.16559/) (2025)
- [Guiding Diffusion Models with Reinforcement Learning for Stable Molecule Generation](https://mlsanity.com/paper/2508.16521/) (2025)
- [Manipulating 3D Molecules in a Fixed-Dimensional E(3)-Equivariant Latent Space](https://mlsanity.com/paper/2506.00771/) (2025)
- [InertialAR: Autoregressive 3D Molecule Generation with Inertial Frames](https://mlsanity.com/paper/2510.27497/) (2025)
- [NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation](https://mlsanity.com/paper/2512.05844/) (2025)
- [Elign: Equivariant Diffusion Model Alignment from Foundational Machine Learning Force Fields](https://mlsanity.com/paper/2601.21985/) (2026)
- [Predictive Feature Caching for Training-free Acceleration of Molecular Geometry Generation](https://mlsanity.com/paper/2510.04646/) (2025)
- [Neural Proposals, Symbolic Guarantees: Neuro-Symbolic Graph Generative Modeling](https://mlsanity.com/paper/2602.16954/) (2026)
- [Predicting and generating antibiotics against future pathogens with ApexOracle](https://mlsanity.com/paper/2507.07862/) (2025)
- [NS-Pep: De novo Peptide Design with Non-Standard Amino Acids](https://mlsanity.com/paper/2510.03326/) (2025)
- [A Conformation-Centric Generative Foundation Model for Linear Polymer Modeling and Design](https://mlsanity.com/paper/2510.16023/) (2025)
- [3D-GSRD: 3D Molecular Graph Auto-Encoder with Selective Re-mask Decoding](https://mlsanity.com/paper/2510.16780/) (2025)
- [VEDA: 3D Molecular Generation via Variance-Exploding Diffusion with Annealing](https://mlsanity.com/paper/2511.09568/) (2025)
- [Chemistry-Enhanced Diffusion-Based Framework for Small-to-Large Molecular Conformation Generation](https://mlsanity.com/paper/2511.12182/) (2025)
- [Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation](https://mlsanity.com/paper/2603.10093/) (2026)
- [Generating Physically Consistent Molecules with Energy-Based Models](https://mlsanity.com/paper/2605.18381/) (2026)
- [Uncertainty-Calibrated Diffusion for Reliable 3D Molecular Graph Generation](https://mlsanity.com/paper/2606.01595/) (2026)
- [OctMole: Multiscale 3D Molecular Generation with Hierarchical Spatial Flows](https://mlsanity.com/paper/2610.16586/) (2026)

## Abstract (arXiv)

Molecular generation and molecular property prediction are both crucial for drug discovery, but they are often developed independently. Inspired by recent studies, which demonstrate that diffusion model, a prominent generative approach, can learn meaningful data representations that enhance predictive tasks, we explore the potential for developing a unified generative model in the molecular domain that effectively addresses both molecular generation and property prediction tasks. However, the integration of these tasks is challenging due to inherent inconsistencies, making simple multi-task learning ineffective. To address this, we propose UniGEM, the first unified model to successfully integrate molecular generation and property prediction, delivering superior performance in both tasks. Our key innovation lies in a novel two-phase generative process, where predictive tasks are activated in the later stages, after the molecular scaffold is formed. We further enhance task balance through innovative training strategies. Rigorous theoretical analysis and comprehensive experiments demonstrate our significant improvements in both tasks. The principles behind UniGEM hold promise for broader applications, including natural language processing and computer vision.