Variational autoencoders and pathwise training

Let $(\mathcal{X},{\mathcal{A}}_X,{\mu }_X)$ and $(\mathcal{Z},{\mathcal{A}}_Z,{\mu }_Z)$ be sigma-finite measure spaces; that is, each has a countable measurable cover by sets of finite measure. Let $\mathrm{\Theta }\subseteq {\mathbb{R}}^{d_{\theta }}$ and $\mathrm{\Phi }\subseteq {\mathbb{R}}^{d_{\varphi }}$ be parameter sets of finite dimension. A latent-variable model with parameter $\theta \in \mathrm{\Theta }$ consists of a measurable prior density $r_{\theta }(z)\ge 0$ on $\mathcal{Z}$ and a decoder density $d_{\theta }(x\mid z)\ge 0$ on $\mathcal{X}$ for each $z$, with ${\int }_{\mathcal{Z}}{r}_{\theta }(z)d{\mu }_Z(z)=1,{\int }_{\mathcal{X}}{d}_{\theta }(x\mid z)d{\mu }_X(x)=1.$ The decoder is jointly measurable in $(x,z)$, and the joint model density is $p_{\theta }(x,z)={\gamma }_{\theta }(z;x):=r_{\theta }(z)d_{\theta }(x\mid z)$ relative to ${\mu }_X\otimes {\mu }_Z$. The decoder is also called the generative model: after drawing $z$ from the prior, it assigns a conditional density to the observation $x$.

For a fixed observation $x$, the evidence density and posterior density are $Z_{\theta }(x):=p_{\theta }(x):={\int }_{\mathcal{Z}}{\gamma }_{\theta }(z;x)d{\mu }_Z(z),{\pi }_{\theta }(z\mid x):=\frac{{\gamma }_{\theta }(z;x)}{Z_{\theta }(x)}\mathrm{.}$ Use this posterior only when $0<Z_{\theta }(x)<\mathrm{\infty }$. A variational encoder is a jointly measurable family of functions $q_{\varphi }(z\mid x)\ge 0$ on $\mathcal{Z}$, parameterized by $\varphi \in \mathrm{\Phi }$ and normalized by ${\int }_{\mathcal{Z}}{q}_{\varphi }(z\mid x)d{\mu }_Z(z)=1$ for every $(\varphi ,x)$; the same parameter $\varphi$ is used for each observation. For an admissible triple $(\theta ,\varphi ,x)$ require $0<Z_{\theta }(x)<\mathrm{\infty }$ and $q_{\varphi }(z\mid x)=0{\mu }_Z\text{-almost everywhere on }\{z:{\pi }_{\theta }(z\mid x)=0\},{\int }_{\mathcal{Z}}{q}_{\varphi }(z\mid x)(\mathrm{\mid }\log {\gamma }_{\theta }(z;x)\mathrm{\mid }+\mathrm{\mid }\log q_{\varphi }(z\mid x)\mathrm{\mid })d{\mu }_Z(z)<\mathrm{\infty }\mathrm{.}$ Terms multiplied by $q_{\varphi }=0$ are defined to be zero. The evidence lower bound is ${\mathrm{ELBO}}_{\theta ,\varphi }(x):={\int }_{\mathcal{Z}}{q}_{\varphi }(z\mid x)(\log {\gamma }_{\theta }(z;x)-\log q_{\varphi }(z\mid x))d{\mu }_Z(z)\mathrm{.}$

A pathwise representation of the encoder at $x$ is a probability space $(\mathcal{E},\mathcal{F},\rho )$ whose law $\rho$ is independent of $\theta$ and $\varphi$, together with a measurable map $g_{\varphi }(ϵ,x)$ such that the law of $g_{\varphi }(ϵ,x)$ for $ϵ\sim \rho$ is the measure with density $q_{\varphi }(\cdot \mid x)$ relative to ${\mu }_Z$. This representation defines how to generate encoder samples; it does not by itself assert that differentiation may pass through an expectation.

The generic admissibility conditions and ELBO convention.

The posterior-kernel formula on positive finite evidence fibres.

Kingma and Welling, Auto-Encoding Variational Bayes, arXiv v1.

Blei, Kucukelbir, and McAuliffe, Variational Inference: A Review for Statisticians.

Zheng, Hu, and Deng, 10-708 Lecture 13: Variational Inference II.

