用于大规模材料平衡采样的可扩展Boltzmann生成器

《Nature Communications》：Scalable Boltzmann generators for equilibrium sampling of large-scale materials

　　

生成结构平衡系综对于分子与材料的建模至关重要，然而分子动力学等传统模拟方法存在采样效率受限的问题。Boltzmann生成器（BGs）提出了一步式深度学习进行平衡采样的概念，但将其扩展至大规模系统始终是主要挑战。本研究中，研究人员通过一种能够对大型材料系统进行建模的BG架构克服了这一规模限制。该研究方法结合增强耦合流（augmented coupling flows）与图神经网络（graph neural networks, GNN）以利用局部环境信息，实现了基于能量的训练与快速推断。与先前设计相比，该架构训练速度更快、资源消耗更少，并达到了更优的采样效率。至关重要的是，该方法能够迁移至更大的系统规模，使得对包含超过一千个原子的模拟胞材料进行高效采样成为可能。研究人员在Lennard-Jones晶体、单原子水（monatomic water, mW）冰相以及硅相图上展示了该方法的能力，产生了准确的平衡系综与自由能结果，其尺度范围已超越有限尺寸效应消失的范围。

生成模型正在兴起为传统采样技术如分子动力学（molecular dynamics, MD）和蒙特卡罗（Monte Carlo, MC）模拟的有前景替代方案，用于生成多粒子系统的平衡系综。与MD或MC通过时间上的顺序更新来探索构型空间不同，生成模型旨在学习从基分布到目标平衡分布的直接映射，提供对非关联样本的即时访问。此类模型的关键优势在于实现统计独立平衡样本的一步式生成，从而绕过传统方法中固有的长相关时间和慢混合问题。因此，生成模型有望显著降低分子模拟的计算成本，尤其对于热力学系综平均和自由能的计算。基于Noé等人在Boltzmann生成器（BGs）方面的工作，归一化流（normalizing flows）等生成模型与统计重加权方法的结合已应用于分子与凝聚态科学的多个领域，涵盖从蛋白质、小分子到凝聚相系统如液体、原子固体和分子晶体的平衡构型直接生成，以及模拟分布的转换以探索不同势能或热力学条件。

尽管BGs具有 promising 特性，但将其扩展至大规模系统仍面临重大挑战，主要体现在高昂的训练成本和低采样效率。因此，大多数已有研究集中于概念验证问题，如少于100个原子的Lennard-Jones团簇和小分子。在材料科学领域，迄今研究的最大系统为包含约500个粒子的原子固体，然而这些模型尽管需要近一个GPU年的训练，仍表现出极低的采样效率，其计算成本超过基于MD的模拟数个数量级，难以应用于实际系统。这一局限性在凝聚相系统中尤为关键，因为更大的系统规模对于确保统计收敛和物理意义明确的预测至关重要。

本研究中，研究人员针对这一挑战开发了可扩展至大型晶体材料的BGs。与先前针对凝聚相系统的方法不同，研究人员采用局部性假设，将生成过程基于局部环境而非完整构型。该目标通过结合图神经网络（GNN）与增强耦合流框架实现，后者支持基于能量的训练而无需目标分布的样本，并在推断过程中保持计算效率。研究人员展示了这种局部架构不仅允许模型达到更高的采样效率，还显著降低了训练成本。该架构的关键特征在于其跨系统尺寸的迁移性，使得在小系统上训练的模型可应用于大得多的系统，从而大幅降低与直接在大型系统上训练相比的计算成本。对热力学状态或原子相互作用势等外部参数的条件化进一步分摊了训练成本。

**大规模Boltzmann生成器的缩放**

在多点体系统背景下缩放任何架构的核心在于基于局部结构特征表述学习问题。遵循这一原则，研究人员旨在训练一个基于局部环境学习坐标变换的流，以实现小系统上的高效训练和对大系统的无缝迁移。在连续流的情境下，构型更新可基于局部结构环境定义，但评估这些模型的密度需要计算Jacobian迹，这在高维度下缩放性较差。另一缺点是连续流的训练需要目标分布的样本，而这些正是该方法旨在生成的量。尽管降低Jacobian迹评估成本的研究以及无需目标样本的连续模型研究是活跃的研究领域并取得了一些初步进展，但这些方法对于本工作相关系统规模尚不够成熟。因此，研究人员转而聚焦于耦合流。

耦合流通过将输入按x=(x^A,x^B)划分为两个通道，并依据一个通道条件更新另一个通道，支持高效的密度估计，产生的三角Jacobian可解析求值。耦合流的关键特征在于其与基于似然的优化方案的兼容性，允许仅使用目标分布的势能函数进行训练，而无需任何来自该分布的样本。由于耦合流需要跨粒子或跨空间维度或两者结合来分割坐标，条件器C无法访问完整的三维原子间距离，因为其输入只能包含坐标的子集。因此，先前架构往往依赖系统中所有原子的绝对坐标，本研究中将此方法称为"全局"架构。

