8.2 Variational Autoencoders

How does this change to the parameterization of the recognition model change the optimization? Since the recognition distribution is still normal, we can still use Eq. 8.1 for the joint relative entropy, and Eqs. 8.3 and 8.4 for its gradients. Beginning with the parameters $\bm{\phi}_{\bm{\nu}}$ of the mean function:

Likewise, for the covariance, we follow Eq. 8.4, although for notational simplicity we write the derivative with respect to only a single parameter, $\phi_{\mathbf{{\Upsilon}}}^{i}$

Computing the Hessian $\partial^{2}{E_{\text{post}}}/\partial{\bm{\hat{x}}}\partial{\bm{\hat{x}}}^{\text{T}}$ naïvely is expensive ($\mathcal{O}({K}^{3})$) but there are clever, $\mathcal{O}({K}^{2})$, alternatives [41].

Now, for sufficiently expressive generative models, the conditional expectations under the recognition distribution will in general be difficult to carry out (although cf. Eqs. 8.10 and 8.11). Instead we might replace the conditional expectation with a sample average, since it is trivial to draw samples from Eq. 8.12. For example, the generative model could be yet more normal distributions,

but with the mean and covariance of the emission, like the recognition distribution, flexible parameterized functions—indeed, neural networks—operating on their “inputs,” in this case $\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}}$. Then Eqs. 8.13 and 8.14 simplify slightly [41]:

with the emission energy $E_{\text{emiss}}$ equal to the energy of the Gaussian distribution on the right in Eq. 8.15,

To optimize these parameters, we can descend these gradients—but discussion of how precisely this is to be carried out we defer until after deriving the E step.

Starting again with the joint relative entropy under Gaussian recognition distributions, Eq. 8.1, assuming the Gaussian generative model of Eq. 8.15, and differentiating with respect to the parameters of the generative model, $\bm{\theta}$, we find

Again we propose to descend this gradient iteratively.

Computing the gradients on the left-hand sides of Eqs. 8.16 and 8.17 evidently requires acquiring the gradients appearing in the right-hand sides. The gradients of ${\bm{\nu}(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{y}{},\bm{\phi}_{\bm{\nu}})}$ and ${\mathbf{{\Upsilon}}(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{y}{}{},\bm{\phi}_{\mathbf{{\Upsilon}}})}$, for their part, obviously pass backwards through the recognition neural networks. But the gradients of the emission energy with respect $\bm{\hat{x}}$ must pass backwards through the generative neural networks, $\bm{\mu}(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}},\bm{\theta})$ and $\mathbf{{\Sigma}}(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}},\bm{\theta})$, for which $\bm{\hat{x}}$ is an input. Along this backward propagating computation, we will compute precisely the quantities required for the gradients $\partial{\bm{\mu}}/\partial{\bm{\theta}}$ and $\partial{\mathbf{{\Sigma}}}/\partial{\bm{\theta}}$ that are needed for the E step, Eq. 8.19.

This suggests that the E and M steps be simultaneous rather than alternating. In particular, a set of ${N}$ observations $\bm{y}_{1},\ldots,\bm{y}_{{N}}$ can be passed “upward” through the recognition distribution, Eq. 8.12, yielding ${N}$ sample latent vectors, $\bm{\check{x}}_{1},\ldots,\bm{\check{x}}_{{N}}$.³³ 3 It might seem that we need multiple samples of ${\bm{\check{X}}}$ for every sample of ${\bm{Y}}$, and indeed we do need (roughly) quadratically more samples to estimate the joint distribution of ${\bm{\check{X}}},{\bm{Y}}$ as opposed to just ${\bm{Y}}$—say, $L^{2}$. But it is more efficient to have $L^{2}$ different values of each of ${\bm{\check{X}}},{\bm{Y}}$ rather than just $L$ different values for ${\bm{Y}}$ and $L^{2}$ values of ${\bm{\check{X}}}$; see Fig. LABEL:fig:XXX. These are then passed on through the generative model to produce values for $\bm{\mu}(\bm{\check{x}}_{n},\bm{\theta})$ and $\mathbf{{\Sigma}}(\bm{\check{x}}_{n},\bm{\theta})$. Notice that this entire process looks like a forward pass through a single neural network that maps each value to itself—i.e., an autoencoder—or, more precisely, to a mean and variance for that sample, encoding the network’s “best guess” for that sample and its (inverse) confidence in that guess.

The gradient is therefore evaluated by initializing a backward pass through the emission energy Eq. 8.18 at the samples $\bm{y}_{1},\ldots,\bm{y}_{{N}}$, working backwards through $\bm{\mu}(\bm{\check{x}}_{n},\bm{\theta})$ and $\mathbf{{\Sigma}}(\bm{\check{x}}_{n},\bm{\theta})$ toward the “inputs” ${\bm{\hat{X}}}$. The gradient at each layer of these neural networks is evaluated at the values computed under the preceding forward pass. Averaging these gradients under all ${N}$ samples yields the right-hand side of Eq. 8.19, the “E-step” gradients. But rather than stop at the parameters of the first layer, the gradient computation is continued into the inputs ${\bm{\hat{X}}}$. The chain of backward computations then proceeds through the parameters of the recognition distribution according to Eqs. 8.16 and 8.17, finally terminating at the first layers of ${\bm{\nu}(\bm{y}_{n},\bm{\phi}_{\bm{\nu}})}$ and ${\mathbf{{\Upsilon}}(\bm{y}_{n}{},\bm{\phi}_{\mathbf{{\Upsilon}}})}$. Averaging these gradients under the ${N}$ samples yields the “M-step” gradients for updating the recognition distribution.

