MixFlows: principled variational inference via mixed flows

Consider a set $\mathcal{X}\subseteq\mathbb{R}^{d}$ and a target probability distribution $\pi$ on $\mathcal{X}$ whose density with respect to the Lebesgue measure we denote $\pi(x)$ for $x\in\mathcal{X}$. In the setting of Bayesian inference, $\pi$ is the posterior distribution that we aim to approximate, and we are only able to evaluate a function $p(x)$ such that $p(x)=Z\cdot\pi(x)$ for some unknown normalization constant $Z>0$. Throughout, we will assume all distributions have densities with respect to the Lebesgue measure on $\mathcal{X}$, and will use the same symbol to denote a distribution and its density; it will be clear from context what is meant.

Variational inference involves approximating the target distribution $\pi$ by minimizing the Kullback-Leibler (KL) divergence from $\pi$ to members of a parametric family $\{q_{\lambda}:\lambda\in\Lambda\}$, $\Lambda\subseteq\mathbb{R}^{p}$, i.e.,

The two objective functions in Equation 2 differ only by the constant $\log Z$. In order to be able to optimize $\lambda$ using standard techniques, the variational family $q_{\lambda}$, $\lambda\in\Lambda$ must enable both i.i.d. sampling and density evaluation. A common approach to constructing such a family is to pass draws from a simple reference distribution $q_{0}$ through a measurable function $T_{\lambda}:\mathcal{X}\to\mathcal{X}$; $T_{\lambda}$ is often referred to as a flow when comprised of repeated composed functions (Tabak & Turner, 2013; Rezende & Mohamed, 2015a; Kobyzev et al., 2021). If $T_{\lambda}$ is a diffeomorphism, i.e., differentiable and has a differentiable inverse, then we can express the density of $X=T_{\lambda}(Y)$, $Y\sim q_{0}$ as

In this case the optimization in Equation 2 can be rewritten using a transformation of variables as

One can solve the optimization problem Equation 4 using unbiased stochastic estimates of the gradient¹¹1We assume throughout that differentiation and integration can be swapped wherever necessary. with respect to $\lambda$ based on draws from $q_{0}$ (Salimans & Knowles, 2013; Kingma & Welling, 2014),

Variational flows are often constructed from a flexible, general-purpose parametrized family $\{T_{\lambda}:\lambda\in\Lambda\}$ that is not specialized for any particular target distribution (Papamakarios et al., 2021); it is the job of the KL divergence minimization Equation 4 to adapt the parameter $\lambda$ such that $q_{\lambda}$ becomes a good approximation of the target $\pi$. However, there are certain functions—in particular, those that are both measure-preserving and ergodic for $\pi$—that naturally provide a means to approximate expectations of interest under $\pi$ without the need for tuning. Intuitively, a measure-preserving map $T$ will not change the distribution of draws from $\pi$: if $X\sim\pi$, then $T(X)\sim\pi$. And an ergodic map $T$, when applied repeatedly, will not get “stuck” in a subset of $\mathcal{X}$ unless it has probability either 0 or 1 under $\pi$. The precise definitions are given in Definitions 2.1 and 2.2.

If a map $T$ satisfies both Definitions 2.1 and 2.2, then long-run averages resulting from repeated applications of $T$ will converge to expectations under $\pi$, as shown by Theorem 2.3. When $\mathcal{X}$ is compact, this result shows that the discrete measure $\frac{1}{N}\sum_{n=0}^{N-1}\delta_{T^{n}x}$ converges weakly to $\pi$ (Dajani & Dirksin, 2008, Theorem 6.1.7).

Based on this result, one might reasonably consider building a measure-preserving variational flow, i.e., $X=T(X_{0})$ where $X_{0}\sim q_{0}$. However, it is straightforward to show that measure-preserving bijections $T$ do not decrease the KL divergence (or any other $f$-divergence, e.g., total variation and Hellinger (Qiao & Minematsu, 2010, Theorem 1)):

Define a reference distribution $q_{0}$ on $\mathcal{X}$ for which i.i.d. sampling and density evaluation is tractable, and a collection of measurable functions $T_{\lambda}:\mathcal{X}\to\mathcal{X}$ parametrized by $\lambda\in\Lambda$. Then the MixFlow family generated by $q_{0}$ and $T_{\lambda}$ is

where $T^{n}_{\lambda}q_{0}$ denotes the pushforward of the distribution $q_{0}$ under $n$ repeated applications of $T_{\lambda}$.

