Stochastic Normalizing Flows

Hao Wu; Jonas Köhler; Frank Noé

Source-linked AI summary

Stochastic Normalizing Flows

Hao Wu, Jonas Köhler, Frank Noé

arXiv:2002.06707v3Feb 16, 2020stat.MLcs.LGphysics.chem-phphysics.data-an

TL;DR

Sampling target distributions known only up to normalization is challenging because stochastic methods can mix slowly and invertible flows face expressivity constraints. The paper introduces Stochastic Normalizing Flows, which combine both mechanisms and train them jointly. SNFs compute exact path-based importance weights, improve sampling efficiency on benchmarks, and support asymptotically unbiased sampling, including molecular applications.

Problem
Generating asymptotically unbiased samples from normalized-unknown target distributions is important, but MCMC can mix slowly and deterministic flows face expressivity limitations.
Method
SNFs combine stochastic sampling blocks and deterministic invertible transformations, with joint training based on nonequilibrium path probabilities and exact path importance weights.
Results
SNFs improve expressivity and sampling efficiency over deterministic flows or pure stochastic sampling, while supporting asymptotically unbiased sampling on benchmark and molecular tasks.
Takeaways & Limitations
Trainable flow blocks transport large-scale probability mass while stochastic dynamics sample local details, combining their complementary roles for target-density sampling.
Takeaways & Limitations
Reweighting reduces bias at the expense of higher variance, and flow-based samplers may be non-ergodic without guarantees of asymptotic target sampling.

Abstract

from arXiv · show

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.

1 Introduction

Sampling unnormalized target distributions is important but difficult: MCMC can mix slowly, while invertible flows can struggle with separated modes. SNFs combine stochastic dynamics with trainable invertible transformations to address both issues and support unbiased sampling.

Motivation: Unnormalized energy-model distributions require asymptotically unbiased sampling in machine learning, statistics, and physics.
Motivation: MCMC is asymptotically unbiased but may remain trapped in local energy minima and fail to converge in practice.
Limitations of existing methods: Invertible normalizing flows can map simple priors to targets exactly, but topological and Bi-Lipschitz constraints hinder separated multimodal distributions.
Contributions: SNFs interweave stochastic sampling blocks with arbitrary sequences of deterministic invertible transformations, improving expressivity and sampling efficiency through joint optimization.
Contributions: Exact importance weights can be computed for SNF samples without marginalizing over stochastic paths, enabling asymptotically unbiased sampling.
Applications: Applications include molecular structures, multimodal torsion-angle distributions, benchmark datasets, and VAE posterior sampling.

2 Stochastic normalizing flows

An SNF alternates deterministic invertible transformations and stochastic dynamics along forward or backward paths. Path-probability ratios make training and asymptotically unbiased reweighting possible without evaluating an intractable marginal path integral.

Architecture: An SNF samples from a prior through a sequence of stochastic and deterministic transformations, producing a forward path and target-space proposal.
Architecture: Each stochastic block has forward and backward sampling densities, while deterministic steps reduce to ordinary invertible transformations.
Path probabilities: The marginal probability of an SNF output requires integrating over all paths ending at that output, unlike in deterministic normalizing flows.
Path probabilities: Independent path samples receive importance weights proportional to forward-backward probability ratios, avoiding direct evaluation of the intractable path integral.
Training: SNF parameters are trained by minimizing forward or backward path divergences, equivalently maximizing corresponding path weights.
Asymptotically unbiased sampling: Path-weight reweighting or Neural MCMC yields asymptotically unbiased target samples, although variational path training does not explicitly compute the marginal target density.

3 Implementing SNFs via Annealed Importance Sampling

SNFs combine deterministic flow layers with stochastic MCMC or Langevin-dynamics blocks arranged across interpolated potentials. In the implementation described here, stochastic blocks anneal between potentials while intervening flow layers approximate the partial density transformations.

Stochastic blocks: SNFs use MCMC or Langevin-dynamics blocks to make local stochastic updates that asymptotically sample an intermediate density.The intermediate density is proportional to exp(−uλ(y)).
Annealed construction: The implementation interpolates between adjacent potentials and places deterministic flow layers between stochastic sampling steps.The flow layers approximate partial density transformations, while stochastic blocks anneal with respect to the intermediate potential.
Langevin dynamics: Overdamped Langevin dynamics use an Euler discretization with Gaussian noise and a step parameter determined by the time step, friction, and mass.The inverse temperature is set to 1 in the described physical-system parameterization.
Langevin dynamics: The backward Langevin step uses the backward noise realization, allowing the corresponding path probability ratio to be computed.The paper states that the log path probability ratio is given in supplementary material.
MCMC: MCMC proposals are required to satisfy detailed balance with respect to the interpolated density, including Metropolis-Hastings and Hamiltonian Monte Carlo moves.When the backward sampling density equals the forward proposal density, the paper gives the log path probability ratio for these moves.

