清华大学智能产业研究院（AIR）马维英教授、刘菁菁教授及周浩副教授课题组与上海药物所的研究团队推出 BFN 系列工作，在先前的 Diffusion、Flow Matching 等生成模型框架之外提出了首个统一分子模态的全新药物设计模型，在分子结构建模方面取得了突破性进展，在无条件和给定靶点分子的生成任务上均表现出卓越的优势。
GeoBFN：课题组设计了基于参数空间的几何分子图生成框架 GeoBFN，用于无条件分子生成。课题组首次使用统一的方式对分子图数据的不同模态变量进行建模，其分子数据分布拟合结果接近理论极限。这一工作为 AI4Science 中的生成任务奠定了新的模型基础，论文获评 ICLR 2024 Oral（TOP 1.2%）。
•论文标题：Unified Generative Modeling of 3D Molecules via Bayesian Flow Networks
•GitHub 开源地址：https://github.com/AlgoMole/GeoBFN
•论文链接：https://arxiv.org/abs/2403.15441
MolCRAFT：课题组将上述成果延展到给定靶点分子生成这一重要应用领域，详细分析了过往生成式建模中存在的问题，并提出基于参数空间的贝叶斯流网络 MolCRAFT 作为解决方案。在给定蛋白质口袋结构的 3D 小分子药物设计任务上，生成构象的亲和力首次达到药物分子水平，在结构合理性（能量、RMSD）等各个维度均超越现有 SBDD 模型，且生成效率提升数十倍以上，论文获得 ICML 2024 接收。
•论文标题：MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space
•GitHub 开源地址：
https://github.com/AlgoMole/MolCRAFT
•论文链接：
https://arxiv.org/abs/2404.12141v4
•Demo 网站：
http://120.240.170.153:10990/
在 AI4Science 领域，针对蛋白质、小分子等分子图数据的建模仍然是一个正在被积极探索的领域。不同于连续的图像数据、离散的文本数据，分子图同时包含了连续原子坐标与离散原子类型等多个模态的变量。对分子图数据的有效表征与合理建模，是一切后续任务的基础。斯坦福大学 Stefano Ermon 组提出了 GeoLDM，即适用于小分子生成的隐空间扩散模型（latent diffusion），然而如何获得合理的小分子隐变量表示仍然是悬而未决的问题。
分子图数据建模难点
分子图数据（Molecular data） 的建模难点主要在于：（1）多模态（Multi-modality），分子点云包含了连续坐标与离散类型等变量；（2）噪声敏感（Noise sensitivity），对坐标或类型的扰动容易带来分子结构较大的信号损失, 因此难以构造信息逐步减少的训练轨迹，因此影响了以往方法如扩散模型（diffusion models）的训练。这一点在分子和图片的对比中更为直观。具体而言，扩散过程在对图像进行加噪的过程中，带噪图像对人而言仍然是可识别的，然而带噪分子由于结构迅速遭到破坏，其对应的属性能量信息难以定义。因此作用在噪声样本空间的扩散模型和流匹配模型（Flow Matching）等对于分子生成并不是最优的选择。
贝叶斯流网络
课题组使用贝叶斯流网络（Bayesian Flow Network, BFN），一个通过在连续可微的参数空间（parameter space）建模统一了连续与离散变量的全新生成模型，来应对 molecular data 的建模挑战。
BFN 区分了传送者（sender）和接收者（receiver）两个角色，生成过程被描述为二者之间的信息交换。Sender 逐步发送带噪分
子
，sender distribution 信噪比逐渐上升。Receiver 在收到噪声分子信息之后，将会通过各维度上独立的贝叶斯更新（Bayesian update）修正自身的 belief
，
得到一个关于分子的后验分布（posterior distribution），即参数分布。注意对参数分布而言，连续变量的参数是高斯分布的均值与方差，离散变量则是 Categorical 分布的概率密度，均为连续实值。
每个中间步在输出分子时，receiver 使用参数化的神经网络去建模不同变量间的相互关系，输入参数分子得到最终的分子输出分布（output distribution）。
课题组首次推导得出满足平移旋转等变性的 BFN 框架，探索了连续原子坐标、离散原子类型与离散化原子电荷量等不同分子特征，统一了各个模态的分子概率建模。
实验结果
在无条件分子生成任务上，QM9、GEOM-DRUG 等数据集的评估显示，GeoBFN 在原子 / 分子稳定性等多个分布指标上均超过之前最好结果，逼近数据集理论极限，由于参数空间方差极小的特点，GeoBFN 能以显著减少的采样步数达到很好的效果。
GeoBFN 的成功显示出在分子图数据生成式建模上的巨大潜力，为 AI4Science 中的生成任务开创了新的模型基础。
基于结构的药物设计（Structure-based Drug Design, SBDD）模型致力于从头（de novo）生成至少满足三方面要求的小分子：（1）与靶点亲和力高，（2）自身类药性好、易于合成，（3）分子构象合理。其中（3）是基于结构的生成模型在生成分子时首先需要满足的前提条件。
然而，课题组分析发现，现有的 SBDD 模型在（3）构象合理性上面临着较大挑战，常常生成扭曲的构想与不合理的结合模式，依赖于分子对接（docking）进行后处理，即使生成的分子亲和力高，这一结果仍然是可怀疑的。课题组通过深入探索问题成因，认为 autoregressive 生成中不合理的顺序要求可能导致模式坍塌（Mode Collapse），以及不一致的连续-离散空间（Hybrid Continuous-Discrete Space）使得 diffusion 过程中样本极易落在 manifold 之外。
课题组探索了 BFN 作为可行的解决方案。作为 non-autoregressive generation，BFN 能够捕捉到多样化的分子结构，同时，在连续的参数空间建模连续-离散变量的联合分布又有助于克服 multi-modality inconsistency。
进一步地，课题组设计了一个参数空间内的采样策略，避免了原始方法中不断从 parameter space 到 sample space 的采样，进一步降低了 variance，使得生成过程中向目标分布的变换更加平滑，并且能够以明显更少的采样步数达到更好的效果。
实验结果
在给定靶点分子生成任务上，从结合能力、构象稳定性、类药性质等方面评估显示，MolCRAFT 在多个指标上均达到 SOTA 水平，进一步说明了连续参数空间建模分子的优势。值得一提的是，在使用 docking 软件进行最优结合构象搜索并打分时，MolCRAFT 生成的构象与 redocked 构象差距最小，RMSD < 2 埃的比例最高，接近 CrossDocked2020 数据集理论极限（～50%），在分子构象稳定性（strain energy, SE）上也是最接近测试集水平的模型。这表明 MolCRAFT 成功学到了蛋白质-小分子结合的相互作用机制。
得益于课题组设计的新颖采样策略，MolCRAFT 在大幅提升生成样本质量的同时，显著降低了采样时间至其他方法的几十甚至几百分之一，为 AI 辅助药物设计增添了新的可能。
目前，基于 MolCRAFT 的 demo 正在试运行中，欢迎大家访问 http://120.240.170.153:10990，后续正式版将在 http://gensi-thuair.com:10990 发布。课题组正在开发基于 MolCRAFT 的药物设计框架，希望在未来真正能够助推临床管线，为 AI for Drug Design 的实际应用添砖加瓦。
