🎯核心问题:分子构象采样的拓扑困境

统计力学的核心洞见是:宏观现象——蛋白质折叠、Ising 模型磁化、晶体结构形成——都由平衡态微观状态的系综决定。这个平衡分布就是玻尔兹曼分布 μ_target(x) ∝ exp(−E(x)),其中 E(x) 是构象 x ∈ ℝ^{n×3} 的无量纲势能。计算挑战在于如何高效地从这个目标分布中抽取统计独立的样本。

分子动力学(MD)模拟是主流方法,但存在严重的时间尺度问题:MD 的时间步长在飞秒(10⁻¹⁵ s)量级,而跨越高能势垒的模式混合需要微秒到秒(10⁻⁶ ~ 10⁰ s)量级。这意味着绝大多数 MD 计算都花在模拟局部极小值内的高频振动上,而非探索全局能量景观。各种加速 MD 方案(增强采样、元动力学等)仍难以根本解决这个混合问题。

玻尔兹曼生成器(BG)由 Noé 等人于 2019 年提出,旨在绕过这一困境。BG 学习一个生成模型 p_θ(x) 来提出重要性采样样本,利用精确的模型似然 p_θ(x) 和目标能量 E(x) 计算重要性权重。这允许并行生成独立且一致的样本,无需在模式之间穿越,使 BG 成为大规模分子系统平衡采样的理想框架。

然而,为了满足可处理似然的要求,BG 几乎完全依赖归一化流(Normalizing Flows)。这带来了根本性的拓扑约束:

离散时间 NF 是微分同胚(diffeomorphism),同时也是同胚(homeomorphism),因为先验是单个高斯分布。这意味着生成模型必须保持定义域的拓扑——它无法建模与先验具有不同拓扑的目标分布,例如不连通的支撑集、不同数量的连通分量,或"空洞"。而分子构象的平衡分布恰好具有这些挑战性拓扑:多个亚稳态被高能势垒区域分隔。为了将高斯的单个连通模式变形到这些有效分离的状态,流被迫执行极端形变——拉伸空间跨越细桥、压缩进入模式——导致高度非光滑的映射、爆炸的 Lipschitz 常数和病态的雅可比矩阵。

连续时间 NF(CNF)虽然不受离散时间 NF 的架构约束,且现代流匹配训练策略可以大幅稳定优化,但病态性在推理时显现:学习到的向量场变得高度非光滑,导致刚性动力学,需要大量函数评估才能获得准确似然。这造成了一个不可避免的权衡:离散时间 NF 似然高效但表达力有限;CNF 更具表达力但推理代价昂贵。

🧬从玻尔兹曼生成器到 ARBG

ARBG 的核心思路简洁而深刻:用自回归模型替代归一化流作为 BG 的提案分布。自回归模型将联合密度分解为一系列条件密度的乘积:

