Stochastic Normalizing Flows

Introduction

A common problem in machine learning and statistics with important applications in physics is the generation of asymptotically unbiased samples from a target distribution defined up to a normalization constant by means of an energy model $u(\mathbf{x})$:

Sampling of such unnormalized distributions is often done with Markov Chain Monte Carlo (MCMC) or other stochastic sampling methods . This approach is asymptotically unbiased, but suffers from the sampling problem: without knowing efficient moves, MCMC approaches may get stuck in local energy minima for a long time and fail to converge in practice.

Normalizing flows (NFs) combined with importance sampling methods are an alternative approach that enjoys growing interest in molecular and material sciences and nuclear physics . NFs are learnable invertible functions, usually represented by a neural network, pushing forward a probability density over a latent or “prior” space $Z$ towards the target space $X$. Utilizing the change of variable rule these models provide exact densities of generated samples allowing them to be trained by either maximizing the likelihood on data (ML) or minimizing the Kullback-Leibler divergence (KL) towards a target distribution.

Let $F_{ZX}$ be such a map and its inverse $F_{XZ}=F_{ZX}^{-1}$. We can consider it as composition of $T$ invertible transformation layers $F_{0},...,F_{T}$ with intermediate states $\mathbf{y}_{t}$ given by:

By calling the samples in $Z$ and $X$ also $\mathbf{z}$ and $\mathbf{x}$, respectively, the flow structure is as follows:

We suppose each transformation layer is differentiable with a Jacobian determinant $\left|\det\mathbf{J}_{t}(\mathbf{y})\right|$. This allows to apply the change of variable rule:

As we often work with log-densities, we abbreviate the log Jacobian determinant as:

The log Jacobian determinant of the entire flow is defined by $\Delta S_{ZX}=\sum_{t}\Delta S_{t}(\mathbf{y}_{t})$ and correspondingly $\Delta S_{XZ}$ for the inverse flow.

Unbiased sampling is particularly important for applications in physics and chemistry where unbiased expectation values are required . A Boltzmann generator utilizing NFs achieves this by (i) generating one-shot samples $\mathbf{x}\sim p_{X}(\mathbf{x})$ from the flow and (ii) using a reweighing/resampling procedure respecting weights

turning these one-shot samples into asymptotically unbiased samples. Reweighing/resampling methods utilized in this context are e.g. Importance Sampling or Neural MCMC .

Training NFs.

NFs are trained in either “forward” or “reverse” mode, e.g.:

Density estimation – given data samples $\mathbf{x}$, train the flow such that the back-transformed samples $\mathbf{z}=F_{XZ}(\mathbf{x})$ follow a latent distribution $\mu_{Z}(\mathbf{z})$, e.g. $\mu_{Z}(\mathbf{z})=\mathcal{N}(\mathbf{0},\mathbf{I})$. This is done by maximizing the likelihood – equivalent to minimizing the KL divergence $KL\left[\mu_{X}\|p_{X}\right]$.

Sampling of a given target density $\mu_{X}(\mathbf{x})$ – sample from the simple distribution $\mu_{Z}(\mathbf{z})$ and minimize a divergence between the distribution generated by the forward-transformation $\mathbf{x}=F_{XZ}(\mathbf{z})$ and $\mu_{X}(\mathbf{x})$. A common choice is the reverse KL divergence $KL\left[p_{X}\|\mu_{X}\right]$.

We will use densities interchangeably with energies, defined by the negative logarithm of the density. The exact prior and target distributions are:

with generally unknown normalization constants $Z_{Z}$ and $Z_{X}$. As can be shown (Suppl. Material Sec. 1) minimizing $KL\left[p_{X}\|\mu_{X}\right]$ or $KL\left[\mu_{X}\|p_{X}\right]$ corresponds to maximizing the forward or backward weights of samples drawn from $p_{X}$ or $\mu_{X}$, respectively.

Topological problems of NFs.

