作者 | 刘圣超

编辑 | 凯霞

在上一篇系列文章分享了 Mila 刘圣超博士介绍的【Geometry 相关的多模态任务】相关研究工作，在这里，将接着介绍关于【Textual description 相关的多模态任务】系列研究。

1 GraphCG 基于图的分子编辑(可控生成)

该研究以《GraphCG: Unsupervised Discovery of Steerable Factors in Graphs》为题，发表在 NeurIPS GLFrontiers Workshop 2022 (oral) 上。

论文链接：https://openreview.net/forum?id=BhR44NzeK_1

GraphCG 是我们第一个使用深度学习来进行可控生成的工作。目前已有的可控生成主要是对于图片的操作，而我们认为，对于分子的可控生成，可以很好地用于 lead optimization 等药物发现的重要任务中。问题的出发点是说，现有的lead optimization 都高度依赖于专家的经验，因此比较耗费人力并且有一定的主观性。而当我们如果能够成功利用深度学习的方法，挖掘出图生成模型中的可控因子，那么就能够提供另外一种 lead optimization 方法，从而帮助药物专家进行研发。

问题的设定是给定一个已经训练完成的图生成模型，然后我们主要进行了两个步骤的操作：

首先我们验证了已有的图生成模型，都在一定程度上是高度耦合的 (entangled)。因为高度耦合，所以要实现可控生成就比较有挑战性。

那么对于高度耦合的图生成模型 (已经训练完毕)，如果我们想对它们进行可控生成，是否可行？答案是可以的，我们对于此提出了 GraphCG。

• 在latent space中，我们有不同的semantic direction，而每一个semantic direction都有各自对应的特定的可控因子。
• 如果我们想学习到这些semantic direction，我们先有一个假设：不同的图，在latent space中如果能够沿着一个semantic direction移动，那么它们就会有对应的可控因子改变 (比如某一个图子结构变多)。这种假设可以通过最大化互信息来实现。这就回到了我们的公式1。具体求解的时候，我们利用了NCE的求解方案。这样我们就利用最大化互信息学到了这些semantic direction。
• 在inference过程中，我们只需要把每一个图对应的latent representation 沿着 semantic direction 进行移动，然后进行解码，这个解码之后的图就是我们希望的某一个因子改变的图。

图 1：GraphCG的学习和推理流程图。

具体定量结果可以参考原文，这里我们主要展示定性结果。

首先在分子图，我们利用了已经训练好的HierVAE，在它的latent space上训练得到了10个semantic direction。然后我们发现其中四个semantic direction都能对应到专门的分子官能团。这个很好地符合了我们的expectation，并且也能够帮助药物专家进行lead optimization设计：比如希望对原始分子图增加/减少更多的 halogen 基团，见图 2 (a)。

图 2：GraphCG在分子图上编辑的结果。每一行的五个格子对应了一个分子的五个编辑步长：从左到右依次为{-3, -1.8, 0, 1.8, 3}。其中0（也就是中心格子）表示原始分子。我们展示了四个编辑方向，对应于a到d。

其次我们还展示了基于PointFlow的点云图的可控结果。下图中的a和b是一个方向，并且我们能看到往左往右两个不同的方向，对应的飞机引擎数目会分别减少和增加。此外在车子和椅子的外形也会随着可控因子，有对应的改变。

图 3：GraphCG在点云上编辑的结果。每一行的五个格子对应了一个分子的五个编辑步长：从左到右依次为{-3, -1.8, 0, 1.8, 3}。其中0（也就是中心格子）表示原始点云。我们展示了三个编辑方向，对应于三个不同的预训练生成模型。

GraphCG的初步尝试给我们带来了非常大的信心，让我们对于分子的可控生成有了更大的把握。它也引导了我们后续的几个工作。

2 MoleculeSTM 基于文本的分子编辑

该研究以《Multi-modal Molecule Structure-text Model for Text-based Editing and Retrieval》为题，发表在 arXiv上。

