今天给大家分享的是来自加拿大蒙特利尔学习算法研究院 (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

• 
  

图 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）这个平台上。