A major caveat of sampling with exactly invertible functions for physical problems are topological constraints. While these can be strong manifold results, e.g., if the sample space is restricted to a non-trivial Lie group , another practical problem are induced Bi-Lipschitz constraints resulting from mapping uni-modal base distributions onto well-separated multi-modal target distributions. For example, when trying to map a unimodal Gaussian distribution to a bimodal distribution with affine coupling layers, a connection between the modes remains (Fig. 1a). This representational insufficiency poses serious problems during optimization – in the bimodal distribution example, the connection between the density modes seems largely determined by the initialization and does not move during optimization, leading to very different results in multiple runs (Suppl. Material, Fig. S1). More powerful coupling layers, e.g., , can mitigate this effect. Yet, as they are still diffeomorphic, strong Bi-Lipschitz requirements can make optimization difficult. This problem can be resolved when relaxing bijectivity of the flow by adding noise as we show in our results. Other proposed solutions are real-and-discrete mixtures of flows or augmentation of the bases space at the cost of losing asymptotically unbiased sampling.

Contributions.

We show that NFs can be interwoven with stochastic sampling blocks into arbitrary sequences, that together overcome topological constraints and improve expressivity over deterministic flow architectures (Fig. 1a, b). Furthermore, NSFs have improved sampling efficiency over pure stochastic sampling as the flow’s and sampler’s parameters can be optimized jointly.

Our main result is that NSFs can be trained in a similar fashion as NFs and exact importance weights for each sample ending in $\mathbf{x}$ can be computed, facilitating asymptotically unbiased sampling from the target density. The approach avoids explicitly computing $p_{X}(\mathbf{x})$ which would require solving the intractable integral over all stochastic paths ending in $\mathbf{x}$.

We apply the model to the recently introduced problem of asymptotically unbiased sampling of molecular structures with flows and show that it significantly improves sampling the multi-modal torsion angle distributions which are the relevant degrees of freedom in the system. We further show the advantage of the method over pure flow-based sampling / MCMC by quantitative comparison on benchmark data sets and on sampling from a VAE’s posterior distribution.

Code is available at github.com/noegroup/stochastic_normalizing_flows

Stochastic normalizing flows

A SNF is a sequence of $T$ stochastic and deterministic transformations. We sample $\mathbf{z}=\mathbf{y}_{0}$ from the prior $\mu_{Z}$, and generate a forward path $(\mathbf{y}_{1},\ldots,\mathbf{y}_{T})$ resulting in a proposal $\mathbf{y}_{T}$ (Fig. 2). Correspondingly, latent space samples can be generated by starting from a sample $\mathbf{x}=\mathbf{y}_{T}$ and invoking the backward path $(\mathbf{y}_{T-1},\ldots,\mathbf{y}_{0})$. The conditional forward / backward path probabilities are

denote the forward / backward sampling density at step $t$ respectively. If step $t$ is a deterministic transformation $F_{t}$ this simplifies as

In contrast to NFs, the probability that an SNF generates a sample $\mathbf{x}$ cannot be computed by Eq. (4) but instead involves an integral over all paths that end in $\mathbf{x}$:

denotes the forward-backward probability ratio of step $t$, and corresponds to the usual change of variable formula in NF for deterministic transformation steps (Suppl. Material Sec. 3). These weights allow asymptotically unbiased sampling and training of SNFs while avoiding Eq. (10). By changing denominator and numerator in (11) we can alternatively obtain the backward weights $w(\mathbf{x}\to\mathbf{z})$.

As in NFs, the parameters of a SNF can be optimized by minimizing the Kullback-Leibler divergence between the forward and backward path probabilities, or alternatively maximizing forward and backward path weights as long as we can compute $\Delta S_{t}$ (Suppl. Material Sec 1):

Variational bound.

Minimization of the reverse path divergence $J_{KL}$ minimizes an upper bound on the reverse KL divergence between the marginal distributions:

Asymptotically unbiased sampling.

As stated in the theorem below (Proof in Suppl. Material. Sec. 2), SNFs are Boltzmann Generators: We can generate asymptotically unbiased samples of $\mathbf{x}\sim\mu_{X}(\mathbf{x})$ by performing importance sampling or Neural MCMC using the path weight $w(\mathbf{z}_{k}\rightarrow\mathbf{x}_{k})$ of each path sample $k$.

