Reparameterized variational training

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.