论文链接：https://arxiv.org/abs/2212.10789

最近随着大模型、多模态的应用，一个很自然的想法就是我们是否也可以将这些技术用到药物发现上？并且这些自然语言的文本描述，是否对于药物发现这个有挑战性的问题带来新的视角？答案是肯定并且乐观的。

具体到方法上。MoleculeSTM的核心思路非常简单直接：分子的描述有内部化学结构和外部功能描述两大类，而我们这里利用了multi-modal learning的思路，将两种类型的信息进行联系，并且基于此我们设计了种类丰富的下游任务来验证其有效性。这里预训练的思路还是通过求解公式1来给两个模态（对应的表征函数）进行链接。

图 4：MoleculeSTM的流程图。

下面我们主要强调几个insight。

2.1 自然语言和大语言模型有什么特点能够帮助scientific discovery？

这个是我们在MoleculeSTM中首先提出的问题。在MoleculeSTM这个工作里，我们是利用了自然语言的open vocabulary和compositionality特性。

• Open vocabulary是说我们对于药物的描述可以非常的多样，比如一个极端例子是“开发一个药物能够治疗明年的突发的xxx疾病”这样的text prompt，只不过这种text prompt是难以验证，并且对于目标的描述过于模糊。
• Compositionality是说我们有一些比较复杂的多目标任务，但是如果在自然语言中，它更容易组合。一个例子是多属性的lead optimization，也就是让一个分子优化到同时符合多个属性；传统做法会需要对每一个属性训练一个分类器，而MoleculeSTM仅仅只需要把两个属性用自然语言描述然后通过“and”连接即可。
• 在我们最近的工作ChatDrug中，我们又挖掘了自然语言和大语言模型的对话特性。这个会在下面展开。

2.2 自然语言和大语言模型的特点能够帮助什么类型的scientific discovery任务？

现有的language-vision task可以认为是艺术相关的任务 (比如生成图片、文字)，也就是说它们的结果是可以多样和不确定。但是scientific discovery是科学问题，通常有着比较明确的结果，比如生成有某个功效的小分子。这个在任务的设计上带来了更大的挑战。

在MoleculeSTM中 (Appendix B)，我们提出了两个准则：

• 首先我们考虑的任务是能够进行计算模拟得到结果。将来会考虑能够有wet-lab验证的结果，但这并不在目前这个工作的考量范畴内。
• 其次我们只考虑有着模糊性描述的问题（问题的描述是相对模糊的，但是答案是相对确定）。具体例子比如让某个分子的水溶性或者穿透性变强。而有一些问题有明确结果，比如在分子的某一个位置加入某一个官能团，我们认为这类任务对于药物、化学专家来说更加简单直接，而DL的帮助比较有限；所以它可以将来当作某一个proof-of-concept任务，但是并不会成为主要的任务目标。

2.3 定性结果展示

MoleculeSTM的文章中我们引入了三类下游任务，来验证MoleculeSTM的有效性。这里我们主要想强调一下zero-shot text-guided molecule editing的定性结果。

这个task就是同时输入一个分子和自然语言描述(比如额外的属性)，然后希望能够输出复合语言文本描述的新的分子。这也是text-guided lead optimization，一种新的lead optimization范式。

具体的方法就是利用已经训练好的分子深成模型和我们预训练好的MoleculeSTM，通过学习二者的latent space的alignment，从而进行 latent space interpolation，再经过解码生成目标分子。流程示意图如下。

图5：zero-shot text-guided molecule editing 两阶段流程示意图。

这里我们展示了几组分子编辑的定性结果 (其余下游任务的结果细节可以参考原论文)。主要我们考虑五类分子编辑任务：