Let $O$ be a function over $X$. An asymptotically unbiased estimator is given by

Implementing SNFs via Annealed Importance Sampling

In this paper we focus on the use of SNFs as samplers of $\mu_{X}(\mathbf{x})$ for problems where the target energy $u_{X}(\mathbf{x})$ is known, defining the target density up to a constant, and provide an implementation of stochastic blocks via MCMC / LD. These blocks make local stochastic updates of the current state $\mathbf{y}$ with respect to some potential $u_{\lambda}(\mathbf{y})$ such that they will asymptotically sample from $\mu_{\lambda}(\mathbf{y})\propto\exp(-u_{\lambda}(\mathbf{y}))$. While such potentials $u_{\lambda}(\mathbf{y})$ could be learned, a straightforward strategy is to interpolate between prior and target potentials

similarly as it is done in annealed importance sampling . Our implementation for SNFs is thus as follows: deterministic flow layers in-between only have to approximate the partial density transformation between adjacent $\lambda$ steps while the stochastic blocks anneal with respect to the given intermediate potential $u_{\lambda}$. The parameter $\lambda$ could again be learned – in this paper we simply choose a linear interpolation along the SNF layers: $\lambda=t/T$.

Overdamped Langevin dynamics, also known as Brownian dynamics, using an Euler discretization with time step $\Delta t$, are given by :

where $\boldsymbol{\eta}_{t}\sim\mathcal{N}(0,\mathbf{I})$ is Gaussian noise. In physical systems, the constant $\epsilon_{t}$ has the form $\epsilon_{t}=\Delta t/\gamma m$ with time step $\Delta t$, friction coefficient $\gamma$ and mass $m$, and $\beta$ is the inverse temperature (here set to $1$). The backward step $\mathbf{y}_{t+1}\rightarrow\mathbf{y}_{t}$ is realized under these dynamics with the backward noise realization (Suppl. Material Sec. 4 and ):

The log path probability ratio is (Suppl. Material Sec. 4):

We also give the results for non-overdamped Langevin dynamics in Suppl. Material. Sec. 5.

Markov Chain Monte Carlo.

Consider MCMC methods with a proposal density $q_{t}$ that satisfies the detailed balance condition w.r.t. the interpolated density $\mu_{\lambda}(\mathbf{y})\propto\exp(-u_{\lambda}(\mathbf{y}))$:

We show that for all $q_{t}$ satisfying (21), including Metropolis-Hastings and Hamiltonian MC moves, the log path probability ratio is (Suppl. Material Sec. 6 and 7):

Results

We first illustrate that SNFs can break topological constraints and improve the representational power of deterministic normalizing flows at a given network size and at the same time beat direct MCMC in terms of sampling efficiency. To this end we use images to define complex two-dimensional densities (Fig. 3a-c, “Exact”) as target densities $\mu_{X}(\mathbf{x})$ to be sampled. Note that a benchmark aiming at generating high-quality images would instead represent the image as a high-dimensional pixel array. We compare three types of flows with 5 blocks each trained by samples from the exact density (details in Suppl. Material Sec. 9):

Normalizing flow with 2 swapped coupling layers (RealNVP or neural spline flow) per block

Non-trainable stochastic flow with 10 Metropolis MC steps per block

SNF with both, 2 swapped coupling layers and 10 Metropolis MC steps per block.

The pure Metropolis MC flow suffers from sampling problems – density is still concentrated in the image center from the prior. Many more MC steps would be needed to converge to the exact density (see below). The RealNVP normalizing flow architecture has limited representational power, resulting in a “smeared out” image that does not resolve detailed structures (Fig. 3a-c, RNVP). As expected, neural spline flows perform significantly better on the 2D-images than RealNVP flows, but at the chosen network architecture their ability to resolve fine details and round shapes is still limite (See dog and small text in Fig. 3c, NSF). Note that the representational power for all flow architectures tend to increase with depth - here we compare the performance of different architectures at fixed depth and similar computational cost.