Fix an observation $x$ in a VAE with latent reference measure ${\mu }_Z$. Let ${\gamma }_{\theta }(\cdot ;x)$ be its nonnegative measurable joint-density slice, let $Z_{\theta }(x)=\int {\gamma }_{\theta }(z;x)d{\mu }_Z(z)\in (0,\mathrm{\infty })$, and let ${\pi }_{\theta }(z\mid x)={\gamma }_{\theta }(z;x)\mathrm{/}{Z}_{\theta }(x)$. Let $q_{\varphi }(\cdot \mid x)$ be a measurable probability density satisfying $q_{\varphi }=0$ ${\mu }_Z$-almost everywhere on $\{{\pi }_{\theta }=0\}$ and $\int q_{\varphi }(z\mid x)(\mathrm{\mid }\log {\gamma }_{\theta }(z;x)\mathrm{\mid }+\mathrm{\mid }\log q_{\varphi }(z\mid x)\mathrm{\mid })d{\mu }_Z(z)<\mathrm{\infty }\mathrm{.}$ Then the choice-free density-ratio identity is $\log Z_{\theta }(x)={\mathrm{ELBO}}_{\theta ,\varphi }(x)+\int q_{\varphi }(z\mid x)\log \text{\hspace{0.17em}⁣}\left(\frac{q_{\varphi }(z\mid x)}{{\pi }_{\theta }(z\mid x)}\right)d{\mu }_Z(z),$ where the ratio integrand is zero when $q_{\varphi }=0$. Under the Axiom of Choice, the final integral is the library reverse-KL divergence $D_{\mathrm{K}\mathrm{L}}(q_{\varphi }{\mu }_Z\mathrm{\parallel }{\pi }_{\theta }{\mu }_Z)$ and is nonnegative; hence ${\mathrm{ELBO}}_{\theta ,\varphi }(x)\le \log Z_{\theta }(x)$.

The VAE density, evidence, posterior, and ELBO definitions.

The generic choice-free density-ratio identity and AC-dependent KL branch.

The Axiom of Choice.

Blei, Kucukelbir, and McAuliffe, Variational Inference: A Review for Statisticians.

Zheng, Hu, and Deng, 10-708 Lecture 13: Variational Inference II.

Let $X=(x_1,\dots ,x_N)$ be a fixed nonempty finite data set, let $\mathrm{\Lambda }$ be an open subset of a finite-dimensional Euclidean parameter space for $\eta =(\theta ,\varphi )$, and fix ${\eta }_0=({\theta }_0,{\varphi }_0)\in \mathrm{\Lambda }$. For each $i$, assume the VAE at $(\eta ,x_i)$ is admissible and has a pathwise representation $g_{\varphi }(ϵ,x_i)$ under one parameter-independent probability law $\rho$. Assume there is one $\rho$-null set, common to every $i$ and $\eta \in \mathrm{\Lambda }$, outside which $0<q_{\varphi }(g_{\varphi }(ϵ,x_i)\mid x_i)<\mathrm{\infty }$ and $0<{\gamma }_{\theta }(g_{\varphi }(ϵ,x_i);x_i)<\mathrm{\infty }$; define the payoff below to be zero on that null set. Define $H_i(\eta ,ϵ):=\log {\gamma }_{\theta }(g_{\varphi }(ϵ,x_i);x_i)-\log q_{\varphi }(g_{\varphi }(ϵ,x_i)\mid x_i),$ with terms on a common $\rho$-null set defined to be zero. Assume each $H_i$ is jointly measurable and integrable for every $\eta \in \mathrm{\Lambda }$. Outside one $\rho$-null set, each $H_i(\cdot ,ϵ)$ is continuously differentiable on $\mathrm{\Lambda }$ with continuous gradient, and ${\mathrm{\nabla }}_{\eta }{H}_i(\eta ,\cdot )$ is measurable for every $\eta$. For every compact neighborhood $K\subset \mathrm{\Lambda }$ of ${\eta }_0$, assume one nonnegative $M_{i,K}\in L^1(\rho )$ satisfies $\mathrm{\parallel }{\mathrm{\nabla }}_{\eta }{H}_i(\eta ,ϵ){\mathrm{\parallel }}_2\le M_{i,K}(ϵ)$ for every $\eta \in K$ outside one common null set, and assume ${\mathbb{E}}_{\rho }\mathrm{\parallel }{\mathrm{\nabla }}_{\eta }{H}_i({\eta }_0,ϵ){\mathrm{\parallel }}_2^2<\mathrm{\infty }$.

Let $I$ be uniform on $\{1,\dots ,N\}$ and independent of $ϵ\sim \rho$. Then the finite-average objective $\bar{\mathcal{L}}(\eta ):=\frac{1}{N}{\sum }_{i=1}^N{\mathrm{ELBO}}_{\theta ,\varphi }(x_i)$ is differentiable at ${\eta }_0$, and its one-example pathwise gradient estimate $G({\eta }_0,I,ϵ):={\mathrm{\nabla }}_{\eta }{H}_I({\eta }_0,ϵ)$ satisfies $\mathbb{E}[G({\eta }_0,I,ϵ)]={\mathrm{\nabla }}_{\eta }\bar{\mathcal{L}}({\eta }_0)$ and has finite coordinate variances. A parameter-independent Gaussian base law alone does not imply the differentiability, envelope, or second-moment hypotheses.

The VAE ELBO and fixed-law encoder representation.