If $T_{\lambda}:\mathcal{X}\to\mathcal{X}$ is a diffeomorphism with Jacobian $J_{\lambda}:\mathcal{X}\to\mathbb{R}$, we can express the density of $q_{\lambda,N}$ by using a transformation of variables formula on each component in the mixture:

This density can be computed efficiently using $N-1$ evaluations of $T^{-1}_{\lambda}$ and $J_{\lambda}$ each (Algorithm 2). For sampling, we can obtain an independent draw $X\sim q_{\lambda,N}$ by treating $q_{\lambda,N}$ as a mixture of $N$ distributions:

The procedure is given in Algorithm 1. On average, this computation requires $\mathbb{E}[K]=\frac{N-1}{2}$ applications of $T_{\lambda}$, and at most it requires $N-1$ applications. However, one often takes samples from $q_{\lambda,N}$ to estimate the expectation of some test function $f:\mathcal{X}\to\mathbb{R}$. In this case, one can use all intermediate states over a single pass of a trajectory rather than individual i.i.d. draws. In particular, we obtain an unbiased estimate of $\int f(x)q_{\lambda,N}(\mathrm{d}x)$ via

We refer to this estimate as the trajectory-averaged estimate of $f$. The trajectory-averaged estimate of $f$ is preferred over the naïve estimate based on a single draw $f(X),X\sim q_{\lambda,N}$, as its cost is of the same order ($N-1$ applications of $T_{\lambda}$) and its variance is bounded above by that of the naïve estimate $f(X)$, as shown in Proposition 3.1. See Section E.2, and Figure 6 in particular, for empirical verification of Proposition 3.1.

We can minimize the KL divergence from $q_{\lambda,N}$ to $\pi$ by maximizing the ELBO (Blei et al., 2017), given by

The trajectory-averaged ELBO estimate is thus

The naïve method to compute this estimate—sampling $X_{0}$ and then computing the log density ratio for each term—requires $O(N^{2})$ computation because each evaluation of $q_{\lambda,N}(x)$ is $O(N)$. Algorithm 3 provides an efficient way of computing $\widehat{\text{ELBO}}(\lambda,N)$ in $O(N)$ operations, which is on par with taking a single draw from $q_{\lambda,N}$ or evaluating $q_{\lambda,N}(x)$ once. The key insight in Algorithm 3 is that we can evaluate the collection of values $\{q_{\lambda,N}(X_{0}),q_{\lambda,N}(T_{\lambda}X_{0}),\dots,q_{\lambda,N}(T_{\lambda}^{N-1}X_{0})\}$ incrementally, starting from $q_{\lambda,N}(X_{0})$ and iteratively computing each $q_{\lambda,N}(T_{\lambda}^{n}X_{0})$ for increasing $n$ in constant time. Specifically, in practice one computes $q_{\lambda,N}(x)$ with a complexity of $O(N)$ and iteratively updates $q_{\lambda,N}(T_{\lambda}^{k}x)$ and the Jacobian for $k=1,2,\dots N-1$. Each update requires only constant cost if one precomputes and stores $T^{-N+1}_{\lambda}x,\dots,T^{N-1}_{\lambda}x$ and $J_{\lambda}(T_{\lambda}^{-N+1}x),\dots,J_{\lambda}(T_{\lambda}^{N-1}x)$. This then implies that Algorithm 3 requires $O(N)$ memory; Algorithm 5 in Appendix B provides a slightly more complicated $O(N)$ time, $O(1)$ memory implementation of Equation 15.

Ergodic MixFlow families are those where $\forall\lambda\in\Lambda$, $T_{\lambda}$ is measure-preserving and ergodic for $\pi$. In this setting, Theorems 4.1 and 4.2 show that MixFlows converge to the target weakly and in total variation as $N\to\infty$ regardless of the choice of $\lambda$ (i.e., regardless of parameter tuning effort). Thus, ergodic MixFlow families provide the same compelling convergence result as MCMC, but with the added benefit of unbiased ELBO estimates. We first demonstrate setwise and weak convergence. Recall that a sequence of distributions $q_{n}$ converges weakly to $\pi$ if for all bounded, continuous $f:\mathcal{X}\to\mathbb{R}$, $\lim_{n\to\infty}\int f(x)q_{n}(\mathrm{d}x)=\int f(x)\pi(\mathrm{d}x)$, and converges setwise to $\pi$ if for all measurable $A\subseteq\mathcal{X}$, $\lim_{n\to\infty}q_{n}(A)=\pi(A)$.