In contrast, SNFs achieve high-quality approximations although they simply combine the same deterministic and stochastic flow components that fail individually in the SNF learning framework (Fig. 3a-c, RNVP+Metropolis and NSF+Metropolis). This indicates that the SNF succeeds in performing the large-scale probability mass transport with the trainable flow layers and sampling the details with Metropolis MC.

Fig. 3d-e quantifies these impressions by computing the KL divergence between generated densities $p_{X}(\mathbf{x})$ and exact densities $\mu_{X}(\mathbf{x})$. Both normalizing flows and SNFs improve with greater depth, but SNFs achieve significantly lower KL divergence at a fixed network depth (Fig. 3d). Note that both RealNVP and NSFs improve significantly when stochasticty is added.

Moreover, SNFs have higher statistical efficiency than pure Metropolis MC flows. Depending on the example and flow architecture, 1-2 orders of magnitude more Metropolis MC steps are needed to achieve similar KL divergence as with an SNF. This demonstrates that the large-scale probability transport learned by the trainable deterministic flow blocks in SNFs significantly helps with the sampling.

Importantly, adding stochasticity is very inexpensive. Although every MCMC or Langevin integration step adds a neural network layer, these layers are very lightweighted, and have only linear computational complexity in the number of dimensions. As an example, for our SNF implementation of the examples in Fig. 3 we can add 10-20 stochastic layers to each trainable normalizing flow layer before the computational cost increases by a factor of 2 (Suppl. Material Fig. S2).

SNFs as asymptotically unbiased samplers.

We demonstrate that SNFs can be used as Boltzmann Generators, i.e., to sample target densities without asymptotic bias by revisiting the double-well example (Fig. 1). Fig. 4 (black) shows the free energies (negative marginal density) along the double-well coordinate $x_{1}$. Flows with 3 coupling layer blocks (RealNVP or neural spline flow) are trained summing forward and reverse KL divergence as a joint loss using either data from a biased distribution, or with the unbiased distribution (Details in Suppl. Material Sec. 9). Due to limitations in representational power the generation probability $p_{X}(\mathbf{x})$ will be biased – even when explicitly minimizing the KL divergence w.r.t. the true unbiased distribution in the joint loss. By relying on importance sampling we can turn the flows into Boltzmann Generators in order to obtain unbiased estimates. Indeed all generator densities $p_{X}(\mathbf{x})$ can be reweighted to an estimate of the unbiased density $\mu_{X}(\mathbf{x})$ whose free energies are within statistical error of the exact result (Fig. 4, red and green).

The differences between multiple runs (see standard deviations of the uncertainty estimate) also reduce significantly, i.e. SNF results are more reproducible than RealNVP flows, confirming that the training problems caused by the density connection between both modes (Fig. 1, Suppl. Material Fig. S1) can be reduced. Moreover, the sampling performance of SNF can be further improved by optimizing MC step sizes based on loss functions $J_{KL}$ and $J_{ML}$ (Suppl. Material Table S1).

Reweighting reduces the bias at the expense of a higher variance. Especially in physics applications, a small or asymptotically zero bias is often very important, and the variance can be reduced by generating more samples from the trained flow, which is relatively cheap and parallel.

Alanine dipeptide.

We further evaluate SNFs on density estimation and sampling of molecular structures from a simulation of the alanine dipeptide molecule in vacuum (Fig. 5). The molecule has 66 dimensions in $\mathbf{x}$, and we augment it with 66 auxiliary dimensions in a second channel $\mathbf{v}$, similar to “velocities” in a Hamiltonian flow framework , resulting in 132 dimensions total. The target density is given by $\mu_{X}(\mathbf{x},\mathbf{v})=\exp\left(-u(\mathbf{x})-\frac{1}{2}\left\|\mathbf{v}\right\|^{2}\right)$, where $u(\mathbf{x})$ is the potential energy of the molecule and $\frac{1}{2}\left\|\mathbf{v}\right\|^{2}$ is the kinetic energy term. $\mu_{Z}$ is an isotropic Gaussian normal distribution in all dimensions. We utilize the invertible coordinate transformation layer introduced in in order to transform $\mathbf{x}$ into normalized bond, angle and torsion coordinates. RealNVP transformations act between the $\mathbf{x}$ and $\mathbf{v}$ variable groups Details in Suppl. Material Sec. 9).

