单卡可跑18万原子!分子模拟的“规模焦虑”该终结了

AITNT-国内领先的一站式人工智能新闻资讯网站
# 热门搜索 #
单卡可跑18万原子!分子模拟的“规模焦虑”该终结了
9175点击    2026-08-26 14:48

人类距离真正理解生命还有多远?


离子穿过通道、药物分子穿过细胞膜、蛋白质折叠成特定构象……生命体从来不是一张静态的蓝图,要探究其中的奥秘,实验之外,计算模拟是另一只眼睛。


近年来,AI模型的介入正在不断刷新分子动力学模拟的边界,让精度与规模兼得的目标越来越近。


但这还远远不够。从“跨过鸿沟”到“可以日常使用”,仍需进一步降低训练成本、提高泛化能力、提升推理速度。


至知创新研究院UBio团队发布的UBio-MolFM v1.5,在这个方向上又往前迈进了一大步。


UBio-MolFM是以量子化学精度、在真实生理环境下对生物大分子进行分子动力学模拟的机器学习原子间势函数。在推理时能以接近经典力场量级的速度,给出接近量子化学精度的预测,而且不需要针对每一个新体系重新调参,同一套权重就能用在不同的分子和环境里,让第一性原理生物分子动力学模拟,不再是“奢侈品”。


先说说背景:精度与规模的取舍


过往分子动力学模拟长期面临一个基本矛盾:精度与规模不可兼得。


量子化学方法(如DFT)能够准确描述电子结构和化学键的断裂与生成,但计算成本高昂,通常被限制在数百个原子,也难以纳入显式的溶剂与膜环境。


经典力场走的是另一条路。它通过预先设定固定不变的原子电荷和键长参数,将模拟规模推至几十万乃至上百万原子,速度也快得多。但代价同样明确:这些参数是对真实物理势能面的近似,误差在具体体系中难以先验评估,而生物分子中微小的能量差,有时恰好是构象转换、配体结合、离子选择性的决定区间,往往落在这些误差的覆盖范围之内。


为了解决这个矛盾,近年来机器学习的出现已经在重新定义这道选择题的边界。基于深度学习的势函数能够在保持量子力学精度的同时,将模拟尺度提升数个数量级。“精度与规模”的取舍不再是绝对的,这道题正在持续被解开与突破,UBio-MolFM的出现,成为这条探索路径上一个新的里程碑。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了

分子动力学方法在精度与规模上的经典取舍。


如图,量子化学(左上)物理严谨,但规模受限于几百个原子;经典力场(右下)能处理生物尺度的体系,但物理是近似的;UBio-MolFM想同时满足两者。


两个实测案例告诉你,UBio-MolFM意味着什么


在介绍整体技术方案之前,先看两个具体的例子。可以看到,在同一个真实体系上,UBio-MolFM和经典力场两者给出的答案有明显差别。


  • 实测案例1:KcsA钾离子通道


KcsA是研究最充分的钾离子通道之一。它最关键的结构是选择性过滤器(selectivity filter),即一段由氨基酸羰基氧排列成的狭窄孔道,能让K⁺高速通过,却几乎不让Na⁺通过,选择性高达10⁴:1。通过分子动力学模拟,可以描述和研究离子通道内部的精细相互作用。


在完全相同的体系(108964个原子,包含真实脂双层与生理盐水环境)和完全相同的采样协议下,分别采用UBio-MolFM、以及两种经典力场(12-6-4与更常规的12-6),分别进行分子动力学模拟,对比结果如下:


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了

KcsA选择性过滤器中钾离子行为的对比。左:UBio-MolFM;右:经典固定电荷力场。标注的是腔内K⁺到最近羰基氧的距离,以及该位点的羰基配位数。


在UBio-MolFM的五条独立轨迹里,滤器深处相邻的两个钾离子会形成一种无水的直接接触,即中间不存在水分子,五条轨迹的平均距离是3.32 Å。其中四条轨迹从头到尾都维持着这一接触,第五条在700皮秒之前也是如此,此后才失去接触。相比之下,两种经典力场共十条独立轨迹里,这一直接接触从未出现过(0%)


钾通道对钾离子的高选择性,被认为与滤器内离子和羰基氧的直接配位密切相关。如果经典力场系统性地低估了这种直接接触,那么建立在这类力场上的一部分机制性结论,可能需要重新审视。


  • 实测案例2:环孢素A的构象自由能


构象转变能垒的准确预测一直是药物设计中的难题。很多候选药物因为穿膜效率低而在临床前被淘汰。环孢素A是一种临床使用的免疫抑制剂,它要穿过细胞膜,必须先完成一次构象收缩。当环孢素A从水溶液中舒展的状态,收拢成一个能穿膜的紧凑构象,这一收缩要付出多少自由能代价,直接决定了这个分子的口服生物利用度。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了

环孢素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Å的有效感受野,且计算复杂度接近线性。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了

UBio-MolFM的模型架构(上)与三阶段训练课程(下)


训练数据:1.6亿条标签,三个层次


训练语料共包含约1.6亿条量子化学标签,来自两个互补的来源:


  • 约1.4亿条来自通用小分子化学数据集OMol25,覆盖广泛的化学空间,但以小分子为主;
  • 约1900万条来自本工作自建的数据集UBio-Mol26,采用两种互补的采样策略:自下而上枚举溶剂化的构筑单元,以及自上而下从AlphaFold预测的蛋白质结构中截取局部区域,连同周围水分子一并采样。