Using Theorem 4.1 as a stepping stone, we can obtain convergence in total variation. Recall that a sequence of distributions $q_{n}$ converges in total variation to $\pi$ if $\mathrm{D_{TV}}\mathchoice{\left({{{{q_{n},\pi}}}}\right)}{({{{{q_{n},\pi}}}})}{({{{{q_{n},\pi}}}})}{({{{{q_{n},\pi}}}})}=\sup_{A}\left|q_{n}(A)-\pi(A)\right|\to 0$ as $n\to\infty$. Note that similar nonasymptotic results exist for the ergodic average law of Markov chains (Roberts & Rosenthal, 1997), but as previously mentioned, these results do not apply to deterministic Markov kernels.

In practice, it is rare to be able to evaluate a measure-preserving map exactly; but approximations are commonly available. For example, in Section 5 we will create a measure-preserving map using Hamiltonian dynamics, and then approximate that map with a discretized leapfrog integrator. We therefore need to extend the result in Section 4.1 to apply to approximated maps.

Suppose we have two maps, $T$ and $\hat{T}$, with corresponding MixFlow families $q_{n}$ and $\hat{q}_{n}$ (suppressing $\lambda$ for brevity). Theorem 4.3 shows that the error of the MixFlow family $\hat{q}_{N}$ reflects an accumulation of the difference between the pushforward of each $q_{n}$ under $\hat{T}$ and $T$.

Suppose $T$ is measure-preserving and ergodic for $\pi$ and $q_{0}\ll N$. Then Theorem 4.2 implies that $\mathrm{D_{TV}}\mathchoice{\left({{{{q_{N},\pi}}}}\right)}{({{{{q_{N},\pi}}}})}{({{{{q_{N},\pi}}}})}{({{{{q_{N},\pi}}}})}\to 0$, and for the second term, we (very loosely) expect that $\mathrm{D_{TV}}\mathchoice{\left({{{{Tq_{n},\hat{T}q_{n}}}}}\right)}{({{{{Tq_{n},\hat{T}q_{n}}}}})}{({{{{Tq_{n},\hat{T}q_{n}}}}})}{({{{{Tq_{n},\hat{T}q_{n}}}}})}\approx\mathrm{D_{TV}}\mathchoice{\left({{{{\pi,\hat{T}\pi}}}}\right)}{({{{{\pi,\hat{T}\pi}}}})}{({{{{\pi,\hat{T}\pi}}}})}{({{{{\pi,\hat{T}\pi}}}})}$ for large $n$. Therefore, the second term should behave like $O(N\cdot\mathrm{D_{TV}}\mathchoice{\left({{{{\pi,\hat{T}\pi}}}}\right)}{({{{{\pi,\hat{T}\pi}}}})}{({{{{\pi,\hat{T}\pi}}}})}{({{{{\pi,\hat{T}\pi}}}})})$, i.e., increase linearly in $N$ proportional to how “non-measure-preserving” $\hat{T}$ is. This presents a trade-off between flow length and approximation quality: better approximations enable longer flows and lower minimal error bounds. Our empirical findings in Section 6 generally confirm this behavior. Corollary 4.4 further provides a more explicit characterization of this trade-off when $\hat{T}$ and its log-Jacobian $\log\hat{J}$ are uniformly close to their exact counterparts $T$ and $\log J$, and $\log q_{n}$ and $\log J$ are uniformly (over $n\in\mathbb{N}$) Lipschitz continuous. The latter assumption is reasonable in practical settings where we observe that the log-density of $q_{n}$ closely approximates $\pi$ (see the experimental results in Section 6).

This result states that the overall map-approximation error is $O(N\epsilon)$ when $\epsilon$ is small. Theorem 1 of Butzer & Westphal (1971) also suggests that $\mathrm{D_{TV}}\mathchoice{\left({{{{q_{N},\pi}}}}\right)}{({{{{q_{N},\pi}}}})}{({{{{q_{N},\pi}}}})}{({{{{q_{N},\pi}}}})}=O(1/N)$ in many cases, which hints that the bound should decrease roughly until $N=O(1/\mathchoice{{\hbox{$\displaystyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=4.30554pt,depth=-3.44446pt}}}{{\hbox{$\textstyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=4.30554pt,depth=-3.44446pt}}}{{\hbox{$\scriptstyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=3.01389pt,depth=-2.41113pt}}}{{\hbox{$\scriptscriptstyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=2.15277pt,depth=-1.72223pt}}})$, with error bound $O(\mathchoice{{\hbox{$\displaystyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=4.30554pt,depth=-3.44446pt}}}{{\hbox{$\textstyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=4.30554pt,depth=-3.44446pt}}}{{\hbox{$\scriptstyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=3.01389pt,depth=-2.41113pt}}}{{\hbox{$\scriptscriptstyle\sqrt{\epsilon\,}$}\lower 0.4pt\hbox{\vrule height=2.15277pt,depth=-1.72223pt}}})$. We leave a more careful investigation of this trade-off for future work.