We compare deterministic normalizing flows using 5 blocks of 2 RealNVP layers with SNFs that additionally use 20 Metropolis MC steps in each block totalling up to 100 MCMC steps in one forward pass. Fig. 5a shows random structures sampled by the trained SNF. Fig. 5b shows marginal densities in all five multimodal torsion angles (backbone angles $\phi$, $\psi$ and methyl rotation angles $\gamma_{1}$, $\gamma_{2}$, $\gamma_{3}$). While the RealNVP networks that are state of the art for this problem miss many of the modes, the SNF resolves the multimodal structure and approximates the target distribution better, as quantified in the KL divergence between the generated and target marginal distributions (Table 2).

Variational Inference.

Finally, we use normalizing flows to model the latent space distribution of a variational autoencoder (VAE) , as suggested in . Table 3 shows results for the variational bound and the log likelihood on the test set for MNIST and Fashion-MNIST . For a 50-dimensional latent space we compare a six-layer RNVP to MCMC using overdamped Langevin dynamics as proposal (MCMC) and a SNF combining both (RNVP+MCMC). Both sampling and the deterministic flow improve over a naive VAE using a reparameterized diagonal Gaussian variational posterior distribution, while the SNF outperforms both, RNVP and MCMC. See Suppl. Material Sec. 8 for details.

Related work

The cornerstone of our work is nonequilibrium statistical mechanics. Particularly important is Nonequilibrium Candidate Monte Carlo (NCMC) , which provides the theoretical framework to compute SNF path likelihood ratios. However, NCMC is for fixed deterministic and stochastic protocols, while we generalize this into a generative model by substituting fixed protocols with trainable layers and deriving an unbiased optimization procedure.

Neural stochastic differential equations learn optimal parameters of designed stochastic processes from observations along the path , but are not designed for marginal density estimation or asymptotically unbiased sampling. It has been demonstrated that combining learnable proposals/transformations with stochastic sampling techniques can improve expressiveness of the proposals . Yet, these contributions do not provide an exact reweighing scheme based on a tractable model likelihood and do not provide efficient algorithms to optimize arbitrary sequences of transformation or sampling steps end-to-end efficiently. These methods can be seen as instances of SNFs with specific choice of deterministic transformations and / or stochastic blocks and model-specific optimizations - see Suppl. Material Table S2) for a categorization. While our experiments focus on nontrainable stochastic blocks, the proposal densities of MC steps can also be optimized within the framework of SNFs as shown in Suppl. Material Table S1.

An important aspect of SNFs compared to trainable Monte-Carlo kernels such as A-NICE-MC is the use of detailed balance (DB). While Monte-Carlo frameworks are usually designed to use DB in each step, SNFs rely on path-based detailed balance between the prior and the target density. This means that SNFs can also perform nonequilibrium moves along the transformation, as done by Langevin dynamics without acceptance step and by the deterministic flow transformations such as RealNVP and neural spline flows.

More closely related is which uses of stochastic flows for density estimation and trains diffusion kernels by maximizing a variational bound of the model likelihood. Their derivation using stochastic paths is similar to ours and this work can be seen as a special instance of SNFs, but it does not consider more general stochastic and deterministic building blocks and does not discuss the problem of asymptotically unbiased sampling of a target density. Ref. proposes a learnable stochastic process by integrating Langevin dynamics with learnable drift and diffusion term. This approach is in a spirit similar as our proposed method, but requires variational approximation of the generative distribution and it has not been worked out how it could be used as a building block within a NF. The approach of combines NF layers with Langevin dynamics, yet approximates the intractable integral with MC samples which we can avoid utilizing the path-weight derivation. Finally, propose a stochastic extension to neural ODEs which can then be trained as samplers. This approach to sampling is very general yet requires costly integration of a SDE which we can avoid by combining simple NFs with stochastic layers.

Conclusions

