跳转至

Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps

会议: NeurIPS 2026
arXiv: 2606.29110
论文: https://arxiv.org/abs/2606.29110
代码: https://github.com/oxai/scallop
领域: 物理/科学计算
关键词: 玻尔兹曼生成器, 流映射蒸馏, 条件散度匹配, 似然估计, 分子构象采样

一句话总结

针对连续流模型似然估计依赖高方差 Hutchinson 迹估计或高阶雅可比行列式的计算瓶颈，本文提出无迹估计且可向量化的条件散度匹配目标 SCALLOP，实现分子玻尔兹曼生成器在数步函数评估下的高效采样与精确似然协同推断。

研究背景与动机

在计算化学、生物物理与材料科学中，从非归一化的玻尔兹曼分布中采样以精确计算宏观热力学可观测量（如自由能差）是核心任务。传统方法多采用玻尔兹曼生成器（Boltzmann Generators）作为提议分布，通过自归一化重要性采样（SNIS）对生成构象进行重加权。这一计算范式不仅要求生成模型能够输出高质量且多样化的原子空间结构，更对模型在单次前向中快速、精确评估样本概率密度（即模型似然）提出了极为严苛的要求。

然而，现有的生成模型在表达能力与似然可计算性之间存在深刻矛盾。经典的离散归一化流（Normalizing Flows）受限于严格的双射网络架构与自回归变换，在处理具有复杂内部对称性（如 SE(3) 等变性）的多原子分子体系时容量严重不足；而无约束的连续正规化流（CNF）与扩散模型虽然生成质量优异，但其瞬时变量替换公式要求沿积分轨迹对速度场的散度进行全空间积分。在实际工程中，散度计算通常依赖 Hutchinson 迹估计器，其方差随数据维度呈二次方急剧膨胀，不仅极大拖慢训练收敛速度，更在有限步推理时带来严重估计偏差。近期基于一致性流映射（Flow Maps）的蒸馏方案虽将采样压缩至极少步数，但要么如 FALCON 需付出 \(\mathcal{O}(d^3)\) 的雅可比行列式计算代价，要么如 F2D2 在蒸馏目标中依然深陷 Hutchinson 估计的高方差泥潭。

针对该计算瓶颈，本文从概率流动 ODE 的全导数与连续性方程出发，探讨能否完全绕过高阶微分与随机迹估计，构造出无偏、低方差且易于并行向量化的条件化似然蒸馏目标。核心 idea：基于全导数链式法则与时空得分条件化恒等式，提出免 Hutchinson 估计的向量化条件散度匹配目标（CDM-v），将雅可比对角元回归解构为闭式高斯条件项，与流映射自蒸馏联合训练实现数步内的样本与对数似然共评估。

方法详解

整体框架

SCALLOP 构建在似然流映射（Likelihood Flow Map）算子之上，将空间样本输运与沿轨对数密度累积联合建模。模型接收噪声样本与时间区间作为输入，通过共享参数的主干网络提取特征，并在输出端分别解耦为空间位移速度预测头与沿时间维度的对数密度变化预测头。在训练阶段，SCALLOP 彻底摒弃了依赖雅可比-向量积（JVP）的 Hutchinson 迹估计，引入具有闭式解的向量化条件散度匹配损失（CDM-v）与流匹配损失联合约束瞬时端点行为，并借助拉格朗日自蒸馏（LSD）损失将多步积分一致性内化至单步前向传播中。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["先验高斯噪声输入 x_0 与时间对 (t, s)"] --> B["共享 DiT 主干特征抽取网络"]
    B --> C["1. 双头跳跃连接似然流映射参数化<br/>空间位移 u_θ 与向量化对数密度头 D_θ"]
    C --> D["2. 免迹估计条件散度匹配 (CDM-v)<br/>全导数展开解耦为解析高斯条件目标"]
    D --> E["3. 雅可比对角线向量化回归<br/>将标量散度升维至 d 维对角元监督"]
    E --> F["4. 分子 SE(3) 软对称性保持<br/>质心归零平移等变与随机旋转数据增强"]
    F --> G["5. 数步自归一化重要性采样推断<br/>Euler 迭代输出平衡态构象与精确密度"]

关键设计

1. 双头跳跃连接似然流映射参数化：协同生成构象坐标与伴随密度

