JCIM｜RJT-RL：一种基于化合物的可逆树表示和深度强化学习的分子设计方法

2022年8月12日，东京大学Ryuichiro Ishitani研究团队在Journal of Chemical Information and Modeling期刊上发表论文《Molecular Design Method Using a Reversible Tree Representation of Chemical Compounds and Deep Reinforcement Learning》，提出了RJT-RL，这是一个基于强化学习（reinforcement learning, RL）的分子生成和优化模型。

在模型中，作者首先定义了可逆连接树（可以从树表示转换回化合物的化学结构表示），然后将分子设计形式化为一个生成可逆连接树的强化学习任务，之后通过神经网络将可逆连接树编码为隐表示，建立强化学习模型，逐步生成可逆连接树，得到优化后的分子。作者使用简单的分子设计任务作为基准测试了该方法，并将其应用于涉及多目标评分函数的药物发现中更现实的任务。

1 摘要

自动化设计具有特定化学和生物化学性质的分子是材料信息学和计算药物发现中的一个重要过程。在本研究中，作者设计了一种新的粗粒度树表示（可逆连接树，Reversible Junction Tree, RJT）方法，它可以在没有外部信息的情况下反向转换回原始分子。利用这种表示，作者进一步形式化分子设计和优化问题为使用深度强化学习（RJT-RL）生成可逆连接树。在该方法中，强化学习的所有中间和最终状态都可转化为有效分子，这可以有效地指导简单基准任务中的优化过程。作者进一步研究了使用RJT-RL的强化学习模型的多目标优化和微调，证明了该方法适用于更现实的药物发现任务。

2 方法

2.1 分子的可逆连接树表示

在原始的连接树(junction tree, JT)表示中，使用JT算法将分子图转为树表示（可使用RDKit实现）。简要地说，连接树中的节点表示分子图中的一个化学键或一个环[这些化学键和环已经预先从数据库中筛选出来定义好。如图S1所示，第三行中矩形框表示连接树中的化学键节点（对应分子图中的一个化学键）；圆圈表示连接树中的环节点(对应分子图中的环)；红色圆圈为单原子节点，其被三个化学键节点连接，详细的介绍可以参考原始论文]，如果连接树中的两个节点映射回分子图中存在公共的原子，那么这两个节点在连接树表示中就存在一条连边。连接树中的节点已经被预先定义好，并被视为word且分配了ID：，粗粒度的连接树表示使得从这些预先定好的满足化学规则的片段(word)组装分子称为可能，但是它是不可逆的，也就是说在不借助外部信息的情况下很难从树表示转换回分子表示。因此，作者提出了可逆的连接树表示(the reversible JT representation )。

图S1 本研究中使用的树分解算法的一个示例。第一行的分子图被分解为片段，作为第二行中的节点。如果第二行中的节点间存在公共的原子，那么两者之间在连接树表示中存在连边（第三行）。

为了消除节点连接的任意性（即确定连接树表示中连接两个节点之间的公共原子），作者记录了两个节点和的ID，并称之为“site information” ，定义如图1所阐释的那样。原始的连接树表示中的边包括两类：两节点之间只有一个公共原子[称为“type-1 edge”（类型1边）]和两个节点之间有两个公共原子[称为“type-2 edge”（类型2边）]。类型1边又可细分为三种情况：(1)化学键节点连接单原子节点；(2)两个化学键节点相连；(3)化学键节点和环节点相连。图1A上示例了两个化学键节点相连，图1A下示例了化学键节点和环节点相连。算法1展示了从可逆连接树转换为分子的化学结构的过程。

图1 使原始分子结构可逆转换的site information的定义。(A) 类型1边的site information（有一个公共原子）。在左图中，圆圈和箭头分别代表连接树表示中的节点和边。箭头附近的数字表示site information。这里展示了利用site information，可以唯一确定分子的原始结构（红色虚线框起来的）。如（A）上，两个连接树节点的公共原子是0号位的N和0号位的N；而（A）下，两个连接树节点的公共原子是2号位的C和0号位的C。(B) 类型2边的site information（有两个公共原子）。两个节点公共原子是0号位的两个C和3号位的两个C。

2.2 基于可逆连接树的神经网络

