🤖 machine learning

Controllable Molecular Generative Foundation Models

The paper introduces CoMole, a controllable molecular generative foundation model that utilizes a motif-aware graph diffusion pipeline and reinforcement learning to overcome atom-level action space challenges, achieving superior performance and validity across diverse materials and drug discovery tasks without requiring post-hoc filtering.

Original authors: Yihan Zhu, Yuhan Liu, Weijiang Li, Tengfei Luo, Meng Jiang

Published 2026-05-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yihan Zhu, Yuhan Liu, Weijiang Li, Tengfei Luo, Meng Jiang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot chef to cook a specific dish, like a "spicy tomato soup," but you only give it instructions one grain of salt or one drop of oil at a time. If the robot adds a grain of salt, then a drop of oil, then another grain, it might accidentally create a salty, oily mess that isn't soup at all. It has to make thousands of tiny, risky decisions before it even knows if the dish is edible. This is exactly the problem scientists face when trying to use AI to design new molecules (the building blocks of drugs and materials).

This paper introduces a new AI system called CoMole that solves this problem by changing how the robot thinks.

The Problem: The "Atom-by-Atom" Trap

Current AI models for designing molecules work like that clumsy robot chef. They try to build a molecule one single atom at a time (like adding one carbon, then one hydrogen).

The Risk: Because there are so many tiny steps, the AI often makes a tiny mistake early on (like connecting two atoms in a way that breaks the laws of chemistry). Once that happens, the whole molecule becomes invalid, and the AI has to start over.
The Result: The AI gets stuck. It tries to learn from rewards (like "this molecule is a good drug"), but by the time it realizes it made a mistake, it's too late. The paper calls this "trajectory collapse"—the AI falls off the cliff of chemical possibility before it can climb back up.

The Solution: The "Lego Block" Strategy

The authors realized that chemists don't think about individual atoms; they think about groups of atoms that work together, like a benzene ring or a specific functional group. They call these groups "motifs."

Think of it like building with LEGOs:

Old Way (Atom-level): You try to build a castle by gluing individual plastic grains together. One wrong grain, and the wall collapses.
CoMole's Way (Motif-level): You build with pre-made LEGO bricks (walls, windows, towers). You snap these big, stable blocks together.

CoMole is a "Foundation Model" (a smart, pre-trained base) that learns to speak the language of these LEGO blocks. Instead of deciding "add a carbon here," it decides "attach a benzene ring here."

How It Works: The Three-Step Training

The paper describes a three-stage training process to make this work:

Pre-training (The Library): First, the AI reads millions of existing molecules to learn what valid "LEGO blocks" (motifs) look like. It builds a mental library of stable chemical structures.
Fine-tuning (The Recipe Book): Next, it learns to follow specific instructions. If you say "I need a molecule that is good at filtering oxygen," it learns to arrange its LEGO blocks to match that goal.
Reinforcement Learning (The Taste Test): Finally, the AI plays a game. It builds a molecule, checks if it meets the goal (e.g., "Is it a good drug?"), and gets a score.
- The Magic: Because it's building with big LEGO blocks instead of tiny grains, it doesn't crash as often. It can learn from its successes and failures much faster and more reliably.

The Results: Better, Faster, and Safer

The authors tested CoMole on three different challenges: designing new materials for gas filters, creating polymers with specific electronic properties, and finding new drug candidates.

Control: On all nine different tasks they tested, CoMole was the best at hitting the target. For example, when asked to design a polymer with a specific gas permeability, it was nearly 50% more accurate than the next best method.
Validity: The most impressive part is that CoMole produces valid molecules 94% of the time without needing a human to fix mistakes afterward. Other methods often produce broken molecules that need to be thrown away.
Generalization: Even when asked to design for a property it had never seen before (like a specific type of electrical energy), CoMole could still do a great job just by understanding the "recipe" of the task, without needing to relearn everything from scratch.

The Bottom Line

The paper claims that by stopping the AI from trying to build molecules atom-by-atom and instead letting it build with pre-approved, chemically meaningful "chunks" (motifs), we can create a much more reliable and controllable system. It's the difference between trying to build a house by stacking individual grains of sand versus snapping together pre-fabricated walls and roofs. CoMole is the first system to successfully apply this "LEGO" approach to a general-purpose AI foundation model for molecular design.

Technical Summary: Controllable Molecular Generative Foundation Models (CoMole)

Problem Statement

Molecular inverse design—the generation of structures with specific functional properties—remains a central challenge in scientific discovery. While recent graph diffusion methods have advanced molecular generation, they largely operate as task-specific models with limited controllability over target properties. The field lacks a unified framework capable of handling heterogeneous design tasks (spanning materials and drug discovery) with reliable controllability.

