MolFlow: 高效3D分子生成方法

DRUGAI

今天为大家介绍的是来自查尔姆斯理工大学的Simon Olsson团队的一篇论文。最近，3D药物设计的生成模型因其在蛋白质口袋中直接设计配体的潜力而获得了广泛关注。然而，目前的方法通常存在采样时间非常慢或生成分子的化学有效性差的问题。为了解决这些限制，作者提出了Semla，一个可扩展的E(3)-等变消息传递架构。作者进一步介绍了一个分子生成模型MolFlow，该模型使用流匹配和尺度最优传输进行训练，这是等变最优传输的一种新扩展。作者的模型在基准数据集上仅需100个采样步骤就能产生最先进的结果。关键是，MolFlow在不牺牲性能下只需20个步骤就能采样出高质量分子，相比于现有技术实现了两个数量级的速度提升。最后，作者比较了MolFlow与当前方法在生成高质量样本方面的能力，进一步展示了其强大性能。

3D药物设计的生成模型最近因其在蛋白质口袋中直接设计配体的潜力而引起了极大关注。尽管一些新提出的模型试图在蛋白质口袋内直接生成配体，但更深入的分析显示，许多这些模型生成的配体具有不现实的结合构象。其他模型尝试训练无条件的3D分子生成器作为起点，尤其是那些应用扩散模型到分子坐标上的模型已经变得特别流行。然而，这些模型也存在显著的实际限制；即它们几乎都需要数百甚至数千次前向传播才能生成分子，这使得它们在大多数下游应用中变得不切实际。当应用于药物类分子数据集时，许多模型还会生成化学不现实或质量较差的样本。

对于将分子表示为字符串或2D图的分子生成器，针对特定蛋白质口袋的微调已被证明非常有效，并且在该领域已成为标准做法。通常，这些模型使用强化学习（RL）引导进入优化的化学空间。这种方法虽然非常有效，但要求能够非常快速地采样高质量分子。现有的3D分子生成器，使用完全连接的消息传递，表现出非常差的扩展性，无法处理更大的分子和更大的模型尺寸。最先进的无条件生成器在采样单批分子时需要数分钟，使其不适合基于RL的微调。

在这项工作中，作者从两个方向解决这个问题。首先，作者引入了一种新型的等变架构用于3D分子生成，该架构比现有方法具有显著更好的效率和可扩展性。即使MolFlow的模型参数数量是现有最先进模型的三倍，它处理一批分子的速度仍是现有模型的两倍多。其次，作者引入了一种新型的最优传输方法，称为尺度最优传输，它扩展了等变最优传输以考虑分子的不同尺寸。最后，作者使用该架构并结合尺度最优传输，通过流匹配训练了一个无条件分子生成器MolFlow。

图 1

图1概述了作者使用尺度最优传输的分子生成方法。

模型部分

流匹配

流匹配旨在学习一个生成过程，该过程将样本从噪声分布pnoise转移到数据分布pdata的样本。条件流匹配(CFM)作为训练流匹配模型的一种有效方法以不同形式出现，它在噪声和数据样本x1 ∼pdata(x1 )之间插值这种无模拟的方式下有效。为此，定义了一个时间依赖的条件流pt∣1 (·∣x1 )，从中可以导出条件向量场ut (xt∣x1 ) 。训练模型vθt (xt )来回归向量场，而其他公式则训练模型以估计分布pθ1∣t(·∣xt)，该分布从噪声数据中重建干净数据。然后可以使用期望构造向量场：

然后通过使用任意ODE求解器积分向量场生成样本。

Semla架构

图 2

如图2所示，作者提出了Semla——一种可扩展的等变模型，它使用多头潜在图注意力机制，消息传递在压缩的潜在表示上进行。这种扩展能够在不显著增加计算成本的情况下扩展节点特征的维度和可学习的模型参数数量。

与之前的方法类似，Semla使用E(3)不变和等变特征。强制实施群对称性提供了强大的归纳偏差，并改善了样本复杂性。然而，与之前的分子生成模型不同，Semla并不区分分子坐标和等变特征向量，而是将它们视为单一的可学习表示（称之为坐标集）。坐标集类似于其他等变架构中的3D几何等变向量集。作者认为这种表示相对于之前的分子生成方法有两个主要优势。首先，可学习的坐标集提供了比仅存储每个分子一组坐标的模型更强的表达能力。其次，等变特征的联合表示允许更简单的更新机制——只需应用线性投影（无偏差）来创建和更新坐标集，同时能够保持等变性。

