图生成扩散模型综述:算法与在分子和蛋白质建模上应用

2023-06-30 18:29:00
刘大牛
转自文章
229

论文简要回顾了扩散模型在图数据上的算法及相关应用的若干研究。

图片

论文链接: https://arxiv.org/pdf/2302.02591.pdf

相关文献信息: https://github.com/ChengyiLIU-cs/Generative-Diffusion-Models-on-Graphs

图数据(Graph-based Data)可以保存现实世界实体(节点)之间丰富多样的关系信息,包括实体间的关联联系、属性特征、以及拓扑结构,已经在社交网络分析、推荐系统、生物信息学等领域有广泛的应用。 图生成模型旨在理解和学习现有的图数据分布,并合成新的图样本。这对于研究图数据中潜在的图结构关系,理解现有数据中的模式、关联和隐藏的信息具有重要的意义。生成模型可以用于探索图数据不同尺度的关系、发现社区结构、预测节点属性等。主要的图生成范式分为两类:自回归生成(Autoregressive Generation)和一次性生成(One-shot Generation)。自回归生成模型将图的生成过程建模为一个递归的条件概率分布,每个节点和边的生成都依赖于之前生成的节点和边;而一次性生成模型则直接生成完整的图样本。由于图数据的复杂特征,基于图的生成模型仍面临着3大挑战 :

1.离散性 (Discreteness):图的数据结构是离散的,由节点和边组成。这增加了图生成模型的训练和优化的复杂性(例如梯度的计算)。这使得广泛应用的优化算法难以直接应用于反向传播训练。

2.复杂的依赖关系 (Complex Intrinsic Dependencies):图数据并非符合独立同分布(independent and identically distributed, i.i.d.),其复杂的结构依赖信息为模型生成带来了巨大的挑战。尤其对于传统的自回归模型,因为生成是逐步进行的。

3.排列不变性 (Permutation Invariance):由于节点在图数据中是无序的,所以对于相同的图最多有种等价的表示方式。然而,传统的生成模型对节点的顺序是敏感的,不同的排列会导致不同的生成结果。这要求生成模型对节点排列的变化具有鲁棒性。

随着神经网络(Neural Networks)的不断发展,深度生成模型,例如变分自编码器(Variational Autoencoder, VAE),生成对抗网络(Generative Adversarial Networks, GAN),和流模型(Normalizing Flows),也被应用于图生成任务。尽管这些模型在一些应用中表现出色,但它们仍然存在一些局限性。例如,Graph-VAE通过训练图神经网络(Graph Neural Networks, GNNs)作为编码器和解码器来学习图形分布。编码器将输入数据映射到低维的连续空间, 而解码器通过采样重建样本。然而变分自编码器基于似然生成数据, 这要求进行大规模的图匹配 (Graph Matching)过程, 或者对每种可能的节点排序进行显式的估计以实现排列不变性。巨大的运算量增加了大规模图生成任务的困难。

Mol-GAN 希望通过对抗训练隐式的生成图样本, 其框架包括生成器和鉴别器。生成器网络通过学习现有图数据的分布特征, 尝试生成新的图样本。同时, 判别器网络与生成器进行对抗, 判断生成的图样本与真实图样本之间的差别。但是由于图数据的离散型, 生成对抗网络易于陷入模式崩溃-一只生成少数几个样本并且缺乏多样性。另外, 基于流的生成模型通过一系列可逆变换将简单的先验分布映射到目标数据分布。每个变换由一个可逆函数和其反函数组成, 确保了数据的双向转换。由于对特定架构的约束, 流模型很难完全学习图数据的结构信息。因此, 对于图的深度生成技术, 如何应用新颖的生成模型吸引力广泛的关注。

图片

图1 深度生成模型

扩散模型(Diffusion Model)是一种新兴的生成式模型,其首先向数据分布中逐步添加随机噪声到预设的先验分布,然后通过学习其逆过程来重建新的数据样本。自2019年第一个扩散模型范式被提出以来,其强大的生成能力引发了研究热潮。通常而言,扩散模型具有三种生成范式Score Matching with Langevin Dynamics (SMLD), Denoising Diffusion Probabilistic Model (DDPM)以及 Score-based Generative Model (SGM).

1.Score Matching with Langevin Dynamics [1]

