分享自:

Next-Mol:结合3D扩散模型与1D语言模型用于3D分子生成

期刊:ICLR

报告:《NEXT-MOL:3D扩散模型与1D语言模型融合用于3D分子生成》

本文为您介绍一篇发表于ICLR 2025会议的前沿研究工作,题为《NEXT-MOL: 3D Diffusion Meets 1D Language Modeling for 3D Molecule Generation》。这项研究由来自新加坡国立大学、中国科学技术大学、香港中文大学和北海道大学的研究团队共同完成,主要作者包括刘志远、罗彦辰、黄涵、张恩智、李思航、方峻峰、施耀睿和王翔(通讯作者)等。该工作旨在解决三维分子生成领域的关键挑战,通过创新性地融合一维(1D)语言模型和三维(3D)扩散模型的优势,提出了一个名为NEXT-MOL的基础模型,在多个基准测试中取得了领先性能。

一、 学术背景与目标

研究领域是计算化学与人工智能的交叉领域,具体聚焦于三维分子生成。这是药物发现和材料设计中的核心任务,旨在同时生成分子的二维化学图(原子和键)及其三维空间构象(原子坐标)。准确预测分子的三维几何结构对于理解其生物活性、物理化学性质至关重要。

在NEXT-MOL之前,主流方法主要依赖3D扩散模型来生成连续的3D坐标,这类方法在建模3D几何方面表现出色,但存在一个显著缺陷:它们无法保证生成分子的100%化学有效性,偶尔会违背化学价键规则,产生“不可能存在”的分子结构。这阻碍了模型对有效分子结构分布的精准学习。另一方面,基于一维字符串表示(如SELFIES)的语言模型(Language Model, LM)在生成二维分子时,能够凭借SELFIES表示的内在鲁棒性,确保生成分子的100%有效性,并且可以充分利用规模庞大(数十亿级)的一维分子序列数据集。然而,纯1D模型无法直接生成3D构象。

因此,该研究的核心目标很明确:将1D语言模型保证分子有效性和数据规模大的优势,与3D扩散模型擅长建模连续空间结构的优势结合起来,创建一个更强大、更可靠的三维分子生成基础模型。他们旨在解决三个关键挑战:1)开发一个高效的自回归1D分子语言模型;2)设计一个强大的3D扩散模型用于精确构象预测;3)探索并实现从丰富的1D数据到稀缺的3D数据的知识迁移学习,以提升3D任务的性能。

二、 详细研究流程

NEXT-MOL采用两步式生成流程:第一步,使用一个名为MolLama的大型语言模型,根据SELFIES表示自回归地生成一个一维分子序列(即确定了原子的二维化学结构)。第二步,使用一个名为扩散分子变换器(Diffusion Molecule Transformer, DMT)的扩散模型,基于第一步生成的1D序列(转化成的2D图信息)来预测其对应的三维分子构象。 此外,研究引入了第三个关键组件:跨模态迁移学习,利用预训练的MolLama模型的知识来增强DMT的3D预测能力。

流程一:1D分子生成与MolLama模型 首先,研究团队从ZINC-15数据库中收集了18亿个分子,将其转换为SELFIES字符串,经过滤后形成了一个包含约900亿个token的庞大数据集。随后,他们从头开始预训练了一个拥有9.6亿参数的、基于Llama-2架构的解码器式语言模型,并将其命名为MolLama。预训练任务是最简单的“下一个token预测”,旨在让模型学会分子序列的语法和化学规则。在后续针对特定数据集(如GEOM-Drugs)进行微调以生成分子时,研究采用了 “随机化SELFIES增强” 技术。由于一个分子的2D图可以通过不同顺序遍历得到多个等价的SELFIES字符串,在训练时随机采用不同的遍历顺序可以增加数据多样性,防止模型过拟合到单一的顺序模式,从而提升生成分子的多样性和新颖性。

流程二:3D构象预测与DMT模型 对于第二步的3D构象预测,研究团队提出了一个新颖的扩散模型架构——扩散分子变换器(DMT)。DMT是一个连续时间扩散模型,其任务是:给定一个分子的原子特征(如原子类型)、成对特征(如化学键)以及被噪声扰动的3D初始坐标,学习去噪并还原出正确的3D坐标。

