分子生成领域的stable diffusion - GEOLDM

一、关于stable diffusion

很多人都知道stable diffusion，stable diffusion的出现改变了机器生成领域，让AI技术第一次无比的接近正常人。大语言模型，AIGC概念于是兴起。基于stable diffusion 大家开发了lora， hyperwork等微调技术，使得我们可以精确的控制生成图片，于是我们可以生成如下一些图片，以假乱真：

​​​​​​​

神奇的stable diffusion模型的架构如下：

stable diffusion模型架构简单理解：首先使用一个编码器，将图片编码到隐式空间的潜向量，然后基于该潜向量实现diffusion model的增噪和去噪过程，生成新的潜向量，然后再使用解码器，解码潜向量，生成新图片。关于stable diffusion的介绍，全网到处都是，感兴趣的可以仔细研究研究。

那么问题来了，在分子生成领域，在药物设计领域，是不是可以有一个stable diffusion 模型？实现真正的可控分子生成？让药物设计、分子设计领域的大语言模型真正的兴起，实现Drug-AIGC？要实现这一点的难点在于：药物分子是具有空间几何特征的（即旋转不变性和空间平移不变性），且不连续稀疏的。而图片更简单，是平面的且连续的。因此，基机器视觉开发的stable diffusion是不能直接使用于分子设计、药物设计等领域的，需要专门开发一个能处理空间几何与稀疏不连续的stable diffusion模型架构。

新发表的GEOLDM 模型实现了这一点。也许从此刻开始，变天了。Drug-AIGC真的要来了。分子设计领域有了stable diffusion模型，那么分子设计领域的lora还远吗？

二、分子生成领域的stable diffusion - GeoLDM

GEOLDM 模型是发表在2023年 ICML（全称：International Conference on Machine Learning）会议上的会刊 文章。第一作者是斯坦福大学计算机系的Minkai Xu。文章链接为：https://arxiv.org/pdf/2305.01140.pdf

GEOLDM模型可以生成3D的分子，在有条件或者无条件下都可以生成。

GEOLDM的全称是Geometric Latent Diffusion Models，几何隐式扩散模型。GEOLDM 是分子几何领域的第一个隐式扩散模型（简称：DM），由将结构编码为连续隐式向量的自动编码器和在隐式空间中运行的扩散模型组成。 文章的关键创新在于，对 3D 分子几何进行建模，通过构建具有不变标量和等变张量的点结构隐式空间来捕获其关键的旋转平移等变约束。

在GEOLDM模型中，分子中原子坐标和特征，即x,h会，被编码器编码成隐式向量，即z_x和z_h。然后在隐式空间内，对z_x和z_h进行T步的加噪过程，使用深度学习网络预测噪音，反向去噪，生成隐式向量。解码器根据去噪后的隐式向量生成分子。

注意，GEOLDM的扩散过程是操作在隐式向量空间的。这一点与之前的分子生成模型，EDM，Difflinker，Targetdiff等不同。之前的分子生成模型的扩散过程是直接操作在原子空间的，即添加噪音在坐标x和原子特征h上。而GEOLDM是将扩散过程操作在坐标x和原子特征h，生成的嵌入隐空间z上。这一点完全与机器视觉中的stable diffusion对齐了。

训练GEOLDM模型的伪代码：

首先，需要按照常规的变分自动编码器一样，在没有加噪和去噪的扩散过程下，先训练变分自动编码器，即First stage。然后，固定自动变分编码器，加入扩散的加噪和去噪过程，训练模型学习去噪。所以，可以简单的理解GEOLDM模型由编码器，扩散去噪，解码器，三个AI网络组成。

使用GEOLDM模型生成分子的伪代码如下：

三、GeoLDM模型的性能

作者在QM9和DRUG数据上进行了测试，考察了分子生成的质量（有效性，唯一性）。作者对每个模型都进行了三次分子生成，每次生成1w个分子。测试结果如下：

与之前的方法相比，同时具有不变变量和等变变量的潜在空间使GEOLDM在大分子生成的有效性方面实现了高达7%的改进。

stable diffusion模型厉害之处，主要是其非常强大的可改造能力，最简单的，就是条件控制生成。这里作者做了尝试。作者尝试让GEOLDM模型生成HOMO等更低的分子，结果发现，GEOLDM确实要好于我们之前常用的EDM模型。如下图：

另一方面，在条件控制生成中，GEOLDM，可以对生成的条件的程度进行控制。作者以生成极化的分子为例，使用了不同的极化控制参数，如下图：

四、总结

整体感觉GEOLDM这个模型非常好模拟了stable diffusion模型的结构，将扩散过程的加噪去噪过程操作在隐式向量空间内。大大降低了，在分子生成领域的计算成本。同时，因为在潜向量空间内，也打开了分子生成领域的优化可能，我们可以CLIP， LORA， Controlnet 等技术，微调模型，实现精确控制分子生成。

看来Drug-AIGC的日子不远了，真的不远了。距离作者正式发表GEOLDM有几个月了，距离作者开始在GitHub上传代码有10个月了。我们将很快看到很多的相关文章出现。老了，脑子跟不上了。