The published pathwise simulation-gradient theorem with a common integrable derivative envelope.

Dominated convergence used in the difference-quotient proof.

Kingma and Welling, Auto-Encoding Variational Bayes, arXiv v1.

Blei, Kucukelbir, and McAuliffe, Variational Inference: A Review for Statisticians.

Zheng, Hu, and Deng, 10-708 Lecture 13: Variational Inference II.

Inputs are a fixed nonempty finite data set $X=(x_1,\dots ,x_N)$; a VAE model and admissible encoder from the preceding definition; a parameter-independent base law $\rho$ and pathwise map $g_{\varphi }$ for every $x_i$; an initial parameter ${\eta }_0=({\theta }_0,{\varphi }_0)$; a finite positive integer $T$; and positive learning rates ${\alpha }_0,\dots ,{\alpha }_{T-1}$. Require the hypotheses of the dominated pathwise-gradient theorem at every parameter value reached by the procedure, and require each update below to remain in its open admissible parameter domain. Let $\mathcal{N}$ be the union of the theorem's common null set for finite positive log values and its null set on which the sample-path $C^1$ condition may fail; thus $\mathcal{N}$ is a common $\rho$-null set for all data indices and parameter values.

Set ${\eta }_0=({\theta }_0,{\varphi }_0)$. For each $t=0,\dots ,T-1$:

1. Draw $I_t$ uniformly from $\{1,\dots ,N\}$ and draw ${ϵ}_t$ from $\rho$, with $I_t$ and ${ϵ}_t$ independent and the pair $(I_t,{ϵ}_t)$ independent of all earlier random draws.
2. If ${ϵ}_t\in \mathcal{N}$, set $h_t=0$. Otherwise set $z_t=g_{{\varphi }_t}({ϵ}_t,x_{I_t})$ and evaluate $h_t=\log {\gamma }_{{\theta }_t}(z_t;x_{I_t})-\log q_{{\varphi }_t}(z_t\mid x_{I_t})\mathrm{.}$
3. If ${ϵ}_t\in \mathcal{N}$, set $G_t=0$. Otherwise, holding ${ϵ}_t$ and $I_t$ fixed, compute the Euclidean gradient $G_t={\mathrm{\nabla }}_{(\theta ,\varphi )}[\log {\gamma }_{\theta }(g_{\varphi }({ϵ}_t,x_{I_t});x_{I_t})-\log q_{\varphi }(g_{\varphi }({ϵ}_t,x_{I_t})\mid x_{I_t})]{\mid }_{(\theta ,\varphi )={\eta }_t}\mathrm{.}$
4. Update ${\eta }_{t+1}={\eta }_t+{\alpha }_t{G}_t$.

Return $({\eta }_0,\dots ,{\eta }_T)$ and $(h_t,G_t{)}_{0\le t<T}$. This finite procedure specifies stochastic ascent on the finite-average ELBO; it makes no claim that the iterates converge or reach a global maximum.

The VAE model, ELBO, and fixed-law encoder sampler.

The conditions under which these pathwise updates estimate the finite-average ELBO gradient.

At a fixed current parameter satisfying the gradient theorem's hypotheses, a fresh independent draw has the unbiased pathwise-gradient and finite-coordinate-variance guarantees proved there. In round $t$, this is a conditional guarantee given the preceding draws, which fix ${\eta }_t$; no unconditional finite-variance claim over the random trajectory is made. The Axiom of Choice identifies this ELBO as a lower bound on log evidence under the corollary's hypotheses.

If the prior and encoder are Gaussian and the expectation of their log-density difference is available analytically, that term can be evaluated exactly instead of sampled. This removes Monte Carlo variation from that term only; the decoder log-likelihood may still be sampled, and no universal variance or training guarantee follows.

The finite-average pathwise-gradient theorem.

The AC-dependent ELBO/KL identity.

Kingma and Welling, Auto-Encoding Variational Bayes, arXiv v1.

Under the full input and admissible-domain hypotheses of reparameterized variational training, let $T\ge 1$ be its finite integer horizon. For every sequence of its prescribed draws satisfying the stated update-domain requirement, all recorded payoffs $h_t$, gradient vectors $G_t$, and parameter vectors ${\eta }_t$ are finite. The mathematical procedure completes exactly $T$ rounds and returns the full trajectories $({\eta }_0,\dots ,{\eta }_T)$ and $(h_t,G_t{)}_{0\le t<T}$, including when one or more noise draws lie in its common null set.

Termination here counts the prescribed draws, branch decisions, pointwise evaluations and vector updates as the operations of the stated mathematical procedure. It asserts neither an implementation or cost bound for arbitrary supplied functions and samplers nor convergence or optimality of the returned parameters.

Reparameterized variational training

Dominated pathwise differentiation for VAE training

Kingma and Welling, Auto-Encoding Variational Bayes, §2.4 Algorithm 1