• 单一属性编辑：对单一属性进行编辑，比如水溶性、穿透性、氢键施主与受主个数。
• 复合属性编辑：同时对多个属性进行编辑，比如水溶性和氢键施主个数。
• 药物相似性编辑：(Appendix D.5)是让输入分子与目标分子药物长得更加接近。
• 专利药物的邻居搜索：对于已经申请到专利的药物，往往会把中间过程的药物一起报道。我们这里就是那中间药物配合自然语言描述，看是否能够生成最终的目标药物。
• Binding affinity编辑：我们选择几个ChEMBL assay作为target，目标是让输入分子和target有更高的binding affinity。

图 6: zero-shot text-guided molecule editing 结果展示。

3 ProteinDT 基于文本的蛋白质生成和编辑

该研究以《A Text-guided Protein Design Framework》为题，发表在 arXiv上。

论文链接：https://arxiv.org/abs/2302.04611

图 7：ProteinDT流程图。

MoleculeSTM的初步探索让我们感受到文本描述的潜在用途是能够提供更多解决问题的视角。这个工作我们就很自然的将这个想法拓展到蛋白质领域上。

首先在预训练上，ProteinDT主要分为三步：

1、第一步ProteinCLAP，是将text sequence和protein sequence通过contrastive learning来学习alignment。

2、第二步Facilitator，是进一步增强alignment。

3、前两个步骤都是通过latent space操作，而第三步deconder则是将latent representation解码到data space。需要提到的是我们这里考虑了两类decoder：auto-regressive和denoising diffusion(如下图)。

• 这里我们还想要强调的是，Transformer并不是生成模型，而仅仅是一个深度学习模块。我们的两种生成模型都考虑了用Transformer作为核心模块的情况。

图 8：ProteinDT中考虑到的两大类条件生成模型。

3.1 Text-to-Protein Generation

当我们有了上述的流程图之后，就可以进行text-to-protein generation的生成(如下图)。并且我们还进行了消融实验证明了facilitator模块的必要性。

图 9：text-to-protein structure生成过程。

3.2 Text-guided Protein Editing

这里基于ProteinCLAP，我们提出了两种protein editing思路。

• Latent interpolation是直接在latent space进行控制interpolation，然后直接解码到protein sequence space。
• Latent optimization则是专门训练一个token-level的解码器，然后利用ProteinCLAP直接优化得到optimal latent，再通过一个预训练的解码器，解码到protein sequence space。

图 10：两种text-guided protein editing方法流程示意图和比较。

图 11：text-guided protein editing inference流程示意图。

整个text-guided protein editing的inference流程图如图11。更多定量实验结果可以参考原文。

4 ChatDrug 基于ChatGPT 对话增强的小分子、多肽、蛋白质编辑

该研究以《ChatGPT-powered Conversational Drug Editing Using Retrieval and Domain Feedback》为题，发表在 arXiv 上。

论文链接：https://arxiv.org/abs/2305.18090

在2022年的11月，ChatGPT推出。紧接着很多领域、应用都开始尝试使用这个基于大模型的工具，并且发现了它的确能够优雅地解决很多问题。

这个工作中，我们就尝试将ChatGPT用于药物发现。首先一个比较有挑战性的瓶颈还是如何设计任务。我们跟随MoleculeSTM和ProteinDT的insight，主要对标drug editing系列任务，并且我们在这里考虑了三种类型的药物：小分子、多肽、蛋白质。

图 12：ChatDrug流程图。

根据ChatGPT的特性，我们提出了ChatDrug，它主要有三个模块：

• PDDS模块是针对domain task设计的prompt。
• ReDF模块是利用retrieval和domain feedback对prompt进行更新。
• Conversation module是让用户和ChatDrug进行交互，从而不断更新结果。

下面我们列出主要的定性结果。关于更详细的任务设定和定量结果，感兴趣的朋友可以看文章细节。

ChatDrug 小分子编辑定性结果：