作者将可逆连接树中的节点和边以one-hot编码作为初始表示，并使用基于树的LSTM网络结构将可逆连接树编码为隐表示，可形式化表示为：

RJTNN表示Reversible Junction Tree Neural Network，可逆连接树神经网络。

2.3 使用可逆连接树表示的强化学习

作者使用强化学习将分子设计形式化为一个可逆连接树生成的任务。将可逆连接树的隐表示视为状态（state），采取action来修改可逆连接树，并将action定义为以下4个组件（示例见图2B）：

1. 选择需要添加新节点的节点（Node）
2. 从预定义节点库中选择一个节点（Word）进行添加
3. 预测两个节点之间的site information，将两个节点连接起来
4. 预测是否结束（stop）

这些action的概率大多使用MLP计算得出。图2A示例了强化学习的整体框架，包括：预定义的可逆连接树节点库、可逆连接树的状态、神经网络计算可逆连接树中节点的隐向量、策略分布和采样行为。作者亦会用数据集预训练模型，这样才能得到才是优化后的分子。

图2 RJT-RL的网络结构和智能体(agent)采取的行为(action)。（A）action的计算过程和价值函数(value function)估计。表示状态的RJT通过RJTNN网络被转换为隐向量，这个隐向量是后续计算的一个重要媒介，并使用后续网络计算策略分布(policy distribution)和价值函数(value function)。智能体所采取的action由该策略分布中的采样确定。树节点中的数字表示预先定义的word ID。（B)从策略分布采样action和对表示状态的RJT的修改的示例。采取action之前的状态表示（A)中所示的分子。从（B)中可以看到，agent选择了22号节点，添加word 22，然后确定了公共原子（依据site information)，然后选择了stop策略。

3 结果

作者进行了广泛的实验，且进行了深入的分析，主要包括以下五个实验：

1. 优化化学属性penalized Log P score来评估模型性能；
2. 根据给定的结构生成一个类似的优化后的分子，相似度由分子指纹的Jaccard相似度评估；
3. 基于结构的骨架跃迁，在已知靶标蛋白和化合物的重要相互作用的情况下设计新的化合物；
4. 对Log P和合成可及性(synthetic accessibility)分数的多目标优化；
5. 微调不同的奖励函数(reward function)，即先在预训练，然后再迁移到多目标优化奖励函数中去。

4 结论

在这项研究中，作者引入了可逆连接树，一种分子的粗粒度表示，可以直接转换为原始化学结构。作者利用这种表示法进行药物发现，将分子设计任务用强化学习形式化为生成一个可逆连接树。在简单的基准测试任务中，所提出的方法比最先进的方法表现出更好或差不多的性能。结果表明了逐步评估分子属性的潜力。因为奖励分数需要多次评估以生成单个化合物，因此当分数的计算成本较低时，逐步优化是有优势的。作者进一步证明，本方法适用于现实世界中的任务，如具有多目标奖励函数的基于结构的骨架跃迁任务。所提出的RTJ-RL模型的另一个优点是策略和价值函数模型的微调，这在基于规则的片段增加的方法中是不可能的。此功能可用于调整复合设计过程中的评分函数。

实验结果也表明了该方法存在的一些问题。例如，考虑到涉及对接模拟的3D生成，RJT-RL模型的有效训练需要正确处理生成化合物的立体异构体。本研究中枚举了可能的立体异构体，并使用暴力法寻找最佳立体异构体。然而，搜索是在单个CPU上执行的，这降低了整体性能（表5）。三维构象生成和对接仿真的并行化可以提高性能。这个问题的另一个解决方案是扩展可逆连接树表示中的site information，以正确处理原子的手性标志，从而使智能体可以生成分子，包括立体异构体。最后，在本研究中，作者只考虑了从相对较小的起始片段尺寸开始的化合物的从头设计。为了将所提出的方法应用于具有更大和更复杂结构的先导化合物优化，有必要将action扩展到包括节点缺失和突变，以便能够修饰起始骨架。

参考资料

Ishitani R, Kataoka T, Rikimaru K. Molecular Design Method Using a Reversible Tree Representation of Chemical Compounds and Deep Reinforcement Learning[J]. Journal of Chemical Information and Modeling, 2022.

--------- End ---------