首先,SMLD通过一系列递增的高斯噪声扰动原始的数据分布,然后在逆过程中学习数据分布的梯度, 通过最小化分数匹配(Score Matching)的损失函数,优化生成模型的参数,并使用朗之万动力学(Langevin Dynamics)采样生成样本。这种方法通过限制噪声分布在低维流形(low dimensional manifold )上的运动和向数据分布低密度区域添加大尺度噪声提高了分数匹配的准确性。

2.Denoising Diffusion Probabilistic Model [2]

DDPM构建了两个参数化的马尔可夫链(Markov Chain),用预定义的噪声在离散的时间步上扩散数据,并从噪声中重建所需的样本。基于马尔可夫链的特性和重参数化技巧,DDPM可以直接通过原始数据和预定义的噪声对任意的正向扩散步进行样本采样。而反向马尔可夫过程由用神经网络参数化的高斯分布组成。神经网络将被训练以优化负对数似然的变分上界 (ELBO)。DDPM通过学习到的反向链递归地生成样本数据。

3.Score-based Generative Model [3]

SGM使用随机微分方程(SDE)来建模扩散过程,把离散的扩散过程推广到连续的时间和空间。SGM利用标准维纳过程(Wiener Process)添加噪声到原始样本中得到扰动样本,然后通过计算扰动样本的分数梯度来估计样本的概率密度函数(PDF)以生成逼近真实数据分布的样本。

目前,扩散模型在图像生成任务上已经击败了最先进的GAN(对抗式生成模型),并被广泛的应用于计算机视觉、自然语言处理、图数据建模等领域。凭借坚实的理论的基础,扩散模型在图生成任务上取得了巨大的成功。这篇文章全面的回顾了三种扩散模型范式在图生成任务上的 算法和通用架构,并特别介绍了其在 分子建模(Molecule Modelling),蛋白质生成和结构预测(Protein Generation and Structure Prediction) 领域的应用。

一、基于图的生成扩散模型:方法

1.1 SMLD on Graphs

EDP-GNN是第一个基于分数匹配的无向图生成扩散算法,发表于AISTATS 2020 [4]。通过利用神经网络对添加不同尺度噪声的邻接矩阵扩散轨迹建模,EDP-GNN隐式的定义了图的排列不变性分布,以学习图数据分布的分数梯度特征。EDP-GNN采用了与SMLD相似的退火朗之万动态采样,并引入了多通道GNN层(Multi-channel GNN layer), 利用消息传递机制获取节点特征。通过训练包含噪声条件项的MLP输出层避免了在每个噪声尺度上单独训练分数网络(Score Network)。

图片

图2 三层EDP-GNN模型[4]

1.2 DDPM on Graphs

基于离散空间的扩散模型可以确保图数据在每个中间步骤保持离散,以更好的保存机构信息。将DDPM应用于图数据的重点在于设计适当的马尔可夫链过渡核(Transition Kernel)。Haefeli等人初次尝试了利用DDPM生成图的邻接矩阵[5]。该工作以one-hot形式对邻接矩阵进行编码,并利用双随机矩阵(Double Stochastic Matrix)加噪。实验表明离散时间步的噪声可以提高采样生成的效率。DiGress进一步将DDPM算法扩展到生成具有分类节点和边属性的图样本[6]。DiGress提出了一种新的马尔可夫噪声模型,在扩散过程中保持节点和边的边际分布,并在每个扩散步骤中添加从噪声图派生的辅助图论特征。利用交叉熵(Cross-entropy)损失函数,该模型将图的分布学习简化为一系列的分类任务,在分子和非分子的数据集上都取得了先进的效果,并展现了大规模图生成的潜力。

另外,E(3) 等变扩散模型(EDM)可以处理图的连续特征(原子坐标)和分类特征(原子类型),用于在三维空间中生成具有等变性的分子。由能量函数引导的等变随机微分方程(EEGSDE)对EDM模型进行了扩展。该框架利用了三维分子构象中的几何对称性,并通过线性组合相应的能量函数来生成具有多个目标属性的分子。虽然一般而言,扩散模型是针对一次生成范式设计,但GRAPHARM提出了一种自回归生成扩散模型,通过顺序预测邻接矩阵的每一行来生成图样本。

图片

图3 DiGress Overview[6]