为克服耦合流的这一架构局限性，研究人员利用增强耦合流框架。增强流引入辅助变量a∈R^3N，使得分割可在物理变量和辅助变量之间进行，从而基于a条件更新物理变量x，反之亦然。重要的是，该方案在物理空间和辅助空间内保留完整的三维坐标，因此允许计算原子间距离。研究人员将辅助系统建模为物理系统的噪声副本，并定义增强系统的联合基分布为q(x,a)=q(x)N(a;x,η_q²I)。在此构造中，辅助系统保留物理变量的结构信息，可被流变换利用。

利用辅助变量，研究人员通过x_i′=g(x_i|h_i^a)定义粒子i的流变换，其中g由嵌入h_i^a参数化，该嵌入使用GNN计算的每个辅助粒子局部邻域信息聚合得到。此处局部方法的关键特征在于其随系统规模的线性缩放，通过固定邻居数量实现。这与基于注意力机制等学习所有粒子间相互作用的全局架构形成鲜明对比，后者随系统规模呈二次方缩放。

为确保双射器（实现为有理二次样条）适用于更大系统规模，其输入量级不随系统规模缩放至关重要。研究人员遵循的可能性之一是建模偏离理想晶格的位移而非绝对位置，这使输入与尺寸无关并保持模型对任意系统规模的泛化能力。

对于作用于增强空间的流，增强系统的变量变换写为q_θ(x′,a′)=q(x,a)|detJ_?θ(x,a)|^?1。流被训练以优化联合目标分布p(x,a)=p(x)N(a;x,η_p²I)，其中p(x)为感兴趣系综的平衡分布。本研究聚焦于正则（或NVT）系综，其平衡分布为Boltzmann分布。流通过最小化生成分布与目标分布在联合空间中的Kullback-Leibler（KL）散度进行训练。

在增强设置中，仅联合生成密度q_θ(x′,a′)可被精确评估。然而，物理系统的边缘分布可通过积分出辅助自由度获得，可近似为q_θ(x′)≈(1/M)∑_m=1^M[q_θ(x′,a_m)/π(a_m|x′)]，其中a_m～π(·|x′)，M为每生成一个物理样本所抽取的辅助样本数。值得注意的是，虽然边缘分布的评估需要每个生成样本通过流进行M次额外的逆向传递，但不需要目标势能的额外评估。

如"方法"部分详述，与分布μ相关的（约化）Helmholtz自由能f_μ=F_μ/k_BT_μ=?logZ_μ的计算需要配分函数Z_μ的估计。虽然该量无法直接计算，训练后的流可用于有针对性的自由能扰动（targeted free energy perturbation, TFEP）来估计基态与目标态之间的自由能差Δf_qp=f_p?f_q。在增强流框架内，生成联合分布的配分函数依赖于物理和辅助变量。由于辅助系统的目标和基分布均为高斯分布，配分函数可分解，增强系统的自由能由f_μ^aug=f_μ+f_μ^aux给出。若η_q=η_p，则Δf_qp^aux=0，Δf_qp^aug=Δf_qp。

基于流的方法还允许计算对应于NPT系综的Gibbs自由能，其中模拟盒形状h∈R^3×3（V=det h）在恒定压力P下变化。尽管先前工作研究了同时建模形状和构型分布的流，研究人员当前的方法通过Legendre变换计算Gibbs自由能：G=min_h[F(h)+det(h)P]。为获得F(h)的可靠估计，研究人员以形状条件化方式训练流，优化流以生成不同盒尺寸的样本。

**尺寸可迁移的训练与评估**

研究人员通过将局部BGs应用于更大系统来展示尺寸迁移性的有效性。对于立方mW冰，在N=216粒子训练的模型应用于512、1000和1728粒子系统；对于FCC LJ，在N=256训练的模型应用于500、864、1000和1372粒子系统。

径向分布函数（radial distribution function, RDF）和势能能量分布显示，BG产生的结果与MD几乎无法区分，远超训练系统的尺寸范围。虽然基分布的RDF已捕捉目标的主要特征，但基分布与MD系综的能量分布存在显著差异，表明基分布构型未反映目标势能中粒子位置间的正确关联。相比之下，局部BG准确再现了两个系统的结构和能量特征，在未重加权情况下即获得了能量分布的极佳一致性。这些结果明确表明，基于局部环境的坐标变换对于生成准确的平衡结构完全足够，并能可靠地迁移至更大的系统尺寸。

