10.2 Variational Autoencoders

As usual, the M step amounts to fitting the parameters $\bm{\theta}$ of the generative model by minimizing the joint relative entropy. We will not yet impose a generative model, so we write the gradient generically:

This gradient can be descended iteratively.

Parameterizing the recognition model with neural networks, rather than fixed parameters (one per observation), only mildly changes the loss gradient. In particular, since the recognition distribution is still normal, we can still use Eq. 10.1 for the joint relative entropy, and Eqs. 10.3 and 10.4 for its gradients. The only addition is the Jacobian of the moment parameters with respect to the neural-network parameters $\bm{\phi}$, which accrues via the chain rule of differentiation:

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

(For brevity, some parameter dependence has been suppressed.) Computing the Hessian $\partial^{2}{\hat{U}_{\text{post}}}/\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{\hat{x}}}}\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{\hat{x}}}}^{\text{T}}$ naïvely is expensive ($\mathcal{O}({K}^{3})$) but there are clever, $\mathcal{O}({K}^{2})$, alternatives [61]. For sufficiently expressive generative models, the conditional expectations under the recognition distribution in Eqs. 10.14 and 10.15 will be impossible to carry out (although cf. Eqs. 10.10 and 10.11), so we will typically resort to sample averages. We can then descend these gradients iteratively.

Computing the gradients on the left-hand sides of Eqs. 10.14 and 10.15 evidently requires acquiring the gradients appearing in the right-hand sides. To compute the gradients of the recognition networks, ${\bm{\nu}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{},\bm{\phi}_{\bm{\nu}})}$ and ${\mathbf{{\Upsilon}}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{}{},\bm{\phi}_{\mathbf{{\Upsilon}}})}$, we can pass observations $\bm{y}$ forward through them, and then backpropagate gradient values. To compute the gradient of the generative log probability, $\log{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}{},{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{};\bm{\theta}}\right)}$, we must pass samples $\bm{\check{x}}$ from the recognition model forward through the generative neural network, and backpropagate gradient values all the way into their inputs. Indeed, we can combine these two pairs of forward/backward passes: We can pass a set of observations $\bm{y}_{1},\ldots,\bm{y}_{{N}}$ through the recognition networks; draw a set of latent samples $\bm{\check{x}}_{1},\ldots,\bm{\check{x}}_{{N}}$, one for each observation³³ 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. 10.1.; pass these on through the generative neural network; and finally propagate gradients backward through the generative and recognition models. Note that backpropagation can simply ignore the sampling step.

But now note that along the backward pass through the generative network, we will compute precisely the quantities required for the M step gradient, $\partial{\log\hat{p}}/\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{\theta}}}$ (recall Eq. 10.13 as well as our discussion of backprop, Section 7.1.3). Thus, we can compute all our gradients estimates, for both the E and M steps, in a single forward/backward pass. Notice that this entire process looks like a forward pass through a single neural network that maps each observation to itself—i.e., an autoencoder^††margin: autoencoder —or, more precisely, to the log probability of that observation.

This suggests (although it does not require) updating the parameters of the generative and recongnition models together, rather than in alternation. This is the standard way to train such models.

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 8.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 8.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 [39] 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 8.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({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}{};\bm{\theta}}\right)}={\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}{}}\right)}$, 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({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{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({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}{}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{};\bm{\phi}}\right)}$, matches the generative posterior, ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}{}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{};\bm{\theta}}\right)}$, and the generative distribution, ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{};\bm{\theta}}\right)}$, matches the data, ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}{}}\right)}$, then the generative prior, ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{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: the joint relative entropy, Eq. 10.16, this time in terms of this relative entropy on “priors” [30]:

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. 10.16—does vanish in the optimal model.