A Unified Approach to Variational Autoencoders and Stochastic Normalizing Flows via Markov Chains

1 Introduction

Deep generative models for approximating complicated and often high-dimensional probability distributions became a rapidly developing research field. Variational autoencoders (VAEs) were originally introduced by Kingma and Welling

[21] and have seen a large number of modifications and improvements for a huge number of quite different applications. For some overview on VAEs, we refer to [22]. Recently, diffusion normalizing flows arising from the Euler discretization of a certain stochastic differential equation were proposed by Zhang and Chen in [39]. On the other hand, finite normalizing flows inclusive residual neural networks (ResNets) [3, 4, 17], invertible neural networks (INNs) [2, 8, 15, 20, 26, 29] and autoregessive flows [7, 9, 18, 27] are a popular class of generative models. To overcome topological constraints and improve the expressiveness of normalizing flow architectures, Wu, Köhler and Noé introduced stochastic normalizing flows [38] which combine deterministic, learnable flow transformations with stochastic sampling methods. In [14], we considered stochastic normalizing flows from a Markov chain point of view. In particular, we replaced the transition densities by general Markov kernels and established proofs via Radon-Nikodym derivatives. This allowed to incorporate deterministic flows or Metropolis-Hasting flows which do not have densities into the mathematical derivation.

The aim of this tutorial paper is to propose the straightforward and clear framework of Markov chains to combine deterministic normalizing and stochastic flows, in particular VAEs and diffusion normalizing flows. More precisely we establish a pair of Markov chains having some special properties. This provides a powerful tool for coupling different architectures. We want to highlight the advantages of the Markov chain approach that can handle distributions with and without densities in a sound mathematical way. We are aware that relations between normalizing flows and other approaches as VAEs were already mentioned in the literature and we point to corresponding references at the end of Section 5.

The outline of the paper is as follows: in the next Section 2 we recall the notation of Markov kernels. Then, in Section 4, we use them to explain normalizing flows. Stochastic normalizing flow were introduced as a pair of Markov chains in Section 4. Afterwards, we show how VAEs fit into the setting of stochastic normalizing flows in Section 5. Related references are given at end of the section. Finally, we demonstrate in Section 6 how diffusion normalizing flows can be seen as stochastic normalizing flows as well.

2 Markov Kernels

In this section, we introduce the basic notation of Markov chains, see, e.g., [24].

Let $(Ω, A, P)$

be a probability space. By a probability measure on

R^{d}

we always mean a probability measure defined on the Borel

σ

-algebra

B (R^{d})

. Let

P (R^{d})

denote the set of probability measures on

R^{d}

. Given a random variable

X : Ω \to R^{d}

, we use the push-forward notation

PX=X#P\coloneqqP∘X−1

for the corresponding measure on $R^{d}$ . A Markov kernel $K : R^{n} \times B (R^{d}) \to [0, 1]$ is a mapping such that

$K (\cdot, B)$ is measurable for any $B \in B (R^{d})$ , and
$K (x, \cdot)$ is a probability measure for any $x \in R^{n}$ .

For a probability measure $μ$ on $R^{n}$ , the measure $μ \times K$ on $R^{n} \times R^{d}$ is defined by

(μ×K)(A×B)\coloneqq∫AK(x,B)dμ(x).

(1)

Note that this definition captures all sets in $B (R^{n} \times R^{d})$ since the measurable rectangles form a $\cap$ -stable generator of $B (R^{n} \times R^{d})$ . Then, it holds for all integrable $f$ that

\int_{R^{n} \times R^{d}} f (x, y) d (μ \times K) (x, y) = \int_{R^{n}} \int_{R^{d}} f (x, y) d K (x, \cdot) (y) d μ (x) .

In the following, we use the notion of the regular conditional distribution of a random variable $X$ given a random variable $Y$ which is defined as the $P_{Y}$ -almost surely unique Markov kernel $P_{Y | X = \cdot} (\cdot)$ with the property

P_{X} \times P_{Y | X = \cdot} (\cdot) = P_{(X, Y)} .

(2)

We will use the abbreviation $P_{Y | X} = P_{Y | X = \cdot} (\cdot)$ if the meaning is clear from the context. A sequence $(X_{0}, \dots, X_{T})$ , $T \in N$ of $d_{t}$ -dimensional random variables $X_{t}$ , $t = 0, . . ., T$ , is called a Markov chain, if there exist Markov kernels

K_{t} = P_{X_{t} | X_{t - 1}} : R^{d_{t - 1}} \times B (R^{d_{t}}) \to [0, 1]

in the sense (2) such that

P_{(X_{0}, . . ., X_{T})} = P_{X_{0}} \times P_{X_{1} | X_{0}} \times \dots \times P_{X_{T} | X_{T - 1}} .

(3)

