[论文解读] Grammars and reinforcement learning for molecule optimization

Egor Kraev|arXiv (Cornell University)|Nov 27, 2018

Machine Learning in Materials Science参考文献 8被引用 5

一句话总结

该论文提出了一种新颖的方法，结合自定义的上下文无关语法（CFG）用于SMILES，以及基于Transformer模型的强化学习，以生成化学上有效的分子并优化其性质。通过规则掩码强制执行价键规则和环状结构约束，该方法确保了分子的有效性，同时实现了高效且高质量的生成，在每原子的模型步数显著减少的情况下，实现了分子性质优化的最先进性能。

ABSTRACT

We seek to automate the design of molecules based on specific chemical properties. Our primary contributions are a simpler method for generating SMILES strings guaranteed to be chemically valid, using a combination of a new context-free grammar for SMILES and additional masking logic; and casting the molecular property optimization as a reinforcement learning problem, specifically best-of-batch policy gradient applied to a Transformer model architecture. This approach uses substantially fewer model steps per atom than earlier approaches, thus enabling generation of larger molecules, and beats previous state-of-the art baselines by a significant margin. Applying reinforcement learning to a combination of a custom context-free grammar with additional masking to enforce non-local constraints is applicable to any optimization of a graph structure under a mixture of local and nonlocal constraints.

研究动机与目标

• 通过深度学习自动化设计具有所需化学性质的分子。
• 解决仅依赖复杂架构或大规模预定义词汇表来生成化学上有效SMILES字符串的挑战。
• 通过减少每原子的模型步数，提高分子生成的样本效率。
• 在局部（价键）和非局部（环闭合、长度）约束下，实现复杂分子性质的优化。
• 证明一种更简单的、基于语法引导的强化学习方法可在分子优化任务中超越最先进结果。

提出的方法

• 设计了一种新的SMILES上下文无关语法（CFG），通过构造确保原子价态正确，并支持复杂的环状结构。
• 引入额外的掩码逻辑以强制执行非局部约束：终端距离掩码确保在模型步数限制内终止，循环ID传播确保环闭合和长度约束。
• 生成过程使用基于Transformer的策略网络，利用完整上下文注意力机制，在每一步做出明智决策。
• 将分子性质优化建模为强化学习问题，采用批量最佳策略梯度方法，奖励函数结合性质得分、SA得分以及对不良特征的惩罚。
• 使用标准CFG解析器（如NLTK）将真实分子分解为生成规则序列，用于训练和验证。
• 该方法通过将有效性验证与性质优化解耦，实现了大分子的高效、可扩展生成。

实验结果

研究问题

• RQ1 能否设计一种自定义的SMILES上下文无关语法，以保证化学上有效的分子，同时支持分支和芳香环等复杂结构特征？
• RQ2 在仅使用局部注意力机制的情况下，如何有效强制执行非局部约束（如环闭合和长度限制）？
• RQ3 基于Transformer的策略结合批量最佳策略梯度方法，是否能在每原子更少的模型步数下，优于先前基于RNN或图神经网络的模型在分子性质优化中的表现？
• RQ4 基于语法规则的生成结合规则掩码，在多大程度上可减少对大规模预定义词汇表或复杂消息传递架构的依赖？
• RQ5 不同的奖励函数设计策略（如锚定、SA惩罚、芳香环惩罚）对生成分子的多样性与类药物性质有何影响？

主要发现

• 该方法在ZINC基准测试中取得了8.46的分子得分，显著超过Jin等人（2018）报告的先前最先进得分7.16。
• 在强锚定和SA惩罚下，生成的最优分子展现出结构多样性，包括六元芳香环、含S和N的五元杂环，且仅含一个氯原子，与真实类药物分子相似。
• 与先前方法相比，该模型每原子所需的模型步数显著减少，从而实现了大分子的高效生成。
• 即使未显式引入环惩罚，当应用强锚定策略时，该方法仍生成了芳香环更少、原子多样性更均衡的分子。
• 尽管架构更简单且基于语法设计，该方法仍优于Kusner等人（2017）的基线模型以及Jin等人（2018）的复杂图神经网络模型。
• 通过使用带循环ID的规则掩码，成功实现了对复杂约束（如最小/最大环长）的强制执行，同时保持了语法正确性和可解析性。