三个DFT精度层级、各自的原子数范围与用途见Figure 3。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了

UBio-Mol26的构建流程:两种互补的采样策略,三个DFT精度层级,与OMol25合并后约1.6亿条量子化学标签。


比原子数范围更能说明问题的,是两项成分分析。它们确认了这份数据确实覆盖了目标化学空间,而不只是「更多的小分子」。


  • 第一项看化学基团。碳原子化学环境的统计显示,UBio-Mol26明显富集亚甲基与酰胺基团,这正是蛋白质骨架、脂类等生物大分子区别于一般小分子化学的特征基团,而OMol25以芳香环为主。
  • 第二项看几何尺度。原子对距离分布在5–6Å之外仍然保持宽泛,而典型的小分子数据集在这一距离之后迅速衰减为零。这个尺度恰好是上一节提到的感受野需要延伸到的范围。


此外,约6.4万条数据来自周期性凝聚相构型,数量虽小,却是模型能够正确重现水密度等宏观热力学性质的关键,用于校准模型对压力和密度的响应。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了


三阶段课程学习:从收敛到泛化


三个阶段共享同一套骨干网络,总计约1000 GPU-day


  • 第一阶段:在通用化学数据上训练。力由一个独立参数化的力头直接给出,这是更容易收敛的起点。
  • 第二阶段:换用自动微分,对预测能量求梯度得到力(F=-∇E),退役独立力头。这一自洽的力预测方式,是最终模型继承下来的核心训练目标,因为长时间轨迹模拟要求力本身来自一个自洽的能量面,而不是两个分开训练、互不担保一致的输出头。
  • 第三阶段:是唯一引入生物大分子数据的一轮,也是产出最终发布模型的一轮。它从四个数据分支采样,训练一个共享的输出头。其中,除了OMol25数据,其他数据分支的损失函数只监督力向量本身,完全不监督能量。因为决定长轨迹是否可靠的,是力的准确性而非能量的绝对值。这个设计还带来一个额外的好处:完全剔除了不同DFT设定带来的系统性能量偏移。


UBio-MolFM:误差降低45%,速度提升8倍


  • 性能与精度


在微观尺度上,UBio-MolFM在精度和效率上,与目前主流模型进行了对比。


下表对比了UBio-MolFM与三个代表性基线——UMA-S-1p2、MACE-OMol、DPA-4(OMol输出头)——的受力误差(单位:meV/Å,数值越低越准)。表中所有数字均可在已发布的技术报告中查到并复核。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了


在基线模型自身的训练分布上(OMol-Bio-10k,第一列),UMA-S-1p2的误差最低,UBio-MolFM没有在这一列上超过它,但是误差也是第一梯队的,这一列衡量的是「在为它们量身定制的分布上表现如何」。


真正有区分度的是后两列:held-out的生物大分子片段,以及体系规模超出所有基线训练上限的场景。在这两个更难的场景下,UBio-MolFM的受力误差比表现最好的基线低约40%–45%,具体到每一类化学环境(蛋白质、DNA、RNA、脂类等),降低幅度在21%到64%之间。


  • 速度与显存(前向+受力,单张141GB H20)


以下是四个模型在同一测试流程下的速度与显存对比(前向+受力,单张141GB H20)。测试体系为标准密度的立方水盒子,这更接近真实生物分子模拟的设定,而非各模型官方发布基准测试时使用的体系。因此,这里比较的是“同一套流程下的实测表现”,而非各自论文中公布的跑分。


单卡可跑18万原子!分子模拟的“规模焦虑”该终结了


在标准模式下,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


文章来自于"量子位",作者 "IQuest Research 团队"。

AI转型,免费服务,就找AITNT
AITNT资源拓展
根据文章内容,系统为您匹配了更有价值的资源信息。内容由AI生成,仅供参考
1
AI数据分析

【开源免费】DeepBI是一款AI原生的数据分析平台。DeepBI充分利用大语言模型的能力来探索、查询、可视化和共享来自任何数据源的数据。用户可以使用DeepBI洞察数据并做出数据驱动的决策。

项目地址:https://github.com/DeepInsight-AI/DeepBI?tab=readme-ov-file

本地安装:https://www.deepbi.com/

【开源免费airda(Air Data Agent)是面向数据分析的AI智能体,能够理解数据开发和数据分析需求、根据用户需要让数据可视化。

项目地址:https://github.com/hitsz-ids/airda

2
RAG

【开源免费】graphrag是微软推出的RAG项目,与传统的通过 RAG 方法使用向量相似性作为搜索技术不同,GraphRAG是使用知识图谱在推理复杂信息时大幅提高问答性能。

项目地址:https://github.com/microsoft/graphrag

【开源免费】Dify是最早一批实现RAG,Agent,模型管理等一站式AI开发的工具平台,并且项目方一直持续维护。其中在任务编排方面相对领先对手,可以帮助研发实现像字节扣子那样的功能。

项目地址:https://github.com/langgenius/dify


【开源免费】RAGFlow是和Dify类似的开源项目,该项目在大文件解析方面做的更出色,拓展编排方面相对弱一些。

项目地址:https://github.com/infiniflow/ragflow/tree/main


【开源免费】phidata是一个可以实现将数据转化成向量存储,并通过AI实现RAG功能的项目

项目地址:https://github.com/phidatahq/phidata


【开源免费】TaskingAI 是一个提供RAG,Agent,大模型管理等AI项目开发的工具平台,比LangChain更强大的中间件AI平台工具。

项目地址:https://github.com/TaskingAI/TaskingAI