本文主要是介绍ICLR2023 | RetMol：基于检索的可控分子生成，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
题目：Retrieval-based Controllable Molecule Generation
文章地址：https://openreview.net/pdf?id=vDFA1tpuLvk
代码地址：GitHub - NVlabs/RetMol: A new retrieval-based framework for controllable molecule generation.
主题：今天我们为大家分享一篇来自英伟达和莱斯大学研究团队的发表于“ICLR2023”的一篇分子生成的文章，作者提出了一种基于检索的分子生成框架“RetMol”，它可以使用仅满足部分目标属性的少量分子生成出满足所有所需属性的分子。
一、摘要
药物发现是一个复杂的多目标问题，生成出满足多种性质的分子在药物发现过程中非常重要。现有方法需要使用大型数据集进行的训练/微调，这在真实世界的生成任务中通常是不现实的。在本研究中，作者提出了一种新的基于检索的可控分子生成框架。不需要特定任务的微调，在数据稀缺的情况下（活性分子有限）容易推广到各种生成任务：仅仅需要使用一小组示例分子来引导模型就能生成出满足给定设计目标的分子。作者设计了一个新的自监督目标训练的检索机制，融合了检索库中的样本分子和输入分子，并预测输入分子附近的相似分子；提出了一个迭代改进过程，动态更新检索数据库以获得更好的泛化性能。从结果上看，在各种任务中，从简单的设计标准到用于设计与SARS-CoV-2主蛋白相结合的引物化合物的现实场景，作者提出的方法比之前的方法具有更好的性能和更广泛的适用性。
二、方法
图1：RetMol框架。该框架结合了检索模块（分子检索器和信息融合）和预训练的生成模型（编码器Encoder和解码器Decoder）。上图示例展示了如何在各种设计条件下优化现有潜在药物法匹拉韦（Favipiravir）的结合力，以更好地治疗COVID-19病毒（SARS-CoV-2主要蛋白质，PDB ID：7L11）。
如上图所示，RetMol框架将一个轻量级检索机制Retriever插入预训练的编码器-解码器生成模型中。对于每个任务，作者首先构建了一个满足部分设计标准的样本分子组成的检索数据库Retrieval database。在给定待优化的输入分子Input molecule之后，检索器模块Retriever会从数据库中检索少量与之相关的示例分子，然后由预先训练好的生成模型的编码器Encoder将它们与输入分子转换为Retrieved embeddings 和 Input embedding。接下来，信息融合模块Information fusion将示例分子的嵌入与输入分子嵌入融合Fused embedding，以引导生成（通过预训练生成模型中的解码器）满足所需性质的分子Output molecule。
信息融合模块Information fusion是RetMol中唯一需要训练的部分。为了训练，作者提出了一个新的自监督目标（即预测输入分子的最近邻），以更新融合模块，使RetMol能够在没有特定任务训练/微调的情况下适用于各种生成任务。在模型的推理阶段，作者提出了一种通过迭代的推理过程，动态更新生成的分子和检索数据库，从而提高了生成质量和多样性。
1）编码器-解码器生成模型主干
预训练分子生成模型构成了RetMol的骨干，它连接了连续嵌入和原始分子表示。编码器将传入的分子编码成嵌入embedding，解码器从嵌入embedding中生成新分子。RetMol不关心底层编码器和解码器架构的设计，使其能够与各种生成模型和分子表示一起使用。在这项工作中，作者使用SMILES字符串表示分子和ChemFormer模型作为主干，其中ChemFormer模型是在十亿级别的ZINC数据集上训练的BART变体。
2）检索数据库（Retrieval database）
检索数据库Retrieval database包含满足各种所需性质的分子，可以引导输入分子向设计标准方向发展，从而生成出具有目标性质的分子，因此对于可控生成非常重要。了解哪些分子符合设计标准以及如何选择满足设计标准的分子在RetMol中发挥着重要作用。此外，作者发现仅有少量分子（例如23个）的数据库就已经可以提供强有力的控制信号，这使得作者的方法可以通过快速检索数据库来轻松适应各种任务。此外，在推理过程中可以动态更新检索数据库，即新产生的分子可以扩充检索数据库以实现更好的泛化效果。
3）分子检索
虽然在生成过程中可以使用整个检索数据库，但出于计算原因（例如内存和效率），选择少量最相关的分子来提供更准确的引导更为可行。因此作者设计了一个基于启发式的检索器，以检索出最适合给定控制设计标准的分子：首先构建一个“可行”分子集合，其中包含满足所有给定约束条件的分子；如果此集合大于常数K（希望检索的最大分子数，本文使用K=10），则选择具有最佳平均属性得分的K个分子，否则，通过一次移除一个约束条件来构造松弛可行集，直到松弛可行集大于K，选择最近删除的属性约束的最佳得分的K个分子。无论上述哪种情况，检索器都会检索比输入更理想的性质的例分子，并引导生成朝着给定的设计标准发展。一般来说，更复杂的检索器设计对生成更好性质的分子可能会更有效，作者把它留作未来的工作。
4）信息融合（Information fusion）
该模块可以使检索的分子将输入分子朝着目标设计标准变化，它通过一个轻量级的可训练标准交叉注意力机制来合并输入分子和检索到的分子的嵌入得到融合嵌入：
融合嵌入包含从检索到的样本分子和输入分子中提取的所需属性的信息，该信息用作解码器的输入。
如何训练信息融合模块?
- 给定一个输入分子，我们使用来自检索数据库的最近邻分子作为预测目标，这是基于CDDD编码器嵌入空间中的余弦相似度进行确定的。【通过输入分子嵌入空间中的余弦相似度确定“最近邻分子”】
- 解码器接收信息融合模块中融合的嵌入（“输入的分子”和“最近邻分子”），生成新的分子。
- 如公式（2）所示，我们计算解码输出与最近邻分子之间的交叉熵距离作为训练目标。
这个动机是因为：来自检索数据库的其他类似分子（即剩余的K - 1个最近邻）通过融合模块可以为将输入转换为其最近邻提供良好的指导（例如，如何扰动分子以及扰动的大小）。因此，通过这个辅助任务，融合模块可以得到有效更新。
5）通过预测输入分子最近的邻居进行自监督训练
传统的训练目标是重建输入分子，但在作者的情况下这种目标不适用，因为重建输入分子不会依赖于检索到的样本分子，这会导致输入分子无法利用到样本分子的优质属性信息。为了使RetMol学习使用样本分子进行可控生成，作者提出了一种新的自监督训练方案，其目标是预测输入分子的最近邻：
【\mathcal{L}(\theta)=\sum_{i=1}^{\mathcal{B}} \operatorname{CE}\left(\operatorname{Dec}\left(f_{\mathrm{CA}}\left(\boldsymbol{e}_{\mathrm{in}}^{(i)}, \boldsymbol{E}_{r}^{(i)} ; \theta\right)\right), x_{1 \mathrm{NN}}^{(i)}\right)】
其中CE是交叉熵损失函数，【x_{1NN}】代表输入分子的最近邻，B是批处理大小，i对输入分子进行索引。检索到的样本分子集合由输入分子的其余K-1个最近邻组成。在训练期间，作者冻结了预训练编码器和解码器的参数，只更新信息融合模块
【f_{CA}】中的参数，从而使训练轻量级和高效。此外，作者使用整个训练数据集作为检索数据库，并通过与输入分子的最佳相似度来检索样本分子，因此训练不是针对特定任务的，而是迫使融合模块
参与以相似性为代理标准的可控生成。在推理时，使用上述训练目标和使用相似度的代理准则训练的模型能够推广到具有不同设计准则的不同生成任务，并且比使用传统重建目标训练的模型性能更好。为了更加有效的训练，作者用高效的近似KNN算法预先计算了所有分子的嵌入及其两两之间的相似性。
采样过程：在本文的所有实验中，我们都在解码器中使用“贪心采样”：当我们需要从给定(融合)嵌入的解码器中采样多个分子时，我们首先用标准差为1的独立随机各向同性高斯对(融合)嵌入进行扰动）（independent random isotropic Gaussian with standard deviation of 1），然后使用解码器从每个扰动(融合)嵌入生成一个样本。
6）通过迭代改进策略进行推理
作者提出了一种迭代改进策略，以在利用样本分子的隐式指导进行可控生成时获得改进的结果。该策略通过替换输入分子并更新检索数据库以获取新生成的分子来工作。这种迭代更新在许多其他可控生成方法中是常见的（例如基于GA的方法，基于潜码优化的方法），此外，如果在推理过程中固定而不进行更新检索数据库，则作者提出的方法会受到检索数据库中分子最佳性质值的限制，这极大地限制泛化能力以及在某些生成场景下的表现，例如无约束优化，因此作者建议在迭代中动态更新检索数据库。
首先通过添加高斯噪声来随机扰动融合嵌入（fused embedding）M次，并从每个扰动嵌入中贪心地解码一个分子，以获得M个生成的分子，根据任务特定的设计标准对这些分子进行评分。然后，如果此集合中最佳分子的分数优于输入分子的分数，则将输入分子替换为最佳分子。同时，如果它们的分数比检索数据库中的最低分数更好，则将其余分子添加到检索数据库中。如果没有生成的分子的分数优于输入分子或检索数据库中最低分数，则输入分子和检索数据库保持不变，以供下一次迭代使用。如果达到最大允许的迭代次数或生成符合所需标准的成功分子便会停止。
三、实验结果
1）QED和Penalized logP优化
QED优化：这是一个带约束的分子优化问题，目标是在满足与目标分子相似度大于一定阈值的条件下，优化目标分子的QED性质成功率。如图2（a）所示，RetMol的成功率远高于其他方法。
Penalized logP优化：这是一个带约束的分子优化问题，目标是在满足与目标分子相似度大于一定阈值的条件下，优化目标分子的Penalized logP性质。如图2（b）所示，RetMol能有效提高分子的此项性质。
表格1：（a）在相似性
【\delta】下满足QED条件的分子生成成功率；（b）在相似性
2）多目标优化
多目标优化：这是一个带约束的分子优化问题，目标是在满足分子的QED性质和SA性质符合阈值内的情况下提高生成分子对GSK3β和JNK3两个靶点的亲和性。
图2：在Guacamol基准测试上进行多属性优化（MPO）任务时与SOTA的比较。左图：QED（↑）与平均基准性能（↑）比较。右图：SA（↓）与平均基准性能（↑）之比较。RetMol在提高基准性能的同时保持所生成分子的可合成性（SA）和药物样性（QED）方面实现了最佳平衡。
表格2：在优化四个属性（QED、SA以及两种与GSK3β和JNK3靶点的结合亲和力）的任务中生成的分子的成功率、新颖性和多样性。
3）新冠药物优化
新冠药物优化：为了展示基于检索的框架在药物发现前沿的应用性，作者将其应用于改善对SARS-CoV-2主蛋白酶（Mpro，PDB ID：7L11）的现有弱抑制剂的抑制作用。作者使用23种已知抑制剂作为检索数据库，并选择8个最弱的抑制剂作为输入，设计了一个优化任务，以最大化生成分子与Mpro之间的docking score。具体结果如上表格所示，可以看到RetMol相比于对比方法，药物优化成功率以及优化力度都要高很多。
表格3：在COVID-19药物优化任务中，作者通过QED、SA和相似性约束来改善所选分子与SARS-CoV-2主要蛋白酶的结合亲和力（通过docking进行估计）。在更严格的相似性条件下，RetMol成功率更高（5/8与3/8相比）。在较轻的相似性条件下，RetMol实现了更高的改进（2.84与1.67平均结合亲和力改进值相比）。
图3：比较RetMol和Graph GA方法在优化与SARS-CoV-2主要蛋白酶结合的原始抑制剂Bromhexine的δ=0.6情况下的三维可视化。RetMol优化的抑制剂具有更多的极性接触（红色线），与原化合物相比也具有更不同的结合模式，这与定量结果相一致。
四、结论：
作者提出了一种基于检索的可控分子生成框架RetMol。通过将检索模块与预训练的生成模型结合，RetMol利用从特定任务的检索数据库中检索的示例分子来引导生成模型生成满足所需设计标准的新分子。RetMol非常灵活，不需要任务特定的微调。作者展示了RetMol在各种基准任务和SARS-CoV-2病毒的真实抑制剂设计任务中的有效性，在每种情况下与现有方法相比均取得了最先进的性能。
问题：
数据集：
检索库中的分子是什么？
训练：
The experiments in Sections 3.1 and 3.3 use the ZINC250k dataset（For the ZINC250k dataset, we follow the train/validation/test splits in (Jin et al., 2019) 【jtnn】and train on the train split. ）， while the experiments in Sections 3.2 and 3.4 use the CheMBL dataset.（For the ChemMBL dataset, we train on the entire dataset without splits. ）
ZINC250k dataset：3.1、IMPROVING QED AND PENALIZED LOGP UNDER SIMILARITY CONSTRAINT + 3.3 GUACAMOL BENCHMARK MULTIPLE PROPERTY OPTIMIZATION
CheMBL dataset：3.2 OPTIMIZING GSK3β AND JNK3 INHIBITION UNDER QED AND SA CONSTRAINTS + 3.4 OPTIMIZING EXISTING INHIBITORS FOR SARS-COV-2 MAIN PROTEASE
推理：
将一个fused embedding perturb M次（by adding Guassian noises） --> 使用贪心搜索算法decoder出M个分子
evaluation：
We perform the evaluation on the test split of the ZINC250k dataset. For each molecule in the test set, we generate 10 molecules by first randomly perturb the input to the decoder, i.e., an encoded (or fused) embedding matrix, 10 independent times using a isotropic random Gaussian with standard deviation of 1.
- Then, for validity, we compute the percentage of the 10 generated molecules that are valid according to RDKit, averaged over all test molecules.
- For novelty, we compute the percentage of the 10 generated molecules that are not in the training split of the ZINC250k dataset, averaged over all molecules in the test set.
- For uniqueness, we compute the percentage of the 10 generated molecules that are unique, averaged over all molecules in the test set.
plogp\QED实验：
- 检索数据库：For both experiments, we use the top 1k molecules with the best property values from the ZINC250k dataset as the retrieval database.
- 推理优化过程：We select 800 molecules with QED in the range [0.7, 0.8] as inputs to be optimized.
- 推理优化过程：We select 800 molecules that have the lowest penalized logP values in the ZINC250k dataset as inputs to be optimized.
是已经经过预先计算plogp\QED...的ZINC/CHEMBL分子，
怎么进行的多目标优化？
信息融合后得到一个分子embedding后，扰动潜码得到多个分子，decoder这些分子后，使用性质预测方法/直接使用rdkit求logp、QED得到这些分子对应的属性值【单目标优化就是预测一个属性、多目标就是进行多个属性预测】，将top的分子替换掉检索库中性质最不好的分子，然后再进行检索、融合、decoder....
训练的时候和推理的时候不一样
训练的时候检索库中没有使用带有标签的数据，推理的时候使用了带有标签的数据，训练的时候使用的是KNN求最近邻，并没有迭代更新检索数据库，推理的时候进行了更新
推理时：
这里的多目标优化使用的是多个损失求和，为了简单，本文将每个损失函数的权重设置为了1
损失函数和 encoder、decoder
where CE is the cross entropy loss function ，we use the BART model as our encoder and decoder【其实是使用的Chemformer，Chemformer使用的是BART】
Loss（decoder(输入分子的embedding和检索示例（retrieved exemplar）融合后的), 输入分子的最近邻）：计算的是一个批次中每个输入分子的最近邻smiles和输入分子的embedding和检索示例（retrieved exemplar）融合后的再Decoder后的smiles直接的损失
备注：我去看了本文的代码，其中的损失函数并没有求decoder后的fushion输出和最近邻的损失，而是求的decoder后的fushion输出和decoder的损失，其中decoder的pred_target是“reconstruction”，但是其实应该是 nearestn才是正确的
Chemformer: a pre-trained transformer for computational chemistry.
本文使用了很多Chemformer的代码：
code：GitHub - MolecularAI/Chemformer
paper：ShieldSquare Captcha
代码相关
交叉注意力机制（融合模块）
MolBART/megatron_molbart/megatron_bart.py 1050行：class RetrievalInformationFusion(MegatronModule):
这篇关于ICLR2023 | RetMol：基于检索的可控分子生成的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
