OmegaFold、EquBind、RELATION、BIMODAL…你都掌握了吗？一文总结生物制药必备经典模型（二）

二、分子生成

1、 BIMODAL

循环神经网络（RNN）能够使用简化的分子输入线输入系统（SMILES）字符串表示的化学结构来生成新的分子设计。基于RNN的结构生成通常是单向进行的，通过从左到右增长SMILES字符串。然而，小分子没有自然的起点或终点，SMILES字符串本质上是分子图的非单点表示。这些特性促使了双向结构的生成。这篇文章介绍了用于基于SMILES的分子设计的双向生成性RNNs，实现了两种既定的双向方法，并引入了一种用于SMILES字符串生成和数据增强的新方法：双向分子设计交替学习法（BIMODAL）。作者将这三种双向策略与用于SMILES字符串生成的单向正向RNN方法进行了比较，内容包括(i)计算机生成的分子的新颖性，(ii)支架的多样性，(iii)化学-生物学相关性。

图4 (a) SMILES字符串，从分子图表示中获得，每个原子用其元素符号表示，而分支和连接性用符号或小写字母表示（例如，"( )"、"="和 "c "分别表示分支、双键和芳香族碳）。图中给出代表药物布洛芬的三个SMILES字符串的示例；用于SMILES字符串生成的起始原子用灰色数字表示。(b) 带有一个递归神经元层的前向RNN的简化方案。RNNs是一个动态系统的模型，其中任何一个时间点t的网络状态都取决于当前的观察（x_t）和之前的状态（t - 1），并被用来预测输出（y_t）

以SMILES字符序列（"token"）作为输入，RNN模型根据序列的前一部分和概率估计，每次学习预测一个token（图4b）。可以从学到的概率分布抽样新的SMILES字符串。RNNs通常被训练成以 "向前 "的方式读取和生成SMILES字符串，即从左到右。然而，SMILES表示可以从任何非氢原子开始，按任何方向进行生成（图4a）。与自然语言不同，小分子没有唯一定义的起点和终点。非单向性和非方向性提供了探索双向序列生成的机会，即在前向和后向都能读取和生成SMILES字符串的方法。然而，纯粹的从头开始（"端到端"）的双向SMILES生成至今还没有被探索过。BIMODAL就是一种双向生成性RNNs。

给定一个输入序列，生成式RNNs被训练成通过预测下一个序列标记来扩展这个序列，定义为y_t = x_t+1。使用带有LSTM单元的RNNs以解决由长序列和大型网络结构引起的梯度消失和梯度爆炸问题。在任何给定的第t个时间步长，这样的网络由以下一组方程描述：

最常见的用于序列生成的RNNs版本从左到右进行（前向），即从t=1到t=L，其中，L是SMILES序列的长度。在训练过程中，输入的第一个位置被填入一个序列开始的token，而输入的最后一个位置被填入一个序列结束的token。一旦RNN模型被训练好，新的序列就会通过(i)输入起始token("G")，(ii)允许模型逐步选择下一个token，给定各自的前一个token序列，直到生成结束token("E")(图5a)。在每一个时间步长t，每一个第k个符号跟随生成的字符串的前一部分的概率是用一个softmax函数计算的

图5 基于RNN的SMILES字符串生成方法。SMILES生成从起始token "G "开始，按预定方向进行。(a) 前向RNN。从起始token "G "开始，从左到右添加新的token。(b) BIMODAL方法：在每个时间步长(t)交替生成token。该模型使用整个序列（前向和后向）来生成下一个token。(c) 前向-后向模型。从 "G "token开始，每个时间步长预测两个token，两边各一个。(d) NADE方法：缺失的 "假 "token（"M"）被替换为有效的SMILES字符，可以向字符串的中心或以随机方式替换

在任何第t个时间步长，BIMODAL沿前向（x_m→x_t）和后向（x_t←x_m）方向读取x={x_m, x_m+1, ..., x_t}，在两个方向上生成SMILES序列。然而，通过同时使用从左到右（向前）和从右到左（向后）的信息，每一步只有一个token被交替预测。BIMODAL由两个RNN组成，每个方向（前向和后向）都有一个读取序列，然后结合起来提供一个联合预测（y_t）：

在SMILES的生成设置中，BIMODAL在每个时间步长t的前向和后向都会读取序列（图5b）。然后，它在前向或后向生成一个新token：

2、 GF-VAE

GF-VAE是一种用于分子图生成的基于流的变分自动编码器(VAE)模型。该模型在原来VAE的基础上增加了Flow模型解码器。其中，编码器主要是加速解码的训练，而解码器则依次优化编码器的性能。由于流模型的可逆性，生成过程很容易通过反转解码器来完成。因此，GF-VAE继承了VAE和基于流的方法的优点。给定𝑅键的类型和𝑀原子的类型，分子图可以用𝐺=（𝐴，𝑋）表示，其中，𝐴∈{0, 1}^𝑁×𝑁×𝑅是𝑁原子的邻接张量，𝑋∈{0,1}^𝑁×𝑀为一个显示𝑁原子节点类型的特征矩阵。我们的主要目标是从给定的图集G中学习生成模型𝑝G(𝐺)，使从分布𝑝G中抽取的样本是一个有效的分子图。GF-VAE的架构如图6所示。

图6 GF-VAE的结构（训练）

