NETS: A Non-Equilibrium Transport Sampler

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Albergo, Michael S., Vanden-Eijnden, Eric
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929673271771136
author Albergo, Michael S.
Vanden-Eijnden, Eric
author_facet Albergo, Michael S.
Vanden-Eijnden, Eric
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_02711
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle NETS: A Non-Equilibrium Transport Sampler
Albergo, Michael S.
Vanden-Eijnden, Eric
Machine Learning
Statistical Mechanics
High Energy Physics - Lattice
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.
title NETS: A Non-Equilibrium Transport Sampler
topic Machine Learning
Statistical Mechanics
High Energy Physics - Lattice
url https://arxiv.org/abs/2410.02711