# Controllable Molecular Generative Foundation Models

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 29 references that are in the map.
- First posted on arXiv 14 May 2026.

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

- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [GuacaMol: Benchmarking Models for De Novo Molecular Design](https://mlsanity.com/paper/1811.09621/) (2018)
- [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations](https://mlsanity.com/paper/2202.02514/) (2022)
- [Fr'echet ChemNet Distance: A metric for generative models for molecules in drug discovery](https://mlsanity.com/paper/1803.09518/) (2018)
- [Deep Reinforcement Learning in Large Discrete Action Spaces](https://mlsanity.com/paper/1512.07679/) (2015)
- [Sample Efficiency Matters: A Benchmark for Practical Molecular Optimization](https://mlsanity.com/paper/2206.12411/) (2022)
- [MARS: Markov Molecular Sampling for Multi-objective Drug Discovery](https://mlsanity.com/paper/2103.10432/) (2021)
- [Exploring Chemical Space with Score-based Out-of-distribution Generation](https://mlsanity.com/paper/2206.07632/) (2022)
- [DeFoG: Discrete Flow Matching for Graph Generation](https://mlsanity.com/paper/2410.04263/) (2024)
- [Hit and Lead Discovery with Explorative RL and Fragment-based Molecule Generation](https://mlsanity.com/paper/2110.01219/) (2021)
- [Molecule Generation by Principal Subgraph Mining and Assembling](https://mlsanity.com/paper/2106.15098/) (2021)
- [Conditional Diffusion Based on Discrete Graph Structures for Molecular Graph Generation](https://mlsanity.com/paper/2301.00427/) (2023)
- [Graph Diffusion Transformers for Multi-Conditional Molecular Generation](https://mlsanity.com/paper/2401.13858/) (2024)
- [TransPolymer: a Transformer-based language model for polymer property predictions](https://mlsanity.com/paper/2209.01307/) (2022)
- [Genetic algorithms are strong baselines for molecule generation](https://mlsanity.com/paper/2310.09267/) (2023)
- [Graph Diffusion Transformers are In-Context Molecular Designers](https://mlsanity.com/paper/2510.08744/) (2025)

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

- [Graph Forward Distribution Matching for Molecular Inverse Design](https://mlsanity.com/paper/2609.32056/) (2026)

## Abstract (arXiv)

Despite the success of foundation models in language and vision, molecular graph generation still lacks a unified framework for heterogeneous design tasks with reliable controllability. While reinforcement learning (RL) offers a natural post-training mechanism for task-specific optimization, applying it to graph generative models is hindered by the vast atom-wise action spaces and chemically invalid intermediate states. We propose Controllable Molecular Generative Foundation Models (CoMole), built with a unified motif-aware graph diffusion pipeline. By learning a motif-aware graph space, CoMole transfers pretrained structural priors into controllable generation, where RL optimizes conditional reverse policies over chemically meaningful decisions. We theoretically characterize the bottleneck of atom-level RL and justify motif-aware policy optimization. Across three heterogeneous benchmarks spanning materials and drug discovery, CoMole ranks first in controllability on all nine targets, reduces MAE by up to 48.2% relative to the strongest baselines, and maintains validity above 0.94 without rule-based correction or post-hoc filtering. We further show that CoMole transfers controllability to unseen properties by optimizing only task embeddings with the generator frozen, achieving performance competitive with strong task-specific baselines.