Our main convergence results (Theorems 4.1 and 4.2) require that $\pi$ dominates the reference $q_{0}$, and that $T$ is both measure-preserving and ergodic. It is often straightforward to design a dominated reference $q_{0}$. Designing a measure-preserving map $T$ with an implementable approximation $\hat{T}\approx T$ is also often feasible. However, verifying the ergodicity of $T$ conclusively is typically a very challenging task. As a consequence, most past work involving measure-preserving transformations just asserts the ergodic hypothesis without proof; see the discussions in ver Steeg & Galstyan (2021, p. 4) and Tupper (2005, p. 2-3).

But because MixFlows provide the ability to estimate the Kullback-Leibler divergence up to a constant, it is not critical to prove convergence a priori (as it is in the case of MCMC, for example). Instead, we suggest using the results of Theorems 4.1, 4.2, 4.3 and 4.4 as a guiding recipe for constructing principled variational families. First, start by designing a family of measure-preserving maps $T_{\lambda}$, $\lambda\in\Lambda$. Next, approximate $T_{\lambda}$ with some tractable $T_{\lambda,\epsilon}$ including a tunable fidelity parameter $\epsilon\geq 0$, such as that $\epsilon\to 0$, $T_{\lambda,\epsilon}$ becomes closer to $T_{\lambda}$. Finally, build a MixFlow from $T_{\lambda,\epsilon}$, and tune both $\lambda$ and $\epsilon$ by maximizing the ELBO. We follow this recipe in Section 5 and verify it empirically in Section 6.

We construct $T_{\lambda}$ by composing the following three steps, which are all measure-preserving for $\bar{\pi}$:

In practice, we cannot simulate the dynamics in step (1) perfectly. Instead, we approximate the dynamics in Equation 21 by running $L$ steps of the leapfrog method, where each leapfrog map $\hat{H}_{\epsilon}:\mathbb{R}^{2d}\to\mathbb{R}^{2d}$ involves interleaving three discrete transformations with step size $\epsilon>0$,

Denote the map $T_{\lambda,\epsilon}$ to be the composition of the three steps with Hamiltonian dynamics replaced by the leapfrog integrator; Algorithm 6 in Appendix C provides the pseudocode. The final variational tuning parameters for the MixFlow are the step size $\epsilon$—which controls how close $T_{\lambda,\epsilon}$ is to being measure-preserving for $\bar{\pi}$—the Hamiltonian simulation length $L\in\mathbb{N}$, and the flow length $N$. In our experiments we tune the number of leapfrog steps $L$ and the step size $\epsilon$ by maximizing the ELBO. We also tune the number of refreshments $N$ to achieve a desirable computation-quality tradeoff by visually inspecting the convergence of the ELBO.

Density evaluation (Algorithm 2) and ELBO estimation (Algorithm 3) both involve repeated applications of $T_{\lambda}$ and $T_{\lambda}^{-1}$. This poses no issue in theory, but in a computer—with floating point computations—one should code the map $T_{\lambda}$ and its inverse $T^{-1}_{\lambda}$ in a numerically precise way such that $(T^{K}_{\lambda})\circ(T^{-K}_{\lambda})$ is the identity map for large $K\in\mathbb{N}$. Figure 8 in Section E.2 displays the severity of the numerical error when $T_{\lambda}$ and $T^{-1}_{\lambda}$ are not carefully implemented. In practice, we check the stability limits of our flow by taking draws from $q_{0}$ and evaluating $T^{K}_{\lambda}$ followed by $T^{-K}_{\lambda}$ (and vice versa) for increasing $K$ until we cannot reliably invert the flow. See Figure 7 in the appendix for an example usage of this diagnostic. Note that for sample generation specifically (Algorithm 1), numerical stability is not a concern as it only requires forward evaluation of the map $T_{\lambda}$.

We begin with a qualitative examination of the i.i.d. samples and the approximated targets produced by MixFlow initialized at $q_{0}=\mathcal{N}(0,1)$ for three one-dimensional synthetic distributions: a Gaussian, a mixture of Gaussians, and the standard Cauchy. We excluded the pseudotime variable $u$ here in order to visualize the full joint density of $(x,\rho)$ in 2 dimensions. More details can be found in Section E.1. Figures 1(a), 1(b) and 1(c) show histograms of $10{,}000$ i.i.d. $x$-marginal samples generated by MixFlow for each of the three targets, which nearly perfectly match the true target marginals. Figures 1(a), 1(b) and 1(c) also show that $\log q_{\lambda,N}$ is generally a good approximation of the log target density.