为了在少步函数评估（NFE）下同时获取分子三维坐标及其对应的生成概率，SCALLOP 定义了增广状态流映射算子 \(\Psi: \mathbb{R}^d \times [0, 1]^2 \to \mathbb{R}^{d+1}\)。该算子通过带尺度的跳跃连接结构进行参数化： $\(\Psi_\theta(x_t, t, s) = \begin{bmatrix} x_t \\ 0 \end{bmatrix} + (s - t) \begin{bmatrix} u_\theta(x_t, t, s) \\ D_\theta(x_t, t, s) \end{bmatrix}\)$ 其中主干网络采用扩散 Transformer（DiT）模块，前 6 层作为空间几何与时序嵌入的共享表示提取器，后接各 3 层的专用 DiT 读出头分别输出速度场 \(u_\theta\) 与密度变化率 \(D_\theta\)。当积分步长趋近于零（\(s \to t\)）时，网络输出自然退化为瞬时速度场 \(v_t(x_t)\) 与瞬时负散度 \(-\nabla \cdot v_t(x_t)\)；当 \(s > t\) 时，网络直接拟合有限步时间跨度内的积分位移与总对数似然增量 \(\log p_s(x_s) - \log p_t(x_t)\)，彻底消除了传统 ODE 求解器数十至数百步细粒度数值步进的需求。

2. 免迹估计条件散度匹配：全导数链式法则消除高阶微分方差

传统流模型在拟合瞬时负散度时被迫采用 Hutchinson 估计器 \(\mathbb{E}_\epsilon [\epsilon^\top \nabla v_t(x_t) \epsilon]\)，该采样方差随体系自由度维度急速上升。SCALLOP 从时间流动轨迹上对数密度的全导数出发，利用链式法则建立等式： $\(\frac{\mathrm{d}}{\mathrm{d}t} \log p_t(x_t) = \nabla \log p_t(x_t) \cdot v_t(x_t) + \partial_t \log p_t(x_t)\)$ 结合线性插值路径 \(x_t = (1-t)x_0 + t x_1\)（其中 \(x_0 \sim \mathcal{N}(0, I)\) 为高斯先验基底），在给定终点样本 \(\xi = x_1\) 的条件下，残存的高斯噪声源使得空间得分 \(\nabla \log p_t(x_t|\xi)\) 与时间得分 \(\partial_t \log p_t(x_t|\xi)\) 均可推导为简单的闭式解析解。在此基础上构建的条件散度匹配（CDM）目标： $\(\mathcal{L}_{\mathrm{CDM}}(\theta; t) = \mathbb{E}_{x_0, x_1} \left[ \left| D_\theta(x_t, t, t) + \frac{x_0}{1-t} \cdot \left( v_t(x_t) - (x_1 - x_0) \right) - \frac{d}{1-t} \right|^2 \right]\)$ 不仅在期望意义下与标准散度匹配损失的参数梯度完全严格等价，而且在整个训练回路中无需进行任何雅可比-向量积运算，彻底摆脱了 Hutchinson 随机投影的方差放大效应。

3. 雅可比对角线向量化回归：将一维标量信号升维为多维稠密监督

标量散度作为雅可比矩阵迹的求和，直接回归极易因标量聚合导致多维空间局部梯度信号相互抵消、数据利用效率低下。为此，SCALLOP 提出向量化条件散度匹配技术（CDM-v），将密度变化预测头 \(D_\theta\) 的输出维度从标量 \(\mathbb{R}\) 扩展至与系统自由度相同的向量空间 \(\mathbb{R}^d\)。网络学习的目标被重新定义为拟合雅可比矩阵的负对角线元素向量 \(-\mathrm{diag}(\nabla v_t(x_t))\)（记作 \(-\nabla \odot v_t(x_t)\)）： $\(\mathcal{L}_{\mathrm{CDM\text{-}v}}(\theta; t) = \mathbb{E}_{x_0, x_1} \left[ \left\| D_\theta(x_t, t, t) + \frac{x_0}{1-t} \odot \left( v_t(x_t) - (x_1 - x_0) \right) - \frac{\mathbf{1}_d}{1-t} \right\|^2 \right]\)$ 该设计不仅为模型提供了维度丰富的信息监督流，与空间坐标预测头具备同构的张量结构，而且在推理阶段只需对向量输出执行内积求和（\(\mathbf{1}_d \cdot D_\theta\)）即可精确恢复全系统标量散度，从根本上加快了网络在微观自由度上的训练收敛速度。

4. 分子 SE(3) 软对称性保持：平移等变与旋转不变的物理几何对齐

