今天给大家分享的是来自加拿大蒙特利尔学习算法研究院 (Mila) 唐建教授团队刘圣超博士的研究工作总结。刘圣超专注于迁移学习(多任务、预训练),多模态学习和统计生成学习研究。
在计算生物、化学、材料领域,伴随着深度学习(DL)的广泛使用,分子的表征(representation)已经成为了最基础的研究问题。一个鲁棒的分子表征,能够支撑丰富的任务。我们团队从研究小分子表征开始研究,是因为它包含了非常丰富的多模态信息,如下图展示的六个模态。
这六种模态又可以分为两大类:3D geometry、2D topology、1D String 和1D Fingerprint 都是关于小分子的“内部”化学结构表示;而 bio KG 和 textual description 则更多的是关于分子的“外部”功能描述。这两大类的模态(即内部和外部模态)是可以相互补充信息。
此外还需要强调的是,对于小分子的内部模态,2D topology 和 1D String、1D Fingerprint 在信息层面等价,主要区别是关于分子图的不同数据结构以及对应的 DL 表征。但是 3D geometry和 2D topology 在信息层面是很大的区别,并且它们的信息能够互相补充。
围绕以上两点,我们将简单介绍围绕这两种信息互享,介绍两个系列工作:
GraphMVP: Pre-training Molecular Graph Representation with 3D Geometry, ICLR 2022. GeoSSL: Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching, ICLR 2023. MoleculeSDE: A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining, ICML 2023. Geom3D: Symmetry-Informed Geometric Representation for Molecules, Proteins, and Crystalline Materials, arXiv 2023.
GraphCG: Unsupervised Discovery of Steerable Factors in Graphs, NeurIPS Workshop 2022. MoleculeSTM: Multi-modal Molecule Structure-text Model for Text-based Editing and Retrieval, arXiv 2022. ProteinDT: A Text-guided Protein Design Framework, arXiv 2023. ChatDrug: ChatGPT-powered Conversational Drug Editing Using Retrieval and Domain Feedback, arXiv 2023.
1 Geom3D 分子的几何表征 benchmark
该研究以《Symmetry-Informed Geometric Representation for Molecules, Proteins, and Crystalline Materials》为题,发表在 arXiv上。
从数据结构上,分子的本质(稳定态)是一个结构稳定的 3D 点云。他的表征主要挑战就是如何保证对于旋转、平移等变。现有的工作都是利用了群表示论来进行描述,而我们在 Geom3D 这个平台中对于现有的工作进行了进一步的整理和总结。现有的 geometric modeling 可以大致分为如下三类:
Equivariant modeling with spherical frame basis 是把相对位置的vector 投影到 spherical harmonics frame 上进行的 modeling。它的优势是一种非常泛化的框架,比如可以 model 更高阶的 particle feature (物理问题中使用更多);而缺点则是需要进行 tensor product,这个计算复杂度非常高。 Equivariant modeling with vector frame basis 则更加针对 3D coordinates:它是把 coordinates 投影到了一个 vector frame 上,然后再对其进行 modeling。它的优点是效率比较高,而缺点就是无法对高阶的粒子进行建模。
16 个 geometric modeling 模型 14 个 geometric pretraining 模型 46 个不同 geometric tasks,包含了小分子、蛋白质、和材料 此外 Geom3D 还包含了 7 个 1D model 和 11 个 2D GNN model。关于2D topology pretraining,我们团队也有一个前序工作,MolGraphEval(https://arxiv.org/abs/2206.08005)。感兴趣的朋友欢迎查阅。
下面我们就重点介绍在单一模态和多模态情况下,如何进行 geometric pretraining。
2 GraphMVP & MoleculeSDE: 2D-3D Pretraining
2.1 GraphMVP的结构化数据预训练框架:从最大化互信息到条件概率求和
该研究以《Pre-training Molecular Graph Representation with 3D Geometry》为题,发表在 ICLR 2022 上。
对于小分子的多模态,我们最先考虑到的就是 2D topology 和 3D geometry。同时对于 2D 和 3D 进行预训练,而预训练的思路非常简单,就是最大化2D topology 和 3D geometry 之间的互信息(MI)。但这里有一个难点就是如何对于结构化数据进行 MI 计算,而 GraphMVP 最大的贡献是提出来一个关于最大化 MI 的一个等价形式:
这就把 MI maximization 问题变成了 summation of two conditional log likelihoods (公式1);而且他在这里还有更具体的意义:2D 生成 3D 的条件概率 + 3D 生成 2D 的条件概率。这个形式对于结构化的数据非常友好,因为可以引入 EBM 这种兼具泛化性、灵活性和强表达能力的概率模型进行估计。
这里我们将结合下面几个工作进行一些延展补充。注:GraphMVP 的推导是基于离散情况,GeoSSL 还提供了连续版本的推导,且二者最终优化目标一样。感兴趣的朋友可以参考比对两篇文章的附录。
GraphMVP 利用 NCE 求解,我们叫做 EBM-NCE。我们发现 EBM-NCE 和 Jensen-Shannon divergence 联系密切。二者的目标函数一样,只是求解的过程和思路不同。而 EBM-NCE 和其他contrastive self-supervised learning 的思路本质都一样:通过构造 positive 和 negative pairs,然后增大 positive pair 的 similarity,并且见效 negative pair 的 similarity。 此外 EBM 还有其他的求解思路,比如 score matching (SM),我们当时在做 GraphMVP 的时候已经意识到了它也是 do-able 的路,但是第一篇工作来不及详细展开。这个也是指导我们后面做 GeoSSL 和 MoleculeSDE 的方法论。
GraphMVP 就首先采用了 VAE 的形式,提出了 variation representation reconstruction (VRR)。VRR 是在 representation space 进行 reconstruction (而不是 data space),从而有了对 ELBO 的估计。并且我们发现,non-contrastive self-supervised learning (比如 BYOL、SimSiam) 就是VRR 的一种特殊情况。 DDPM 也是在优化 ELBO,并且它和 denoising score matching (DSM) 是非常类似的。它们的区别可以通过 Stochastic Differential Equation (SDE)一个统一框架下的两种变形体现。这一点我们在 MoleculeSDE 中进行了更加详细的解释。
第一类是 NCE (包含了 EBM-NCE、InfoNCE、GAN),因为它们本质思路就是把概率估计问题转换为了分类问题,也就是 contrastive learning。它是基于 data pair 来进行 distribution estimation。 第二类是类似 DSM、VRR,目标任务是为了重构某一个 data 或者data 的 representation。本质是把概率估计问题转换为了重构问题,也就是 generative learning 或者 reconstruction learning。它是直接基于每一个单独的 data point 来进行 distribution estimation。 2.2 基于latent space的GraphMVP
![图片](https://image.jiqizhixin.com/uploads/editor/cadcb9eb-86cc-458b-852d-43ac9e319706/640.png)
图 3:GraphMVP的流程图。
当我们有了公式 1 引出的一系列求解思路之后就非常直接。GraphMVP 是完全基于 latent space,使用了如下两个目标函数:一个 contrastive loss,也就是 EBM-NCE;一个 generative loss,也就是 VRR。除此以外,GraphMVP 还有两个 variant,GraphMVP-C 和 GraphMVP-G,分别考虑到了如何加入contrastive 和 generative 2D SSL。
2.3 基于 data space 的 MoleculeSDE
该研究以《A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining》为题,发表在 ICML 2023 上。
MoleculeSDE 是 GraphMVP 的 follow-up 工作。在 GraphMVP 中的VRR,我们是利用了 VRR 来对 ELBO 进行估计,但是这种估计会造成信息损失。这里我们提出了更加严格的概率估计,也就是直接在 data space (geometry 和 topology) 进行重构。
但这里又有一个挑战,就是从 2D 到 3D 的条件概率 (也就是 2D 生成 3D)需要遵照 SE(3)等变,也就是对于旋转、平移等变,并且对于对称反对称。为了实现这个目标,我们基于 vector frame basis 引入了 SE(3)-equivariant and reflection-antisymmetric SDE。这个思路本质上是利用 score matching 或者diffusion 去求解公式 1。
另外我们想强调的是,对于 downstream task,除了常规的 2D 和 3D proeprty prediction,MoleculeSDE 也让我们有了更多样化的选择。主要是对于小分子 2D 到 3D 的生成,这个 task 的意义并不仅仅在于能够有 conformation generation,而是能够有基于生成的 3D coordinates 进一步进行modeling,如下图 5(3)所示:
3 GeoSSL: 3D pretraining
该研究以《Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching》为题,发表在 ICLR 2023 上。
GraphMVP 和 MoleculeSDE 都是考虑模态之间的预训练。同时我们还想强调仅仅考虑 3D geometry 的预训练。这个工作的推出是在 GraphMVP 和MoleculeSDE 这两个工作之间,而原因则是因为预训练的数据集。GraphMVP 的预训练是在 GEOM 数据上(当时取了 250K data,是已知较大的小分子 3D 数据集)。而 2021 年暑假开始,陆续有几个比较大的数据集相继推出,比如 Molecule3D 和 PCQM4Mv2。GeoSSL 就是在 Molecule3D 上进行预训练。
GeoSSL 是仅仅考虑到 geometry 的 single-modality pretraining。我们首先需要定义 view。这里的出发点是在计算或者模拟中,分子的 3D coordinates 有一定的误差,并且分子的 3D geometry 哪怕是稳定态(势能面),也会在一个小区域内进行运动。由此,我们定义了两个 view:original geometry 和 perturbed geometry,如下图所示。
基于这两个 view,我们又可以利用公式 1 引出的一系列方法来最大化 MI。已有的方法 (比如 EBM-NCE、InfoNCE、RR)都在 GeoSSL 中进行了 benchmark。此外我们还利用了 geometry 数据的特性,也就是连续的 3D coordinates,提出了利用 denoising score matching 的方法,来进行 denoising distance matching,从而帮助了参数估计,具体推导过程可以看论文。大概流程则如下图所示:
小结
分子 geometry representation learning 本身已经是一件挑战性同时非常重要的任务,因为 geometry 是这些物理粒子最本质的并且很复杂的数据结构;而geometry预训练任务的复杂性更加大一些。我们团队的工作一直在探索相关问题。从最开始的 benchmark,一直到预训练,所有的代码除了每一个工作各自的 github repo 开源,也已经整合到了 Geom3D (https://github.com/chao1224/Geom3D)这个平台上。
关于作者:刘圣超博士现在是 Mila 的第四年博士,将于 2023 年秋季加入 UC Berkeley 和 Caltech 联合博后,导师是 Prof. Jennifer Chayes, Prof. Christian Borgs, 和 Prof. Anima Anandkumar。(个人主页https://chao1224.github.io)
参考内容:https://zhuanlan.zhihu.com/p/639560302