DMT的神经网络架构是其创新核心。它摒弃了许多3D生成模型中常用的、具有内置等变性的复杂网络结构,转而采用了一个更通用、可扩展的Transformer架构。其关键改进在于引入了 “关系多头自注意力”(Relational Multi-Head Self-Attention, RMHA) 模块。与标准自注意力不同,RMHA同时维护并迭代更新两套表示:原子表示(每个原子的特征)和成对表示(每对原子间的相互作用,如化学键信息)。在注意力计算中,成对表示的查询和值会通过元素级乘法动态地调制原子表示的查询和值,这使得模型能够充分、灵活地利用2D分子图的完整结构信息(原子和键),而不仅仅是损失性的简化表示。此外,DMT采用了自适应层归一化(AdaLN)来条件化扩散过程的时间步和可选的其他化学性质条件。研究还应用了随机旋转增强技术,在训练时对输入和目标的3D坐标施加相同的随机旋转,迫使模型学会对旋转的等变性,从而提升了鲁棒性。

流程三:1D到3D的跨模态迁移学习 为了利用海量1D数据预训练获得的知识来辅助数据相对稀缺的3D任务,研究设计了 “跨模态投影器” 和相应的三阶段训练策略。其核心思想是将预训练好的MolLama模型作为“化学知识库”,将其对分子序列的理解注入到DMT中。 1. 跨模态投影器:由于MolLama基于因果自注意力,每个token只能看到前面的token,且SELFIES token与原子并非一一对应,因此需要一个桥梁。该投影器首先将MolLama的SELFIES token表示通过一个双向自注意力层,使其获得全局上下文。然后通过一个编程式的 “SELFIES到原子映射” ,将多个对应同一个原子的token表示池化,并为没有对应token的氢原子使用可学习的嵌入。最终,处理后的表示通过一个多层感知机(MLP)后,与DMT原有的原子表示拼接,共同用于3D预测。 2. 三阶段训练策略: * 阶段一:在3D构象数据集上独立训练一个基础DMT模型直至收敛。 * 阶段二:将冻结参数的MolLama(使用LoRA低秩适应技术以减少内存占用)和随机初始化的跨模态投影器连接到冻结的预训练DMT上,进行10个epoch的“预热”训练,仅更新投影器和LoRA的参数,防止随机梯度破坏DMT已有的知识。 * 阶段三:解冻整个集成模型(DMT、投影器、MolLama的LoRA参数),进行端到端的微调直至收敛。在此阶段,研究发现使用规范的(Canonical)SELFIES比随机化的效果更好,可能因为固定表示有助于更快地弥合1D和3D之间的模态差距。

三、 主要实验结果

研究在GEOM-Drugs、GEOM-QM9和QM9-2014等多个标准数据集上,对NEXT-MOL进行了全面的评估,涵盖三个任务:从头3D分子生成、条件性3D分子生成(给定目标化学性质)和3D构象预测

结果一:从头3D分子生成性能领先 在最具药学相关性的GEOM-Drugs数据集上,NEXT-MOL展现了全面优势。在评估2D分子图质量的指标上,其弗雷歇化学网络距离(FCD)得分从之前最好基线(MolGen)的0.655大幅提升至0.334(相对改进49%),并且在原子稳定性、分子稳定性、有效性、独特性、新颖性以及片段/骨架相似度等几乎所有指标上均达到最佳或接近最佳。这证明了MolLama在捕获有效且多样的1D/2D分子分布方面的卓越能力。在评估3D结构质量的指标上,NEXT-MOL的3D FCD为14.69,最接近训练集本身的13.73,显著优于所有基线(如EDM的31.29,JoDO的19.99)。在几何相似性指标(键长、键角、二面角的最大平均差异MMD)上,NEXT-MOL也取得了最低误差,表明其生成的3D构象在几何上与真实分布高度一致。在QM9-2014数据集上的结果同样验证了其在处理较小分子时的稳健性。

结果二:条件性3D分子生成取得显著改进 在基于QM9-2014数据集、以量子化学性质(如偶极矩、极化率、HOMO-LUMO能隙等)为条件生成分子的任务中,NEXT-MOL再次超越了所有对比的扩散模型基线(如EDM, JoDO, GeoLDM等)。在六个目标性质的预测平均绝对误差(MAE)上,NEXT-MOL相比当前最优的JoDO模型取得了平均13%的相对提升。这一成功部分归功于DMT模型能够根据目标性质条件,为MolLama生成的1D分子寻找最匹配该性质的3D构象。