在多肽分子的构象空间中，能量面关于全局刚体旋转平移具有刚性物理不变性。SCALLOP 针对分子坐标预先执行质心对齐（Center-of-Mass Zeroing）消除平移自由度，并在训练循环每次采样批次构象时施加三维随机旋转矩阵。空间位移头 \(u_\theta\) 与向量化密度头 \(D_\theta\) 均被约束为旋转等变映射，而在计算标量总密度变化时对 \(D_\theta\) 各原子坐标分量求和，数学上严格满足对标量对数似然的 SE(3) 旋转不变性，避免了非物理刚体变换引起的概率密度震荡。

损失函数 / 训练策略

SCALLOP 的总训练损失综合了瞬时边界损失、向量化条件散度匹配损失与拉格朗日自蒸馏损失（LSD）。针对时间端点 \(t \to 1\) 处条件方差趋于无穷的数学奇异性，算法设置截止截断区间 \([0, 1 - \mathrm{eps}]\)（\(\mathrm{eps} = 10^{-5}\)）并乘以 \((1-t)\) 线性阻尼因子。整体训练目标函数定义为： $\(\mathcal{L}_{\mathrm{SCALLOP}}(\theta) = \int_0^1 \Bigl[ e^{-w_\psi(t,t)} \left( \mathcal{L}_v(\theta; t) + (1-t) \mathcal{L}_{\mathrm{CDM\text{-}v}}(\theta; t) \right) + w_\psi(t,t) \Bigr] \mathrm{d}t + \int_0^1 \int_0^s \Bigl[ e^{-w_\psi(t,s)} \mathcal{L}_{\mathrm{SD}}(\theta; t, s) + w_\psi(t,s) \Bigr] \mathrm{d}t \mathrm{d}s\)$ 其中 \(\psi\) 为可学习的自适应方差加权小网络。自蒸馏项采用拉格朗日约束 \(\mathcal{L}_{\mathrm{LSD}}\)，在反向传播中使用停止梯度算子（Stop-Gradient）稳固教师轨迹目标。模型采用 AdamW 优化器（初始学习率 \(10^{-4}\)，权重衰减 \(10^{-4}\)，梯度截断阈值 10.0），配合前 \(2.5\%\) 步数的预热与平方根退火调度，利用指数移动平均（EMA，衰减率 0.999）跟踪模型参数。

实验关键数据

主实验

论文在四种不同分子长度的丙氨酸多肽系统（ALA-2 至 ALA-6，构象自由度维度从 \(d=66\) 扩展至 \(d=189\)）上开展了系统性评估。通过生成 \(2 \times 10^5\) 个分子构象并执行自归一化重要性采样（SNIS），对比有效样本量（ESS，越高越好）、能量分布 2-Wasserstein 距离（\(\mathcal{E}\text{-}\mathcal{W}_2\)，反映局部键长/键角精细松弛，越低越好）以及主链扭转角分布 2-Wasserstein 距离（\(\mathbb{T}\text{-}\mathcal{W}_2\)，反映宏观构象亚稳态分布，越低越好）。

体系 (维度) 指标 SCALLOP (本文) F2D2 FALCON FALCON-A SBG-IS NFE
ALA-2 (\(d=66\)) \(\mathrm{ESS} \uparrow\) 0.088 ± 0.001 0.066 ± 0.001 0.067 ± 0.013 0.097 ± 0.007 0.030 ± 0.012 4
\(\mathcal{E}\text{-}\mathcal{W}_2 \downarrow\) 0.555 ± 0.116 0.958 ± 0.114 0.225 ± 0.104 0.512 ± 0.038 0.873 ± 0.338 4
\(\mathbb{T}\text{-}\mathcal{W}_2 \downarrow\) 0.310 ± 0.027 0.383 ± 0.022 0.402 ± 0.021 0.180 ± 0.005 0.439 ± 0.129 4
ALA-3 (\(d=99\)) \(\mathrm{ESS} \uparrow\) 0.088 ± 0.001 0.079 ± 0.001 0.077 ± 0.004 0.104 ± 0.004 0.052 ± 0.013 8
\(\mathcal{E}\text{-}\mathcal{W}_2 \downarrow\) 1.596 ± 0.028 2.155 ± 0.033 0.544 ± 0.013 1.385 ± 0.182 0.758 ± 0.506 8
\(\mathbb{T}\text{-}\mathcal{W}_2 \downarrow\) 0.212 ± 0.001 0.219 ± 0.001 0.452 ± 0.011 0.343 ± 0.004 0.502 ± 0.016 8
ALA-4 (\(d=126\)) \(\mathrm{ESS} \uparrow\) 0.067 ± 0.001 0.053 ± 0.001 0.055 ± 0.003 0.094 ± 0.007 0.046 ± 0.014 8
\(\mathcal{E}\text{-}\mathcal{W}_2 \downarrow\) 1.545 ± 0.033 4.209 ± 0.066 0.686 ± 0.047 2.929 ± 0.068 1.068 ± 0.495 8
\(\mathbb{T}\text{-}\mathcal{W}_2 \downarrow\) 1.399 ± 0.039 1.435 ± 0.023 0.858 ± 0.077 1.094 ± 0.034 0.969 ± 0.067 8
ALA-6 (\(d=189\)) \(\mathrm{ESS} \uparrow\) 0.042 ± 0.001 0.033 ± 0.001 0.060 ± 0.017 0.077 ± 0.007 0.034 ± 0.015 16
\(\mathcal{E}\text{-}\mathcal{W}_2 \downarrow\) 5.845 ± 0.101 10.140 ± 0.065 0.892 ± 0.311 1.211 ± 0.105 1.021 ± 0.239 16
\(\mathbb{T}\text{-}\mathcal{W}_2 \downarrow\) 1.120 ± 0.014 1.143 ± 0.039 1.256 ± 0.132 1.163 ± 0.112 1.431 ± 0.085 16

