编辑丨王多鱼
排版丨水成文
人类距离真正理解生命还有多远?
离子穿过通道、药物分子穿过细胞膜、蛋白质折叠成特定构象……生命体从来不是一张静态的蓝图,要探究其中的奥秘,实验之外,计算模拟是另一只眼睛。
近年来,AI 模型的介入正在不断刷新分子动力学模拟的边界,让精度与规模兼得的目标越来越近。
但这还远远不够。从“跨过鸿沟”到“可以日常使用”,仍需进一步降低训练成本、提高泛化能力、提升推理速度。
近日,至知创新研究院 UBio 团队发布的UBio-MolFM v1.5,在这个方向上又往前迈进了一大步。
UBio-MolFM是以量子化学精度、在真实生理环境下对生物大分子进行分子动力学模拟的机器学习原子间势函数。在推理时能以接近经典力场量级的速度,给出接近量子化学精度的预测,而且不需要针对每一个新体系重新调参,同一套权重就能用在不同的分子和环境里,让第一性原理生物分子动力学模拟,不再是“奢侈品”。
先说说背景:精度与规模的取舍
过往的分子动力学模拟长期面临一个基本矛盾:精度与规模不可兼得。
量子化学方法(例如 DFT)能够准确描述电子结构和化学键的断裂与生成,但计算成本高昂,通常被限制在数百个原子,也难以纳入显式的溶剂与膜环境。
经典力场走的是另一条路,它通过预先设定固定不变的原子电荷和键长参数,将模拟规模推至几十万乃至上百万原子,速度也快得多。但代价同样明确:这些参数是对真实物理势能面的近似,误差在具体体系中难以先验评估,而生物分子中微小的能量差,有时恰好是构象转换、配体结合、离子选择性的决定区间,往往落在这些误差的覆盖范围之内。
为了解决这个矛盾,近年来机器学习的出现已经在重新定义这道选择题的边界。基于深度学习的势函数能够在保持量子力学精度的同时,将模拟尺度提升数个数量级。“精度与规模”的取舍不再是绝对的,这道题正在持续被解开与突破,UBio-MolFM的出现,成为这条探索路径上一个新的里程碑。
分子动力学方法在精度与规模上的经典取舍
如图,量子化学(左上)物理严谨,但规模受限于几百个原子;经典力场(右下)能处理生物尺度的体系,但物理是近似的;而 UBio-MolFM 想同时满足两者。
两个实测案例告诉你,UBio-MolFM 意味着什么
在介绍整体技术方案之前,先看两个具体的例子。可以看到,在同一个真实体系上,UBio-MolFM 和经典力场两者给出的答案有明显差别。
实测案例1:KcsA 钾离子通道
KcsA 是研究最充分的钾离子通道之一,它最关键的结构是选择性过滤器(selectivity filter),即一段由氨基酸羰基氧排列成的狭窄孔道,能让 K⁺ 高速通过,却几乎不让 Na⁺ 通过,选择性高达 10⁴:1。通过分子动力学模拟,可以描述和研究离子通道内部的精细相互作用。
在完全相同的体系(108964 个原子,包含真实脂双层与生理盐水环境)和完全相同的采样协议下,分别采用 UBio-MolFM、以及两种经典力场(12-6-4 与更常规的 12-6),分别进行分子动力学模拟,对比结果如下:
KcsA 选择性过滤器中钾离子行为的对比。左:UBio-MolFM;右:经典固定电荷力场。标注的是腔内 K⁺ 到最近羰基氧的距离,以及该位点的羰基配位数。
在 UBio-MolFM 的五条独立轨迹里,滤器深处相邻的两个钾离子会形成一种无水的直接接触,即中间不存在水分子,五条轨迹的平均距离是 3.32 Å。其中四条轨迹从头到尾都维持着这一接触,第五条在 700 皮秒之前也是如此,此后才失去接触。相比之下,两种经典力场共十条独立轨迹里,这一直接接触从未出现过(0%)。
钾通道对钾离子的高选择性,被认为与滤器内离子和羰基氧的直接配位密切相关。如果经典力场系统性地低估了这种直接接触,那么建立在这类力场上的一部分机制性结论,可能需要重新审视。
实测案例2:环孢素 A 的构象自由能
构象转变能垒的准确预测一直是药物设计中的难题。很多候选药物因为穿膜效率低而在临床前被淘汰。环孢素A是一种临床使用的免疫抑制剂,它要穿过细胞膜,必须先完成一次构象收缩。当环孢素 A 从水溶液中舒展的状态,收拢成一个能穿膜的紧凑构象,这一收缩要付出多少自由能代价,直接决定了这个分子的口服生物利用度。
环孢素 A 构象自由能地形的对比。左:UBio-MolFM 解出清晰的漏斗状地形;右:经典力场给出的地形近乎平坦。
在同一套二维伞形采样协议下,UBio-MolFM 解出的自由能地形落差为 5.51 kcal/mol,其中穿膜构象本身对应约 3.5 kcal/mol 的能量代价;经典力场在相同采样窗口内的落差只有 1.51 kcal/mol,地形近乎平坦,看不出明确的能量最小值在哪。
这个势能面是否真的达到了量子化学精度,经过了直接验证:从采样窗口中抽取 100 个去相关构象逐个做单点 DFT 计算,与模型预测逐一比对。在模型能量头的标定基准上,两者偏差为 0.44 kcal/mol,比三个 DFT 参考方法彼此之间的分歧(最小 1.18)还要小。受力误差(两个基准为 7.2 meV/Å)也落在下方精度表格中对应层级的水平上,说明这个大环分子的描述质量与基准测试预测的一致。
这两个案例指向了共同的逻辑:当体系的物理行为由具体的相互作用细节决定时(直接配位 vs 水介导,能垒高度 vs 平均结构),经典力场的局限已经超越了定量误差的范畴,可能导向定性层面的误判。
UBio-MolFM 在架构、数据和训练三个方面进行创新
骨干网络:三个约束,三重设计
架构的设计必须同时满足三个条件:旋转等变性、达到非共价距离的感受野、以及能在一张 GPU 上运行十万原子级体系的硬件效率。这三个约束分别对应三层设计。
旋转等变性由骨干网络本身承担。本工作沿用团队此前发表于 ICML 2026 的 E2Former-V2 架构——一个等变(equivariant)Transformer。它满足一个最基本的物理要求:分子旋转后,模型预测的力向量按相同角度旋转,而能量保持不变。这是任何试图学习物理规律、而不是记住特定朝向的模型都须满足的约束。
E2Former-V2 架构自带三项工程优化,是模型能在合理时间内完成训练、并在单张 GPU 上完成十万原子推理的前提:
SO(3)→SO(2) 的基变换:把稠密的张量收缩运算换成按局部坐标轴重新索引的稀疏运算,卷积阶段提速约 6 倍。
以节点、而非以原子对为中心的 Wigner-6j 张量分解:把计算开销从数量远多于原子本身的原子对,转移到原子上。
一套融合的 Triton 流式注意力核心:计算时无需在显存中保存完整的邻居张量,激活显存占用与节点数成正比,而不是与原子对数成正比,算力吞吐最高提升 20 倍。
在此基础上,本工作对骨架网络做了进一步的架构调整,在骨干之上叠加了四层混合截断半径堆叠:
三层短程等变注意力(半径5Å)覆盖全部原子(含氢),承载局部多体相互作用与氢键几何。相比原架构的 6Å 注意力半径,5Å 的截断使邻居数量减少 40% 以上。第四层将截断半径扩展到 8Å,同时在构建邻域时排除氢原子。三层短程层的输出经残差连接,与第四层输出在一个融合节点汇合,拼接通道后投影回统一维度,送入能量头。
这样设计的目标是:把感受野延伸到第二水合层、盐桥这类短程非共价接触,同时避免溶剂化体系里氢-氢原子对数量二次方增长导致的算力爆炸。三层短程加一层中程,最终给出约 23Å 的有效感受野,且计算复杂度接近线性。
UBio-MolFM的模型架构(上)与三阶段训练课程(下)
训练数据:1.6 亿条标签,三个层次
训练语料共包含约 1.6 亿条量子化学标签,来自两个互补的来源:
约 1.4 亿条来自通用小分子化学数据集 OMol25,覆盖广泛的化学空间,但以小分子为主;
约 1900 万条来自本工作自建的数据集 UBio-Mol26,采用两种互补的采样策略:自下而上枚举溶剂化的构筑单元,以及自上而下从 AlphaFold 预测的蛋白质结构中截取局部区域,连同周围水分子一并采样。
三个DFT精度层级、各自的原子数范围与用途见Figure 3。
UBio-Mol26 的构建流程:两种互补的采样策略,三个 DFT 精度层级,与 OMol25 合并后约 1.6 亿条量子化学标签。
比原子数范围更能说明问题的,是两项成分分析。它们确认了这份数据确实覆盖了目标化学空间,而不只是「更多的小分子」。
第一项看化学基团。碳原子化学环境的统计显示,UBio-Mol26 明显富集亚甲基与酰胺基团,这正是蛋白质骨架、脂类等生物大分子区别于一般小分子化学的特征基团,而 OMol25 以芳香环为主。
第二项看几何尺度。原子对距离分布在 5-6Å 之外仍然保持宽泛,而典型的小分子数据集在这一距离之后迅速衰减为零。这个尺度恰好是上一节提到的感受野需要延伸到的范围。
此外,约 6.4 万条数据来自周期性凝聚相构型,数量虽小,却是模型能够正确重现水密度等宏观热力学性质的关键,用于校准模型对压力和密度的响应。
三阶段课程学习:从收敛到泛化
三个阶段共享同一套骨干网络,总计约 1000 GPU-day。
第一阶段:在通用化学数据上训练。力由一个独立参数化的力头直接给出,这是更容易收敛的起点。
第二阶段:换用自动微分,对预测能量求梯度得到力(F=-∇E),退役独立力头。这一自洽的力预测方式,是最终模型继承下来的核心训练目标,因为长时间轨迹模拟要求力本身来自一个自洽的能量面,而不是两个分开训练、互不担保一致的输出头。
第三阶段:是唯一引入生物大分子数据的一轮,也是产出最终发布模型的一轮。它从四个数据分支采样,训练一个共享的输出头。其中,除了 OMol25 数据,其他数据分支的损失函数只监督力向量本身,完全不监督能量。因为决定长轨迹是否可靠的,是力的准确性而非能量的绝对值。这个设计还带来一个额外的好处:完全剔除了不同 DFT 设定带来的系统性能量偏移。
UBio-MolFM:误差降低 45%,速度提升 8 倍
性能与精度
在微观尺度上,UBio-MolFM 在精度和效率上,与目前主流模型进行了对比。
下表对比了 UBio-MolFM 与三个代表性基线——UMA-S-1p2、MACE-OMol、DPA-4(OMol 输出头)——的受力误差(单位:meV/Å,数值越低越准)。表中所有数字均可在已发布的技术报告中查到并复核。
在基线模型自身的训练分布上(OMol-Bio-10k,第一列),UMA-S-1p2 的误差最低,UBio-MolFM 没有在这一列上超过它,但是误差也是第一梯队的,这一列衡量的是「在为它们量身定制的分布上表现如何」。
真正有区分度的是后两列:held-out 的生物大分子片段,以及体系规模超出所有基线训练上限的场景。在这两个更难的场景下,UBio-MolFM 的受力误差比表现最好的基线低约 40%-45%,具体到每一类化学环境(蛋白质、DNA、RNA、脂类等),降低幅度在 21%-64% 之间。
速度与显存(前向+受力,单张141GB H20)
以下是四个模型在同一测试流程下的速度与显存对比(前向+受力,单张 141GB H20)。测试体系为标准密度的立方水盒子,这更接近真实生物分子模拟的设定,而非各模型官方发布基准测试时使用的体系。因此,这里比较的是“同一套流程下的实测表现”,而非各自论文中公布的跑分。
在标准模式下,DPA-4,MACE-Omol 以及 UMA-S-1p2 的速度接近,约 8-9K 原子/s,最大体系在 10000-15000原子,UBio-MolFM 约 37K 原子/s,最大支持约 65000 原子。UMA 与 UBio-MolFM 均支持「激活重计算」运行模式,也就是通过重新计算部分激活值以达到降低显存的效果,而代价则是速度的下降。在激活重计算模式下,两个模型均可支持 12 万原子,在这个规模下,UMA 已经用满了 141GB 显存,UBio-MolFM 仍有约三分之一余量,吞吐是 25.1 对 3.0 千原子/秒,这正是速度提升 8.3 倍这个数字的确切来源。
真正决定这个差距的,是显存随体系增大的增长速度:UBio-MolFM 每增加 1000 个原子,激活重计算模式下多占用约 0.73GB 显存,标准模式下约 1.97GB;UMA 在标准模式下这个数字是约 17.5GB,接近 UBio-MolFM 的 9 倍,这也是它在标准模式下体系刚过九千原子就耗尽显存的原因。按这个斜率计算,一张 141GB 显卡理论上能装下约 19.3 万原子,实测已经在一张卡上跑到 18 万原子,留有安全余量。把同样的显存开销铺到一个 8 卡节点上,可以到 144 万原子,相当于一整个卫星烟草花叶病毒。
在基线模型自身的训练分布范围内(通常是数百原子的小分子体系),UMA-S-1p2 仍具有精度优势,UBio-MolFM 并未在这一区间提出更优的主张。它真正的适用范围,是这一区间之外:从需要显式溶剂与生物大分子环境的体系,直至百万原子规模。
现已全面开源
UBio-MolFM 的模型权重、代码与技术报告已经开源,欢迎复现、检验,也欢迎把它用在尚未测试过的体系上。
模型:https://huggingface.co/IQuestLab/IQuest-UBio-MolFM-V1.5
代码:https://github.com/IQuestLab/UBio-MolFM
技术报告:https://huggingface.co/IQuestLab/IQuest-UBio-MolFM-V1.5/blob/main/MolFM-1p5-Technical-Report.pdf
热门跟贴