We observed above that this optimization procedure suggests thinking of the generative and recognition models together as a single, feedforward neural network (Fig. LABEL:fig:VAE). To emphasize this intuition, let us again reorganize the joint relative entropy:

The second term, ${\text{H}_{(p\check{p})\hat{p}}{\mathopen{}\mathclose{{}\left[{\bm{Y}}{}\middle|{\bm{\check{X}}}{};\bm{\phi},\bm{\theta}}\right]}}$, can be interpreted as a reconstruction cost (see Section 6.3.2): the number of bits⁴⁴ 4 If a base-2 logarithm is used. required to encode the observations under the generative model, given latent states inferred with the recognition model. It measures the “lossiness” of the autoencoder. By itself, this term enforces a discriminative rather than generative cost. Note that, without further constraint, the reconstruction cost could be either greater or lower than the marginal entropy, ${\text{H}_{p}{\mathopen{}\mathclose{{}\left[{\bm{Y}}{}}\right]}}$; i.e., the final two terms could together be either positive or negative.⁵⁵ 5 For example, let ${Y}$ be a biased coin (entropy less than 1 bit). A model could simply copy through the observations (reconstruction cost of 0 bits), or again predict heads and tails with equal probability (reconstruction cost of 1 bit).

We have previously (Section 6.3.2) interpreted the first term as the sum of a “code cost,” ${\text{H}_{(p\check{p})\hat{p}}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}{};\bm{\phi},\bm{\theta}}\right]}}$, required to encode the latent states inferred by the recognition model; and $-{\text{H}_{(p\check{p})}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}{}\middle|{\bm{Y}}{};\bm{\phi}}\right]}}$, a regularizer that prevents the optimization from decreasing reconstruction and code costs merely by making recognition more confident. Alternatively, $\operatorname*{\text{D}_{\text{KL}}}\mathopen{}\mathclose{{}\left\{{\check{p}\mathopen{}\mathclose{{}\left({\bm{\check{X}}}\middle|{\bm{Y}};\bm{\phi}}\right)}\middle\|{\hat{p}\mathopen{}\mathclose{{}\left({\bm{\check{X}}};\bm{\theta}}\right)}}\right\}$ can be interpreted [27] as a regularizer that encourages the recognition posterior to resemble the generative prior. Intuitively, optimizing reconstruction cost obliges the emission density to “play well” with the recognition model (as well as to fit the observations); but to complete the generative model, we also need it to play well with the generative prior! This requires the recognition model and the generative prior to resemble each other. Indeed, the regularization that enforces this resemblance is the only place that the generative prior enters the loss.

It is important to realize that we do not expect this relative entropy to vanish in the optimal model, which will embody a compromise between this regularizer and the reconstruction cost (cf. Section 6.3.2). Instead, the joint relative entropy $\mathcal{L}_{\text{JRE}}$ vanishes at the optimum, so the most we can say is

This equation tells us something intuitive: Since the left-hand side is non-negative, so must the right-hand side be; or in other words, the optimal reconstruction cost never exceeds the entropy of the data distribution.⁶⁶ 6 Returning to our previous example of a biased coin, we find that the optimal model should correctly identify heads-vs.-tails at least as frequently as the bias of the coin.

In the simplest VAEs, the generative priors are not parameterized, ${\hat{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}}{};\bm{\theta}}\right)}={\hat{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}}{}}\right)}$, as in Eq. 8.15, in which case the KL regularizer pushes the recognition model toward the generative prior but not vice versa. For the generative model considered in isolation, such a truly structureless prior distribution is desirable. The ideal latent variables are akin to the fundamental particles of physics: they should be as independent as possible, since any remaining structure is something that still needs to be explained. We prefer to let the emission, ${\hat{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{y}}{}\middle|\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}}{};\bm{\theta}}\right)}$, capture all structure.

However, the recognition model makes additional demands on the generative prior: A simple calculation shows that if the recognition model, ${\check{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\check{x}}{}\middle|\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{y}{};\bm{\phi}}\right)}$, matches the generative posterior, ${\hat{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}}{}\middle|\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{y}}{};\bm{\theta}}\right)}$, and the generative distribution, ${\hat{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{y}}{};\bm{\theta}}\right)}$, matches the data, ${p\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{y}{}}\right)}$, then the generative prior, ${\hat{p}\mathopen{}\mathclose{{}\left(\leavevmode\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\bm{\hat{x}}{};\bm{\theta}}\right)}$, must match the “aggregated posterior,”

In brief, for optimal generative and recognition models,

must hold. To understand this requirement, we yet again rewrite^††margin: Exercise LABEL:ex:JREELBOsurgery the joint relative entropy, Eq. 8.20, this time in terms of this relative entropy on “priors” [17]:

The penalty on information transmission prevents the reconstruction loss from being driven to zero, and vice versa. But the relative entropy in this equation—unlike the relative entropy in Eq. 8.20—does vanish in the optimal model.