图像域高维扩展与消融实验

为了进一步验证条件散度匹配在复杂高维自然数据上的泛化能力，论文在 CelebA-64 无条件图像生成基准（维度 \(3 \times 64 \times 64 = 12288\)）上对比了不同 Euler 步数下的样本质量（FID）与负对数似然估计（BPD，每维度比特数）。参考教师流匹配模型在 1024 步时达到 BPD 1.75，在 200 步时达到 FID 2.48。其中 \(|\Delta|\) 表示相对于教师 BPD 的绝对偏差，\(\mathrm{Err}\) 表示样本级 BPD 估计的平均绝对误差。

| 步数 | 模型 | NLL (BPD) | 绝对偏差 \(|\Delta| \downarrow\) | 逐样本误差 \(\mathrm{Err} \downarrow\) | 样本质量 \(\mathrm{FID} \downarrow\) | |---|---|---|---|---|---| | 1 步 | Flow Map (无密度头) | -69.83 | 71.58 | 71.64 | 12.96 | | | F2D2 (Hutchinson 蒸馏) | 1.64 | 0.11 | 0.53 | 6.94 | | | SCALLOP (CDM-v 蒸馏) | 1.60 | 0.15 | 0.17 | 6.91 | | 2 步 | Flow Map (无密度头) | -32.72 | 34.47 | 34.46 | 6.32 | | | F2D2 (Hutchinson 蒸馏) | 1.73 | 0.02 | 0.33 | 3.86 | | | SCALLOP (CDM-v 蒸馏) | 1.69 | 0.06 | 0.12 | 3.81 | | 4 步 | Flow Map (无密度头) | -14.89 | 16.64 | 16.59 | 4.04 | | | F2D2 (Hutchinson 蒸馏) | 1.75 | 0.00 | 0.24 | 2.75 | | | SCALLOP (CDM-v 蒸馏) | 1.75 | 0.00 | 0.12 | 2.70 | | 8 步 | Flow Map (无密度头) | -6.78 | 8.53 | 8.45 | 3.33 | | | F2D2 (Hutchinson 蒸馏) | 1.64 | 0.11 | 0.18 | 2.41 | | | SCALLOP (CDM-v 蒸馏) | 1.68 | 0.07 | 0.10 | 2.37 |

关键发现

  • 训练方差暴降与稳定性提升：在 ALA-2 动力学训练监测中，由于完全移除了 Hutchinson 迹估计中的随机投影噪声，SCALLOP 的梯度信噪比（SNR）相比 F2D2 获得显著提升，损失梯度方差降低高达近 100 倍，不仅收敛更平稳，而且在更少的 GPU 训练时长内达到了更优的 ESS 极限。
  • 高维扩展与逐样本似然保真度：在 CelebA-64 高维视觉实验中，SCALLOP 计算单步损失的耗时仅为 F2D2 的 83.1%。在 1 至 8 步的极简采样推断下，SCALLOP 将单样本对数似然误差 \(\mathrm{Err}\) 从 F2D2 的 0.53/0.33/0.24 压缩至 0.17/0.12/0.12，证明向量化条件散度匹配能够精准捕捉局部高维概率密度曲率。
  • 构象宏观保真度与推断吞吐优势：在分子主链扭转角分布指标 \(\mathbb{T}\text{-}\mathcal{W}_2\) 上，SCALLOP 在 ALA-3 和 ALA-6 上均斩获全基准最优成绩（分别达到 0.212 与 1.120）。推断测速表明，FALCON 由于依赖昂贵的高阶全雅可比行列式计算，在 ALA-6 上仅能支持 100 的微小批次；而 SCALLOP 支持全量 \(10^4\) 大批次并行前向推断，实现了相比 FALCON 100 倍的推断加速，并比前沿归一化流基准 SBG 快 10 倍以上。

