Controllable Molecular Generative Foundation Models

TL;DR

CoMole introduces a motif-aware graph diffusion framework for controllable molecular generation, effectively integrating pretrained priors and reinforcement learning to optimize chemically meaningful structures across diverse benchmarks.

Contribution

It presents a unified motif-aware graph diffusion model that enables controllable molecular generation with reinforcement learning, outperforming existing methods in validity and controllability.

Findings

CoMole ranks first in controllability across all benchmarks.

Reduces MAE by up to 48.2% compared to baselines.

Maintains validity above 0.94 without post-hoc filtering.

Abstract

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 \textbf{Co}ntrollable \textbf{Mole}cular 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.