p_θ(x) = ∏_j p_θ(x_j | x_{

这个看似简单的改变带来三个关键优势:

1. 突破表达力瓶颈。AR 模型不需要可逆架构,不要求建模微分同胚。它可以自然地建模不连续跳跃和分离模式——恰好是复杂多模态目标密度的特征。离散时间 NF 的拓扑约束在 AR 框架中完全消失。

2. 精确且高效的似然。Eq. 1 允许在单次前向传播中计算精确似然,无需雅可比行列式或昂贵的 Neural ODE 数值求解器。这直接满足 BG 中 SNIS(自归一化重要性采样)的核心需求。

3. 享受 LLM 投资红利。AR 模型在离散生成建模中已取得巨大成功——从文本(Gemini 2.5)到图像(ViT)。ARBG 展示了这些在大语言模型中验证的缩放行为,现在在分子构象采样领域同样成立:模型规模和推理样本数都呈现可预测的缩放律。

ARBG 的两个核心需求——快速准确的似然和可扩展性——都可以被自回归模型满足,但不一定被流模型满足。这是范式转换的根本动机。

📦条件分布参数化:从混合密度到均匀分箱

自回归模型在连续空间上的核心技术挑战是:如何参数化条件分布 p_θ(x_j | x_{?ARBG 探索了多条路径。

混合密度网络路径

第一种方法借鉴 PixelCNN++ 的先驱工作,将分子构象的条件分布建模为混合分布。具体做法是先将输入 x_j(例如一个原子的某个空间维度)离散化为 BΔ,区间为 I = [C_min, C_max)。然后假设一个潜在分配 kK 分量混合中采样,每个分量用逻辑分布(Mixture of Logistics, MoL)或高斯分布(GMM)参数化。

离散化概率质量通过 CDF 差分计算:p(X = x_j | x_{ 等于逻辑分布在 x̃_j + Δ/2x̃_j - Δ/2 处的 CDF 之差的混合加权和。这产生了 MoL-PixelCNN++ 和 GMM-PixelCNN++ 两种参数化方案。

然而,混合密度网络存在一个已知问题:混合权重 π_k 的模式坍缩。在训练中,混合分量容易坍缩到小子集,导致在更难的系统上性能次优。

均匀分箱参数化:分子 token

ARBG 的关键创新是完全抛弃混合模型,引入可能最简单的参数化方案:直接预测分箱中心 b_l 作为分类分布(Categorical distribution)。这允许直接使用自回归模型——与 LLM 的 next-token 预测完全相同——但现在用于分子数据。

训练时,模型预测每个维度的分箱索引,用交叉熵损失优化。推理时,从采样的分箱中心恢复连续坐标:x_j = b_l + u_l,其中 u_l ~ Unif(-Δ/2, Δ/2)。这诱导出一个分段常数条件密度:

p_θ(x̃_j | x_{

Δ 对所有分箱相同(均匀分箱)时,条件对数密度有一个常数偏移 log Δ,随分箱数增加而消失。这个参数化与 LLM 训练高度对齐,揭示了可预测的缩放行为。

命题 1量化了分箱内均匀去量化引入的内在误差:KL 散度的最低可达误差等于真实条件密度在每个分箱内与均匀分布的 KL 散度,按真实分箱质量加权求和。直觉上,这精确刻画了真实密度在每个分箱内不恰好均匀所导致的不可约误差。通过选择合适的分箱数,这个误差可以被有效控制。

⚙️自回归建模:残基级排序与Transformer架构

AR 模型需要一个维度排序。ARBG 采用残基级排序:每个残基的侧链原子紧跟其骨架原子之后。输入 x ∈ ℝ^{n×3} 被展平为 d = 3n 维向量,模型按此顺序逐维度生成。

这个设计选择有几个考量:

物理意义。残基是蛋白质的基本结构单元,按残基排序使得模型可以在残基边界处进行物理有效性检查(如空间位阻检测),为后续的 SMC 干预提供自然粒度。

与力场兼容。残基级粒度允许使用标准肽体力场进行部分能量评估,为扭曲 SMC 中的中间能量函数 E_s(·) 提供基础。

Transformer 架构。ARBG 使用与 LLM 相同的 Transformer 架构,包括 FlashAttention-2 加速和 Muon 优化器。模型规模从 8M 到 126M 参数,在 Chignolin 系统上展示了清晰的缩放律:更大的模型在相同 FLOPs 预算下达到更低的训练损失。

🔄扭曲 SMC:子结构级推理时干预

ARBG 的自回归分解解锁了一整套流模型无法实现的推理时干预工具。因为生成被分解为每个条件维度的离散步骤,我们可以在完整分子实现之前分析、纠正或丢弃部分构象。

核心实现是自回归扭曲序贯蒙特卡洛(Autoregressive Twisted SMC)。定义一系列扭曲函数 ψ_j(x_j),构造中间密度 η_j(x_j) := p_θ(x_j | x_{。目标不是直接从 p_θ(x_j | x_{ 采样,而是从扭曲的中间分布 η_j 采样。

扭曲函数使用部分能量评估定义:

ψ_j(x_j) = exp(-E_s(r_s(x_{≤j}))) / p_θ(r_s(x_{≤j}) | r_{s-1}(x_{≤j}))

其中 r(x_{≤j}) 定义了 x_{≤j} 中包含 x_j 的最大(截断)残基子集,E_s(·) 是原始能量函数 E(·) 的版本,作用于长度 s < d 的肽链。

这个扭曲函数允许我们在残基级别注入物理有效性约束。系统重采样在每个残基的原子坐标生成完毕后触发,使用标准肽体力场进行部分能量评估。这与先前基于流的重采样方法(如 SBG)根本不同:SBG 必须生成完整候选才能重采样,而 ARBG 允许子结构级引导——一旦检测到错误就立即纠正生成过程。

实验观察到,扭曲 SMC 仅边际优于 SNIS,这归因于 ROBIN 基础提案的高质量——模型已充分学习玻尔兹曼分布,SMC 校正收益递减。但在 8 肽系统上,约 7% 的样本最终部分能量超过 100 + E_min,表明在更大或更复杂的系统中,提案精度下降时 SMC 的价值将更加显著。

📊实验结果:从丙氨酸二肽到 Chignolin

实验在四个单肽系统上评估:三丙氨酸(AL3, 3残基)、四丙氨酸(AL4, 4残基)、六丙氨酸(AL6, 6残基)和 Chignolin(GYDPETGTWG, 10残基)。使用三个互补的 Wasserstein 指标评估:

E-W2(能量距离):测量生成与参考能量分布的一致性,对局部物理准确性敏感。

T-W2(扭转空间距离):在扭转角空间计算,捕捉全局构象差异和缺失模式。

TICA-W2:在时间滞后独立成分分析的低维空间中比较,覆盖最慢动力学模式。

核心发现

在温度调优后,ARBG 在所有系统的 E-W2 和 T-W2 上全面超越所有基线方法。即使不调优温度(T=1),ARBG 仍是整体最佳方法,尤其在更大系统(AL6 和 Chignolin)上优势明显。

在 10 残基 Chignolin 系统上——因 β-发夹二级结构而特别具有挑战性——ARBG 的表现尤为突出:

方法E-W2 ↓T-W2 ↓
SBG(先前 SOTA)10.819 ± 7.2063.778 ± 0.440
FALCON
ARBG (T=1)7.942 ± 1.0532.780 ± 0.105
ARBG (调优 T)1.723 ± 0.0752.632 ± 0.044

能量直方图显示,ARBG 的重加权能量分布与 MD 模拟数据高度吻合。Ramachandran 图表明模型准确捕捉了 Chignolin 几乎所有的构象模式。

MDN 基线(GMM-PixelCNN++ 和 GIVT)虽然略逊于 ARBG,但展现出有前景的结果,进一步验证了 AR 模型在 BG 中的整体潜力——相对于流模型的优势是系统性的。

🚀ROBIN:可迁移的零样本生成

ROBIN 是用 ARBG 框架训练的可迁移模型,132M 参数,通过额外条件信息实现零样本迁移到未见肽链。在 ManyPeptidesMD 数据集上评估,覆盖 30 条长度为 4 和 8 残基的不同序列。

8 肽系统结果(2×10⁵ 能量评估)

模型E-W2 ↓T-W2 ↓TICA-W2 ↓
TarFlow0.9290.7760.498
Prose(先前 SOTA)0.6460.6070.349
ROBIN0.5310.6490.379

ROBIN 在 8 肽系统上将 E-W2 从 0.646 降低到 0.531,改善超过 60%。TICA 图显示 ROBIN 对未见八肽 CGSWHKQR 的零样本预测与真实 MD 分布高度一致。

推理缩放行为

ROBIN 在推理样本数方面展现出有利的缩放行为:在 T-W2 指标上,ROBIN 用比 Prose 少一个数量级的样本、比 MD 少三个数量级的样本达到相同性能。这表明 ARBG 不仅在模型规模上缩放,也在推理预算上缩放——与 LLM 的推理时缩放特性一致。

🔥消融实验:分箱分辨率与采样温度

分箱分辨率

增加分箱数会提高模型生成坐标的粒度。在 AL3 系统上的消融实验验证:增加分箱数单调改善重采样 E-W2 性能。分箱数的上限实验(先离散化到固定分箱数,再用均匀噪声重建)确认了性能上界随分箱数增加而提升。

采样温度

推理时对 Transformer 采样温度进行消融。能量分布随温度系统性变化:低温时概率质量集中于高似然(低能量)模式,可能抑制或完全遗漏其他模式;高温时分布变平,鼓励多样性但可能过度采样模式并降低生成质量。温度扫描揭示了 AL4 的最优采样温度约为 T=1.02。更大的分子系统倾向于更低的最优温度。

💡技术启示与局限

核心启示

1. 微分同胚不是必要条件。ARBG 证明了玻尔兹曼生成器不需要可逆架构。自回归分解天然绕过了流模型的拓扑约束,为分子构象采样开辟了全新路径。

2. LLM 范式可直接迁移到连续分子空间。均匀分箱参数化将连续坐标转化为离散"分子 token",使 next-token 预测范式直接适用。这意味着 LLM 领域的技术积累——FlashAttention、Muon 优化器、推理时缩放——可以直接惠及分子建模。

3. 子结构级干预是流模型无法实现的。自回归的逐步生成天然支持在残基级别进行物理有效性检查和纠正,这为更大更复杂系统的高效采样提供了新工具。

4. 缩放律在分子领域同样成立。从 8M 到 126M 参数,ARBG 展示了与 LLM 一致的缩放行为——更大模型在相同计算预算下性能更好,更多推理样本带来更好的分布覆盖。

局限

维度排序。AR 模型要求特定排序,而分子本身没有自然排序。这个选择可能影响性能,尤其在更小的分子上。

分箱精度。均匀分箱将模型精度限制在 Δ 级别,在更大系统上可能面临更尖锐能量剖面的挑战。

信息先验。流模型可以利用信息先验(如 TFEP),AR 模型如何利用类似先验是未来工作方向。

🔬补充:TrajCast — 自回归等变网络的无力场分子动力学

在 ARBG 用自回归模型革新平衡态采样的同时,IBM Research 团队在 Nature Machine Intelligence(2026年5月)发表了 TrajCast——一个同样基于自回归范式但面向动力学轨迹生成的框架。两者从不同角度证明了自回归模型在分子模拟中的巨大潜力:ARBG 采样平衡分布,TrajCast 生成时间有序的轨迹。

核心思路:跳过力计算,直接预测下一状态

传统 MD 的计算瓶颈有两个:(1) 精确计算原子间力(需要求解薛定谔方程的近似);(2) 数百万步的数值积分(时间步长 δt 通常仅 0.5–1.0 fs)。机器学习原子间势(MLIP)如 MACE 加速了力的计算,但仍受限于小时间步长。

TrajCast 采取了根本不同的路径:使用等变消息传递神经网络(MPNN)直接预测时间 t + Δt 的原子位置和速度,完全跳过力计算和数值积分。预测步长 Δt 比传统 MD 时间步长大 10–30 倍。轨迹通过自回归展开生成——每一步的输出作为下一步的输入,保持了轨迹的马尔可夫性质和时间有序性。

架构设计:等变 MPNN 与动量守恒

TrajCast 的 MPNN 接收原子位置 r(t)、速度 v(t) 和原子类型 Z 作为输入,通过 T 个消息传递块逐步精炼特征。关键设计包括:

等变消息构建。消息通过将潜在特征与可学习径向基函数和球谐展开的边向量进行卷积构建,确保 SE(3) 等变性。

速度感知更新。在更新块中,来自邻居的消息通过张量积与速度向量在可学习的径向和球谐基中组合,使模型能利用当前速度信息预测下一状态——这是 TrajCast 与其他生成式方法(如 Timewarp、MDGen)的关键区别,后者仅传播原子位置而不保留动量。

守恒律约束。在 Readout 阶段后,通过调整位移和速度来强制总线性动量和角动量守恒,确保物理一致性。

恒温器集成。训练在 NVE 系综上进行,推理时通过 CSVR 恒温器在每步 MPNN 前向传播后重标度速度,实现 NVT 系综采样。

实验验证:从小分子到晶体到液体

TrajCast 在三个代表性系统上验证:

扑热息痛(小分子)。训练仅用 10,000 个 NVE 构型(<1 ns 轨迹),预测步长 7 fs(14× MD 步长)。生成 7 ns 稳定轨迹,振动密度态(VDOS)重叠分数 ~0.95,势能分布与 MD 高度吻合。自由能面(FES)准确捕捉了顺式/反式构型,值得注意的是训练集中不包含 ∣ψ₂∣ > π/2 的构型,但 TrajCast 在展开中正确生成了翻转状态并恢复了相应自由能势垒。

α-石英(晶体)。训练仅用 5,000 个样本(162 原子系统),预测步长 30 fs(30× MD 步长)。生成 15 ns 轨迹(50万步),VDOS 重叠分数 >0.95,能量分布几乎完全一致。在 3×3×3 超胞(~4,300 原子)上每天生成近 20 ns 轨迹——相比之下,MACE MLIP 对 1,000 原子系统每天约 1 ns。

液态水(液体)。训练 5,000 个样本(64 分子),预测步长 5 fs(10× MD 步长)。VDOS 重叠分数 ~0.99,扩散系数 2.18 vs 1.84 × 10⁻⁹ m²/s(TrajCast vs MD),径向分布函数(RDF)在所有元素对上高度一致。在 3×3×3 超胞(>5,000 原子)上每天生成 >1 ns。

零样本泛化:超淬火玻璃态水

TrajCast 最令人印象深刻的实验是零样本泛化:仅用室温液态水训练的模型,直接用于预测从 300 K 冷却到 180 K 的超淬火玻璃态水(HGW)形成过程。四条独立轨迹(每条 12 ns、240万步)全部保持稳定,扩散系数在全温度范围内与 MD 吻合,O-O RDF 几乎不可区分,结构序参数(τ-Q 平面)遵循与 MD 相同的演化路径。这表明 TrajCast 能外推到训练分布之外未见过的构型和速度空间。

数据效率与缩放

TrajCast 的数据效率令人瞩目:学习曲线遵循幂律,缩放指数 |s| 高达 MLIP 的 4 倍。这归因于将速度纳入输入——每个构型携带更多信息。仅需几百皮秒的轨迹数据即可训练出准确模型,而 Timewarp、MDGen 等方法需要数百纳秒训练数据。

ARBG 与 TrajCast:自回归分子建模的两翼

ARBG 和 TrajCast 共同勾勒出自回归模型在分子模拟中的两大应用方向:

维度ARBGTrajCast
目标平衡态分布采样时间有序轨迹生成
自回归对象原子坐标维度时间步
架构Transformer(LLM 范式)等变 MPNN(GNN 范式)
物理先验玻尔兹曼能量 + SNISSE(3) 等变 + 动量守恒
输出无序平衡构象位置 + 速度(完整相空间)
零样本能力ROBIN:未见肽链未见温度区间(玻璃态)
时间信息不保留保留(可计算动力学性质)
缩放律模型规模 + 推理样本训练数据量(幂律)

两者互补而非竞争:ARBG 适合需要平衡系综采样的任务(如自由能计算、构象搜索),TrajCast 适合需要动力学信息的任务(如扩散系数、振动谱、非平衡过程)。两者都证明了自回归模型可以绕过传统分子模拟的根本瓶颈——ARBG 绕过拓扑约束,TrajCast 绕过小时间步长约束——共同指向自回归分子建模的新范式。

TrajCast 论文信息

标题:Force-free molecular dynamics through autoregressive equivariant networks

作者:Fabian L. Thiemann, Thiago Reschützegger, Massimiliano Esposito, Tseden Taddese, Juan D. Olarte-Plata & Fausto Martelli

机构:IBM Research Europe (Daresbury), Microsoft Research AI for Science (Cambridge), IBM Research (Rio de Janeiro), STFC Hartree Centre, CNR-Istituto dei Sistemi Complessi (Rome)

发表:Nature Machine Intelligence, Volume 8, May 2026, pp. 764–776

DOI10.1038/s42256-026-01227-7

📄论文信息

标题:Autoregressive Boltzmann Generators

作者:Danyal Rehman, Charlie B. Tan, Yoshua Bengio, Avishek Joey Bose, Alexander Tong

机构:Mila – Québec AI Institute, Broad Institute of MIT & Harvard, Aithyra, Université de Montréal, University of Oxford, Imperial College London

发表:ICML 2026(Proceedings of the 43rd International Conference on Machine Learning, Seoul, South Korea, PMLR 306)

代码github.com/danyalrehman/autobg

arXiv:2606.27361v1