图 13：ChatDrug在小分子编辑的结果。任务是关于小分子单/多属性优化，具体描述在对应表格中。蓝色、红色和绿色分别表示编辑过程的初始状态、中间状态和最终状态。这里展示的中间状态的分子并不能满足我们的预期，而通过ChatDrug对话形式，我们能将其进一步优化。

ChatDrug 多肽编辑定性结果：

图 14：ChatDrug在多肽编辑的结果。任务是希望优化多肽对目标任务有更高的binding affinity。这里展示的是输入的多肽、优化的多肽和目标任务对应多肽的motif。这里motif是由PWM测量。

ChatDrug 蛋白质编辑定性结果：

图 15：ChatDrug在蛋白质二级结构编辑的结果。任务是希望蛋白质有更多的alpha或者beta结果。对应的改变部分都用蓝色圆圈进行了标记。

总结

我们通过前面两个系列从ML技术和domain解释两个角度介绍了分子的多模态任务。此外我们团队还有一些其他的探索，和前面两个系列的工作一起在下图展示。

图 16：我们团队关于分子多模态表示的探索。

整个AI for molecule/drug discovery的发展其实还是比较初步，但是已经吸引到学术界、工业界的大量关注。回首十年前深度学习刚开始随着GPU的发展影响学术圈，而十年后的现在深度学习在艺术领域已经取得了非常大的突破。我们团队也是对深度学习在科学领域的前景很有信心，但这两个领域之间目前也存在着非常大的gap：

• 对于domain科学家，往往把深度学习当做一个可以直接使用解决问题的一个黑盒模块，但是这个操作往往忽略了重要的优化过程。这个可以直接通过理解优化过程(目前有一些物理统计的尝试，但还是非常困难)，或者目前更加可操作的就是更严格的计算上的控制变量实验。
• 而人工智能领域的科研工作者往往考虑一些比较简单的评估方法和任务。这个就可以多去和domain专家沟通、理解domain上的问题。

这些挑战是交叉领域发展的自然过程，下面的一个方向就需要数学家、物理学家、统计学家、化学家、生物学家、工程师等各个领域/技能的科研学者进行更多深入的交流合作，构造起一个更加严格的pipeline。一个非常好的例子就是AlphaFold，它在兼顾对解决folding问题有突破性进展的同时，又有非常深厚的数学物理支撑，并且对solvable的问题敏锐程度非常精准。希望将来能有更多这样的工作出现。

[1] GraphMVP: Pre-training Molecular Graph Representation with 3D Geometry, ICLR 2022.

[2] GeoSSL: Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching, ICLR 2023.

[3] MoleculeSDE: A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining, ICML 2023.

[4] Geom3D: Symmetry-Informed Geometric Representation for Molecules, Proteins, and Crystalline Materials, arXiv 2023.

[5] GraphCG: Unsupervised Discovery of Steerable Factors in Graphs, NeurIPS Workshop 2022.

[6] MoleculeSTM: Multi-modal Molecule Structure-text Model for Text-based Editing and Retrieval, ArXiv 2022.

[7] ProteinDT: A Text-guided Protein Design Framework, arXiv 2023.

[8] ChatDrug: ChatGPT-powered Conversational Drug Editing Using Retrieval and Domain Feedback, arXiv 2023.

[9] N-Gram Graph: Simple Unsupervised Representation for Graphs, with Applications to Molecules, NeurIPS 2019

[10] AWARE: Attentive Walk-Aggregating Graph Neural Networks, TMLR 2022

[11] SGNN-EBM: Structured Multi-task Learning for Molecular Property Prediction, AISTATS 2022

[12] GIMLET: A Unified Graph-Text Model for Instruction-Based Molecule Zero-Shot Learning, arXiv 2023

[13] MolGraphEval: Evaluating Self-Supervised Learning for Molecular Graph Embeddings, arXiv 2022

[14] D3G: Leveraging Domain Relations for Domain Generalization, arXiv 2023

参考内容：https://zhuanlan.zhihu.com/p/639560302