结构准确性虽是正确自由能估计的必要条件，但不足以保证高采样效率或精确的自由能计算。更严格的度量是有效样本量（effective sample size, ESS）。随着系统规模增加，ESS自然下降，因为即使广泛的势能中恒定的每粒子误差也会导致Boltzmann分布中指数放大的误差。但即使对于N=1728的立方mW冰和N=1372的FCC LJ等大系统，局部BG的ESS仍足以确保可靠的系综平均统计量。

重要的是，大系统的ESS显著高于全局方法获得的结果，而计算成本显著降低。对于N=512的mW冰，文献报道的全局BGs需要超过330 GPU天的训练直至收敛，但ESS仅约0.2%。相比之下，研究人员N=216的局部BG在相同硬件上约4 GPU天即达到收敛，而512粒子系统的ESS比全局BG高一个数量级。此外，全局架构扩展至更大系统的计算成本令人望而却步，使超过500的粒子数因过高训练成本和恶化采样效率而实际上不可能实现。

图3进一步展示了局部架构相对于全局架构的改进训练性能。在所有系统中，局部BG在任意给定训练时间下均持续达到比全局BG更高的ESS。这一点在较大系统尺寸下尤为明显，展示了局部BG随粒子数的优越缩放性。局部BG还表现出ESS的显著更低方差，表明即使样本量较小也可获得可靠估计，并随样本量增加保持稳定。

研究人员注意到BG在不同系统上表现不一。例如，尽管粒子数仅适度增加，局部和全局BG在立方mW冰上均达到比FCC LJ显著更高的采样效率。研究人员还观察到同一势能内某些晶体结构可表现出相当不同的采样效率，这需要本工作范围之外的进一步研究。

**自由能估计**

训练后BG的关键优势在于能够在TFEP框架内直接评估自由能差。利用局部架构，在相对小系统上训练的BG允许对更大系统进行准确的绝对Helmholtz和Gibbs自由能评估，与MBAR等传统自由能估计器相比大幅降低计算量。

立方冰相的绝对自由能随粒子数的变化显示，联合和边缘自由能估计与参考值达到极佳一致，直至1000粒子偏差低于10^?3k_BT每粒子。对于最大系统尺寸（N=1728），联合估计略高估真实值，而边缘评估对此部分修正，获得与表1中报告ESS一致的更准确结果。研究人员基于所考虑示例，0.1%或更高的平均ESS足以确保使用5×10⁴样本时自由能估计偏差低于10^?3k_BT每粒子。FCC LJ晶体观察到类似行为。

不同晶相之间准确自由能差的获取尤为具有挑战性但对探索相图至关重要。达到立方与六角mW冰之间自由能差的收敛值需要高达1728粒子的系统尺寸。联合和边缘估计的Δf与参考值非常一致，联合估计的良好表现可能源于两相之间的误差抵消。收敛自由能差所需的大量粒子凸显了评估扩展系统规模能力的重要性，而这对于全局架构是不可行的。

为量化尺寸可迁移自由能估计带来的计算节省，研究人员将能量评估次数与标准MD+MBAR方法进行比较。MBAR设置需要约50个中间态，每态10³非关联样本，总计约5×10⁷次能量评估，加上MBAR评估额外的2.5×10⁶次评估。流模型训练直至收敛涉及可比的评估次数（约6×10⁷），但仅在小的系统规模上进行。训练后，用BG评估自由能仅需大系统的数万样本及相应能量评估。相比之下，MD+MBAR每个系统规模需要数千万能量评估，对于成对势能计算成本二次方缩放，对于密度泛函理论等量子力学方法更陡峭。

虽然BG模型在所需能量评估数量上比MD+MBAR更高效，但这一优势对于单次自由能估计并不直接转化为挂钟时间。对于N=1728的立方冰系统，完整MD+MBAR自由能估计仅需数小时，而将流模型训练至收敛约需8 GPU天。然而，这种比较特定于mW等廉价相互作用模型。机器学习原子间势能（当今计算材料科学的主力）通常比经验势能昂贵10³至10⁴倍。在这些情况下，能量评估主导成本，而流模型的优化和神经网络评估开销基本与势能无关。因此，BG方法减少所需能量评估次数预计将使其相对挂钟成本比传统MD+MBAR计算显著更有利，甚至更低。

**包含体积涨落**

为展示方法在等温等压系综中的效率和准确性，研究人员确定了恒定压力下FCC与六角密堆积（hexagonal close-packed, HCP）LJ晶体之间的Gibbs自由能差。与六角和立方冰相类似，这些结构具有极小的自由能差，需要大系统规模收敛。关键的是，FCC和HCP的相对稳定性依赖于LJ势能的截断半径，小截断甚至可导致稳定性的定性错误预测。