结果三:DMT在3D构象预测任务中达到新高度 在“给定2D图预测其3D构象”这一独立任务上,提出的DMT模型展现了卓越性能。在GEOM-Drugs数据集上,参数量为5500万的DMT-B模型在平均覆盖率(COV-R,衡量预测构象覆盖真实构象的能力)和平均最小RMSD(AMR-R,衡量预测构象与最近真实构象的平均距离)上均超越了所有同规模及更大规模的基线,包括参数量为2.42亿的MCF-L模型。这验证了DMT架构(特别是RMHA模块)在有效利用完整2D图信息方面的设计优势。当将DMT-B扩展到1.5亿参数的DMT-L时,性能得到进一步提升,证明了该架构良好的可扩展性。

结果四:跨模态迁移学习有效提升3D预测 关键性发现:将MolLama的1D表示通过跨模态投影器集成到DMT中,能够显著提升DMT的3D构象预测性能。在GEOM-Drugs上,这为DMT-B带来了约1.3%的COV-R绝对提升。更重要的是,这种提升在“未见过的分子骨架”(Unseen Scaffold)子集上尤为明显。对于训练集中从未出现过的骨架分子,DMT-B单独预测时AMR-R误差显著增大,而结合了MolLama知识后,误差得到了有效缓解。这表明,在大规模1D数据上预训练的MolLama学习到的通用化学启发式知识,能够有效迁移到3D任务,帮助模型更好地泛化到新颖的分子结构。可视化案例也显示,融合MolLama后,DMT对复杂二面角(扭转角)的预测更加准确。

四、 研究结论与价值

该研究成功地构建了NEXT-MOL这一强大的三维分子生成基础模型。其核心结论是:通过一个两阶段的、模块化的框架,将保证有效性的1D自回归语言模型与精确的3D扩散模型相结合,并辅以从1D到3D的跨模态知识迁移,能够实现更有效、更可靠、几何更准确的三维分子生成。 该模型在从头生成、条件生成和构象预测等多个关键任务上确立了新的性能标杆。

这项工作的科学价值在于:它弥合了1D序列建模与3D几何建模之间的鸿沟,为解决3D分子生成中的数据稀缺和有效性约束问题提供了一个新颖且有效的范式。它证明了大规模预训练的1D语言模型可以作为丰富的化学知识源,通过迁移学习来增强3D任务,这为整个计算分子科学领域利用不同模态数据提供了新思路。其应用价值非常直接且重大:作为一款高性能的基础模型,NEXT-MOL可以极大地加速药物先导化合物的发现与优化、新型功能材料的设计等过程,因为它能快速生成大量同时具备正确化学结构和合理三维形态的候选分子。

五、 研究亮点

  1. 范式创新:首次系统性地提出并验证了“1D LM生成结构 + 3D扩散模型预测构象”的两步式3D分子生成框架,巧妙地整合了两种范式的优势。
  2. 模型架构创新
    • MolLama:大规模(9.6亿参数)自回归分子语言模型,利用超大规模1D数据预训练,为下游任务奠定了坚实的化学结构基础。
    • DMT:新颖的扩散分子变换器,采用关系多头自注意力(RMHA)同时建模原子和成对相互作用,在保持Transformer可扩展性的同时,充分保留了2D分子图的完整信息,在构象预测上达到新高度。
  3. 有效的跨模态迁移:设计了跨模态投影器和三阶段训练策略,成功实现了从海量1D数据到稀缺3D数据的知识迁移,显著提升了模型对新颖分子结构的泛化能力,这是本工作的一个关键贡献。
  4. 全面领先的性能:在多个基准测试和任务上的实验结果一致且显著地超越了现有最先进方法,充分证明了NEXT-MOL作为下一代分子生成基础模型的潜力。
  5. 注重化学合理性:从源头(100%有效的SELFIES生成)确保分子化学有效性,并通过几何指标严格评估3D结构的合理性,使模型输出更贴合实际应用需求。

NEXT-MOL研究是一次高质量、多学科交叉的典范,它不仅提出了高性能的模型,更重要的是为三维分子智能生成领域指明了融合不同数据模态和建模范式的新方向,具有重要的学术影响力与应用前景。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com