1.3 SGM on Graphs

尽管EDP-GNN开发了基于得分的生成模型来得到图的邻接矩阵,但分数函数(Score Function)的估计取决于离散步骤中的噪声尺度,这限制了其生成大规模图的能力。GraphGDP使用连续时间步的扩散过程实现了满足排列不变性的图样本生成[7]。该工作利用位置增强的图分数网络(Position-enhanced Graph Score Network, PGSN)学习节点和边的信息以模拟逆时随机微分方程生成样本。GDSS提出了在连续时间步上的随机微分方程系统同时对节点和边的联合分布建模[8]。GDSS在扩散过程中直接向邻接矩阵和节点特征添加高斯噪声,并推导了相应的分数匹配目标以估计每个分量(节点和边)的联合对数密度梯度。GDSS还利用基于分数的马尔可夫链蒙特卡洛(MCMC)估计纠正分数估计的准确性。

值得一提的是,GDSS是第一个能够基于节点,边依赖关系生成整个图的扩散框架,并能够生成与训练分布接近且不违反化学价规则的分子,展示了SDE系统在建模节点-边关系方面的有效性。然而,标准的扩散过程会在几个步骤中消除稀疏图的特征,这可能会导致反向扩散过程中的分数估计信息不足。为了解决这个问题,GSDM对邻接矩阵的特征值矩阵(谱分解)进行扩散操作,运用低秩的高斯噪声扰动数据。这种方法在降低了运算量的同时提高了生成质量。此外,另一份工作SGGM将变分自编码器和扩散模型相融合,把高维的离散的数据映射到低维的潜在空间进行图生成建模。

图片

图4 GDSS:基于分数的图生成框架[8]

二、基于图的生成扩散模型:应用

2.1 分子建模

分子的结构可以自然地用图数据进行表示,例如包含原子和化学键的分子图。这使得图学习技术在分析和研究分子结构方面具有重要的应用价值,并可用于各种下游任务,如药物研发、化学计算、材料科学和生物信息学等。通过对分子图进行建模,我们可以生成具有特定性质的新分子。具体而言,分子建模可以进一步分为分子构象生成和分子对接两个任务。

分子构象生成(Molecule Conformation Generation)

分子的三维构象对其生物学和物理学特性具有显著影响。在对分子构象进行建模时,需要考虑其旋转-平移不变性。GeoDiff利用非平衡热力学模拟扩散过程,使生成的分子逐步靠近目标构象[9]。此外,引入其他领域的知识也可以增强分子表示学习。EEGSDE引入了能量函数(Energy Function)来引导3D分子构象生成, MDM考虑了原子间距对原子作用力的影响,而DiffBridges则为分子生成设计了一个包含物理信息和统计先验的能量函数[10]。

分子对接(Molecular Docking)

分子对接是一项计算任务,用于预测分子与其他分子(通常是蛋白质)结合时的最佳方向。这在药物研发中非常重要,可以帮助找到最适合的小分子进入目标蛋白质的活性位点。TargetDiff结合了扩散模型和SE(3)等变网络学习原子类型和坐标特征,以生成具有所需几何特性的蛋白质靶分子[11]。DiffLink则以片段为基础,利用E(3)等变去噪扩散模型生成给定分子片段的剩余部分。它可以根据原子的位置选择性的生成可以适应蛋白质口袋的分子结构[12]。DiffDock使用扩散模型将对接位置预测问题(Docking Pose Prediction Problem)转化为生成问题,并通过分别以配体和蛋白质作为输入,随机选择初始状态并对其进行排名来执行逆向扩散过程。

图片

图5 分子与蛋白质的扩散建模过程

2.2 蛋白质建模

蛋白质建模旨在生成和预测蛋白质的结构。这项任务有助于理解蛋白质的功能和相互作用,并广泛应用于药物研发和设计具有特定特征的新蛋白质等领域。

蛋白质生成(Protein Generation)

计算蛋白质设计(Computational Protein Design)的目标是自动生成具有特定结构和功能特性的蛋白质。SiamDiff通过在天然蛋白质结构和序列上预训练扩散模型以更好地学习目标蛋白质的分布特征。而ProteinSGM采用条件生成的方法,通过将合理的骨架和功能位点融合到预定长度的结构中生成蛋白质。SMCDiff使用粒子滤波算法对具有特定结构的稳定蛋白质骨架进行条件采样,从而生成稳定的蛋白质骨架结构[13]。在免疫系统蛋白质的生成领域,DiffAntigen基于抗体的框架区域和目标抗原,联合生成抗体的CDR序列和结构[14]。