体积条件化训练后，研究人员将模型应用于N=1080的系统以评估一系列截断半径下的Gibbs自由能。BG预测与MD+MBAR参考值展现极佳一致。随着截断增加，敏感性降低，强调了容纳足够大截断半径所需的大模拟胞.需要。此外，通过Legendre变换从BG获得的粒子密度与NPT MD模拟中观察到的平均密度紧密匹配。BG估计在整个截断半径范围内均产生高度准确的结果，强烈表明局部流模型可使用单个短截断在小系统中训练，同时仍可迁移至大得多的系统并继续生成准确构型。

**条件化于原子类型**

另一重要应用展示是流可条件化于相互作用势参数，使单一模型能够泛化至整个系统类别。研究人员采用Stillinger-Weber（SW）势能，其结合两体势和三体项，三体相互作用强度由因子λ₃控制。当用约化单位表示时，区分不同参数化.的唯一有效参数为三体相互作用强度λ₃。不同λ₃值稳定或 destabilize 不同晶体结构。

为研究变化三体强度和系统尺寸下的相图，研究人员训练了N=216的钻石立方和β-tin结构以及N=128的体心立方（body-centered cubic, BCC）结构的局部BGs。BGs条件化于λ₃∈[15,24]和模拟盒形状，使模型能够泛化至整个SW势能类别并覆盖不同材料范围。小系统和大系统的Gibbs自由能显示，在整个λ₃范围内，BG的自由能估计与MD+MBAR参考值非常一致。虽然立方和β-tin结构之间的转变点基本不随系统尺寸变化，但BCC和β-tin结构之间的转变点观察到向较低三体相互作用强度的轻微偏移，需要大系统尺寸的准确自由能来确定这一微妙效应。

流方法的一个额外优势是模型可在整个λ₃光谱上针对所有结构训练，即使某些结构仅在部分区域稳定。

**硅相图**

聚焦于硅（Si）的SW参数化，研究人员条件化于温度和模拟盒形状训练局部BGs以绘制(T,P)相图。具体目标是确定钻石立方和β-tin相之间的共存线。BGs的训练和评估均在N=216进行。虽然立方结构的尺寸迁移性再次极好，β-tin结构在训练尺寸上已更具挑战性，限制了大系统中自由能估计的准确性。为改进相图预测准确性，β-tin相的Gibbs自由能基于M=200的边缘密度评估。

即使不迁移至更大系统，条件化训练也在热力学状态间提供了大量成本分摊，因为参考MD+MBAR计算需要在至少7×7网格上完全收敛的模拟以覆盖完整(T,P)范围，而局部BGs实现了所有热力学状态下Gibbs自由能的高效评估。结果Si相图展示了广泛条件下的立方和β-tin相竞争，准确预测了相应的共存线。与MD+MBAR自由能估计的验证显示极佳一致，共存线实际上不可区分。

**讨论**

研究人员介绍的局部流架构用于训练生成模型，实现了材料系统的平衡采样，从小训练胞无缝缩放至超过1000粒子的显著更大系统。该方法超越先前全局架构，实现了更快训练时间和更高采样效率。至关重要的是，流模型表现出跨系统尺寸的迁移性，能有效采样短程和长程相互作用势，在各种材料系统中产生高度准确的绝对自由能估计。通过条件化于热力学变量和相互作用参数，训练过程进一步分摊，实现了跨温度、压力和组成的相稳定性直接高效评估。所展示结果凸显了该架构作为强大可扩展工具的潜力，用于广泛热力学条件下复杂材料的研究。

未来方向包括扩展该方法以高效采样液相，以及开发具有化学空间广泛泛化能力的模型。在此背景下，重要的一点是研究为何模型对不同势能-结构组合表现出不同性能，以更好理解和指导生成方法的设计。超越当前原子固体表述的一个特别有趣的扩展将是适用于分子晶体的模型开发。虽然在概念上可在所开发方法框架内实现，但这需要更精细的架构修改以考虑取向自由度，如刚体约束产生的自由度。未来工作的另一重要方向是发展能够捕捉长程相互作用（如静电产生的相互作用）的模型，这无法在任何纯局域表述中表示。这些发展 together 可显著扩展生成建模在材料科学中的范围和影响。

**研究结论**

研究人员引入了用于训练生成模型的局部流架构，实现材料系统的平衡采样，从小训练胞无缝缩放至超过1000粒子的显著更大系统。该方法超越先前全局架构，实现更快训练时间和更高采样效率。流模型表现出跨系统尺寸的迁移性，能有效采样短程和长程相互作用势，在各种材料系统中产生高度准确的绝对自由能估计。通过条件化于热力学变量和相互作用参数，训练过程进一步分摊，实现跨温度、压力和组成的相稳定性直接高效评估。所展示结果凸显了该架构作为强大可扩展工具的潜力，用于广泛热力学条件下复杂材料的研究。