We have introduced stochastic normalizing flows (SNFs) that combine both stochastic processes and invertible deterministic transformations into a single learning framework. By leveraging nonequilibrium statistical mechanics we show that SNFs can efficiently be trained to sample asymptotically unbiased from target densities. This can be done by utilizing path probability ratios and avoiding intractabe marginalization. Besides possible applicability in classical machine learning domains such as variational and Bayesian inference, we believe that the latter property can make SNFs a key component in the efficient sampling of many-body physics systems. In future research we aim to apply SNFs with many stochastic sampling steps to accurate large-scale sampling of molecules.

Broader Impact

The sampling of probability distributions defined by energy models is a key step in the rational design of pharmacological drug molecules for disease treatment, and the design of new materials, e.g., for energy storage. Currently such sampling is mostly done by Molecular Dynamics (MD) and MCMC simulations, which is in many cases limited by computational resources and generates extremely high energy costs. For example, the direct simulation of a single protein-drug binding and dissociation event could require the computational time of an entire supercomputer for a year. Developing machine learning (ML) approaches to solve this problem more efficiently and go beyond existing enhanced sampling methods is therefore of importance for applications in medicine and material science and has potentially far-reaching societal consequences for developing better treatments and reducing energy consumption. Boltzmann Generators, i.e. the combination of Normalizing Flows (FNs) and resampling/reweighting are a new and promising ML approach to this problem and the current paper adds a key technology to overcome some of the previous limitations of NFs for this task.

A risk of the method is that flow-based sampling bears the risk that non-ergodic samplers can be constructed, i.e. samplers that are not guaranteed to sample from the target distribution even in the limit of long simulation time. From such an incomplete sample, wrong conclusions can be drawn. While incomplete sampling is also an issue with MD/MCMC, it is well understood how to at least ensure ergodicity of these methods in the asymptotic limit, i.e. in the limit of generating enough data. Further research is needed to obtain similar results with normalizing flows.

Special thanks to José Migual Hernández Lobarto (University of Cambridge) for his valuable input on the path probability ratio for MCMC and HMC. We acknowledge funding from the European Commission (ERC CoG 772230 ScaleCell), Deutsche Forschungsgemeinschaft (GRK DAEDALUS, SFB1114/A04), the Berlin Mathematics center MATH+ (Project AA1-6 and EF1-2) and the Fundamental Research Funds for the Central Universities of China (22120200276).

References

Supplementary Material

If the target density $\mu_{X}$ is known up to a constant $Z_{X}$, we minimize the forward KL divergence between the generated and the target distribution.

The importance weights wrt the target distribution can be computed as:

Maximum likelihood and backward weight maximization.

maximum likelihood equals log weight maximization:

Proof of theorem 1 (unbiased sampling with SNF importance weights)

Derivation of the deterministic layer probability ratio

In order to work with delta distributions, we define $\delta^{\sigma}(\mathbf{x})=\mathcal{N}(\mathbf{x};\boldsymbol{0},\sigma\mathbf{I})$, i.e. a Gaussian normal distribution with mean $\mathbf{0}$ and variance $\sigma$ and then consider the limit $\sigma\rightarrow 0^{+}$. In the case where $\sigma>0$, by defining

where $p_{t}(y_{t})$ denotes the marginal distribution of $\mathbf{y}_{t}$. By considering

and using the definition of $\Delta S_{t}$ in terms of path probability rations, we obtain:

Derivation of the overdamped Langevin path probability ratio

These results follow . The backward step is realized by

Derivation of the Langevin probability ratio

These results follow . We define constants:

Then, the forward step of Brooks-Brünger-Karplus (BBK, leap-frog) Langevin dynamics are defined as:

Note that the factor $4$ in sqrt is different from – this factor is needed as we employ $\Delta t/2$ in both half-steps. The backward step with reversed momenta, $(\mathbf{x}_{t+1},-\mathbf{v}_{t+1})\rightarrow(\mathbf{x}_{t},-\mathbf{v}_{t})$ is then defined by:

Combining Eqs. (31), (32) and (35), we obtain:

Combining Eqs. (29), (34) and (35), we obtain:

To compute the path probability ratio we introduce the Jacobian

where the Jacobian ratio cancels as the Jacobians are independent of the noise variables.

Derivation of the probability ratio for Markov Chain Monte Carlo

For MCMC, $q_{t}$ satisfies the detailed balance condition

with respect to the potential function $u_{\lambda}$. We have

Derivation of the probability ratio for Hamiltonian MC with Metropolis acceptance

Hamiltonian MC with Metropolis acceptance defines a forward path density

which satisfies the joint detailed balance condition

Considering the velocity $\mathbf{v}$ is independently drawn from $\mathcal{N}(\mathbf{v}|\mathbf{0},\mathbf{I})$, the “marginal” forward path density of $\mathbf{y}_{t}\to\mathbf{y}_{t+1}$ is

Details on using SNFs for variational inference

Here we elaborate on the details of using SNFs as a variational approximation of the posterior distribution of a variational autoencoder (VAE) as presented in our last results section. In contrast to the usual notation used in common VAE literature, we choose $\mathbf{x}$ to indicate the latent variable, while we call the observed variable $\mathbf{s}$. This is due to being consistent with the use of $\mathbf{x}$ as the sampled variable of interest throughout our former discussions.

Together, this defines the joint distribution

Conditioned on a given $\mathbf{s}$, we can utilize a SNF to approximate the posterior distribution

If the SNF consists of only deterministic transformations, $J_{KL}$ is equivalent to $\mathcal{F}$ in .

We estimate $J_{KL}(\mathbf{s})$ on samples as

by sampling $M$ paths $\{(\mathbf{y}_{0}^{(i)},\ldots,\mathbf{y}_{T}^{(i)})\}_{i=1}^{M}$ for each $\mathbf{s}$ and setting $M=5$.

We define a simple base distribution $q(\mathbf{z})=\mathcal{N}(\mathbf{z}|\mathbf{0},\mathbf{I})$, together with a conditional diffeomorphism $F_{LL}(\mathbf{x}|\mathbf{s})$ transforming $\mathbf{z}$ to $\mathbf{x}$ and vice-versa conditioned on $\mathbf{s}$:

We realize such a conditional flow via RealNVP transformations, where coupling layers are additionally conditioned on $\mathbf{s}$ and only $\mathbf{x}/\mathbf{z}$ is transformed during the flow. Together with $q(\mathbf{z})$ this defines the conditional distribution

which we use as variational approximation to the true posterior. We then train $q_{LL}$ by minimizing the KL divergence

Hyper-parameters and other benchmark details

All experiments were run using PyTorch 1.2 and on GTX1080Ti cards. Optimization uses Adam with step-size $0.001$ and otherwise default parameters. All deterministic flow transformations use RealNVP . A RealNVP block is defined by two subsequent RealNVP layers that are swapped such that each channel gets transformed once as a function of the other channel. The affine transformation of each RealNVP layer is given by a fully connected ReLU network. For the NSF layers we substitute the simple affine transformations used in RealNVP by the rational-quadratic (RQ) spline transformation implemented in https://github.com/bayesiains/nflows. As before the width, height and slope of the RQ transformations are given by fully connected ReLU networks. Again a NSF block consists of two subsequent NSF layers with intermediate swap layers.

Both normalizing flow and SNF networks use 3 RealNVP blocks with three hidden layers of dimension 64. The SNF additionally uses 20 Metropolis MC steps per block using a Gaussian proposal density with standard deviation 0.25.

Training is done by minimizing $J_{ML}$ for 300 iterations and $\frac{1}{2}J_{ML}+\frac{1}{2}KL$ for 300 iterations using a batch-size of 128.

“Biased data” is defined by running local Metropolis MC in each of the two wells. These simulations do not transition to the other well and we use 1000 data points in each well for training.

“Unbiased data” is produced by running Metropolis MC with a large proposal step (standard deviation 1.5) to convergence and retaining 10000 data points for training.