4 Results

Across image, double-well, molecular, and variational-inference benchmarks, SNFs improve expressivity, sampling efficiency, reproducibility, and unbiased estimation by combining trainable flow transformations with stochastic sampling blocks.

Two-dimensional densities: SNFs combine deterministic coupling layers with Metropolis MC steps, while comparisons include RealNVP, neural spline flows, and pure Metropolis sampling.The two-dimensional benchmark fixes either 50 MC steps or 10 deterministic transformations when varying the other component.
Two-dimensional densities: SNFs achieve high-quality image-density approximations by transporting large-scale probability mass with trainable flows and sampling details with Metropolis MC.The individual deterministic and stochastic components perform worse than their combined SNF architecture.
Sampling efficiency: 1-2 orders of magnitude more Metropolis MC steps are needed to reach similar KL divergence as an SNF, depending on the example and flow architecture.Adding 10-20 lightweight stochastic layers per trainable flow layer can keep computational cost below twice the original cost in the reported image experiments.
Asymptotically unbiased sampling: Importance sampling reweights deterministic and stochastic flow proposals to free energies within statistical error of the exact double-well result.With 100,000 samples, SNFs with Metropolis MC reduce both bias and uncertainty by half compared with purely deterministic flows.
Sampling efficiency: SNFs reduce run-to-run differences and improve reproducibility compared with RealNVP flows, with further gains possible by optimizing MC step sizes.The reported optimization uses loss functions JKL and JML.
Molecular sampling: For alanine dipeptide, SNFs resolve multimodal torsion-angle structure that deterministic RealNVP networks miss and approximate the target distribution better by KL divergence.The SNF uses up to 100 MCMC steps in one forward pass.
Variational inference: In variational inference, SNF outperforms both RNVP and Langevin-proposal MCMC, while all three improve over a naive diagonal-Gaussian VAE posterior.The comparison uses variational-bound and test-set log-likelihood results on MNIST and Fashion-MNIST.

5 Related work

The paper situates SNFs among stochastic-flow and trainable-sampling approaches, emphasizing their broader building-block flexibility and exact reweighting for asymptotically unbiased sampling.

5 Related work: SNFs replace fixed stochastic protocols with trainable deterministic and stochastic layers optimized within a generative model.This generalizes the nonequilibrium candidate Monte Carlo framework.
5 Related work: Unlike neural stochastic differential equations, SNFs target marginal density estimation and asymptotically unbiased sampling.The cited alternatives learn stochastic processes from observations but are not designed for these objectives.
5 Related work: SNFs use path-based detailed balance between prior and target, permitting nonequilibrium moves along transformations.This differs from frameworks that impose detailed balance at every Monte Carlo step.
5 Related work: Compared with closely related stochastic-flow density estimation, SNFs allow more general stochastic and deterministic blocks and address asymptotically unbiased target sampling.The related approach is described as a special instance of SNFs.
5 Related work: The paper concludes that path probability ratios enable efficient SNF training without intractable marginalization.It identifies future applications to large-scale molecular sampling.

Broader Impact

The paper frames energy-based sampling as important for drug and materials design, while describing reweighting-based estimation and its bias–variance trade-off.

Broader Impact: Energy-model sampling supports rational drug-molecule and energy-storage material design, where conventional simulations can be computationally expensive.A single protein–drug binding and dissociation event may require a supercomputer for a year.
Broader Impact: Reweighting reduces bias at the expense of higher variance, which can be reduced by generating more parallel samples.The paper emphasizes low or asymptotically zero bias for physics applications.
Broader Impact: Forward-KL training minimizes divergence between the generated distribution and the target when the target is known up to a normalization constant.The target density is represented through an energy model.
Broader Impact: Importance weights can be computed for generated samples, supporting normalized estimators of target expectations without knowing the target normalization constant.The estimator uses unnormalized path weights and normalization across samples.
Broader Impact: The normalized estimator converges to the target expectation as the sample count tends to infinity.The convergence is attributed to the law of large numbers.

3. Derivation of the deterministic layer probability ratio

The deterministic-layer derivation develops probability ratios for noisy approximations to delta distributions and combines forward and backward path relations.

3. Derivation of the deterministic layer probability ratio: Delta distributions are represented as Gaussian densities with variance σ, followed by the limit σ → 0+.This provides a finite-variance construction before taking the deterministic limit.
3. Derivation of the deterministic layer probability ratio: The marginal distribution of an intermediate state y_t is related to path probabilities in the derivation.The derivation explicitly introduces p_t(y_t) for this intermediate marginal.
3. Derivation of the deterministic layer probability ratio: The derivation uses the definition of ΔS_t in terms of path-probability ratios to obtain the intermediate result.This step connects stochastic path probabilities to the layer entropy-like quantity.
3. Derivation of the deterministic layer probability ratio: The backward step is specified after deriving the relevant probability relation.The construction follows previously established results.
3. Derivation of the deterministic layer probability ratio: Combining the stated equations yields the deterministic-layer path probability ratio.The derivation explicitly combines Equations (18) and (28).