The Markov kernels $K_{t}$ are also called transition kernels. If the measure $K_{t} (x, \cdot) = P_{X_{t} | X_{t - 1} = x}$ has a density $k_{t} (x, y)$ , and $P_{X_{t - 1}}$ resp. $P_{X_{t}}$ have densities $p_{t - 1}$ resp. $p_{t}$ , then setting $A\coloneqqRdt−1$ in equation (1) results in

p_{t} (y) = \int_{R^{d_{t - 1}}} k_{t} (x, y) p_{t - 1} (x) d x .

(4)

3 Normalizing Flows

In this section, we show how normalizing flows can be interpreted as finite Markov chains. A normalizing flows [28] is often understood as deterministic, invertible transform, which we call $T_{θ} : R^{d} \to R^{d}$ .

For better readability, we likewise skip the dependence of $T_{θ}$ on the parameter $θ$ and write just $T = T_{θ}$ . Normalizing flows can be used to model the density $p_{X}$ of a distribution $P_{X}$ by a simpler distribution $P_{Z}$

, usually the standard normal distribution, by learning

T

such that it holds approximately

P_{X} \approx T_{#} P_{Z}, or equivalently P_{Z} \approx T_{#}^{- 1} P_{X} .

Note that we have by the change of variable formula for the corresponding densities

p_{T_{#} P_{Z}} (x) = p_{Z} (T^{- 1} (x)) | d e t \nabla T^{- 1} (x) | .

(5)

The approximation can be done by minimizing the Kullback-Leibler divergence

	$K L (P_{X}, T_{#} P_{Z})$	$= E_{x \sim P_{X}} [log (\frac{p_{X}}{p_{T_{#} Z}})] = E_{x \sim P_{X}} [log p_{X}] - E_{x \sim P_{X}} [log p_{T_{#} Z}]$		(6)
		$= E_{x \sim P_{X}} [log p_{X}] - E_{x \sim P_{X}} [log p_{Z} \circ T^{- 1}] - E_{x \sim P_{X}} [log \| d e t (\nabla T^{- 1}) \|] .$		(7)

Noting that the first summand is just a constant, this gives the loss function

L_{NF} (θ) = - E_{x \sim P_{X}} [log p_{Z} \circ T^{- 1}] - E_{x \sim P_{X}} [log | d e t (\nabla T^{- 1}) |] .

The network $T$ is constructed by concatenating smaller blocks

T = T_{T} \circ \dots \circ T_{1}

which are invertible networks on their own. Then, the blocks $T_{t} : R^{d} \to R^{d}$ generate a pair of Markov chains $((X_{0}, . . ., X_{T}), (Y_{T}, . . ., Y_{0}))$ by

X_{0} \sim P_{Z}, X_{t} = T_{t} (X_{t - 1}) and Y_{T} \sim P_{X}, Y_{t - 1} = T_{t}^{- 1} (Y_{t}) .

Here, for all $t = 0, . . ., T$ , the dimension $d_{t}$ of the random variables $X_{t}$ and $Y_{t}$ is equal to $d$ . The transition kernels $K_{t} = P_{X_{t} | X_{t - 1}}$ and $R_{t} = P_{Y_{t - 1} | Y_{t}}$ are given by the Dirac distributions

K_{t} (x, \cdot) = δ_{T_{t} (x)} and R_{t} (x, \cdot) = δ_{T_{t}^{- 1} (x)}

which can be seen by (2) as follows: for any $A, B \in B (R^{d})$ it holds

	$P_{(X_{t - 1}, X_{t})} (A \times B)$	$= \int_{R^{d}} 1_{A \times B} (x_{t - 1}, x_{t}) d P_{(X_{t - 1}, X_{t})} (x_{t - 1}, x_{t})$		(8)
		$= \int_{R^{d}} 1_{A} (x_{t - 1}) 1_{B} (x_{t}) d P_{(X_{t - 1}, X_{t})} (x_{t - 1}, x_{t}) .$		(9)

Since $P_{(X_{t - 1}, X_{t})}$ is by definition concentrated on the set ${(x_{t - 1}, T_{t} (x_{t - 1}) : x_{t - 1} \in R^{d}}$ , this becomes

$P_{(X_{t - 1}, X_{t})} (A \times B)$	$= \int_{R^{d} \times R^{d}} 1_{A} (x_{t - 1}) 1_{B} (T_{t} (x_{t - 1})) d P_{(X_{t - 1}, X_{t})} (x_{t - 1}, x_{t})$	(10)
	$= \int_{A} 1_{B} (T_{t} (x_{t - 1})) d P_{X_{t - 1}} (x_{t - 1})$	(11)
	$= \int_{A} δ_{T_{t} (x_{t - 1})} (B) d P_{X_{t - 1}} .$	(12)

Consequently, by (1), the transition kernel $K_{t} = P_{X_{t} | X_{t - 1}}$ is given by $K_{t} (x, \cdot) = δ_{T_{t} (x)}$ . Due to their correspondence to the layers $T_{t}$ and $T_{t}^{- 1}$ from the normalizing flow $T$ , we call the Markov kernels $K_{t}$ forward layers, while the Markov kernels $R_{t}$ are called reverse layers.

4 Stochastic Normalizing Flows

The idea of stochastic normalizing flows is to replace some of the deterministic layers $T_{t}$ from a normalizing flow by random transforms. From the Markov chains viewpoint, we replace the kernels $K_{t}$ and $R_{t}$ with the Dirac measure by more general Markov kernels.

Formally, a stochastic normalizing flow (SNF) is a pair $((X_{0}, \dots, X_{T}), (Y_{T}, . . ., Y_{0}))$ of Markov chains of $d_{t}$ -dimensional random variables $X_{t}$ and $Y_{t}$ , $t = 0, . . ., T$ , with the following properties:

$P_{X_{t}}, P_{Y_{t}}$ have the densities $p_{X_{t}}, p_{Y_{t}} : R^{d_{t}} \to R_{> 0}$ for any $t = 0, . . ., T$ .
There exist Markov kernels $K_{t} = P_{X_{t} | X_{t - 1}}$ and $R_{t} = P_{Y_{t - 1} | Y_{t}}$ , $t = 1, . . ., T$ such that

$P_{(X_{0}, . . ., X_{T})}$ $= P_{X_{0}} \times P_{X_{1} | X_{0}} \times \dots \times P_{X_{T} | X_{T - 1}},$ (13)

$P_{(Y_{T}, . . ., Y_{0})}$ $= P_{Y_{T}} \times P_{Y_{T - 1} | Y_{T}} \times \dots \times P_{Y_{0} | Y_{1}} .$ (14)
For $P_{X_{t}}$ -almost every $x \in R^{d_{t}}$ , the measures $P_{Y_{t - 1} | Y_{t} = x}$ and $P_{X_{t - 1} | X_{t} = x}$ are absolutely continuous with respect to each other.

We say that the Markov chain $(Y_{T}, . . ., Y_{0})$ is a reverse Markov chain of $(X_{0}, \dots, X_{T})$ . In applications, Markov chains usually start with a latent random variable

X_{0} = Z

on $R^{d_{0}}$ , where it is easy to sample from and we intend to learn the Markov chain such that $X_{T}$ approximates a target random variable $X$ on $R^{d_{T}}$ , while the reversed Markov chain is initialized with a random variable

Y_{T} = X

from a data space and $Y_{0}$ should approximate the latent variable $Z$ . As outlined in the previous paragraph, each deterministic normalizing flow is a special case of a SNF. In the following, let $N (μ, Σ)$ denote the normal distribution with density $N (\cdot; μ, Σ)$ .

4.1 Stochastic Layers

In the following, we briefly recall the two stochastic layers which were used in [14, 38]. Another kind of layer arising from VAEs is detailed in the next section. In both cases from [14, 38], we choose as for the deterministic layers

d_{t - 1} = d_{t} = d

and the basic idea is to push the distribution of $X_{t - 1}$ into the direction of some proposal density $p_{t} : R^{d} \to R_{> 0}$

, which is usually chosen as some interpolation between

p_{X}

and

p_{Z}

. For a detailed description of this interpolation, we refer to [14]. As reverse layer, we use the same Markov kernel as the forward layer, i.e.,

R_{t} = K_{t} .

Metropolis-Hastings (MH) Layer: The Metropolis-Hastings algorithm outlined in Alg. 1 is a frequently used Markov Chain Monte Carlo type algorithm to sample from a distribution $P$ with known density $p$ , see, e.g., [30].

Input:

x_{0} \in R^{d}

, proposal density

p_{t} : R^{d} \to R_{> 0}

Repeat

Draw

x^{'}

from

N (x_{k}, σ^{2} I)

and

u

uniformly in

[0, 1]

Compute the acceptance ratio

α(xk,x′)\coloneqqmin{1,p(x′)p(xk)}.

Set

xk+1\coloneqq{x′ifu<α(x′,xk),xk+1\coloneqqxkotherwise.

Output:

{x_{k}}_{k}

Algorithm 1 Metropolis-Hastings Algorithm

Under mild assumptions, the corresponding Markov chain $(X_{k})_{k \in N}$ admits the unique stationary distribution $P$ and $P_{X_{k}} \to P$ as $k \to \infty$ in the total variation norm, see, e.g. [35].

In the MH layer, the transition from $X_{t - 1}$ to $X_{t}$ is one step of a Metropolis-Hastings algorithm. More precisely, let $ξ_{t} \sim N (0, σ^{2} I)$ and $U \sim U_{[0, 1]}$ be random variables such that $(σ (ξ_{t}), σ (U), σ (\cup_{s \leq t - 1} σ (X_{s})))$ are independent. Here $σ (X)$ denotes the smallest $σ$ -algebra generated by the random variable $X$ . Then, we set

X_{t}

\coloneqqXt+1[U,1](αt(Xt−1,Xt−1+ξt))ξt

(15)

where

αt(x,y)\coloneqqmin{1,pt(y)pt(x)}

with a proposal density $p_{t}$ which has to be specified. The corresponding transition kernel $K_{t} : R^{d} \times B (R^{d}) \to [0, 1]$ was derived, e.g., in [34] and is given by

Kt(x,A)\coloneqq∫AN(y;x,σ2I)αt(x,y)dy+δx(A)∫RdN(y;x,σ2I)(1−αt(x,y))dy.

(16)

Note that another kind of MH layers coming from the Metropolis-adjusted Langevin algorithm (MALA) [11, 31] was also used in [14, 38] under the name Markov Chain Monte Carlo (MCMC) layer.
Langevin Layer: In the Langevin layer, we model the transition from $X_{t - 1}$ to $X_{t}$ by one step of an explicit Euler discretization of the overdamped Langevin dynamics [37]. Let $ξ_{t} \sim N (0, I)$ such that $σ (ξ_{t})$ and $σ (\cup_{s \leq t - 1} σ (X_{s}))$ are independent. Again we assume that we are given a proposal density $p_{t} : R^{d} \to R_{> 0}$ which has to be specified. We denote by $ut(x)\coloneqq−log(pt(x))$ the negative log-likelihood of $p_{t}$ and set

Xt\coloneqqXt−1−a1∇ut(Xt−1)+a2ξt,

where $a_{1}, a_{2} > 0$ are some predefined constants. To determine the corresponding kernel, we use the the independence of $ξ_{t}$ of $X_{t}$ and $X_{t - 1}$ to obtain that $X_{t}$ and $X_{t - 1}$ have the common density

$p_{(X_{t - 1}, X_{t})} (x_{t - 1}, x_{t})$	$= p_{X_{t - 1}, ξ_{t}} (x_{t - 1}, \frac{1}{a_{2}} (x_{t} - x_{t - 1} + a_{1} \nabla u_{t} (x_{t - 1}))$	(17)
	$= p_{X_{t - 1}} (x_{t - 1}) p_{ξ_{t}} (\frac{1}{a_{2}} (x_{t} - x_{t - 1} + a_{1} \nabla u_{t} (x_{t - 1})))$	(18)
	$= p_{X_{t - 1}} (x_{t - 1}) N (x_{t}; x_{t - 1} - a_{1} \nabla u_{t} (x_{t - 1}), a_{2}^{2} I) .$	(19)

Then, for $A, B \in B (R^{d})$ , it holds

$P_{(X_{t - 1}, X_{t})} (A \times B)$	$= \int_{A \times B} p_{X_{t - 1}} (x_{t - 1}) N (x_{t}; x_{t - 1} - a_{1} \nabla u_{t} (x_{t - 1}), a_{2}^{2} I) d (x_{t - 1}, x_{t})$	(20)
	$= \int_{A} \int_{B} N (x_{t}; x_{t - 1} - a_{1} \nabla u_{t} (x_{t - 1}), a_{2}^{2} I) d x_{t} p_{X_{t - 1}} (x_{t - 1}) d x_{t - 1}$	(21)
	$= \int_{A} K_{t} (x_{t - 1}, B) p_{X_{t}} (x_{t}) d P_{X_{t - 1}} (x_{t - 1}),$	(22)

where

Kt(x,⋅)\coloneqqN(x−a1∇ut(x),a22I).

(23)

By (1) and (2) this is the Langevin transition kernel $P_{X_{t} | X_{t - 1}}$ .

4.2 Training SNFs

We aim to find parameters of a SNF such that $P_{X_{T}} \approx P_{X}$ . Recall, that for deterministic normalizing flows, it holds that $P_{X_{T}} = T_{#} P_{Z}$ , such that the loss function $L_{NF}$ reads as $L_{NF} = K L (P_{X}, P_{X_{T}})$ . Unfortunately, the stochastic layers make it impossible to evaluate and minimize $K L (P_{X}, P_{X_{T}})$

. Instead, we minimize the KL divergence of the joint distributions

L_{SNF} = K L (P_{(Y_{0}, . . ., Y_{T})}, P_{(X_{0}, . . ., X_{T})}),

which is an upper bound of $K L (P_{Y_{T}}, P_{X_{T}}) = K L (P_{X}, P_{X_{T}})$ . It was shown in [14, Theorem 5] that this loss function can be rewritten as

$L_{SNF} (θ)$	$= K L (P_{(Y_{0}, . . ., Y_{T})}, P_{(X_{0}, . . ., X_{T})})$	(24)
	$= E_{(x_{0}, . . ., x_{T}) \sim P_{(Y_{0}, . . ., Y_{T})}} [log (\frac{p_{X} (x_{T})}{p_{X_{T}} (x_{T})} T \prod t = 1 f_{t} (x_{t - 1}, x_{t}))]$	(25)
	$= E_{(x_{0}, . . ., x_{T}) \sim P_{(Y_{0}, . . ., Y_{T})}} [log (\frac{p_{X} (x_{T})}{p_{Z} (x_{0})} T \prod t = 1 \frac{f_{t} (x_{t - 1}, x_{t}) p_{X_{t - 1} (x_{t - 1})}}{p_{X_{t} (x_{t})}})],$	(26)

where $f_{t} (\cdot, x_{t})$ is given by the Radon-Nikodym derivative $\frac{d P_{Y_{t - 1} | Y_{t} = x_{t}}}{d P_{X_{t - 1} | X_{t} = x_{t}}}$ . Finally, note that by [14, Theorem 6] we have for any deterministic normalizing flow that $L_{NF} = L_{SNF}$ .

5 VAEs as Special SNF Layers

In this section, we introduce variational autoencoders (VAEs) as another kind of stochastic layers of a SNF. First, we briefly revisit the definition of autoencoders and VAEs. Afterwards, we show that a VAE can be viewed as one-layer SNF.

Autoencoders.

Autoencoders (see [12]

for an overview) are a dimensionality reduction technique inspired by the principal component analysis. For

d > n

, an autoencoder is a pair

(E, D)

of neural networks, consisting of an encoder

E = E_{ϕ} : R^{d} \to R^{n}

and a decoder

D = D_{θ} : R^{n} \to R^{d}

, where

θ

and

ϕ

are the neural networks parameters. The network

E

aims to encode samples

x

from a

d

-dimensional distribution

P_{X}

in the lower-dimensional space

R^{n}

such that the decoder

D

is able to reconstruct them. Consequently, it is a necessary assumption that the distribution

P_{X}

is approximately concentrated on a

n

-dimensional manifold. A possible loss function to train

E_{ϕ}

and

D_{θ}

is given by

L_{AE} (ϕ, θ) = E_{x \sim P_{X}} [∥ x - D_{θ} (E_{ϕ} (x)) ∥^{2}] .

Using this construction, autoencoders have shown to be very powerful for reduce the dimensionality of very complex datasets.

Variational Autoenconders via Markov Kernels.

Variational autoencoders (VAEs) [21] aim to use the power of autoencoders to approximate a probability distribution $P_{X}$ with density $p_{X}$ using a simpler distribution $P_{Z}$ with density $p_{Z}$ which is usually the standard normal distribution. Here, the idea is to learn random transforms that push the distribution $P_{X}$ onto $P_{Z}$ and vice versa. Formally, these transforms are defined by the Markov kernels

K (z, \cdot) = N (μ_{θ} (z), Σ_{θ} (z)) % a n d R (x, \cdot) = N (μ_{ϕ} (x), Σ_{ϕ} (x)),

(27)

where

D (z) = D_{θ} (z) = (μ_{θ} (z), Σ_{θ} (z))

is a neural network with parameters $θ$ , which determines the parameters of the normal distribution within the definition of $K$ . Similarly,

E (x) = E_{ϕ} (x) = (μ_{ϕ} (x), Σ_{ϕ} (x))

determines the parameters within the definition of $R$ . In analogy to the autoencoders in the previous paragraph, $D$ and $E$ are called stochastic decoder and encoder. By definition, $K (z, \cdot)$ has the density $p_{θ} (x | z) = N (x; μ_{θ} (z), Σ_{θ} (z))$ and $R (x, \cdot)$ has the density $q_{ϕ} (z | x) = N (z; μ_{ϕ} (x), Σ_{ϕ} (x))$ .

Now, we aim to learn the parameters $θ$ such that it holds approximately

p_{X} (x) \approx \int_{z \in R^{n}} p_{θ} (x | z) p_{Z} (z) d z or equivalently P_{X} (A) \approx \int_{R^{n}} K (z, A) d P_{Z} (z) .

(28)

Assuming that the above equation holds true exactly, we can generate samples from $P_{X}$ by sampling first $z$ from $P_{Z}$ and secondly, sampling $x$ from $K (z, \cdot)$ .

The first idea, would be to use the maximum likelihood estimator as loss function, i.e., maximize

E_{x \sim P_{X}} [log (p_{θ} (x))], p_{θ} (x) = \int_{z \in R^{n}} p_{θ} (x | z) p_{Z} (z) d z .

Unfortunately, computing the integral directly is intractable. Thus, using Bayes’ formula

p_{θ} (z | x) = \frac{p_{θ} (x | z) p_{Z} (z)}{p_{θ} (x)},

we artificially incorporate the stochastic encoder by the computation

$log (p_{θ} (x))$	$= E_{z \sim q_{ϕ} (\cdot \| x)} [log (\frac{p_{θ} (x) p_{θ} (z \| x)}{p_{θ} (z \| x)})]$	(29)
	$= E_{z \sim q_{ϕ} (\cdot \| x)} [log (\frac{p_{θ} (x) p_{θ} (z \| x)}{q_{θ} (z \| x)})] + E_{z \sim q_{ϕ} (\cdot \| x)} [log (\frac{q_{ϕ} (z \| x)}{p_{θ} (z \| x)})]$	(30)
	$= E_{z \sim q_{ϕ} (\cdot \| x)} [log (\frac{p_{θ} (x \| z) p_{Z} (z)}{q_{θ} (z \| x)})] + K L (q_{ϕ} (\cdot \| x), p_{θ} (\cdot \| x))$	(31)
	$\geq E_{z \sim q_{ϕ} (\cdot \| x)} [log (\frac{p_{θ} (x \| z) p_{Z} (z)}{q_{θ} (z \| x)})]$	(32)

Then the loss function given by

(33)

is a lower bound on the so-called evidence $log (p_{θ} (x))$ . Therefore it is called the evidence lower bound (ELBO). Now the parameters $θ$ and $ϕ$ of the VAE $(E_{ϕ}, D_{θ})$ can be trained by maximizing the expected ELBO, i.e., by minimizing the loss function

L_{VAE} (θ, ϕ) = - E_{x \sim P_{X}} [L_{θ, ϕ} (x)] .

(34)

VAEs as one Layer SNFs.

In the following, we show that a VAE is a special case of one layer SNF. Let $((X_{0}, X_{1}), (Y_{1}, Y_{0}))$ be a one-layer SNF, where the layers $K_{1} = K = P_{X_{1} | Z}$ and $R_{1} = R = P_{Y_{0} | X}$ are defined as in (27) with densities $p_{θ} (\cdot | z)$ and $q_{ϕ} (\cdot | x)$ , respectively. Note that in contrast to the stochastic layers from Section 4 the dimensions $d_{0}$ and $d_{1}$ are not longer equal. Now, with $X_{0} = Z$ , the loss function (25) of the SNF reads as

L_{SNF} (θ, ϕ) = E_{(z, x) \sim P_{(Y_{0}, Y_{1})}} [- log (\frac{p_{X_{1}} (x)}{p_{X} (x) f_{1} (z, x)})],

(35)

where $f_{1} (\cdot, x)$ is given by the Radon-Nikodym derivative $\frac{d P_{Y_{0} | Y_{1} = x}}{d P_{X_{0} | X_{1} = x}}$ . Now we can use the fact, that by the definition of $K$ and $R$ the random variables $(Y_{0}, Y_{1})$ as well as the random variables $(X_{0}, X_{1})$ have a joint density to express $f_{1}$ by the corresponding densities $p_{X_{0} | X_{1} = x}$ of $p_{Y_{0} | Y_{1} = x}$ . Together with Bayes’ formula we obtain

f_{1} (z, x) = \frac{p_{Y_{0} | Y_{1} = x} (z)}{p_{X_{0} | X_{1} = x} (z)} = \frac{q_{ϕ} (z | x)}{p_{X_{1} | X_{0} = z} (x)} \frac{p_{X_{1}} (x)}{p_{X_{0}} (z)} = \frac{q_{ϕ} (z | x)}{p_{θ} (x | z)} \frac{p_{X_{1}} (x)}{p_{Z} (z)} .

Inserting this into (35), we get

L_{SNF} (θ, ϕ)

= E_{(z, x) \sim P_{(Y_{0}, Y_{1})}} [- log (\frac{}{p_{θ} (x | z) p_{Z} (z)} q_{ϕ} (z | x) p_{X} (x))]

(36)

and using (2) further

$L_{SNF} (θ, ϕ)$	$= E_{x \sim P_{X}} [E_{z \sim R (x, \cdot)} [- log (\frac{p_{θ} (x \| z) p_{Z} (z)}{q_{ϕ} (z \| x) p_{X} (x)})]]$	(37)
		(38)
		(39)

where $L_{θ, ϕ}$ denotes the ELBO as defined in (33) and $E_{x \sim P_{X}} [log (p_{X} (x))]$ is a constant independent of $θ$ and $ϕ$ . Consequently, minimizing $L_{SNF} (θ, ϕ)$ is equivalent to minimize the negative expected ELBO, which is exactly the loss $L_{VAE} (θ, ϕ)$ for VAEs from (34).

The above result could alternatively be derived via the relation of the ELBO to the the KL divergence between the probability measures defined by the densities $(x, z) \mapsto p_{θ} (x | z) p_{Z} (z)$ and $(x, z) \mapsto q_{ϕ} (z | x) p_{X} (x)$ , see [22, Section 2.7].

Related Combinations of VAEs and Normalizing flows.

There exist several works which model the latent distribution $P_{Z}$ of a VAE by normalizing flows [6, 29], SNFs [38] or sampling-based Monte Carlo methods [36] and often achieve state-of-the-art results. Using the above derivation, all of these models can be viewed as special cases of a SNF, even though some of them employ different training techniques for minimizing the loss function. Further, the authors of [13] modify the learning of the covariance matrices of decoder and encoder of a VAE using normalizing flows. However, analogously as in the previous section, this can be viewed as one-layer SNF.

A similar idea was applied in [25], where the authors model the weight distribution of a Baysian neural network by a noramlizing flow. However, we are not completely sure, how this approach relates to SNFs.

Finally, to overcome the problem of expansive training in high dimensions, some recent papers [5, 23] propose also other combinations of a dimensionality reduction and normalizing flows. However, [5] can be viewed as a variational autoencoder with special structured generator and can therefore be considered as one-layer SNF. In [23] the authors propose to reduce the dimension in a first step by a non-variational autoencoder and the optimization of a normalizing flow in the reduced dimensions in a second step.

6 Diffusion Normalizing Flows as special SNFs

Recently, Song et. al [33] proposed to learn the drift $f_{t} : R^{d} \to R^{d}$ , $t \in [0, S]$ and diffusion coefficient $g_{t} \in R$ of a stochastic differential equation

d X_{t} = f_{t} (X_{t}) d t + g_{t} d B_{t}

(40)

with respect to the Brownian motion $B_{t}$ , such that it holds approximately $X_{S} \sim P_{X}$ for some $S > 0$ and some data distribution $P_{X}$ . The explicit Euler discretization of (40) with step size $ϵ > 0$ reads as

X_{t} = X_{t - 1} + ϵ f_{t - 1} (X_{t - 1}) + \sqrt{ϵ} g_{t - 1} ξ_{t - 1}, t = 1, . . ., T,

where $ξ_{t - 1} \sim N (0, I)$ is independent of $X_{0}, . . ., X_{t - 1}$ . With a similar computation as for the Langevin layers, this corresponds to the Markov kernel

Kt(x,⋅)\coloneqqPXt|Xt−1=x=N(x+ϵft−1(x),ϵg2t−1).

(41)

Song et. al. parametrize the functions $f_{t} : R^{d} \to R^{d}$ by some a-priori learned score network [19, 32] and achieve competitive performance in image generation. Motivated by the time-reversal [1, 16, 10] of the SDE (40), Zhang and Chen [39] introduce the backward layer

R_{t} (x, \cdot) = P_{Y_{t - 1} | Y_{t} = x} = N (x + ϵ (f_{t} (x) - g_{t}^{2} s_{t} (x)), ϵ g_{t}^{2})

and learn the parameters of the neural networks $f_{t} : R^{d} \to R^{d}$ and $s_{t} : R^{d} \to R^{d}$ using the loss function $L_{SNF}$ from (25) to achieve state-of-the-art results. Even though Zhang and Chen call their model diffusion flow, it is indeed a special case of a SNF using the forward and backward layers $K_{t}$ and $R_{t}$ .

On the other hand, not every SNF can be represented as a discretized SDE. For example, the forward layer (16) from the MH layer has not the form (41).

References

[1] B. D. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313–326, 1982.
[2] L. Ardizzone, J. Kruse, C. Rother, and U. Köthe. Analyzing inverse problems with invertible neural networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
[3] J. Behrmann, W. Grathwohl, R. T. Chen, D. Duvenaud, and J.-H. Jacobsen. Invertible residual networks. In
International Conference on Machine Learning
, pages 573–582, 2019.
[4] R. T. Q. Chen, J. Behrmann, D. K. Duvenaud, and J.-H. Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[5] E. Cunningham, R. Zabounidis, A. Agrawal, I. Fiterau, and D. Sheldon. Normalizing flows across dimensions. arXiv preprint arXiv:2006.13070, 2020.
[6] B. Dai and D. P. Wipf. Diagnosing and enhancing VAE models. In International Conference on Learning Representations, 2019.
[7] N. De Cao, I. Titov, and W. Aziz. Block neural autoregressive flow. ArXiv:1904.04676, 2019.
[8] L. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using real NVP. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017.
[9] C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios. Neural spline flows. Advances in Neural Information Processing Systems, 2019.
[10] C. Durkan and Y. Song. On maximum likelihood training of score-based generative models. arXiv preprint arXiv:2101.09258, 2021.
[11] M. Girolami and B. Calderhead. Riemann manifold Langevin and Hamiltonian Monte Carlo methods. J. R. Stat. Soc.: Series B (Statistical Methodology), 73(2):123–214, 2011.
[12] I. Goodfellow, Y. Bengio, and A. Courville. Deep learning. MIT press, 2016.
[13] A. A. Gritsenko, J. Snoek, and T. Salimans.
On the relationship between normalising flows and variational- and denoising autoencoders.
In Deep Generative Models for Highly Structured Data, ICLR 2019 Workshop, 2019.
[14] P. Hagemann, J. Hertrich, and G. Steidl. Stochastic normalizing flows for inverse problems: a Markov Chains viewpoint. arXiv preprint arXiv:2109.11375, 2021.
[15] P. L. Hagemann and S. Neumayer. Stabilizing invertible neural networks using mixture models. Inverse Problems, 2021.
[16] U. G. Haussmann and E. Pardoux. Time reversal of diffusions. The Annals of Probability, pages 1188–1205, 1986.
[17] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In
Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition
, pages 770–778, 2016.
[18] C.-W. Huang, D. Krueger, A. Lacoste, and A. Courville. Neural autoregressive flows. In Proc. of the 35th International Conference on Machine Learning, pages 2078–2087. PMLR, 2018.
[19] A. Hyvärinen and P. Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[20] D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. ArXiv preprint arXiv:1807.03039, 2018.
[21] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[22] D. P. Kingma and M. Welling. An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4):307–392, 2019.
[23] K. Kothari, A. Khorashadizadeh, M. de Hoop, and I. Dokmanić. Trumpets: Injective flows for inference and inverse problems. arXiv preprint arXiv:2102.10461, 2021.
[24] J.-F. Le Gall. Brownian motion, martingales, and stochastic calculus, volume 274 of Graduate Texts in Mathematics. Springer, [Cham], 2016.
[25] C. Louizos and M. Welling. Multiplicative normalizing flows for variational Bayesian neural networks. In Proc. of the 34th International Conference on Machine Learning, pages 2218–2227. PMLR, 2017.
[26] T. Müller, McWilliams, R. B., M. F., Gross, and J. Novák. Neural importance sampling. ArXiv:1808.03856, 2018.
[27] G. Papamakarios, T. Pavlakou, , and I. Murray. Masked autoregressive flow for density estimation. Advances in Neural Information Processing Systems, page 2338–2347, 2017.
[28] D. Rezende and S. Mohamed. Variational inference with normalizing flows. In F. Bach and D. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1530–1538, Lille, France, 07–09 Jul 2015. PMLR.
[29] D. J. Rezende and S. Mohamed. Variational inference with normalizing flows. ArXiv:1505.05770, 2015.
[30] G. O. Roberts and J. S. Rosenthal. General state space Markov chains and MCMC algorithms. Probabability Surveys, 1:20 – 71, 2004.
[31] G. O. Roberts and R. L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, 2(4):341–363, 1996.
[32] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint arXiv:1907.05600, 2019.
[33] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
[34] L. Tierney. A note on Metropolis-Hastings kernels for general state spaces. Annals of Applied Probability, 8(1):1–9, 1998.
[35] D. Tsvetkov, L. Hristov, and R. Angelova-Slavova. On the convergence of the Metropolis-Hastings Markov chains. arXiv 1302.0654v4, 2020.
[36] A. Vahdat, K. Kreis, and J. Kautz. Score-based generative modeling in latent space. ArXiv preprint arXiv:2106.05931, 2021.
[37] M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference on Machine Learning, page 681–688, 2011.
[38] H. Wu, J. Köhler, and F. Noé. Stochastic normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 2020, 2020.
[39] Q. Zhang and Y. Chen. Diffusion normalizing flow. In Conference on Neural Information Processing Systems, 2021.

A Unified Approach to Variational Autoencoders and Stochastic Normalizing Flows via Markov Chains

Stochastic Normalizing Flows for Inverse Problems: a Markov Chains Viewpoint

SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flows

Ladder Variational Autoencoders

Diffusion Priors In Variational Autoencoders

Linearizing Visual Processes with Convolutional Variational Autoencoders

Reversible Markov chains: variational representations and ordering

Predictive Coding, Variational Autoencoders, and Biological Connections

1 Introduction

2 Markov Kernels

3 Normalizing Flows

4 Stochastic Normalizing Flows

4.1 Stochastic Layers

4.2 Training SNFs

5 VAEs as Special SNF Layers

Autoencoders.

Variational Autoenconders via Markov Kernels.

VAEs as one Layer SNFs.

Related Combinations of VAEs and Normalizing flows.

6 Diffusion Normalizing Flows as special SNFs

References

	$P_{(X_{0}, . . ., X_{T})}$	$= P_{X_{0}} \times P_{X_{1} \| X_{0}} \times \dots \times P_{X_{T} \| X_{T - 1}},$		(13)
	$P_{(Y_{T}, . . ., Y_{0})}$	$= P_{Y_{T}} \times P_{Y_{T - 1} \| Y_{T}} \times \dots \times P_{Y_{0} \| Y_{1}} .$		(14)

A Unified Approach to Variational Autoencoders and Stochastic Normalizing Flows via Markov Chains

Related Research

Stochastic Normalizing Flows for Inverse Problems: a Markov Chains Viewpoint

SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flows

Ladder Variational Autoencoders

Diffusion Priors In Variational Autoencoders

Linearizing Visual Processes with Convolutional Variational Autoencoders

Reversible Markov chains: variational representations and ordering

Predictive Coding, Variational Autoencoders, and Biological Connections

1 Introduction

2 Markov Kernels

3 Normalizing Flows

4 Stochastic Normalizing Flows

4.1 Stochastic Layers

4.2 Training SNFs

5 VAEs as Special SNF Layers

Autoencoders.

Variational Autoenconders via Markov Kernels.

VAEs as one Layer SNFs.

Related Combinations of VAEs and Normalizing flows.

6 Diffusion Normalizing Flows as special SNFs

References