为了确保稳定训练，作者在整个模型中使用了归一化层。LayerNorm应用于不变特征，对于等变特征，作者采用MiDi的归一化方案以允许多个坐标集。作者假设这允许坐标集学习不同长度尺度的等变特征，从而帮助克服不同大小的分子被归一化为相同平均向量范数的问题。

尺度最优传输

对于分子分布，原子坐标与分子中心之间的平均长度随着原子数量N增加。在聚合物理论中，已证明自排斥聚合物链中的原子坐标标准偏差随√N缩放，这被称为Flory半径。然而，分子生成模型通常从单位高斯分布中采样坐标噪声，这并不能反映分子大小的差异。与之相反，作者从方差依赖于N的高斯分布中采样，其中N是x1中的原子数量。在训练和推理过程中，作者从噪声分布pπnoise中生成样本xπ0，如下所示：

其中，fπ(x0, x1)是等变最优传输变换，对应于应用一个排列和旋转以最小化x0和x1之间的传输成本（在这种情况下为均方误差）。由于现有的3D生成模型在推理时已经需要知道N，因此使用尺度OT时对采样没有额外的限制。

由于它适用于长聚合物链，作者将Flory半径视为小分子药物的上限，并根据一组类药物分子数据集经验性地发现σN=k log(N)，其中k=0.2是一个合适的值。

实验结果

数据集

作者使用了两个基准数据集QM9和GEOM Drugs来评估MolFlow作为无条件分子生成器的能力。QM9是一个包含多达9个重原子的分子集合，这些分子满足一组预定义的约束。QM9中的分子平均包含大约18个原子。另一方面，GEOM Drugs是一个包含较大药物样本的分子数据集，其平均大小约为44个原子。对于这两个数据集，作者使用了与MiDi和EQGAT-diff相同的数据划分。然而，为了提高训练时间，作者从GEOM Drugs训练集中删除了超过72个原子的分子。这相当于大约1%的训练数据。验证和测试集保持不变。以下所有关于MolFlow的指标均通过首先从测试集中采样分子大小，然后使用模型生成具有采样原子数量的分子来计算。

基线

作者将MolFlow与一些最近提出的3D分子生成模型进行了比较，包括推断基于坐标的键和直接生成键的模型。EDM、GCDM、MUDiff和GFMDiff都是基于扩散的方法，通过原子位置推断键。作者还与使用流匹配和等变最优传输生成原子类型和坐标的EquiFM进行了比较；然后根据这些推断键。最后，作者还比较了MiDi和EQGAT-diff，它们学习分子键并被视为当前的最先进方法。作者使用标准基准评估指标：原子稳定性、分子稳定性、有效性、唯一性和新颖性，这些指标在之前的工作中已被详细描述。

分子生成结果

表 1

表1比较了MolFlow与现有方法在QM9数据集上的性能，除了唯一性之外，MolFlow在所有指标上都能匹敌或超过所有模型，同时使用的采样步骤比MiDi和EQGAT-diff少了5倍。MolFlow还显著优于EquiFM，这是表中唯一的基于流匹配的方法。值得注意的是，由于EquiFM使用自适应ODE求解器，表中的函数评估次数是平均值，对于更大、更复杂的分子，其评估次数可能会显著更高。

表 2

由于QM9只包含非常小的分子，GEOM Drugs作为区分模型性能的更有用的基准。表2比较了作者的模型在该数据集上的表现。作者没有包括推断键的模型，因为这些模型在较大的分子上通常表现非常差，并且通常不提供所有评估指标的结果。MolFlow在所有关键评估指标上都优于现有的最先进模型，同时生成的全新分子稍微少一些。除了显著减少的评估步骤之外，作者的模型只需要更少的计算资源就能够训练。MolFlow在单个Nvidia A100 GPU上训练200个epoch，而EQGAT-diff在4个GPU上训练800个epoch。

采样效率

表 3

为了进一步提高模型效率，作者尝试了不同数量的ODE积分步骤，并在表3中提供了这些结果与EQGAT-diff的一对一比较。作者还测量了每个模型生成5000个分子所需的时间。表中可以看到两个模型之间的运行时间差异显著。MolFlow在采样步骤仅为20步时提供了与EQGAT-diff相当的性能，这对应于采样时间的两个数量级的速度提升。大部分这种改进归因于非常少的采样步骤，作者假设这是由于使用了最优传输，因为这旨在最小化在分布之间传输概率质量的成本。然而，作者进一步看到每次前向传递时间的两倍改进，这证明了设计高效的等变架构在分子生成中的重要性。

编译 | 于洲

审稿 | 曾全晨

参考资料

Irwin R, Tibo A, Janet J P, et al. Efficient 3D Molecular Generation with Flow Matching and Scale Optimal Transport[J]. arXiv preprint arXiv:2406.07266, 2024.