自回归玻尔兹曼生成器:突破流模型限制的高效分子采样新范式

针对统计物理中分子系统热力学平衡采样的难题,研究提出自回归玻尔兹曼生成器(ArBG)。传统方法依赖归一化流,面临表达能力受限或似然计算昂贵的瓶颈。ArBG摒弃流式范式,采用大语言模型架构,通过自回归建模克服拓扑约束并支持序列推理干预。实验显示其在所有基准测试中显著优于流模型,尤其在较大肽系统中表现突出。此外,1.32亿参数的Robin模型在8残基系统上将零样本能量误差降低超60%,刷新最先进水平,为分子模拟提供更具扩展性的解决方案。

在统计物理学与计算化学的交叉领域,如何在热力学平衡状态下高效地对分子系统进行采样,始终是一个核心挑战。这一挑战直接推动了玻尔兹曼生成器(BGs)的发展,其旨在通过结合生成模型、精确似然估计以及重要性采样校正,快速生成无相关性的平衡样本。然而,现有的主流玻尔兹曼生成器大多依赖于归一化流(Normalizing Flows, NFs)。这类方法存在明显的局限性:离散时间的流模型受限于严格的可逆性约束,导致表达能力不足;而连续时间的流模型虽然表达能力更强,但其似然计算往往计算成本高昂。针对这一痛点,本研究提出了自回归玻尔兹曼生成器(ArBG),这是一种全新的自回归建模范式。ArBG 的核心贡献在于彻底脱离了基于流的玻尔兹曼生成器范式,利用自回归模型的特性,不仅规避了流模型固有的拓扑约束,还使得在推理阶段进行序列干预成为可能。更重要的是,ArBG 借鉴了大语言模型中经过验证的高效架构,从而在可扩展性上实现了显著提升,为分子系统的模拟开辟了新的技术路径。在技术方法层面,ArBG 详细阐述了如何利用自回归架构重构分子生成过程。与归一化流必须满足的双射映射不同,自回归模型通过条件概率分解来建模高维分子构型空间,这允许模型逐步生成分子的各个部分,从而自然地处理复杂的拓扑结构。

这种方法避免了流模型中常见的数值不稳定性和计算瓶颈。具体而言,ArBG 采用了类似于大型语言模型中的注意力机制和分层结构,能够捕捉分子内部长程依赖关系。在训练策略上,模型通过最大化数据的对数似然进行优化,并结合玻尔兹曼分布的物理约束进行校正。这种设计不仅简化了似然计算的过程,还使得模型能够更灵活地适应不同规模和复杂度的分子系统。此外,自回归的特性允许在生成过程中引入额外的干预信号,例如固定某些原子位置或调整局部构型,这在传统流模型中难以实现或计算代价极高。通过这种架构创新,ArBG 在保持物理一致性的同时,大幅提升了生成效率和模型的表达能力。为了验证 ArBG 的有效性,研究团队在多个标准基准数据集上进行了广泛的实验评估。结果显示,ArBG 在所有测试基准上均显著优于基于流的模型,特别是在处理较大规模的肽系统时,优势更为明显。以10残基的Chignolin蛋白为例,ArBG 展现出了卓越的采样能力和能量预测精度。

除了性能提升,作者还引入了一个名为 Robin 的1.32亿参数可迁移模型,该模型基于 ArBG 框架训练而成。实验数据显示,Robin 模型在8残基系统上将零样本能量误差(E-W2)降低了超过60%,这一结果大幅超越了之前的最先进水平。消融实验进一步证实了自回归架构在捕捉分子长程相互作用方面的优势,以及重要性采样校正对最终生成质量的关键作用。这些结果不仅证明了 ArBG 在理论上的优越性,也展示了其在实际分子模拟任务中的巨大潜力,特别是在需要高精度能量预测的场景下。从行业意义和潜在影响来看,ArBG 的提出对开源社区和工业落地具有深远影响。首先,该方法打破了归一化流在分子生成领域的垄断,为研究者提供了另一种高效且可扩展的技术选择。其代码的公开(https://github.com/danyalrehman/autobg)将加速相关领域的复现与创新。在工业应用方面,更高效的分子采样意味着药物发现和材料设计过程中的模拟速度可以大幅提升,从而缩短研发周期。此外,ArBG 支持推理时干预的特性,使其在需要精细控制分子构型的场景中具有独特优势,例如在蛋白质折叠预测或分子对接任务中。对于后续研究而言,ArBG 框架为结合物理先验与深度学习提供了新的思路,未来可能扩展到更复杂的生物大分子系统或材料科学领域。随着大模型架构在科学计算中的进一步渗透,此类工作有望推动计算化学向更智能化、自动化的方向演进,促进跨学科的技术融合与创新。

Sources