We propose an algorithm, termed the Non-Equilibrium Transport Sampler (NETS), to sample from unnormalized probability distributions. NETS can be viewed as a variant of annealed importance sampling (AIS) based on Jarzynski's equality, in which the stochastic differential equation used to perform the non-equilibrium sampling is augmented with an additional learned drift term that lowers the impact of the unbiasing weights used in AIS. We show that this drift is the minimizer of a variety of objective functions, which can all be estimated in an unbiased fashion without backpropagating through solutions of the stochastic differential equations governing the sampling. We also prove that some these objectives control the Kullback-Leibler divergence of the estimated distribution from its target. NETS is shown to be unbiased and, in addition, has a tunable diffusion coefficient which can be adjusted post-training to maximize the effective sample size. We demonstrate the efficacy of the method on standard benchmarks, high-dimensional Gaussian mixture distributions, and a model from statistical lattice field theory, for which it surpasses the performances of related work and existing baselines.

翻译：我们提出了一种名为非平衡输运采样器（NETS）的算法，用于从非归一化概率分布中进行采样。NETS可被视为基于雅尔津斯基等式的一种退火重要性采样（AIS）变体，其中用于执行非平衡采样的随机微分方程通过一个额外的学习漂移项得到增强，该漂移项降低了AIS中使用的无偏校正权重的影响。我们证明该漂移项是多种目标函数的最小化解，所有这些目标函数均能以无偏方式估计，而无需对控制采样的随机微分方程的解进行反向传播。我们还证明了其中一些目标函数控制了估计分布与其目标分布之间的Kullback-Leibler散度。NETS被证明是无偏的，并且具有可调的扩散系数，可在训练后进行调整以最大化有效样本量。我们在标准基准测试、高维高斯混合分布以及一个统计晶格场论模型上验证了该方法的有效性，其性能超越了相关工作和现有基线。