亮点与洞察

  • 从动力学全导数视角重构似然匹配：将对数概率的时间演化拆分为空间对流项与显式时间偏导项，并巧妙利用已知高斯先验基底将两者显式解构为无需反向传播求导的封闭代数解析形式，数学形式极为优雅。
  • 雅可比矩阵对角元向量化监督：敏锐地指出了标量散度回归存在的信息瓶颈，创新性地把回归目标升维到雅可比对角向量，不仅统一了位移头与密度头的结构维度，而且极大丰富了多维动力学系统的梯度反传信息。
  • 玻尔兹曼生成器推断范式解耦：彻底破除了玻尔兹曼生成器在“归一化流的严格拓扑受限”与“连续流模型推断受阻于雅可比行列式”之间的两难困境，使扩散与流匹配模型能够以近乎无额外开销的单次前向输出精确的非归一化重加权密度。

局限与展望

  • 对速度场预训练精度的理论依赖：条件散度匹配损失的核心成立条件依赖于真实边缘速度场 \(v_t\) 已被完美拟合（Assumption 1）。当第一阶段教师流模型的速度场存在逼近误差时，由此引发的分布漂移如何理论量化与向后累积，目前仍缺乏严格的泛化误差上界分析。
  • 极少步数与超高维扩展的折中：尽管模型在理论上支持 1 步超快速联合采样与似然输出，但在分子复杂高维多肽或大型生物大分子体系中，实际生成往往仍需 \(8 \sim 16\) 步数方能平滑微观碰撞，完全无损的一步玻尔兹曼采样生成仍需进一步突破。
  • 局部构象微观松弛的后续优化：在高自由度大分子中，SCALLOP 的能量分布 Wasserstein 距离 \(\mathcal{E}\text{-}\mathcal{W}_2\) 虽大幅超越 F2D2，但仍略高于引入精确重构惩罚的 FALCON，未来可结合微观力场局部几何弛豫（Local Relaxation）后处理进一步平抑高能量重叠冲突。

相关工作与启发

  • vs F2D2 (Ai et al., 2026): F2D2 首次建立了似然流映射的双头自蒸馏框架，但其散度瞬时匹配完全依赖随机 Hutchinson 迹估计；本文提出的 SCALLOP 通过全导数展开与条件化技巧实现了免 Hutchinson 的向量化散度匹配，训练梯度方差暴降百倍且训练速度提升近 17%。
  • vs FALCON (Rehman et al., 2026): FALCON 通过正则化约束流映射的可逆性，但推断计算似然时必须显式计算雅可比矩阵行列式（复杂度达 \(\mathcal{O}(d^3)\)），内存占用随原子数激增；SCALLOP 将似然计算解耦为轻量 DiT 读出头的前向求和，推断速度提升高达两个数量级。
  • vs SBG (Tan et al., 2026a): 基于离散归一化流的 SBG 依赖自回归或耦合层架构，模型容量与对称性受限；SCALLOP 释放了连续流映射与 DiT 架构的强大几何表达力，在取得更优扭转角分布精度的同时推断速度提升 10 倍。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用全导数与条件得分规避散度高阶微分，提出免迹估计的向量化蒸馏目标，理论推导严密且立意深刻。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖不同长度的四种丙氨酸多肽系统及高维 CelebA 图像基准，包含详尽的构象几何分布、能量直方图、训练动力学稳定性与吞吐测速分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨连贯，理论推导步步为营，数学证明在附录中交代极为完整，图表展现信息充实。
  • 价值: ⭐⭐⭐⭐⭐ 攻克了玻尔兹曼生成器在复杂物理分子体系中兼顾高表达力与极速精确似然评估的关键瓶颈，具有重要的理论与工程实用价值。

相关论文