Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels¶
会议: NeurIPS 2026
arXiv: 2606.16073
OpenReview: https://openreview.net/forum?id=zB413mSyhL
代码: https://github.com/kkorolev1/stop-the-sampler
领域: 优化/理论
关键词: 采样算法, 非无环生成流网络, 马尔可夫链蒙特卡洛, 自适应提前终止, 朗之万动力学
一句话总结¶
将连续状态空间下的 MCMC 采样过程表述为带非无环结构的生成流网络（Non-Acyclic GFlowNet），通过训练神经网络分类器自适应学习轨迹终止决策与转移漂移修正，在大幅缩短采样步数的同时显著改善复杂多模态分布的混合与覆盖。
研究背景与动机¶
从复杂未归一化概率密度函数中高效采样是贝叶斯推断、统计物理与生成建模的核心基础问题。传统马尔可夫链蒙特卡洛（MCMC）方法，如不调整朗之万算法（ULA）与哈密顿蒙特卡洛（HMC），虽然具备渐近收敛保证，但在高维或多模态能量曲面上极易遭遇混合缓慢与维数灾难。更为棘手的是，传统 MCMC 通常依赖人为预设的固定迭代步长与固定轨迹长度，或者依赖 Gelman-Rubin \(\hat{R}\) 与 Geweke 等被动的后验收敛统计量。这些统计检验不仅往往需要极长的样本窗口才能触发，且阈值激进时极易带来不可逆的截断估计偏差，缺乏主动感知局部几何形态并做出终止决策的能力。
另一方面，基于正规化流（Normalizing Flows）与扩散采样器（Diffusion Samplers）的摊销生成模型通过将计算前置于离线训练来换取实时推理效率，但往往依赖预设的时间网格或固定扩散反演步数。虽然连续流网络（Continuous GFlowNets）为概率密度采样提供了全新的流平衡框架，但经典流网络普遍建立在有向无环图（DAG）与带有离散时间步扩展的状态空间假设上，无法自然容纳具有循环往复探索行为的动力学采样过程。
面对传统采样链与现代流网络之间的割裂，本文切入的角度是将动力学的终止行为本身看作采样器内部可学习的原生组件。核心 idea：将连续采样核的随机轨迹形式化为非无环生成流网络（Non-Acyclic GFlowNet），利用状态相关的神经分类器学习跳出流并进入吸收终态的自适应终止策略，并联合学习前向漂移与反向转移核，实现高密度区域极速停机与低密度区域高效跨模态探索的协同。
方法详解¶
整体框架¶
本文构建了一个非无环连续状态空间 \(\bar{\mathcal{S}} = \{s_0\} \cup \mathcal{S} \cup \{s_f\}\)，其中 \(\mathcal{S} = \mathbb{R}^d\) 为连续样本空间，\(s_0\) 与 \(s_f\) 分别表示唯一的源状态与吸收汇状态。与传统将时间步作为状态一部分的 DAG GFlowNet 不同，本文剥离了显式时间变量，允许粒子在连续空间 \(\mathbb{R}^d\) 内任意游走成环。生成过程自初始分布 \(p_0(s)\) 注入，每一步由前向分类器 \(d_F(s) \in [0, 1]\) 判定是否直接转入终态 \(s_f\)，若不终止则以 \(1 - d_F(s)\) 的概率调用连续转移核 \(q_F(s' \mid s)\) 推进到下一个连续状态。
在多模态与高维几何场景下，单层网络直接从高斯先验拉扯到复杂目标分布容易陷入局部极小。为此，论文进一步扩展为多层退火架构（Multilevel Generalization），设置 \(L\) 个退火层级 \(\ell \in \{1, \dots, L\}\)，每层定义几何插值中间奖励，并在每一层内部通过分类器自主决定是在当前层继续探索还是晋级到下一层，从而形成层次递进的自适应退火采样流。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["源状态 s0"] -->|"初始先验 p0(s)"| B["连续游走状态 s"]
B -->|"前向分类器决策"| C{"是否终止？"}
C -->|"终止概率 dF(s)"| D["汇状态 sf (采样完成)"]
C -->|"继续概率 1 - dF(s)"| E["漂移修正转移核<br/>qF(s' | s)"]
E -->|"循环探索更新"| B
subgraph 多层退火机制
direction TB
F["第 1 层：平滑目标"] -->|"分类器晋级"| G["第 ℓ 层：中间退火态"]
G -->|"分类器晋级"| H["第 L 层：真实目标分布"]
end关键设计¶
1. 非无环流网络与流测度紧确界：建立流守恒与平均轨迹长度的显式对应
为了摆脱离散时间步的有向无环限制，本文将粒子在 \(\mathbb{R}^d\) 内的遍历占据测度形式化定义为关于前向转移核 \(P_F\) 的连续势流测度： $\(F(\{s_0\}) = F(\{s_f\}) = Z, \quad F(A) = Z \sum_{n=0}^\infty P_F^n(s_0, A), \quad \forall A \subseteq \mathbb{R}^d\)$ 在细致平衡条件（Detailed Balance）下，论文证明了由此诱导的流测度严格满足局部流守恒条件 \(F(A) = \int_{\bar{\mathcal{S}} \setminus \{s_f\}} F(\mathrm{d}s) P_F(s, A)\) 与奖励匹配条件 \(r(\mathrm{d}x) = F(\mathrm{d}x) P_F(x, \{s_f\})\)。更为重要的是，不同于以往非无环理论仅给出松散的上界，本文证明了采样器的平均轨迹长度恰好等于非终态总流与源流的比值： $\(\mathbb{E}_{\tau \sim P}[n_\tau] = \frac{F(\mathcal{S})}{F(\{s_0\})}\)$ 这一紧确等式不仅在理论上确立了采样步数与全空间占据流的严格代数关系，更从数学本质上为引入流正则化项以压缩采样轨迹长度提供了唯一无偏的理论依据。
2. 分类器自适应终止机制：揭示最优分类器与目标密度的对偶关系
本文引入前向分类器 \(d_F(s)\) 与反向分类器 \(d_B(s)\)，分别表示在状态 \(s\) 终止并跳入 \(s_f\) 的概率，以及从 \(s\) 逆向回退到源点 \(s_0\) 的概率。通过细致平衡推导，最优前向分类器直接由目标未归一化密度 \(r(s)\) 与状态流密度 \(f(s)\) 的比值决定： $\(f(s) = \frac{r(s)}{d_F(s)} = \frac{Z p_0(s)}{d_B(s)} \implies d_F(s) = \frac{\pi(s)}{p_0(s)} d_B(s)\)$ 当将转移核 \(Q_F\) 固定为几何各态历经的基准 MCMC 核（如 ULA）且平稳分布为 \(\pi_Q\) 时，论文定理 3.6 进一步证明，可解细致平衡的充要条件是流测度密度满足 \(f(s) = Z(u(s) + n_Q \pi_Q(s))\)，其中常数 \(n_Q\) 严格等价于平均轨迹长度，其理论极限下界为 \(n_Q^* = \sup_{s} (\max\{\pi(s), p_0(s)\} - u(s)) / \pi_Q(s)\)。这意味着分类器不仅充当截断门控，其实际扮演了消除固定核转移偏差（de-biasing）的重赋权算子，并在几何高密度区域自动赋予高终止概率，形成天然的目标密度轮廓感知。
3. 可微漂移与方差修正核：攻克传统朗之万跨模态转移受阻瓶颈
若仅训练终止分类器而保持固定 ULA 连续核，粒子在深势阱与高势垒之间依旧受到扩散局限。为此，本文将前向连续转移核 \(q_F(s' \mid s; \theta)\) 参数化为增强型高斯动力学： $\(\mu_F(s; \theta) = s + \gamma \Bigl(u_F(s; \theta) + \bigl(\mathbf{1} + v_F(s; \theta)\bigr) \odot \nabla \log r(s)\Bigr)\)$ $\(\Sigma_F(s; \theta) = 2\gamma \operatorname{diag}\Bigl(\exp\{C_F \tanh(w_F(s; \theta))\}\Bigr)\)$ 其中 \(u_F(s; \theta)\) 为神经网络学习的全局漂移偏置向量，\(v_F(s; \theta)\) 为局部几何分量缩放，而 \(w_F\) 动态调制扩散噪声方差。这一参数化在保持一阶梯度引导的同时，赋予采样动力学跨越平坦低密度区与陡峭势垒的主动动量，彻底摆脱了传统 Langevin 动力学受困于单一模态的局限。对应地，反向核 \(q_B\) 参数化为带可学习漂移衰减与方差调制的 Ornstein-Uhlenbeck 反向破环过程，与前向扩散形成完备闭环。
4. 多层几何退火扩展：平滑高维复杂能量曲面的探索路径
针对维度达到数十乃至数百维的高难度分布，论文构建了由 \(L\) 个层级组成的退火链，状态扩充为 \((s, \ell)\)，各级中间目标由线性插值退火因子 \(\beta_\ell \in [0, 1]\) 给出： $\(r(s, \ell) = p_0(s)^{1 - \beta_\ell} r(s)^{\beta_\ell}, \quad 0 = \beta_0 < \beta_1 < \dots < \beta_L = 1\)$ 网络骨干通过引入正弦位置编码区分层级变量 \(\ell\)。采样器在层级 \(\ell\) 内部通过分类器自适应决定何时达成当前分布平稳并晋级至 \(\ell+1\)。与扩散模型中每步强制推进到下一个固定时间步的机械设定截然不同，该机制允许粒子在陡峭复杂的相变层级驻留更多自适应步数充分混合，而在容易平移的平缓区域快速跨越，兼具退火采样（AIS）的渐进稳定性与流网络的自适应效率。
损失函数 / 训练策略¶
为解决连续流网络中细致平衡（DB）优化极度不稳定、子轨迹平衡（SubTB）计算复杂度随轨迹长度呈二次方膨胀的痛点，模型采用高效的前缀轨迹平衡损失（Prefix Trajectory Balance）。
在训练采集阶段，为保证粒子能够充分遍历全空间，模型主动停用分类器，强制以无截断方式前向生成长度为 \(N_{\max}\) 的长探索轨迹。随后，算法提取每一个前缀子轨迹 \(\tau_{0:i} = (s_0 \to s_1 \dots \to s_i \to s_f)\)，赋予伪终止跳步，并计算加权前缀轨迹平衡损失与流正则化项： $\(\mathcal{L}(\theta) = \frac{1}{B} \sum_{b=1}^B \sum_{i=1}^{N_{\max}} w_i^b \left[ \mathcal{L}_{\mathrm{TB}}(\theta, \tau_{0:i}^b) + \rho \frac{r(s_i^b)}{d_F(s_i^b; \theta)} \right]\)$ 其中前缀权重 \(w_i = \operatorname{sg}[d_F(s_i; \theta) \prod_{t=1}^{i-1}(1 - d_F(s_t; \theta))]\) 为当前策略诱导的截断概率边缘分布，\(\operatorname{sg}\) 为梯度截断算子，使得损失函数聚焦于概率显著的前缀，稳定训练动态。正则化系数 \(\rho > 0\) 直接惩罚总流 \(F(\mathcal{S})\)，驱使模型压低无效游走步数。
训练策略交替执行前向探索（采样先验 \(p_0\) 推进）与基于优先经验回放池（Prioritized Replay Buffer）的反向生成，回放采样比例固定为 \(N_{\text{RR}} = 2\)。同时，算法每隔 \(N_{\text{LS}} = 25\) 步调用 300 步 MALA 局部搜索提炼回放池样本，强效打破模式塌陷。
实验关键数据¶
主实验¶
论文在四个典型基准能量环境（2 维非均衡 9 高斯混合 GMM9、2 维漏斗分布 Funnel、40 维高斯混合 GMM9 以及 40 维高势垒双阱势 ManyWell）上全面评测了该方法。对比基线包括标准固定步长 ULA、基于 Gelman-Rubin 诊断终止的 ULA RG，以及基于 Geweke 谱方差诊断终止的 ULA Geweke。评价指标涵盖 Sinkhorn 距离（SD）、核最大均值差异（MMD）、变分下界与上界（ELBO/EUBO）以及期望轨迹长度 \(\mathbb{E}[n_\tau]\)。
| 目标分布与维度 | 算法配置 | Sinkhorn 距离 (SD) ↓ | MMD ↓ | ELBO ↑ | 期望步数 \(\mathbb{E}[n_\tau]\) ↓ |
|---|---|---|---|---|---|
| GMM9 (\(d=2\)) | ULA (固定步长) | \(0.463 \pm 0.131\) | \(0.002 \pm 0.005\) | — | \(10^4\) |
| ULA RG (Gelman-Rubin) | \(0.463 \pm 0.154\) | \(0.002 \pm 0.001\) | — | \(10^4\) | |
| ULA Geweke | \(3.234 \pm 0.243\) | \(0.010 \pm 0.001\) | — | \(5307.3 \pm 29.7\) | |
| ULA w/ clf (仅学分类器) | \(1.522 \pm 0.051\) | \(0.053 \pm 0.002\) | \(-1.753 \pm 0.045\) | \(15.0 \pm 0.2\) | |
| ULA w/ clf & corr (本文完整模型) | \(0.221 \pm 0.026\) | \(0.027 \pm 0.006\) | \(-0.157 \pm 0.025\) | \(29.6 \pm 0.3\) | |
| Funnel (\(d=2\)) | ULA (固定步长) | \(14.101 \pm 0.243\) | \(0.142 \pm 0.014\) | — | \(500\) |
| ULA RG (Gelman-Rubin) | \(11.686 \pm 0.165\) | \(0.112 \pm 0.002\) | — | \(10^3\) | |
| ULA Geweke | \(16.262 \pm 0.133\) | \(0.216 \pm 0.004\) | — | \(406.6 \pm 0.7\) | |
| ULA w/ clf (仅学分类器) | \(8.849 \pm 0.126\) | \(0.128 \pm 0.003\) | \(-1.264 \pm 0.038\) | \(21.8 \pm 0.3\) | |
| ULA w/ clf & corr (本文完整模型) | \(4.947 \pm 0.097\) | \(0.089 \pm 0.013\) | \(-0.392 \pm 0.023\) | \(36.3 \pm 0.4\) | |
| GMM9 (\(d=40\)) | ULA (固定步长) | \(376.406 \pm 2.165\) | \(0.347 \pm 0.011\) | — | \(500\) |
| ULA RG (Gelman-Rubin) | \(376.652 \pm 2.377\) | \(0.345 \pm 0.002\) | — | \(10^3\) | |
| ULA Geweke | \(376.407 \pm 1.935\) | \(0.346 \pm 0.001\) | — | \(437.2 \pm 1.4\) | |
| ULA w/ clf (仅学分类器) | \(321.803 \pm 2.038\) | \(0.314 \pm 0.017\) | \(-64.581 \pm 0.912\) | \(37.7 \pm 2.6\) | |
| ULA w/ clf & corr (本文完整模型) | \(243.809 \pm 1.984\) | \(0.251 \pm 0.019\) | \(-51.701 \pm 0.778\) | \(27.8 \pm 3.5\) | |
| ManyWell (\(d=40\)) | ULA (固定步长) | \(71.508 \pm 0.876\) | \(0.332 \pm 0.018\) | — | \(1500\) |
| ULA RG (Gelman-Rubin) | \(70.614 \pm 0.189\) | \(0.328 \pm 0.001\) | — | \(2000\) | |
| ULA Geweke | \(73.531 \pm 0.081\) | \(0.338 \pm 0.001\) | — | \(1468.2 \pm 8.5\) | |
| ULA w/ clf (仅学分类器) | \(54.804 \pm 4.203\) | \(0.298 \pm 0.019\) | \(92.507 \pm 10.109\) | \(79.1 \pm 4.2\) | |
| ULA w/ clf & corr (本文完整模型) | \(62.472 \pm 1.237\) | \(0.314 \pm 0.017\) | \(150.008 \pm 8.421\) | \(75.4 \pm 2.9\) |
消融实验¶
论文在 40 维 GMM9、40 维 ManyWell 及 196 维 MNIST Digits（基于预训练 NICE 流模型的真实连续流分布）上系统评估了层级数 \(L \in \{1, 2, 5\}\) 的影响，并与 SOTA 连续扩散采样器 PIS（Path Integral Sampler）和 DDS（Denoising Diffusion Sampler）进行横向对比：
| 数据集与维度 | 模型配置 | Sinkhorn 距离 (SD) ↓ | MMD ↓ | ELBO ↑ | 期望步数 \(\mathbb{E}[n_\tau]\) ↓ |
|---|---|---|---|---|---|
| GMM9 (\(d=40\)) | PIS (扩散基线) | \(329.400 \pm 1.814\) | \(0.186 \pm 0.012\) | \(-278.4 \pm 1.841\) | 45 |
| DDS (扩散基线) | \(180.200 \pm 1.310\) | \(0.181 \pm 0.015\) | \(-7.631 \pm 5.183\) | 45 | |
| 单层 (\(L=1\), 即 ULA w/ clf & corr) | \(243.809 \pm 1.984\) | \(0.251 \pm 0.019\) | \(-51.701 \pm 0.778\) | \(27.8 \pm 3.5\) | |
| 多层级 (\(L=2\)) | \(186.200 \pm 1.723\) | \(0.203 \pm 0.015\) | \(-34.735 \pm 6.192\) | \(44.6 \pm 3.6\) | |
| 多层级 (\(L=5\)) | \(152.347 \pm 2.281\) | \(0.184 \pm 0.011\) | \(-28.421 \pm 2.156\) | \(72.9 \pm 4.2\) | |
| ManyWell (\(d=40\)) | PIS (扩散基线) | \(38.740 \pm 1.329\) | \(0.258 \pm 0.011\) | \(200.9 \pm 7.924\) | 75 |
| DDS (扩散基线) | \(38.810 \pm 1.329\) | \(0.259 \pm 0.013\) | \(201.100 \pm 8.914\) | 75 | |
| 单层 (\(L=1\), 即 ULA w/ clf & corr) | \(62.472 \pm 1.237\) | \(0.314 \pm 0.017\) | \(150.008 \pm 8.421\) | \(75.4 \pm 2.9\) | |
| 多层级 (\(L=2\)) | \(36.330 \pm 0.561\) | \(0.108 \pm 0.021\) | \(143.500 \pm 4.254\) | \(74.2 \pm 2.6\) | |
| 多层级 (\(L=5\)) | \(28.400 \pm 4.843\) | \(0.095 \pm 0.020\) | \(154.267 \pm 3.845\) | \(68.7 \pm 2.7\) | |
| Digits (\(d=196\)) | PIS (扩散基线) | \(1728.000 \pm 10.351\) | \(0.540 \pm 0.023\) | \(-3671.000 \pm 21.931\) | 65 |
| DDS (扩散基线) | \(202.400 \pm 4.189\) | \(0.346 \pm 0.021\) | \(-149.300 \pm 3.194\) | 65 | |
| 单层 (\(L=1\), 即 ULA w/ clf & corr) | \(469.500 \pm 1.532\) | \(0.410 \pm 0.036\) | \(-210.340 \pm 3.591\) | \(71.6 \pm 6.1\) | |
| 多层级 (\(L=2\)) | \(343.600 \pm 0.386\) | \(0.394 \pm 0.042\) | \(-184.391 \pm 8.138\) | \(62.3 \pm 5.1\) | |
| 多层级 (\(L=5\)) | \(196.938 \pm 3.519\) | \(0.304 \pm 0.028\) | \(-121.570 \pm 4.193\) | \(68.3 \pm 5.3\) |
关键发现¶
- 显著压缩轨迹步数并实现精准截断：无论在 2 维还是 40 维任务中，单纯学习前向分类器（ULA w/ clf）即可将 ULA 的迭代步数从 \(500 \sim 10000\) 步急剧压缩至 \(15 \sim 79\) 步（缩短 1 到 2 个数量级），并且分类器在几何高密度核心处精准呈现极高停机概率，直接验证了理论推导中 \(d_F(s) \propto \pi(s) / f(s)\) 的密度映射特性。
- 漂移修正激活跨模态跃迁能力：固定 ULA 核在非对称混合高斯中极易被陷于某一局部峰（Gelman-Rubin 与 Geweke 均无法有效促成模态跳跃）；而完整模型（ULA w/ clf & corr）通过联合学习可变漂移 \(u_F\) 与方差 \(w_F\)，在 GMM9 上将 Sinkhorn 距离从 \(0.463\) 砍半至 \(0.221\)，ELBO 从 \(-1.753\) 跃升至 \(-0.157\)，展现出强大的跨模态传输动力学。
- 多层级退火对于高维曲面至关重要：在 \(d=196\) 的图像流任务中，单层模型 Sinkhorn 距离达 \(469.5\)，而扩充至 \(L=5\) 层后骤降至 \(196.9\)，击败经典扩散采样器 DDS（\(202.4\)），证明在多层非无环退火下各层自适应调配采样步数的机制能极大缓解高维复杂能量面的相变瓶颈。
亮点与洞察¶
- 将 MCMC 终止决策提升为非无环流网络的原生动作：跳出了传统启发式后验收敛诊断（如 \(\hat{R}\)）的外部观测框架，将提前停机形式化为流网络向吸收汇状态 \(s_f\) 的策略转移，建立了最优分类器与未归一化密度的严格细致平衡闭式解。
- 轨迹长度等式与流正则化的紧确对偶：严格证明了非无环连续状态流网络下“期望轨迹长度严格等于非终态总流”，从而使经验惩罚项 \(\rho r(s) / d_F(s)\) 具有严格的最短路径变分几何意义，避免了人工截断造成的密度失真。
- 自适应时间步与动力学修正的双重解耦：将连续扩散核（负责空间遍历）与状态分类器（负责终止判定）正交解耦，使得该框架具备即插即用特性，未来能够轻松迁移至重度依赖路径长度调节的哈密顿蒙特卡洛（HMC）或辛积分器体系。
局限与展望¶
- 训练收敛对超参数平衡较为敏感：分类器与连续核的联合优化极易因流正则化系数 \(\rho\) 或前缀截断权重的微小扰动而在过早停机（模式欠探索）与无限游走（流过大）之间震荡，需要仔细调节梯度截断与先验退火。
- 未包含拒识校正（Metropolis Step）的渐近无偏性损失：当前前向核采用可学习的连续高斯参数化，类似于免调朗之万算法，但由于缺乏 Metropolis-Hastings 接受-拒绝步骤，在大步长或强非线性漂移下可能引入不可消除的渐近系统截断偏差。
- 改进方向：可考虑将不可逆保结构动力学（如 Underdamped Langevin 或 HMC 动量扩展）引入非无环 GFlowNet 框架，并将可微 M-H 检验与分类器截断协同建模。
相关工作与启发¶
- vs 传统 MCMC 收敛诊断（ULA RG / Geweke）: 传统诊断是在生成固定长序列后做离线或滑动窗口的统计矩检验，反应滞后且容易陷入单模态假收敛；本文方法将停止条件作为在线策略在局部状态直接推断，步数缩短近两个数量级且显著改善模态平衡。
- vs 连续流网络与扩散采样器（Continuous GFlowNet / DDS / PIS）: 扩散采样器与经典连续 GFlowNet 强绑定在固定时间步长或有限视界的 DAG 上；本文框架基于非无环图构建，支持连续空间内多轮自适应循环探索，且多层级方案能针对各退火层级自适应分配探索步数。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从非无环 GFlowNet 理论出发解决连续采样核的自适应提前终止问题，理论推导严密且漂亮。
- 实验充分度: ⭐⭐⭐⭐☆ 覆盖低维到高维（\(d=2 \sim 196\)）的多种典型难分布，对比基线丰富，兼顾了传统统计量与扩散生成模型。
- 写作质量: ⭐⭐⭐⭐⭐ 符号系统高度严谨，测度论背景完备，算法框架逻辑链条极其清晰。
- 价值: ⭐⭐⭐⭐⭐ 为解决 MCMC 轨迹长度设定盲区提供了兼具坚实理论保证与卓越经验效率的全新范式。