Abstract:The sampling of probability distributions specified up to a normalization constant is an important problem in both machine learning and statistical mechanics. While classical stochastic sampling methods such as Markov Chain Monte Carlo (MCMC) or Langevin Dynamics (LD) can suffer from slow mixing times there is a growing interest in using normalizing flows in order to learn the transformation of a simple prior distribution to the given target distribution. Here we propose a generalized and combined approach to sample target densities: Stochastic Normalizing Flows (SNF) -- an arbitrary sequence of deterministic invertible functions and stochastic sampling blocks. We show that stochasticity overcomes expressivity limitations of normalizing flows resulting from the invertibility constraint, whereas trainable transformations between sampling steps improve efficiency of pure MCMC/LD along the flow. By invoking ideas from non-equilibrium statistical mechanics we derive an efficient training procedure by which both the sampler's and the flow's parameters can be optimized end-to-end, and by which we can compute exact importance weights without having to marginalize out the randomness of the stochastic blocks. We illustrate the representational power, sampling efficiency and asymptotic correctness of SNFs on several benchmarks including applications to sampling molecular systems in equilibrium.
What problem does this paper attempt to address?
### The problems the paper attempts to solve
The paper aims to solve the important problem of generating asymptotically unbiased samples from probability distributions with unknown normalization constants. Specifically, the paper focuses on how to efficiently sample from the target distribution $\mu_X(x) \propto \exp(-u(x))$ defined by a given energy model $u(x)$ in machine learning and statistical mechanics.
### Background and challenges
1. **Limitations of traditional methods**:
- Classical stochastic sampling methods such as **Markov Chain Monte Carlo (MCMC)** and **Langevin Dynamics (LD)** are asymptotically unbiased, but may get trapped in local energy minima due to slow mixing times, resulting in difficulty in convergence in practical applications.
- **Normalizing Flows (NFs)** are an alternative method that generates samples by learning the transformation from a simple prior distribution to the target distribution. However, NFs are restricted by the reversibility constraint and may not be able to fully represent complex multimodal distributions.
2. **The core of the problem**:
- How to combine the advantages of determinism and randomness, overcome the representational limitations of NFs, and at the same time improve the efficiency and accuracy of sampling.
### Solutions
The paper proposes **Stochastic Normalizing Flows (SNFs)**, a method that combines deterministic reversible transformations and stochastic sampling blocks. The specific contributions are as follows:
1. **Combining determinism and randomness**:
- SNFs overcome the topological constraints of pure NFs through an arbitrary sequence of deterministic reversible functions and stochastic sampling blocks, improving the representational ability.
- The randomness part is introduced through the idea of non - equilibrium statistical mechanics, enabling the model to sample more effectively while remaining asymptotically unbiased.
2. **Efficient training methods**:
- An end - to - end optimization method is proposed, which can optimize the parameters of the sampler and the flow simultaneously and can calculate the exact importance weights without explicitly integrating out the randomness of the stochastic blocks.
3. **Application examples**:
- Through multiple benchmark tests, including equilibrium sampling of molecular systems, the representational ability, sampling efficiency, and asymptotic correctness of SNFs are demonstrated.
### Conclusions
The paper successfully solves the problem of efficiently generating asymptotically unbiased samples from complex probability distributions through SNFs, especially showing significant advantages when dealing with multimodal distributions. This method is not only innovative in theory but also shows a strong performance improvement in practical applications.