Attempts to bridge this gap using Reinforcement Learning (RL) for post-training face a critical bottleneck when applied to atom-level graph diffusion. The vast space of low-level atom-wise actions is poorly aligned with the chemically meaningful substructures (motifs) that chemists use to formulate structure-property rationales. As demonstrated in the paper, atom-level RL often suffers from "trajectory collapse," where a single invalid atomic choice pushes the generation off the feasible chemical manifold, making recovery difficult. This fragility leads to poor validity and unstable credit assignment, preventing the model from reliably associating structural changes with property rewards.

Methodology

The authors propose CoMole (Controllable Molecular Generative Foundation Models), a unified framework built on a motif-aware graph diffusion pipeline. The core innovation is reparameterizing the graph-diffusion RL process from an atom-level decision space to a motif-aware decision space.

1. Motif-Aware Graph Tokenization

CoMole employs a Node Pair Encoding (NPE)-based tokenizer that learns a vocabulary from pretraining data. This tokenizer maps chemically valid molecular graphs into motif graphs ( $z_0$ ), where nodes represent atom-disjoint units:

Singleton atoms.
Preserved ring units.
Data-driven merged substructures (higher-level motifs).

This representation preserves attachment-level information and inter-motif bond labels, allowing for lossless reconstruction while significantly compressing the graph state.

2. Three-Stage Training Pipeline

CoMole follows a foundation-model paradigm involving three stages:

Pretraining (PT): The model learns transferable structural priors over the motif graph space using unconditional diffusion.
Supervised Fine-Tuning (SFT): The model is conditioned on multiple properties (numerical and categorical) to introduce task-specific alignment.
Reinforcement Learning (RL): The conditional reverse diffusion policy is optimized via Proximal Policy Optimization (PPO) to align with terminal target-property rewards. Crucially, RL operates over the motif-aware action space, where actions correspond to chemically meaningful decisions (e.g., attaching a benzene ring) rather than individual atom edits.

3. Theoretical Justification

The paper provides a theoretical analysis characterizing the bottleneck of atom-level RL. By defining the decision complexity as the number of categorical factors in the reverse kernel, the authors show that motif-aware representations reduce the action space size by a factor proportional to the square of the atom-to-motif compression ratio ( $\chi$ ). This reduction stabilizes the KL-regularized policy optimization, allowing rewards to be credited more directly to property-relevant actions and preventing trajectory collapse.

Key Contributions

Unified Foundation Model: CoMole is presented as the first family of controllable molecular generative foundation models designed for heterogeneous inverse design, bridging the gap between current task-specific methods and practical needs.
Motif-Aware RL: The work introduces a novel approach to stabilizing RL for graph generation by shifting the decision space from atoms to chemically meaningful motifs, theoretically justifying why this reduces decision complexity and improves credit assignment.
Transferable Controllability: The framework demonstrates the ability to transfer controllability to unseen property targets by optimizing only task embeddings while keeping the generator frozen, achieving performance competitive with task-specific baselines.

Experimental Results

CoMole was evaluated across three heterogeneous benchmarks: polymer gas permeability, polymer DFT properties, and drug discovery tasks (FreeSolv, BACE, BBBP).

Controllability: CoMole ranked first in controllability on all nine targets.
- On polymer benchmarks, it reduced Mean Absolute Error (MAE) by over 44% relative to the best baselines.
- On drug tasks, it reduced FreeSolv MAE by 13.1% and achieved 1.0 accuracy on BACE classification.
- Compared to the strongest baselines (e.g., Graph DiT, FREED), CoMole reduced MAE by up to 48.2%.
Validity: The model maintained a validity rate above 0.94 across all datasets without rule-based correction or post-hoc filtering. In contrast, atom-level RL variants suffered validity drops (e.g., from 0.95 to 0.07) and failed to maintain controllability.
Generalization: When tested on unseen DFT targets ( $E_i$ , $\epsilon_s$ ) by freezing the generator and learning only new task embeddings, CoMole achieved performance competitive with baselines trained directly on those targets.
Ablation: Replacing the motif tokenizer with an atom-level representation caused a catastrophic failure in both validity and controllability, confirming the necessity of the motif-aware design.

Significance and Claims

The paper claims that CoMole provides a practical foundation for unified, scalable, and controllable molecular inverse design. By transferring pretrained structural priors into controllable generation through motif-aware policy optimization, the model learns transferable structure-property knowledge that supports diverse inverse-design tasks.

The authors emphasize that their design addresses the fundamental fragility of atom-level RL in graph generation. By stabilizing the optimization geometry through chemically meaningful decisions, CoMole achieves high controllability while preserving chemical validity, offering a pathway toward foundation-model-based design systems for molecules and polymers that can bridge the gap between hypothetical generation and practical scientific discovery. The work suggests that such models can help researchers prioritize promising candidates under heterogeneous constraints before committing to expensive simulations or experiments.