今天给大家介绍一篇来自浙江大学侯廷军教授课题组、武汉大学陈曦教授课题组、中南大学曹东升教授课题组和腾讯量子实验室联合发表的一篇文章。该文章提出了多约束分子生成新方法MCMG，通过结合知识蒸馏（Knowledge Distillation）、条件 Transformer（c-Transformer）和强化学习（RL）来生成满足多个约束的分子。c-Transformer用于通过有效学习并将构效关系合并到有偏差的生成过程中来训练分子生成模型。知识蒸馏模型可降低模型的复杂性，便于通过 RL对其进行微调，并增强生成分子的结构多样性。实验证明，MCMG方法可以高效地遍历复杂的化学空间以寻找满足多种性质约束的新型化合物。
1
研究背景
药物分子设计本质上是一个多参数优化问题，其目标是设计和发现具有某些理想属性的新分子，因此如何生成新的分子结构及优化分子的关键属性（如生物活性、成药性、安全性、选择性和可合成性等）是影响药物设计成败的两个关键问题。据估计，类药化学空间内中可探索的分子数高达1023至1060，在如此巨大的化学空间内如何进行分子结构的智能生成并进一步实现结构的快速演化是药物分子设计面临的巨大挑战。近几年，多类基于人工智能（AI）技术的深度分子生成模型在复杂化学空间探索上展现出巨大的潜力，有望成为从头药物设计的关键技术，如循环神经网络（RNN）、变分自动编码器（VAE）、生成对抗网络（GAN）和强化学习（RL）等。
RL可用于微调生成模型的参数，以将自由参数空间引导至具有最佳目标函数值（例如生物活性）的集合。尽管 RL已被证实能够找到一组高度多样的生物活性分子，但它仍然需要大量的优化步骤才能通过奖励训练来学习这些模式。在之前的研究中，为了缓解这个问题并提升RL的效率，在RL操作之前首先使用迁移学习（TL）快速聚焦化学空间中的某些区域。但TL的副作用是可访问化学空间的显著收缩，这使得模型容易陷入局部最优，这意味着生成的分子往往与 TL 训练集中的化合物高度类似。
本研究的目标是通过预处理生成模型而不影响其在多约束任务中输出多样性的同时，提高分子生成模型输出所需分子的效率。为此，作者提出了一种新的分子生成方法，即多约束分子生成方法MCMG。首先，使用c-Transformer来构建生成模型；然后，采用知识蒸馏模型来降低模型的复杂度，并提升生成分子的多样性；最后，通过RL对其进行微调。蒸馏模型最初是为了将从大型模型或多个模型的集合中学到的知识转移到另一个轻量级模型以实现快速部署，但在本研究中该方法还可以大大提高生成分子的结构多样性。在针对多个靶点的模拟实验中，MCMG表现出比传统模型更优的预测效果。
2
MCMG方法
模型框架
MCMG模型的工作流程如图 1 所示。首先训练一个c-Transformer，然后将该c- Transformer蒸馏为RNN，以便于后续与RL的应用。蒸馏后的RNN不仅可以减轻RL时的训练负担，还可以提升生成分子的结构多样性。
图1. MCMG方法的工作流程图
Prior model的架构如图 1B 所示。该模型学习生成具有由一组条件标记编码的特定的分子。Prior model在训练过程中学习了约束和SMILES的联合嵌入；除了约束代码的处理外，其余的训练过程类似于标准的seq2seq训练。由于Transformer模型庞大，直接通过 RL 算法进行微调较为困难，并且这种c-Transformer的重构化学空间过于集中，将产生易于陷入局部最优的潜在问题。为了应对这一挑战，作者采用并比较了两种知识蒸馏方法。第一种是构建具有三层门控循环单元（GRU）的RNN，以学习从prior model中采样的分子子集；第二种是直接使用c-Transformer模型生成100万个所需分子的数据集（给定一组适当的条件标记），然后使用该数据集训练具有上述相同结构的RNN。作者将蒸馏模型和以这种方式训练的最终模型分别命名为蒸馏分子模型（DM）和MCMG分子模型（MCMGM）；此外，作者还提出了两种不同的模型，semi-DM和semi-MCMGM，试图只引入方便计算且可靠的标签，例如QED和SA，避免标记具有预测生物活性的分子。最后，作者采用了REINVENT模型中用到的RL 算法来微调蒸馏模型，并为药物设计中分子生成通常需要的多个目标构建定制的奖励函数。
数据集
训练数据集是REINVENT模型采用的数据集，包含100万个ChEMBL分子。生物活性数据集包括 3 种不同蛋白质靶标的生物活性实验数据，即dopamine type 2 receptor（DRD2）、 c-Jun N-terminal kinase-3（JNK3）和glycogen synthase kinase-3 beta（GSK3β）。DRD2数据集包含10万个阴性和7219个阳性化合物，JNK3数据集包含5万个阴性和2665个阳性化合物，GSK3β数据集包含5万个阴性和740个阳性化合物。
实验任务
两种实验设置用于评估多条件生成模型的质量。第一个评估目标是建立和保存一个最优的分子生成模型，可以在后期重复使用；第二个评估设置是在 RL 辅助微调阶段收集有用的分子。其基本原理是最终目标是获得一组满足要求的生成分子。作者分别在如下两种任务上与几种传统深度生成模型进行了比较：
任务一：DRD2、QED和SA
任务二：JNK3、GSK3β、QED和SA
3
实验结果
Conditional Transformer相对于conditional RNN的优势
鉴于Transformer在自然语言处理（NLP）领域的优越性能，作者将它作为prior model，并验证了c-Transformer能够比c-RNN更好地捕获和利用生成任务的构效关系。为了验证生成分子具有更高的质量，作者通过将条件标记添加到训练中来构建c-RNN。作者比较RNN、c-RNN和c-Transformer，以确定哪个模型更适合用作prior model。计算结果表明，与c-RNN相比，c-Transformer在分子结构的有效性（validity）和成功率（success）方面具有巨大的领先优势，更适合作为prior model。
蒸馏模型的表现
作者分析了两种不同的蒸馏方法（DL和DM）。基于生成分子的MOSES评估指标，DL在各个方面的表现几乎与prior model相同，但DM模型的表现却大不相同，生成的分子满足所有四个约束的成功率急剧下降，但这并不意味着DM模型无法生成我们所需的分子（只是没有那么多同时满足所有约束）。虽然DM模型的成功率明显降低，但其他两个重要属性（Frag和IntDiv）得到了显着改善。作者发现DM模型生成的分子分布与上述其他模型生成的分子分布不同，其重建的化学空间比DL模型大得多。为了证实上述猜想，作者对5000个分子进行了采样，并采用每个模型的平均负对数似然（NLL）进行了化学空间的评估。如图2所示，DL分布的NLL在谱下端附近达到峰值，因此其重构的化学空间最小，接下来依次是DM、semi- DM和RNN。
图2. 平均负对数似然DLL的分布。
MCMG在实验1中优于传统方法
作者首先针对实验1进行了两个任务的基准研究，即构建和保存最佳分子生成模型。该实验的目标是使用RL在化学空间中找到最佳区域，区域的大小和位置可以通过RL进行调整。实验结果表明，MCMG方法与JT-VAE、GCPN、REINVENT、REINVENT2.0和RationaleRL等模型相比，其Success、Novelty、Div、Real success、Unique、Frag、SNN、IntDiv等多个评估参数表现更优；特别是Real success以及Novelty比最优的传统模型高出约30%与20%。
MCMG在实验2中优于传统方法
实验2不是专注于获得最佳模型，而是在RL微调阶段收集有用的分子。无论模型是否收敛，每个模型的训练都在第5000步终止。每个模型收集这5000个步骤中产生的所有成功分子，从池中随机抽取10000个分子作为代表。随后分析了这组具有代表性的10000个分子，以评估模型的性能。总体而言，在任务1中，semi-MCMGM实现了最令人印象深刻的性能，在7个指标中的5个表现出色；在任务2中，MCMGM在5个指标上取得了最佳性能。
作者还分析了与RL优化步骤数量相关的不重复的成功分子的累积数量。一开始MCMGM、MCMGL和REINVENT2.0 迅速积累了大量成功分子，而REINVENT和 semi-MCMGM在第1000步左右才开始加快。原因是REINVENT和semi-MCMGM是从更大化学空间的分子中采样的（没有由于TL或预处理而受到聚焦子空间的影响），并且RL需要更多的迭代来学习所需分子的特征。对于MCMGL，在它输出了一定数量的独特成功分子后，无论再训练多少步，都很难产生更多独特成功的分子，REINVENT2.0也遇到了同样的问题。这是因为这两种模型重建的化学空间很小，成功分子的密度虽大，但是数量很少，RL很容易陷入局部最优，从而导致许多成功分子最终被重复生成。至于其他三个具有更大化学空间的模型，它们的累积数量随着训练步骤的增加而不断增加。总之，在评估设置2中MCMG方法的表现优于其他模型。
图3. 生成的满足要求的分子数量与优化步数之间的关系。
骨架分析
作者对用于实验2中生成的10000个分子进行了骨架分析，提取了所有分子的Murcko骨架并与已知活性化合物的骨架进行了比较，并计算了它们之间的相似性。对于任务1，semi-MCMGM 在骨架新颖性方面比其他方法具有绝对优势。对于任务2，结果相似，MCMGM和semi-MCMGM模型在生成的分子骨架的新颖性方面显示出优势。为了更好地证明MGMGM和semi-MCMGM 具有更好的能力来生成具有独特骨架的分子，选取任务1中部分相似度小于0.2的骨架进行展示（图4），可见模型生成的分子与真实活性分子具有显着的区别。
图4. 生成的分子与真实活性分子的结构示意图。
4
总结
作者将Transformer、knowledge distillation和RL结合提出了MCMG方法，并在平衡分子生成模型的收敛速度和输出多样性的挑战性问题上取得了实质性的进展。作者采用两种评估设置并在多个常见任务下对多个版本的 MCMG模型进行了测试。MCMG在这两项任务上的实际成功率分别为89.26%和70.9%，与其他模型相比有较大的领先优势（REINVENT为72.8%，RationaleRL为51.7%），这表明MCMG可以生成更多的成功分子。总之，MCMG 在多目标分子生成任务中取得了有希望的性能，并提供了一种高效的方法来遍历大型复杂的化学空间以寻找潜在的候选药物。
参考资料
Multi-constraint molecular generation based on conditional transformer, knowledge distillation and reinforcement learning, Nature Machine Intelligence, 2021, 3, 914-922.
