Coarse-Grained Boltzmann Generators¶

会议: ICML 2026
arXiv: 2602.10637
代码: https://github.com/tummfm/cg-bg
领域: 科学计算 / 分子模拟
关键词: Boltzmann Generator, 粗粒化建模, 重要性采样, 平均力势, 归一化流

一句话总结¶

提出 Coarse-Grained Boltzmann Generators (CG-BGs)，在粗粒化坐标空间中结合归一化流生成模型和学到的平均力势 (PMF) 进行重要性采样，以远低于原子级 BG 的计算成本实现渐近正确的分子平衡态采样。

研究背景与动机¶

领域现状：从 Boltzmann 分布中采样平衡态分子构型是统计物理中的核心难题。Boltzmann Generator (BG) 通过精确似然生成模型与重要性采样相结合来应对此问题，能够生成提案样本并进行重加权以得到无偏估计。同时，粗粒化 (CG) 方法通过降低自由度来处理更大的分子系统。

现有痛点：原子级 BG 在维度增长时面临两大瓶颈——(1) 生成分布与目标分布的重叠度下降导致重要性权重方差爆炸，重加权失效；(2) Jacobian 行列式计算随维度急剧增长。另一方面，Boltzmann Emulator 虽通过粗粒化降维提升了可扩展性，但省略了重加权步骤导致无法纠正分布偏差，且依赖难以获得的长时间无偏模拟数据训练。

核心矛盾：BG 有重加权机制但难以扩展到大系统；CG Emulator 可扩展但没有纠偏机制——两者的优势互补但未被整合。

本文目标：在粗粒化坐标空间中实现带重要性采样的生成建模，同时从快速收敛的增强采样数据中学习目标能量函数。

切入角度：粗粒化坐标下的边缘分布 \(p(\mathbf{R})\) 同样可以写成 Boltzmann 形式 \(p(\mathbf{R}) \propto e^{-\beta U(\mathbf{R})}\)，其中 \(U(\mathbf{R})\) 是平均力势 (PMF)。如果能学到 PMF，就可以在低维 CG 空间中复用 BG 的重要性采样框架。

核心 idea：用增强采样力匹配 (ESFM) 从快速收敛的有偏轨迹中学 PMF，用归一化流在 CG 空间生成提案分布，用学到的 PMF 做重要性重加权，形成完整的 CG-BG 框架。

方法详解¶

整体框架¶

输入为原子级分子动力学模拟轨迹（可以是有偏的增强采样数据），通过粗粒化映射 \(\mathbf{R} = \Xi(\mathbf{r})\) 投影到低维 CG 坐标。框架包含两个并行训练的组件：(1) 基于连续归一化流 (CNF) 的提案分布 \(q_\theta(\mathbf{R})\)；(2) 基于神经网络的 PMF \(U_\eta(\mathbf{R})\)。推理时流模型生成 CG 构型，PMF 计算重要性权重 \(w(\mathbf{R}) \propto e^{-\beta U_\eta(\mathbf{R})} / q_\theta(\mathbf{R})\)，通过自归一化重要性采样估计量得到无偏平衡态可观测量。

关键设计¶

增强采样力匹配 (ESFM) 学习 PMF:
- 功能：从快速收敛的有偏模拟数据中学习粗粒化平均力势，不依赖昂贵的无偏平衡态轨迹
- 核心思路：利用纤维分布不变性——在 CG 坐标上施加偏置势 \(V(\mathbf{R})\) 时，给定 \(\mathbf{R}\) 的原子级条件分布不变，即 \(p_V(\mathbf{r}|\mathbf{R}) = p(\mathbf{r}|\mathbf{R})\)。因此投影力的条件均值（即平均力）在有偏采样下不变，力匹配回归目标不受偏置影响。训练损失为 \(\mathcal{L}_{\mathrm{ESFM}}(\eta) = \mathbb{E}_{\mathbf{r} \sim \mathcal{D}_{\mathrm{bias}}}[\|\nabla_{\mathbf{R}} U_\eta(\Xi(\mathbf{r})) + \mathcal{F}_{\mathrm{proj}}(\mathbf{r})\|^2]\)，其中力从无偏原子势重新计算
- 设计动机：标准力匹配需要收敛的无偏数据，而增强采样（如 well-tempered metadynamics）可快速覆盖亚稳态间的过渡区域。ESFM 与标准力匹配有相同全局最优解，且 KL 散度受力误差平方上界约束
CG 空间连续归一化流 (CNF) 提案生成:
- 功能：在低维粗粒化坐标空间中学习提案密度 \(q_\theta(\mathbf{R})\) 以近似目标边缘分布
- 核心思路：使用 Flow Matching 训练神经向量场 \(v_\theta(t, \mathbf{x})\)，通过线性插值路径 \(\mathbf{x}_t = (1-t)\mathbf{x}_0 + t\mathbf{x}_1\) 回归目标向量场。流模型在 CG 空间操作维度远小于原子级（如丙氨酸六肽从 72 原子降至 Core Beta 映射的少量珠子），使得 Jacobian 计算和分布重叠度显著改善
- 设计动机：在低维空间中，生成模型与目标分布的重叠更好，重要性权重方差更低，ESS 更高；同时推理时 Jacobian 计算成本大幅降低
PMF 引导的重要性重加权:
- 功能：将流模型的有偏提案样本纠正为渐近正确的平衡态分布
- 核心思路：对流模型生成的样本 \(\mathbf{R}_i \sim q_\theta\) 计算重要性权重 \(w(\mathbf{R}_i) \propto e^{-\beta U_\eta(\mathbf{R}_i)} / q_\theta(\mathbf{R}_i)\)，用自归一化估计量计算可观测量。通过归一化有效样本量 \(\mathrm{ESS} = (\sum w_i)^2 / (B \sum w_i^2)\) 评估重加权质量，并采用权重截断策略提升对 MLP 外推异常和生成伪影的鲁棒性
- 设计动机：Boltzmann Emulator 直接用 \(q_\theta\) 估计可观测量无法纠偏。引入学到的 PMF 作为目标能量函数恢复了 BG 的重加权能力，且 PMF 可捕获隐式溶剂模型无法表达的溶剂介导效应

