![]()
这项由匹兹堡大学计算与系统生物学系及CMU-Pitt计算生物学博士项目联合完成的研究,以预印本形式于2026年6月发布在arXiv平台,编号为arXiv:2605.31498v2,感兴趣的读者可通过该编号查阅原文。
**一段关于分子"舞步"的困惑**
每一种分子,无论是药物分子还是蛋白质,都在不停地运动——它们扭转、弯曲、振动,就像舞台上永不停歇的舞者。这些舞步的规律,决定了分子能不能和某个蛋白质结合,能不能发挥药效,甚至决定了它在人体内的命运。科学家把所有可能的舞步及其出现概率,称为"玻尔兹曼分布"——这是一张描述分子在特定温度下如何分配自身姿态的概率地图。
问题在于,这张地图极难绘制。传统办法是让计算机模拟分子的每一步运动,就像用超慢动作摄影机记录舞者每一毫秒的动作,再拼出完整的舞蹈画面。这种方法需要以飞秒(千万亿分之一秒)为步长推进,耗时极长,计算量庞大,而且分子很容易陷入某个固定姿势"卡死",怎么也跳不出来。更麻烦的是,只要分子结构稍有改变,就得重新从头模拟,毫无复用可能。
正因如此,研究者们一直在寻找更聪明的替代方案。深度学习生成模型的兴起带来了新希望——如果能训练一个神经网络,让它直接"生成"符合玻尔兹曼分布的分子构型,就能绕开漫长的模拟过程。然而,这里存在一个令人头疼的循环悖论:训练这种模型,需要大量已经处于平衡状态的分子构型数据;而获取这些数据,恰恰就是我们最初想要解决的难题。
匹兹堡大学的研究团队提出了一种名为SITA(可扩展推断时退火,Scalable Inference-Time Annealing)的新方法,试图打破这个死循环。这套方案的核心思路,就是用一种"温度阶梯"策略,让模型像退火工艺一样,从高温逐步冷却到室温,一步步学会在低温下精确描绘分子的舞步分布。
**一、为什么"从高温开始"是个聪明的起点**
退火,是冶金工人的老手艺:把金属加热到高温,再缓慢冷却,让内部结构逐渐排列整齐。分子模拟领域早就借用了这个概念。在高温下,分子拥有更多能量,能够越过各种能量屏障,自由探索所有可能的姿态;而在低温(比如室温300K)下,分子更倾向于蜷缩在能量最低的几个姿态里,几乎无法跳出去。
传统分子动力学模拟在低温下极易"卡死",就是因为分子没有足够能量越过那些能量山丘,去探索其他构型。高温模拟则宽松得多,分子可以畅游整个构型空间,覆盖所有可能的"舞步"。
SITA正是利用这一特点,先在1200K的高温模拟数据上训练一个生成模型,让模型先学会高温下分子的"豪放舞风",再通过一系列精心设计的步骤,逐步把模型引导向300K(室温)下更"保守"、更专注于低能量姿态的精确分布。这个过程,就像先教一个舞者学习自由即兴舞,再一步步约束他学习古典芭蕾的严格规范。
**二、"温度阶梯"究竟怎么走**
SITA的训练过程分为四个循环往复的步骤,整个过程像是爬下一段楼梯,每走一步,模型对低温分布的理解就更深一层。
第一步,研究团队利用当前的生成模型(称为"流模型",可以理解为一个专门生成分子构型的神经网络),在稍低一些的温度下批量产出分子样本。这里有一个精妙的技巧:流模型本身不需要修改,只需改变输入的"噪声"来源的方差,就能让输出的样本更倾向于低能量状态。具体来说,如果当前温度是1200K,目标是755K,那么只需将输入噪声的方差按照两个温度的比值缩小,模型就会自然地生成更符合低温分布的样本。这种只需调整输入、无需重训练模型的温度切换方式,是SITA的一大亮点。
第二步,用这批新生成的样本,训练一个"能量模型"(EBM,Energy-Based Model)。能量模型的作用,是学习这批样本的概率密度——通俗地说,就是学会判断每个分子构型"有多大可能真实存在"。直接计算流模型生成样本的精确概率,在高维度下需要计算一个极其复杂的"雅可比行列式"(可以理解为描述空间如何被变换的数学量),计算量随分子自由度的增加急剧膨胀,完全不可行。能量模型则通过一种叫做BoltzNCE的训练方法,绕开了这个计算瓶颈,用一个近似但高效的方式估计出样本的概率。
BoltzNCE的训练分两个部分:第一部分叫"分数匹配",通过学习样本能量函数的梯度来约束能量的形状;第二部分叫"噪声对比估计",通过让模型区分不同时间点的样本来锚定能量的绝对值。两者合力,使能量模型能够给每个样本打出一个"概率分数",而无需经历高昂的计算。
第三步,利用能量模型估算出的概率,计算每个样本的重要性权重。重要性权重的含义是:这个样本在目标低温分布下,相比在当前高温分布下,更应该被重视多少倍。计算公式是用玻尔兹曼分布给出的真实低温概率,除以能量模型估算的高温概率。权重越高,说明这个样本越符合低温分布的特征。通过按这些权重对样本重新采样,就得到了一批"更像室温分布"的训练数据。为了防止极少数权重极高的样本主导整个重采样结果(这会严重降低样本多样性),研究团队还对权重做了99百分位截断处理。
第四步,用这批重新采样的数据,对流模型进行微调,让它更好地逼近当前目标温度下的分布。微调完成后,整个循环重新开始,目标温度再降一个台阶。研究中用到的温度阶梯是:从1200K出发,依次经过755.95K、555.52K、408.24K,最终到达300K。每一步都让模型对室温分布的理解更进一步。
**三、与对手的正面交锋:SITA凭什么说自己更好**
在这个领域,最直接的竞争对手是一个叫PITA(Progressive Inference-Time Annealing)的方法。PITA的思路与SITA有几分相似,同样是沿温度阶梯逐步退火,但它采用的是基于扩散模型的架构,并依赖一种叫"费曼-卡茨公式"的数学工具来估算样本的重要性权重。这个估算过程,需要沿整个生成轨迹积分计算速度场的散度(可以理解为计算每一步生成过程中空间如何被"压缩"或"膨胀"),计算量随分子自由度的增加呈灾难性增长。对于只有几十个自由度的丙氨酸二肽,这已经很吃力;对于更大的分子系统,基本上是不可行的。
SITA用能量模型替代了这个昂贵的散度积分,从根本上绕开了这个瓶颈。代价是引入了一点近似误差——能量模型毕竟不是精确的概率密度,而是一个近似估计。这个近似会导致重采样后的分布不是精确的目标低温分布,而是一个"倾斜版本",其偏差取决于能量模型与真实密度的差距。研究团队坦率地承认了这一点,并从理论上推导出了倾斜分布的形式。
然而,令人意外的是,这种带有近似偏差的方法,在实验基准上竟然全面超越了理论上更精确的PITA。
**四、用丙氨酸检验功力**
研究团队在两个标准分子基准上测试了SITA:丙氨酸二肽(Alanine Dipeptide,ADP)和丙氨酸三肽(Alanine Tripeptide,ATP)。丙氨酸二肽是计算化学领域最经典的测试案例,就像机器学习领域的MNIST手写数字识别,既足够简单,又能暴露方法的关键缺陷。
评价指标主要有四个。"Rama-KL"衡量的是生成样本的构型分布与真实室温MD模拟之间的KL散度(一种衡量两个概率分布差异的数学量),数值越小说明覆盖越全面、越精确,这个指标尤其对"模式崩溃"(即模型只生成少数几种构型)非常敏感。"Energy-W1"和"Energy-W2"衡量的是生成样本与参考样本在能量分布上的差异,分别对应一阶和二阶Wasserstein距离。"T-W2"则衡量样本在扭转角空间(即拉马钱德兰坐标,描述蛋白质骨架弯曲方式的两个角度)上的二阶Wasserstein距离。
在丙氨酸二肽上,SITA在Rama-KL(0.517±0.013)和Energy-W2(0.939±0.079)两项最重要的指标上拿到了所有方法中的最佳成绩,显著优于PITA(Rama-KL为4.773±0.460,Energy-W2为1.615±0.053)。有一个有趣的细节:另一个叫MD-NF的基准方法,直接用300K的MD数据训练,在Energy-W1这一指标上表现最好,但它的Rama-KL极高(13.533±0.024),说明它其实严重模式崩溃了——只生成了能量较低的几种构型,所以能量均值看起来不错,但覆盖的构型空间极为有限。SITA生成的样本则覆盖了丙氨酸二肽所有主要的构型盆地,与MD参考分布在拉马钱德兰自由能图上高度吻合。
在丙氨酸三肽上,SITA的优势更为突出。PITA在没有进行额外MD松弛(即用短时间的MD模拟修正生成样本)的情况下,Rama-KL高达8.535,能量误差更是飙升到86.270。SITA在完全不做任何后处理的情况下,Rama-KL只有0.361,Energy-W1为1.933。这意味着SITA生成的样本质量,直接超越了PITA加上MD松弛后处理的水平,而后者需要额外消耗3倍的能量计算预算。
两个系统的有效样本比例(ESS,衡量重采样效率的指标,越高说明权重分布越均匀)都处于较低水平——丙氨酸二肽在0.13到0.27之间,丙氨酸三肽在0.04到0.07之间。这说明重采样过程中存在一定的样本浪费,但即便如此,最终结果仍然优于其他方法,说明SITA对样本质量的把握非常准确,即使少量高权重样本也能有效指引模型向正确方向优化。
**五、还能更精准:独立梅特罗波利斯-黑斯廷斯精炼**
对于追求极致精度的场景,研究团队还设计了一个可选的后处理步骤:独立梅特罗波利斯-黑斯廷斯(IMH)算法。
这个算法的原理,可以类比为一个挑剔的艺术评审。每次从流模型生成一个新的分子构型候选,然后评审团(由真实玻尔兹曼权重和能量模型共同组成)比较这个新候选与当前样本的"质量得分比",如果新候选明显更好,就接受它;如果差不多或更差,就以相应概率接受或拒绝。这个机制保证了最终保留下来的样本集合会逐渐向真实分布靠拢。
理论上,如果评审团使用的是精确概率,这个算法会收敛到精确的目标分布。由于SITA使用的是能量模型近似概率,收敛目标实际上是前面提到的那个"倾斜版本"的分布——研究团队在论文中给出了严格的数学证明,说明倾斜程度完全由能量模型与真实密度的偏差决定。
实验结果显示,在丙氨酸三肽上运行50步IMH后,Rama-KL从0.361进一步降至0.313,T-W2从0.798降至0.704,多项指标均有改善。相比之下,另一种简单的重要性重采样后处理(SITA-IS)虽然在能量误差上更低,但Rama-KL出现了恶化,说明样本多样性受损。IMH则在改善能量精度的同时,较好地保持了构型空间的覆盖,代价是额外需要5×10?次能量计算(而整个预训练阶段需要5×10?次,这个开销相对可以接受)。
**六、关于评估方式的一场"较真"**
SITA论文中还专门辟出一节,指出了竞争方法PITA在评估指标上的两处方法论问题,这在学术界算是相当直接的批评。
第一个问题是TICA(时间滞后独立分量分析)评估中的滞后时间选择。TICA是一种分析分子动力学轨迹慢模式的技术,其投影结果对滞后时间参数极为敏感。PITA对丙氨酸三肽使用了滞后时间10,对丙氨酸二肽使用了滞后时间100。过小的滞后时间会让TICA捕捉到高频振动和热噪声,而非真正有物理意义的慢构型转变,从而使基于TICA的评估指标失去区分模式覆盖能力的灵敏度。
第二个问题是MD参考轨迹的降采样方式。PITA取轨迹的前10000帧作为参考分布,但MD轨迹帧之间存在强时间相关性,前10000帧往往集中于轨迹起始的少数构型,无法代表完整的平衡分布。使用这种有偏参考来评估模型,会让那些同样模式崩溃的方法因为与有偏参考"匹配得好"而获得虚假的高分。SITA团队因此改用均匀间隔采样的方式选取参考帧,结果发现所有方法在这种更公平的评估下表现均有变化,而SITA的优势在这种更严格的评估下更加显著。
这一较真细节,体现了研究者在方法论层面的严谨态度——评估方式的细节,同样能决定"谁更好"的结论。
**七、模型架构背后的技术选择**
SITA的流模型使用了一种叫GVP-GNN(几何向量感知图神经网络)的架构,这是一种专门为分子设计的神经网络,能够保持物理系统的基本对称性:无论分子在空间中怎么旋转、平移,模型给出的物理量预测值应当保持不变或协变。这种对称性约束大大减少了模型需要学习的冗余信息,提高了学习效率。
能量模型则采用了Graphormer架构,这是一种为图结构数据设计的Transformer变体。Transformer原本是处理自然语言的神经网络,Graphormer通过在注意力机制中加入分子内原子间距离矩阵作为结构偏置,使其能够有效处理三维分子构型。在SITA中,这个距离信息被直接注入注意力计算,让模型天然具备对分子三维几何的感知能力。
两个模型的训练都使用了三角函数插值(αt=cos(πt/2),βt=sin(πt/2))作为随机插值框架的调度函数,这种选择能让插值路径在时间上更均匀,避免训练信号在某些时间段过于集中。优化器使用Adam,配合Reduce-on-Plateau学习率调度,在损失停止下降时自动降低学习率。所有实验在单张NVIDIA L40 GPU上完成,每张显卡配备46GB显存。
**八、这套方法的局限与未来**
研究团队非常坦率地说明了SITA引入近似的代价。由于能量模型只是流模型分布的近似,重采样后的分布并非精确的目标低温分布,而是受能量模型精度影响的倾斜分布。当前实验中有效样本比例较低(丙氨酸三肽最低只有0.045),说明重采样效率还有较大提升空间。
此外,SITA目前只在小型分子系统上测试,对于更大的蛋白质或药物分子,能量模型的近似精度是否足够、流模型的表达能力是否足够,仍然是开放的问题。研究团队将架构优化、跨分子系统的迁移能力,以及应用到更大分子系统的可行性,列为未来主要研究方向。
归根结底,SITA代表的是一种"以近似换规模"的工程哲学:在精确性上做出可接受的让步,换来在更大系统上的可行性。在分子模拟这个高维度的领域,这往往是通往实用的唯一路径。
当分子生物学家和药物研发者面对一个全新的蛋白质靶点时,如果能用SITA这样的工具在数小时内生成可靠的室温构型分布,而不是等待数周的MD模拟,这种效率上的质变,可能直接影响到下一代药物从实验室走向临床的速度。而这,才是这场关于分子"舞步"的研究,与普通人最真实的关联。
Q&A
Q1:SITA方法与传统分子动力学模拟相比,在计算效率上有多大优势?
A:SITA的主要效率来源于两点。第一,它用神经网络直接生成分子构型,避免了传统MD模拟以飞秒步长逐步推进的高昂开销;第二,它用能量模型替代了精确散度积分,避免了随分子自由度增长而急剧膨胀的计算量。在实验中,整个预训练加退火流程消耗约5.08×10?次能量计算,比对手PITA在丙氨酸三肽上需要的8×10?次少约37%,且无需额外MD松弛后处理。
Q2:SITA中的能量模型近似会不会让生成的分子构型不准确?
A:会引入一定偏差。能量模型估算的是流模型分布的近似密度,而非精确密度,导致重采样后的分布是目标分布的"倾斜版本",偏差大小取决于能量模型与真实密度的差距。研究团队在理论上推导了这个偏差的形式,并指出只有当能量模型与流模型分布完全一致时,才能精确恢复目标分布。尽管存在这种近似偏差,实验结果显示SITA在多个指标上仍优于其他方法。
Q3:SITA的温度阶梯中温度点是怎么选的,能不能直接从1200K跳到300K?
A:温度点的选择影响重要性权重的质量,跨度越大,相邻温度分布差异越大,有效样本比例会急剧下降。SITA采用755.95K、555.52K、408.24K、300K四步过渡,每步约降低约26%的温度,以保证足够高的有效样本比例(丙氨酸二肽约13%~27%)。若直接从1200K跳到300K,两个温度下的玻尔兹曼分布差异极大,几乎所有生成样本的权重都会趋近于零,重采样几乎失效,训练数据质量会严重劣化。





京公网安备 11011402013531号