In Table S1, the sampling results of SNFs with RealNVP blocks and Metropolis MC steps. MC step sizes of the first SNF is fixed to be $0.25$ as before, and all step sizes of the second one are trainable parameters in $[0.01,0.3]$. The other settings are the same as in Table 1.

Two-dimensional image densities in Figure 3

RealNVP and NSF flows both use 5 blocks. All involved transformation parameters (translation/scale in RealNVP layers, width/height/slope in NSF layers) use three hidden layers of dimension 64. For the NSF layers we used 20 knot points in the RQ-spline transformation. Training was done by minimizing $J_{ML}$ for 2000 iterations with batch-size 250.

Purely stochastic flow (column 2) uses five blocks with 10 Metropolis MC steps each using a Gaussian proposal density with standard deviation 0.1.

SNF (column 3/5) uses 5 blocks (RNVP/NSF block and 10 Metropolis MC steps with same parameters as above). Training was done by minimizing $J_{ML}$ for 6000 iterations with batch-size 250.

Alanine dipeptide in Fig. 5

Normalizing flow uses 3 RealNVP blocks with 3 hidden layers and $$nodes in their transformers. Training was done by minimizing$J_{ML}$ for 1000 iterations with batch-size 256.

SNF uses the same architecture and training parameters, but additionally 20 Metropolis MC steps each using a Gaussian proposal density with standard deviation 0.1.

As a last flow layer before $\mathbf{x}$, we used an invertible transformation between Cartesian coordinates and internal coordinates (bond lengths, angles, torsion angles) following the procedure described in . The internal coordinates were normalized by removing the mean and dividing by the standard deviation of their values in the training data.

Training data: We set up Alanine dipeptide in vacuum using OpenMMTools. Parameters are defined by the force field ff96 of the AMBER program . Simulations are run at standard OpenMMTools parameters with no bond constraints, 1 femtosecond time-step for $10^{6}$ time-steps (1 nanosecond) at a temperature of 1000 K in order to facilitate rapid exploration of the $\phi/\psi$ torsion angles and a few hundred transitions between metastable states. $10^{5}$ atom positions were saved as training data.

MNIST and Fashion-MNIST VAE in Table 3

where $[\mathbf{s}]_{i}$, $[D(\mathbf{x})]_{i}$ denote the $i$th pixel of $\mathbf{s}$ and the $i$th output of $D$.

Adam algorithm is used to train all models. Training was done by minimizing $\hat{J}_{KL}$ (see (37)) for 40 epochs with batch-size 128 and step size $10^{-3}$ unless otherwise stated.

In simple VAE, the encoder $E$ consists of 2 fully connected hidden layers, with 1024 nodes and ReLU non-linearities for each hidden layer. The encoder has $100$ outputs, where the activation function of the first $50$ outputs is the linear function and the activation function of the last $50$ outputs is the absolute value function. The transformation from $\mathbf{z}$ to $\mathbf{x}$ is given by

MCMC uses 30 Metropolis MC steps each using a overdamped Langevin proposal, where the interpolated potential are used. The interpolation coefficients and the step size of the proposal are both trained as parameters of the flow.

Normalizing flow uses 6 RealNVP blocks with 2 hidden layers and $$ nodes in their transformers.

SNF uses three units with each unit consisting of 2 RealNVP blocks + 10 Metropolis MC steps, where architectures are the same as the above. During the training procedure, we first train parameters of the 6 RealNVP blocks without the Metropolis MC steps for 20 epochs, and then train all parameters for another 20 epochs. The training step size is $10^{-3}$ for the first 20 epochs and $10^{-4}$ for the last 20 epochs.

Comparison with related sampling methods

A brief comparison of the proposed SNF and selected sampling methods with learnable proposals and transformations is provided in Table S2. Most previous sampling methods are developed based on the detailed balance in each step, except that HVI presented in can perform nonequilibrium sampling steps by using annealed target distributions. Furthermore, some sampling techniques and also improve the sampling efficiency by linear or nonlinear deterministic transformation, where the transformation is performed only once. It can be seen from the comparison that SNF provides a universal framework for sampling, where the deterministic and stochastic blocks can be flexibly designed and combined.

Supplementary Figures