训练策略¶

框架分为两阶段独立训练：(1) PMF 网络通过 ESFM 损失在有偏或无偏原子级轨迹数据上训练；(2) 归一化流通过 Conditional Flow Matching 损失在 CG 坐标数据上训练。两者可并行训练，推理时组合使用。

实验关键数据¶

主实验¶

在丙氨酸二肽 (22 原子)、三肽 (42 原子)、六肽 (72 原子) 上评估，以显式溶剂 MD 模拟为参考标准。

模型	JS 散度 (↓)	PMF 误差 (↓)	ESS (↑)
CG-BG Heavy Atom	0.0048	0.2005	0.5112
CG-BG Heavy Atom (有偏数据)	0.0063	0.2277	0.4115
CG-BG Core Beta	0.0052	0.2210	0.5528
CG-BG Core Beta (有偏数据)	0.0057	0.2093	0.4818
隐式溶剂 GB (OBC1)	0.0157	0.3709	—
隐式溶剂 GB (OBC2)	0.0182	0.4028	—

计算效率对比（丙氨酸二肽，\(10^4\) 样本）¶

CG 映射	训练时间	推理时间	总时间
Core Beta	0.45h	0.95min	0.47h
Heavy Atom	0.80h	3.78min	0.86h
All Atom (仅溶质)	2.55h	14.91min	2.80h

更大体系验证（三肽 & 六肽）¶

模型	三肽 JS (↓)	三肽 PMF (↓)	三肽 ESS (↑)	六肽 JS (↓)	六肽 PMF (↓)	六肽 ESS (↑)
CG-BG Core Beta	0.0060	0.2112	0.4212	0.0100	0.3646	0.1231
CG-BG Heavy Atom	0.0056	0.1957	0.3201	—	—	—
隐式溶剂 GB (OBC2)	0.0932	1.0274	—	0.1652	1.8401	—

关键发现¶

CG-BG 重加权后在所有指标上显著优于隐式溶剂基线，尤其在三肽和六肽等更大体系中差距扩大（六肽 JS 散度 0.0100 vs 0.1652）
粗粒化分辨率存在精度-效率 trade-off：Core Beta 映射 ESS 更高（分布重叠更好）但重加权后精度略低于 Heavy Atom 映射
从 10ns 有偏数据训练的 CG-BG 与 500ns 无偏数据训练版本精度接近，证明了 ESFM 对数据效率的提升
原子级 BG 只能达到隐式溶剂精度上限，而 CG-BG 通过从显式溶剂数据学 PMF 突破了这一限制

亮点与洞察¶

纤维分布不变性的巧妙利用：CG 偏置势不改变给定 CG 坐标下的原子级条件分布，使得力匹配目标在有偏采样下等价——这个理论保证使得用增强采样的快速收敛数据替代昂贵的无偏轨迹成为可能
无模拟 PMF 评估：学好一个提案分布后可以通过切换不同 PMF 的重要性权重一次性评估多个候选 CG 力场，无需为每个模型单独跑 MD——这是对 CG 力场开发工作流的显著加速
粗粒化 + 重加权的互补设计：粗粒化解决了维度问题使 ESS 可控，重加权解决了分布偏差问题保证渐近正确——这种将降维与纠偏正交分解的思路可迁移到其他高维采样问题

局限与展望¶

依赖预定义的集体变量（CG 映射和增强采样的 CV 选择），对复杂体系可能难以确定合适的 CV
当前实验仅在丙氨酸短肽（≤72 原子）上验证，更大的蛋白质体系效果待验
六肽的 ESS 已降至 0.1231，随体系增大重要性采样效率可能进一步下降
未来方向包括结合自动 CV 发现方法、引入可迁移生成模型架构、以及探索基于能量的训练替代 Flow Matching