We then present similar visualizations on four more challenging synthetic target distributions: the banana (Haario et al., 2001), Neal’s funnel (Neal, 2003), a cross-shaped Gaussian mixture, and a warped Gaussian. All four examples have a 2-dimensional state $x\in\mathbb{R}^{2}$, and hence $(x,\rho,u)\in\mathbb{R}^{5}$. In each example we set the initial distribution $q_{0}$ to be the mean-field Gaussian approximation. More details can be found in Section E.2. Figures 1(d), 1(e), 1(f) and 1(g) shows the scatter plots consisting of $1{,}000$ i.i.d. $x$-marginal samples drawn from MixFlow, as well as the approximated MixFlow log density and exact log density sliced as a function of $x\in\mathbb{R}^{2}$ for a single value of $(\rho,u)$ chosen randomly via $(\rho,u)\sim{\sf{Lap}}(0,I)\times{\sf{Unif}}[0,1]$ (which is required for visualization, as $(x,\rho,u)\in\mathbb{R}^{5}$). We see that, qualitatively, both the samples and approximated densities from MixFlow closely match the target. Finally, Figure 9 in Section E.2 provides a more comprehensive set of sample histograms (showing the $x$-, $\rho$-, and $u$-marginals). These visualizations support our earlier theoretical analysis.

Next, we provide a quantitative comparison of MixFlow, NFs, NUTS on 7 real data experiments outlined in Section E.5. We tune each NF method under various settings (Tables 1, 2 and 3), and present the best one for each example. ELBOs of MixFlow are estimated with Algorithm 3, averaging over $1{,}000$ independent trajectories. ELBOs of NFs are based on $2{,}000$ samples. To obtain an assessment for the target marginal distribution itself (not the augmented target), we also compare methods using the kernel Stein discrepancy (KSD) with an inverse quadratic (IMQ) kernel (Gorham & Mackey, 2017). NUTS and NFs use $5{,}000$ samples for KSD estimation, while MixFlow is based on $2{,}000$ i.i.d. draws (Algorithm 1). For KSD comparisons, all variational methods are tuned by maximizing the ELBO (Figures 2 and 17).

In order to tune MixFlow, we simply run a 1-dimensional parameter sweep for the step size $\epsilon$, and use a visual inspection of the ELBO to set an appropriate number of flow steps $N$. Tuning an NF requires optimizing its architecture, number of layers, and its (typically many) parameters. Not only is this time consuming—in our experiments, tuning took 10 minutes to roughly 1 hour (Figure 4)—but the optimization can also behave in unintuitive ways. For example, performance can be heavily dependent on the number of flow layers, and adding more layers does not necessarily improve quality. Figures 16, 1, 2 and 3 show that using more layers does not necessarily help, and slows tuning considerably. In the case of RealNVP specifically, tuning can be unstable, especially for more complex models. The optimizer often returns NaN values for flow parameters during training (see Table 1). This instability has been noted in earlier work (Dinh et al., 2017, Sec. 3.7).

Finally, Figure 4 presents timing results for two of the real data experiments (additional comparisons in Figures 14 and 19). In this figure, we use MixFlow iid to refer to i.i.d. sampling from MixFlow, and MixFlow single to refer to collecting all intermediate states on a single trajectory. This result shows that the per sample time of MixFlow single is similar to NUTS and HMC as one would expect. MixFlow iid is the slowest because each sample is generated by passing through the entire flow. The NF generates the fastest draws, but recall that this comes at the cost of significant initial tuning time; in the time it takes NF to generate its first sample, MixFlow single has generated millions of samples in Figure 4. See Section E.4 for a detailed discussion of this trade-off.

Figures 4 and 19 further show the computational efficiency in terms of ESS per second on real data examples, which reflects the autocorrelation between drawn samples. Results show that MixFlow produce comparable ESS per second to HMC. MixFlow single behaves similarly to HMC as expected since the pseudo-momentum refreshment we proposed (steps (2-3) of Section 5) resembles the momentum resample step of HMC. The ESS efficiency of MixFlow iid depends on the trade-off between a slower sampling time and i.i.d. nature of drawn samples. In these real data examples, MixFlow iid typically produces a high ESS per second; but in the synthetic examples (Figure 14(b)), MixFlow iid is similar to the others.