蛋白质-配体复合物结构预测(Protein-ligand Complex Structure Prediction)

蛋白质-配体复合物的普遍存在使得预测它们的三维结构对于生成新的酶和药物化合物非常有价值。NeuralPLexer通过将生物分子复合物中的多尺度诱导偏差(multi-scale induced bias in biomolecular complexes)与扩散模型相结合,来预测蛋白质-配体复合物的结构。它以分子图作为配体输入,并利用学习到的统计分布生成3D结构。DiffEE 提出了一个基于预训练的蛋白质端到端扩散生成模型。它能够生成具有正确结合位置的多种蛋白质-配体复合物的结构。

三、机遇与挑战

扩散模型在图领域的应用仍有很多值得研究的方向。

3.1 图数据的离散性

如前所述,图数据的离散性使得研究人员很难直接将扩散模型应用于图数据的分析和生成。在这种情况下,一些研究试图通过引入离散概率分布或将离散的数据映射到连续的潜在空间的方法使扩散模型适应离散的图数据。然而,目前还缺乏一种通用且广泛认可的方法来解决这个问题。

3.2 图扩散模型的条件生成

生成符合一定条件的图样本是至关重要的。在生物信息学中,为了生成具有特定属性的分子和蛋白质,对图生成模型的设计设置一定的约束非常关键。因此,将额外的信息作为条件引入图扩散模型已成为一项迫切的研究方向。这些信息包括知识图谱、视觉和文本信息等。

3.3 图扩散模型的可信度

图生成任务可能会在各种实际任务中对用户造成意外的伤害,特别是在药物发现等安全关键领域。例如,基于数据驱动的图扩散模型容易受到恶意攻击者的对抗性攻击。这成为一个重要的安全问题。此外,由于图扩散模型的复杂性,理解和解释图生成的工作机制变得非常困难,这限制了其在实际应用中的可信度。

实现图扩散模型的可信度需要考虑多个关键因素。首先是安全性和稳健性,确保模型对于恶意攻击具有较强的抵抗力。其次是可解释性,即使图扩散模型通常很复杂,也需要能够解释其生成过程和结果。公平性也是一个重要的维度,确保模型在生成图时不产生偏见或不公平的结果。最后是隐私性,保护用户数据和敏感信息的安全和隐私。

3.4 图样本评估

现有的图扩散技术大多用于分子和蛋白质的生成,而在图上的许多应用很少被探索,例如推荐系统,图数据异常检测,因果图生成等等。

参考文献

[1] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. NeurIPS, 2019.

[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NeurIPS, 2020.

[3] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021.

[4] Chenhao Niu, Yang Song, Jiaming Song, Shengjia Zhao, Aditya Grover, and Stefano Ermon. Permutation invariant graph generation via score-based generative modelling. In AISTATS, 2020.

[5] Kilian Konstantin Haefeli, Karolis Martinkus, Nathana ̈el Perraudin, and Roger Wattenhofer. Diffusion models for graphs benefit from discrete state spaces. In LoG, 2022.

[6] Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graphgeneration. In ICLR, 2023

[7] Han Huang, Leilei Sun, Bowen Du, Yanjie Fu, and Weifeng Lv. Graphgdp: Generative diffusion processes for permutation invariant graph generation. In IEEE ICDM, pages 201–210, 2022.

[8] Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In ICML, 2022.

[9] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. InICLR, 2022

[10] Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and qiang liu. Diffusion-based molecule generation with informative prior bridges. In NeurIPS, 2022.

[11] Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equivariant diffusion for target-aware molecule generation and affinity prediction. In ICLR, 2023

[12] Ilia Igashov, Hannes St ̈ark, Clement Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael M Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design. In NeurIPS, 2022

[13] Brian L Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. In ICLR, 2023

[14] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen-specific antibody design and optimization with diffusion-based generative models for protein structures. In NeurIPS, 2022

入门
发表评论
评论通过审核后显示。
文章分类
联系我们
联系人: 透明七彩巨人
Email: weok168@gmail.com