Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels

Technical Summary: Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels

Problem Statement
Sampling from complex, unnormalized probability densities, defined as $\pi(x) = r(x)/Z$, is a fundamental challenge in Bayesian inference and probabilistic modeling. While Markov chain Monte Carlo (MCMC) methods provide asymptotic guarantees, they frequently suffer from slow mixing and high computational costs, particularly in high-dimensional spaces. A critical limitation of standard MCMC is the reliance on fixed or manually tuned trajectory lengths. Existing approaches to determine stopping times often rely on online diagnostics (e.g., Gelman–Rubin, Geweke) computed from generated samples, which do not always provide affirmative stopping criteria and can lead to premature termination or biased estimates. Conversely, amortized approaches (e.g., normalizing flows, diffusion samplers, GFlowNets) shift the computational burden to training but often bypass the exact sampling guarantees of MCMC.

Methodology
The authors propose a novel framework that treats trajectory termination as a learnable component of sampling dynamics, bridging MCMC and amortized sampling. The core methodology involves:

1. Non-acyclic GFlowNet Formulation: The authors frame MCMC with adaptive stopping within the theory of non-acyclic Generative Flow Networks (GFlowNets). They define a structured state space $\bar{S} = \{s_0\} \cup S \cup \{s_f\}$, where $S = \mathbb{R}^d$ represents the continuous state space, $s_0$ is the source, and $s_f$ is an absorbing sink state. Unlike standard acyclic GFlowNets, this setting allows for cycles within $S$.
2. Classifier-Based Stopping: Instead of running a Markov chain for a fixed number of steps, the method employs state-dependent neural classifiers, $d_F(s)$ and $d_B(s)$, to decide when a trajectory should terminate.
   • The forward classifier $d_F(s)$ determines the probability of transitioning directly to the sink state $s_f$ (stopping).
   • The backward classifier $d_B(s)$ controls the probability of moving to the source state $s_0$.
   • The forward transition density is defined as $p_F(s'|s) = (1 - d_F(s))q_F(s'|s)$ for continuous transitions, and $p_F(s_f|s) = d_F(s)$ for termination.
3. Theoretical Connections: The paper establishes a rigorous connection between the learned stopping policies and the target density $\pi$ via detailed balance conditions. Theoretically, the optimal forward classifier is shown to satisfy $d_F(s) = r(s)/f(s)$, where $f(s)$ is the flow measure. This relates the stopping decision to the accumulated discrepancy between the injected mass and the target distribution.
4. Training Algorithm: The authors utilize a prefix trajectory balance (TB) objective, which is computationally efficient for non-acyclic environments. The training process involves:
   • Sampling partial trajectories of fixed maximum length $N_{max}$ with classifiers disabled to encourage exploration.
   • Extending prefixes with terminal transitions to form complete trajectories.
   • Minimizing a loss function that combines the prefix trajectory balance loss with a flow regularization term. The regularization encourages shorter trajectories by penalizing the total flow, effectively minimizing the expected trajectory length.
5. Multilevel Generalization: To handle complex geometries, the framework is extended to a multilevel scheme. The state space is augmented with level variables $\ell \in \{1, \dots, L\}$, allowing the model to switch between regimes. Intermediate rewards $r(s, \ell)$ are defined to interpolate between the initial distribution and the target, stabilizing training and exploration similar to annealed importance sampling.

Key Contributions

1. Theoretical Framework: The authors demonstrate that sampling dynamics in $\mathbb{R}^d$ with the possibility of stopping at any point can be viewed as a non-acyclic GFlowNet environment. They derive the relationship between optimal classifiers and the target density using flow and detailed balance conditions.
2. Practical Algorithms: They introduce practical algorithms combining learned stopping with adaptive forward and backward kernels. This includes a multilevel extension designed to improve training stability and exploration in complex geometries.
3. Empirical Evaluation: The method is evaluated on various benchmark densities (GMM9, Funnel, ManyWell, and high-dimensional Digits tasks). The results demonstrate that the approach significantly reduces average trajectory lengths compared to standard MCMC baselines (like ULA) while improving mode coverage and mixing.

Experimental Results

• Trajectory Length: The classifier-based approach substantially reduces the mean trajectory length ($E[n_\tau]$) compared to fixed-step ULA and ULA with standard convergence diagnostics (Gelman–Rubin, Geweke). For instance, on the GMM9 ($d=2$) task, the learned classifier reduced the trajectory length from ~104 steps (ULA) to ~15 steps.
• Sample Quality: While the classifier alone reduces trajectory length, the full model (including learned forward kernel corrections) significantly improves approximation quality. This is evidenced by lower Sinkhorn distances (SD), lower Maximum Mean Discrepancy (MMD), and tighter Evidence Lower/Upper Bounds (ELBO/EUBO).
• Mode Coverage: Visualizations and quantitative metrics show that the proposed method effectively transitions between distant high-density regions and modes, whereas standard ULA often remains trapped in local regions.
• Multilevel Performance: The multilevel extension ($L=5$) consistently outperforms single-level models and becomes competitive with established diffusion sampling methods (PIS, DDS) on high-dimensional tasks like Digits ($d=196$).

Significance and Claims
The paper claims to provide a unified framework where sampling trajectories are dynamically shortened in high-density regions while remaining exploratory elsewhere. By treating stopping as a learnable component, the method addresses the inefficiency of fixed-length MCMC trajectories without sacrificing the theoretical guarantees associated with flow-based sampling. The authors note that while the approach introduces additional optimization complexity (balancing exploration and regularization), it offers a principled way to de-bias fixed kernels (like ULA) and adapt to the geometry of complex target distributions. The work suggests that future directions could involve improving training objectives for continuous non-acyclic GFlowNets and extending the framework to more expressive dynamics like Hamiltonian Monte Carlo.