Flow-generator由两个部分组成：atom flow用于近似原子的分布，bond flow用于捕捉键的分布。通过向𝐴和𝑋添加一个均匀的随机噪声𝑈[0,1]来对𝐺进行去量化处理：

将得到的𝐺′=(𝐴′,𝑋′)送到流生成器（Flow-generator），学习从原空间到潜在空间的偏射，使潜在空间中映射的图向量遵循高斯分布。由于我们已经获得了高斯分布中的VAE编码向量z，原子流和键流只需要将𝑋′和𝐴′贴近z，而不是进行像完全边际归一化流模型那样多的工作来映射到标准高斯分布。

对于bond flow，以Glow模型为基础。它为z𝐴𝑁(𝐴′)学习了一个可逆映射，它将𝐴′∈R^𝑅×𝑁×𝑁映射到z𝐴∈R^𝑅×𝑁。使用图16中bond flow中从下到上的所有过程，如squeeze+actnorm、可逆1∗1卷积，以及掩码卷积。对于bond coupling layer，类似于图像处理中使用的，通过3×3𝑐𝑜𝑛𝑣2𝑑→𝑐ℎ𝑁𝑜𝑟𝑚→𝑅𝑒𝐿𝑢层并堆叠𝐾𝐵次建立。此外，整个模块被堆叠𝐿𝐵次以提高映射性能。

对于atom flow，通过𝑜𝑚𝜔(𝑋′,𝐴)学习一个可逆的映射函数来得到z𝑋∈R^𝑁×𝑀。如图16的atom flow所示，atom affine coupling layer由𝐾𝐴次叠加𝑐ℎ𝑁𝑜𝑟𝑚→𝑅𝑒𝐿𝑢→𝑅𝑒𝐿𝑢层与多层感知器（MLP）输出层组成。这里𝑔𝑟𝑎𝑝ℎ𝑐𝑜𝑛𝑣也由Relational-GCN实现。由于二维原子矩阵中没有通道，需要对每一行而不是每一个通道进行归一化处理，这与屏蔽卷积过程类似。此外，整个模块被堆叠了𝐿𝐴次，以达到更好的映射性能。由于𝑓𝑎𝑡𝑜𝑚𝜔和𝑓𝑏𝑜 都是可倒置的，给定潜在向量z𝑋和z𝐴来自已知先验分布，给定图形的概率𝐺=（𝐴，𝑋）可以计算为：

图7 GF-VAE中的原子流和键合流

生成过程如图17所示，首先从标准高斯分布中提取随机样本z，并将其分成zA和zX。由于Flow模型的生成器（𝑓𝑎𝑡𝑜𝑚𝜔和𝑜𝑚𝜔）是可逆的，只需将zA放入𝑓𝑎𝑡𝑜𝑚𝜔，zX放入𝑜𝑚𝜔，但顺序相反。值得注意的是，原子流和键流可以在训练中同时进行，因为它们的更新是相互独立的。然而，原子流在生成过程中需要有效的键张量作为输入。因此，需要首先在zA上应用键合流的反向步骤以得到键合张量𝐴。然后，把它和zX一起送入原子流的反向步骤，得到原子张量𝑋。最后，把原子张量和键张量放入有效性修正模块，得到最终生成的分子。

图8 GF-VAE的生成过程

3、 MCMG

MCMG是一种基于知识蒸馏的多约束分子生成方法，该方法将Conditional Transformer和强化学习算法相结合，满足多个约束条件。利用条件Transformer训练分子生成模型，有效地学习并将结构性质关系纳入有偏生成过程中。然后利用知识蒸馏模型降低模型的复杂度，通过强化学习对模型进行有效的微调，增强生成分子的结构多样性。

MCMG由三个部分组成 ：先验模型、基于RNN的蒸馏模型、使用RL微调的专家模型。蒸馏模型可以帮助学习到更加specific的信息，从而平滑不必要的噪音。MCMG模型的工作流程如图 18 所示。首先训练一个c-Transformer，然后将该c-Transformer蒸馏为RNN，以便后续与RL的应用。蒸馏后的RNN不仅可以减轻RL的训练负担，还可以提升生成分子的结构多样性。

图9 MCMG结构

Prior model的架构如图 18B 所示。该模型学习生成具有由一组条件标记编码的特定的分子。Prior model在训练过程中学习了约束和SMILES的联合嵌入；除了约束代码的处理外，其余的训练过程类似于标准的seq2seq训练。由于Transformer模型庞大，直接通过 RL 算法进行微调较为困难，并且这种c-Transformer的重构化学空间过于集中，将产生易于陷入局部最优的潜在问题。为了应对这一挑战，作者采用并比较了两种知识蒸馏方法。第一种是构建具有三层门控循环单元（GRU）的RNN，以学习从prior model中采样的分子子集；第二种是直接使用c-Transformer模型生成100万个所需分子的数据集（给定一组适当的条件token），然后使用该数据集训练具有上述相同结构的RNN。作者将蒸馏模型和以这种方式训练的最终模型分别命名为蒸馏分子模型（DM）和MCMG分子模型（MCMGM）；此外，作者还提出了两种不同的模型，semi-DM和semi-MCMGM，试图只引入方便计算且可靠的标签，例如QED和SA，避免标记具有预测生物活性的分子。最后，作者采用了REINVENT模型中用到的RL 算法来微调蒸馏模型，并为药物设计中分子生成通常需要的多个目标构建定制的奖励函数。