5. Derivation of the Langevin probability ratio

The Langevin-layer derivation defines forward and reverse BBK steps, combines their probability expressions, and simplifies the resulting path ratio using a Jacobian cancellation.

3. Derivation of the Langevin probability ratio: Forward BBK Langevin dynamics are specified as leap-frog updates for the position and velocity variables.The forward step is stated explicitly for the Langevin protocol.
3. Derivation of the Langevin probability ratio: The reverse step uses reversed momenta and accounts for the factor 4 required by using Δt/2 in both half-steps.This factor differs from the cited prior derivation.
3. Derivation of the Langevin probability ratio: The forward and reverse equations are combined to derive the expressions for reverse-step realization and path probabilities.The derivation combines Equations (30–33), then Equations (31), (32), and (35).
3. Derivation of the Langevin probability ratio: A Jacobian is introduced for the path-ratio calculation, and its ratio cancels because the Jacobians do not depend on the noise variables.This cancellation simplifies the Langevin probability ratio.

6. Derivation of the probability ratio for Markov Chain Monte Carlo

The derivation establishes detailed-balance properties for MCMC and Hamiltonian Monte Carlo, including the forward and marginal path densities used for transitions.

MCMC transition densities satisfy the detailed balance condition.
Hamiltonian Monte Carlo with Metropolis acceptance defines a forward path density.
The Hamiltonian Monte Carlo path density satisfies joint detailed balance.
Independently drawing velocity v from N(v|0, I) yields a marginal forward path density for yt → yt+1.
The marginal transition expression is obtained from the preceding detailed-balance relation.

8. Details on using SNFs for variational inference

The paper uses SNFs as conditional variational approximations to VAE posteriors and estimates the resulting objective and evidence with Monte Carlo samples.

A VAE decoder defines an observed variable s conditioned on a latent variable x, with a tractable prior over x.
The joint VAE distribution is defined as PD(x, s) = P(x) · PD(s|x).
Conditioned on s, an SNF approximates the posterior distribution PD(x|s).
When the SNF contains only deterministic transformations, JKL is equivalent to F in [35].
The objective JKL is estimated from sampled paths, while −log PD(s) is approximated by marginalizing x with Monte Carlo sampling using N = 2000.
The conditional flow uses q(z) = N(z|0, I) and a conditional diffeomorphism FLL(x|s) implemented with RealNVP transformations.
The variational approximation is trained by minimizing KL divergence, reaching qLL(x|s) = PD(x|s) at convergence.

9. Hyper-parameters and other benchmark details

The supplementary experiments specify architectures, stochastic-step settings, optimization schedules, and coordinate or latent-space choices across benchmark systems.

Experiments use PyTorch 1.2, GTX1080Ti cards, Adam with step-size 0.001, and RealNVP for deterministic flow transformations.
Double-well experiments use three RealNVP blocks and, for SNFs, 20 Metropolis MC steps per block with Gaussian proposal standard deviation 0.25.
Other flow benchmarks use five blocks, with SNFs combining RealNVP or NSF blocks and 10 Metropolis MC steps using Gaussian proposal standard deviation 0.1.
The molecular benchmark uses three RealNVP blocks for both flows and SNFs, with SNFs additionally applying 20 Metropolis MC steps per block.
The VAE benchmark uses latent dimension 50, a decoder with two 1024-unit hidden layers, and sigmoid output activation.
MCMC uses 30 Metropolis MC steps with overdamped Langevin proposals whose interpolation coefficients and step size are trained parameters.
A further SNF setup uses three units, each containing two RealNVP blocks and 10 Metropolis MC steps, with staged training of flow and all parameters.

10. Comparison with related sampling methods

The supplementary comparison situates SNFs among samplers with learnable proposals or transformations and documents reproducibility and computational-cost comparisons.

Most listed sampling methods are based on detailed balance in each step, while HVI can use nonequilibrium steps with annealed target distributions.
Some prior sampling techniques improve efficiency through linear or nonlinear deterministic transformations.
Figure S1 compares independent runs of three RealNVP blocks with and without 20 birth-death steps before and after the blocks.
Figure S2 measures training time versus stochastic layers per RNVP or NSF layer, normalized to one deterministic layer.
The figure caption states that deterministic flow layers are much more computationally expensive than stochastic flow layers, though timings depend on implementation and computing conditions.