Controllable Molecular Generative Foundation Models

📄 arXiv: 2605.15354v1 📥 PDF

作者: Yihan Zhu, Yuhan Liu, Weijiang Li, Tengfei Luo, Meng Jiang

分类: cs.LG

发布日期: 2026-05-14


💡 一句话要点

CoMole:可控分子生成基础模型,用于异构设计任务。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 分子生成 图扩散模型 强化学习 可控生成 药物发现 材料科学 motif感知

📋 核心要点

  1. 分子图生成缺乏统一框架,难以实现具有可靠可控性的异构设计任务,现有方法在原子级别操作面临巨大搜索空间和无效中间态问题。
  2. CoMole通过motif感知的图扩散管道,学习motif感知的图空间,将预训练的结构先验知识迁移到可控生成中,并利用强化学习优化条件逆向策略。
  3. 实验表明,CoMole在材料和药物发现的多个基准测试中,可控性排名第一,MAE显著降低,同时保持较高的分子有效性,并能将可控性迁移到未见过的属性。

📝 摘要(中文)

尽管基础模型在语言和视觉领域取得了成功,但分子图生成仍然缺乏一个统一的框架,以实现具有可靠可控性的异构设计任务。强化学习(RL)为特定任务优化提供了一种自然的后训练机制,但将其应用于图生成模型受到巨大原子级动作空间和化学上无效的中间状态的阻碍。我们提出了可控分子生成基础模型(CoMole),它建立在统一的motif感知图扩散管道上。通过学习motif感知图空间,CoMole将预训练的结构先验知识转移到可控生成中,其中RL优化了化学上有意义的决策上的条件逆向策略。我们从理论上描述了原子级RL的瓶颈,并证明了motif感知策略优化的合理性。在跨越材料和药物发现的三个异构基准测试中,CoMole在所有九个目标上的可控性排名第一,相对于最强的基线,MAE降低了高达48.2%,并在没有基于规则的校正或事后过滤的情况下,保持了0.94以上的有效性。我们进一步表明,CoMole通过仅优化任务嵌入并冻结生成器,将可控性转移到未见过的属性,实现了与强大的特定任务基线竞争的性能。

🔬 方法详解

问题定义:分子图生成旨在设计具有特定性质的分子结构。现有方法,特别是基于强化学习的方法,在原子级别进行操作,导致巨大的动作空间,难以训练。此外,生成的中间状态可能在化学上无效,需要额外的规则或过滤步骤来保证结果的合理性。

核心思路:CoMole的核心思路是利用motif(分子片段)作为基本构建块,而不是单个原子。通过在motif级别进行操作,可以显著减小动作空间,并确保生成的中间状态更可能在化学上有效。同时,利用图扩散模型学习分子结构的先验知识,并结合强化学习进行微调,以实现对特定性质的可控生成。

技术框架:CoMole包含预训练阶段和强化学习微调阶段。在预训练阶段,使用图扩散模型学习motif感知的图空间。在强化学习微调阶段,利用预训练的生成器作为基础,通过优化条件逆向策略,控制生成过程,使其生成具有目标性质的分子。整体流程包括:1) 从噪声图开始;2) 通过扩散过程逐步添加motif;3) 使用强化学习优化逆向扩散过程,使其生成具有目标性质的分子。

关键创新:CoMole的关键创新在于motif感知的图扩散生成框架。与传统的原子级别操作相比,motif级别操作显著减小了动作空间,提高了生成效率和化学有效性。此外,通过预训练图扩散模型,CoMole能够学习分子结构的先验知识,从而更好地指导生成过程。

关键设计:CoMole的关键设计包括:1) Motif的选择和表示方法;2) 图扩散模型的网络结构和训练方式;3) 强化学习的奖励函数设计,用于指导生成过程,使其生成具有目标性质的分子;4) 任务嵌入的设计,用于将任务信息融入到生成过程中,实现对不同性质的可控生成。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CoMole在三个异构基准测试中,在所有九个目标上的可控性排名第一,相对于最强的基线,MAE降低了高达48.2%,并在没有规则校正或事后过滤的情况下,保持了0.94以上的有效性。此外,CoMole能够将可控性迁移到未见过的属性,实现了与强大的特定任务基线竞争的性能。

🎯 应用场景

CoMole在药物发现和材料科学领域具有广泛的应用前景。它可以用于设计具有特定生物活性或物理化学性质的分子,例如新型药物、高性能材料等。通过控制生成过程,可以加速新分子的发现和优化,降低研发成本,并为解决实际问题提供新的解决方案。

📄 摘要(原文)

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.