【NeurIPS 2024 Oral】用于多条件分子生成的图扩散Transformer

基于扩散模型的逆分子设计在材料和药物发现方面具有巨大的潜力。尽管在无条件分子生成方面取得了一定的成功，但将多个性质（如合成评分和气体渗透性）作为条件约束集成到扩散模型中仍未被探索。我们提出了用于多条件分子生成的图扩散Transformer（Graph DiT）。Graph DiT集成了一个编码器，用于学习数值和分类属性表示，并结合了基于Transformer的去噪器。与之前在前向扩散过程中分别对原子和键添加噪声的图扩散模型不同，Graph DiT通过一种新颖的图相关噪声模型进行训练，能够准确估计分子中的图相关噪声。我们对Graph DiT在多条件聚合物和小分子生成中的表现进行了广泛验证。结果表明，Graph DiT在分布学习和分子性质的条件控制的九项指标上都表现出优越性。通过气体分离的聚合物逆向设计任务并结合领域专家的反馈，进一步证明了其实际应用价值。 https://www.zhuanzhi.ai/paper/6950383d6ff9f64d7042d00bb91db831

1. 引言

分子图的扩散模型对于通过生成分子和聚合物（大分子）进行材料和药物的逆向设计至关重要，因为这些模型能够有效地训练来预测去噪过程中的离散图结构和原子/键类型 [40, 46]。实际的逆向设计考虑了多个因素，如分子的合成评分和各种性质 [15]，这被称为多条件图生成任务。 现有工作将多个条件转化为单个条件，并将任务作为单条件生成来解决 [5, 25]。然而，多属性关系可能无法被适当或明确地定义 [5]。首先，这些属性具有不同的尺度和单位。例如，合成复杂度的范围为1到5 [8]，而气体渗透性则变化很大，超过10,000 Barrier单位 [4]。这种差异使得模型难以平衡这些条件。其次，多条件包括分类属性和数值属性的组合。通常使用的加法 [47] 或乘法 [25] 不足以实现组合。图1(a) 经验性地展示了多条件生成的挑战，即发现符合多种性质的分子。我们使用了一个包含100个数据点的测试集，具有三个属性：合成可行性 (Synth.) [12]、O2 渗透性 (O2Perm) 和 N2 渗透性 (N2Perm) [4]。每个条件使用单条件扩散模型生成多达30个图，总共为三个条件生成了90个图。我们使用聚合物属性Oracle对每组中的30个图进行排序（见附录B.3）。然后，我们检查是否可以在不同条件集中找到符合多属性约束的共享聚合物结构。如果找到，则其排名K（K在1到30之间）表示其在列表中的排名，考虑所有条件集。如果找不到，我们将K设置为30。图1(a) 显示了100个测试案例中K的频率分布。中位数K为30，表明尽管生成了大量图形，但超过一半的测试聚合物未能满足多个属性。为了解决这些挑战，我们通过学习将多属性投射到表示中，从而引导分子生成的扩散过程。我们提出了用于条件下图去噪的图扩散Transformer (Graph DiT)。Graph DiT 具有一个用于属性表示学习的条件编码器和一个图去噪器。条件编码器采用了一种基于聚类的数值属性学习方法，并使用独热编码来表示分类属性，从而学习多属性表示。图去噪器首先将节点和边特征整合为图tokens，然后在Transformer层中通过自适应层归一化（AdaLN）对这些tokens进行去噪 [19, 34]。AdaLN用条件表示中的分子统计数据（均值和方差）替代每个隐藏层中的统计数据，有效地优于其他基于预测器和无预测器的条件方法 [22, 43, 34]，如第4.4节所示。我们观察到现有的前向扩散过程 [43, 22] 分别对原子和键应用噪声，这可能会影响Graph DiT在噪声估计中的准确性。因此，我们提出了一种新颖的图依赖噪声模型，该模型有效地应用了针对图中原子和键之间依赖关系的噪声。图1(b)中的结果显示，Graph DiT生成的聚合物与多属性约束高度一致。对于每个测试案例，我们根据三个属性从Graph DiT生成一个图。Oracle确定该图在每个条件下生成的30个单条件图中的排名。我们发现Synth.、O2 Perm和N2 Perm的中位排名分别为4、9和11，均远高于30。值得注意的是，这30个图的排序集非常具有竞争力，因为模型专门针对特定条件进行了训练。在实验中，我们评估了模型在一个聚合物数据集和三个小分子数据集上的表现。聚合物数据集包含四个用于多条件评估的数值条件。我们的模型具有最低的平均绝对误差（MAE），与最佳基线相比显著减少了17.86%的误差。在小分子任务中，模型同样表现出色，分类条件相关任务的准确率超过0.9，显著超过基线的不到0.6的准确率。我们还考察了模型在用于O2/N2气体分离的聚合物逆向设计中的实用性，并通过领域专家的反馈突出了我们模型在多条件分子设计中的实际应用价值。