Any-Property-Conditional Molecule Generation with Self-Criticism using Spanning Trees

📄 arXiv: 2407.09357v3 📥 PDF

作者: Alexia Jolicoeur-Martineau, Aristide Baratin, Kisoo Kwon, Boris Knyazev, Yan Zhang

分类: cs.LG, q-bio.BM

发布日期: 2024-07-12 (更新: 2025-07-15)

备注: Code: https://github.com/SamsungSAILMontreal/AnyMolGenCritic

💡 一句话要点

提出STGG+模型，利用生成树和自批判机制实现任意属性条件下的分子生成

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 分子生成 条件生成 生成树 Transformer 自批判 药物发现 材料设计

📋 核心要点

  1. 现有分子生成模型常产生大量无效分子，限制了其应用。
  2. STGG+模型通过生成树结构保证分子有效性，并引入自批判机制优化属性条件生成。
  3. 实验表明，STGG+在条件分子生成和奖励最大化方面均达到最优性能。

📝 摘要（中文）

生成新颖分子极具挑战，多数表示方法导致生成模型产生大量无效分子。基于生成树的图生成(STGG)是一种有前景的方法，能确保生成有效分子，在无条件生成方面优于最先进的SMILES和图扩散模型。在实际应用中，我们希望能够根据一个或多个期望属性有条件地生成分子，而不是无条件生成。因此，本文将STGG扩展到多属性条件生成。我们的方法STGG+，结合了现代Transformer架构、训练期间属性的随机掩码(支持对任何属性子集进行条件控制和无分类器指导)、辅助属性预测损失(允许模型自我批判分子并选择最佳分子)以及其他改进。我们证明STGG+在分布内和分布外条件生成以及奖励最大化方面实现了最先进的性能。

🔬 方法详解

问题定义：论文旨在解决多属性条件下的有效分子生成问题。现有方法，如基于SMILES和图扩散的模型，在生成分子时容易产生无效结构，且难以有效控制多个目标属性。

核心思路：论文的核心思路是利用生成树结构来保证生成分子的有效性，并结合Transformer架构和自批判机制来实现对多个属性的条件控制和优化。生成树确保了分子结构的连通性和化学有效性，Transformer用于学习分子结构的表示和生成规则，自批判机制则用于评估生成分子的质量并选择最佳结果。

技术框架：STGG+模型主要包含以下几个模块：1) 基于生成树的图生成器：负责生成有效的分子图结构。2) Transformer编码器-解码器：用于学习分子结构的表示和生成规则，并根据给定的属性条件生成新的分子。3) 属性预测器：用于预测生成分子的属性值，并作为自批判机制的依据。4) 训练策略：包括随机掩码属性和辅助属性预测损失，以提高模型的泛化能力和生成质量。

关键创新：STGG+的关键创新在于以下几个方面：1) 将生成树结构引入到条件分子生成中，保证了生成分子的有效性。2) 采用随机掩码属性的方法，使得模型可以灵活地根据任意属性子集进行条件生成。3) 引入辅助属性预测损失，使得模型可以自我批判生成分子的质量并选择最佳结果。

关键设计：在STGG+模型中，关键的设计包括：1) Transformer架构的选择和参数设置，例如层数、隐藏单元数等。2) 属性预测器的设计，包括网络结构和损失函数。3) 随机掩码属性的比例和策略。4) 辅助属性预测损失的权重。这些设计都对模型的性能和生成质量有重要影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

STGG+在分布内和分布外的条件分子生成任务上均取得了state-of-the-art的性能。在奖励最大化任务中，STGG+也显著优于其他基线模型，表明其能够有效地生成具有期望属性的分子。具体性能数据在论文中有详细展示，证明了该方法的有效性和优越性。

🎯 应用场景

该研究成果可应用于药物发现、材料设计等领域。通过指定目标分子的多个属性，例如溶解度、活性等，可以生成满足特定需求的候选分子，加速新药和新材料的研发过程。此外，该方法还可以用于优化现有分子的属性，提高其性能。

📄 摘要（原文）

Generating novel molecules is challenging, with most representations leading to generative models producing many invalid molecules. Spanning Tree-based Graph Generation (STGG) is a promising approach to ensure the generation of valid molecules, outperforming state-of-the-art SMILES and graph diffusion models for unconditional generation. In the real world, we want to be able to generate molecules conditional on one or multiple desired properties rather than unconditionally. Thus, in this work, we extend STGG to multi-property-conditional generation. Our approach, STGG+, incorporates a modern Transformer architecture, random masking of properties during training (enabling conditioning on any subset of properties and classifier-free guidance), an auxiliary property-prediction loss (allowing the model to self-criticize molecules and select the best ones), and other improvements. We show that STGG+ achieves state-of-the-art performance on in-distribution and out-of-distribution conditional generation, and reward maximization.