Variational Autoencoders (VAEs) and Generative Adversarial Networks (GANs)

In the early 2010s, the most striking successes of deep learning were discriminative models, such as classifiers. Deep generative models had less impact, as their training by maximum likelihood involves many intractable probabilistic computations. Two papers proposed ways around this: Auto-Encoding Variational Bayes in 2013, which introduced the Variational Autoencoder, and Generative Adversarial Networks in 2014.

In the quest to understand generative models in Machine Learning, a previous post looked into Variational Inference as a method to perform generative ML on non-parametric or computationally intractable Bayesian formulations. This post continues the quest and looks at two new methods based on neural networks for distributions that no standard parametric family describes: Variational Autoencoders (VAEs) and Generative Adversarial Networks (GANs). It presents the structure and the loss function of each model, the reparameterization trick of the VAE, and the challenges of training a GAN.

Variational Autoencoder (VAE)

Credit: most content in this section is a summary of the papers Auto-Encoding Variational Bayes, by Diederik P. Kingma and Max Welling at the Universiteit van Amsterdam, and An Introduction to Variational Autoencoders, by the same authors.

The papers introduce Variational Autoencoders, aiming at performing efficient inference and learning in probabilistic models whose latent variables and parameters have intractable posterior distributions, and with large datasets. The SGVB (Stochastic Gradient Variational Bayes) estimator presented can be used for efficient approximation of posterior inference in almost any model with continuous latent variables, and because it is differentiable, it can be optimized with e.g. gradient descent. The Auto-Encoding Variational Bayes (AEVB) algorithm presented uses the SGVB estimator to optimize a model that does efficient approximate posterior inference using simple sampling (from parametric distributions).

The VAE is a probabilistic auto-encoder, where the probabilistic encoder (aka recognition model) and decoder (aka generative model) are represented as neural networks. The AEVB algorithm is a learning/inference algorithm that can be used to find the parameters of the VAE, i.e. to perform approximate variational inference in the directed graphical model of the VAE.

Input datapoints \(x\) are generated as follows:

• a value \(z^{(i)}\) is generated from some prior distribution \(p_{\theta^*}(z)\); and
• a value \(x^{(i)}\) is generated from the likelihood distribution \(p_{\theta^∗}(x \mid z)\).

The assumptions are that both the prior and likelihood functions are parametric differentiable distributions, and that \(\theta^*\) and \(z^{(i)}\) are unknown. There is also an assumption of intractability, as in:

• the integral of the evidence or marginal likelihood \(\int p_\theta(z) p_\theta(x \mid z) dz\) is intractable;
• the true posterior density \(p_\theta (z \mid x)\) is intractable, so the expectation maximization algorithm cannot be used; and
• the integrals required by any mean-field approximation algorithm are also intractable.

This is the level of intractability common to DNNs with nonlinear hidden layers.

The paper proposes an efficient approximation of the Maximum Likelihood (MLE) or Maximum a Posteriori (MAP) estimation of the parameters \(\theta\), of the posterior inference of the latent variable \(z\), and of the marginal inference of the variable \(x\). Similarly to other variational inference methods, the intractable true posterior \(p_{\theta}(z \mid x)\) is approximated by \(q_\phi(z \mid x)\) (the Encoder), whose parameters \(\phi\) are not computed by a closed-form expectation but learnt as the weights of the Encoder DNN instead. \(p_\theta(x \mid z)\) is the Decoder, which, given a \(z\), produces/generates the output, a distribution over the possible values of \(x\). Given a datapoint \(x\), the encoder produces a distribution (e.g. a Gaussian) over the possible values of the code \(z\) from which the datapoint \(x\) could have been generated.

Non-Variational vs Variational AutoEncoders

The VAE proposed includes a DNN decoder and a DNN encoder, with parameters \(\theta\) and \(\phi\) optimized together with the AEVB algorithm, where \(p_\theta(x \mid z)\) is a Gaussian/Bernoulli with distribution parameters computed from \(z\). Therefore the VAE can be viewed as two coupled, independently parameterized models: the encoder/recognition model, and the decoder/generative model (trained together), where the encoder delivers to the decoder an approximation to its posterior over latent random variables.

The difference between Variational Autoencoders and regular Autoencoders (and other encoder-decoders, such as sequence-to-sequence models) is that:

• in regular AEs, the latent representation \(z\) is both the output of the encoder and the input of the decoder;
• in VAEs, the outputs of the encoder are the parameters of the distribution of \(z\) (e.g. the mean and variance in the case of a Gaussian distribution), and \(z\) is drawn from the distribution with those parameters and then passed to the decoder;

In practice, they can be visualized as follows, with the path of the AE at the top and the path of the VAE at the bottom of the dashed frame:

Structure of an autoencoder (AE) and of a variational autoencoder (VAE). Adapted from: Variational Autoencoders, S. Yadav, Data Science Blog

Why VAE instead of Variational Inference

One advantage of the VAE framework, relative to ordinary Variational Inference, is that the encoder is now a stochastic function of the input variables. In VI, each data-case has a separate variational distribution, which is inefficient for large datasets. The encoder uses one set of parameters \(\phi\) for all datapoints, which is called amortized inference.

The reparameterization trick

The authors noticed that the sampling induces noise in the gradients required for learning (or, put differently, that \(z\) is randomly generated and cannot be back-propagated). In practice, randomness disrupts the deterministic computation of gradients, making it impossible to directly optimize the model via gradient descent. To counteract that, the authors propose the reparameterization trick.

The reparameterization trick decouples the sampling process from the stochasticity by introducing a deterministic transformation. Instead of sampling \(z \sim \mathcal{N}(\mu, \sigma^2)\), \(z\) is rewritten as \(z = \mu + \sigma \cdot \epsilon\), where \(\epsilon\) follows the standard Gaussian, i.e. \(\epsilon \sim \mathcal{N}(0,1)\).

Why does this work?

• The randomness is now isolated in \(𝜖\), which does not depend on the model parameters.
• The computation of \(𝑧\) becomes a differentiable function of \(𝜇\) and \(𝜎\) (the outputs of the encoder), allowing gradients to flow through \(𝜇\) and \(𝜎\) during backpropagation, because the transformation \(𝑧 = 𝜇 + 𝜎 ⋅ 𝜖\) is fully differentiable with respect to these parameters.

The following picture illustrates the computation graph of the objective before and after the reparameterization:

Illustration of the reparameterization trick. The variational parameters ϕ affect the objective f through the random variable z ~ q_ϕ(z|x). Optimizing the objective with SGD requires the gradients ∇_ϕf. In the original form (left), f cannot be differentiated with respect to ϕ, because gradients cannot be backpropagated through the random variable z. The reparameterized form (right) externalizes the randomness of z by writing it as a deterministic and differentiable function of ϕ, x and a new random variable ε. Gradients then flow through z, and ∇_ϕf can be computed. Adapted from: An Introduction to Variational Autoencoders

The loss function

The loss function is the negative of the evidence lower bound (ELBO) below, i.e. the sum of two terms:

\[\mathcal{L}(\theta, \phi; \mathbf{x}) = \mathbb{E}_{q_\phi(\mathbf{z} \vert \mathbf{x})} [\log p_\theta(\mathbf{x} \vert \mathbf{z})] - D_{KL} \big( q_\phi(\mathbf{z} \vert \mathbf{x}) \parallel p(\mathbf{z}) \big)\]

The first term is the reconstruction loss (or expected negative log-likelihood of the i-th datapoint). The expectation is taken with respect to the encoder’s distribution over the representations, encouraging the VAE to generate valid datapoints. This loss compares the model output with the model input. It can be any of the losses used in regular autoencoders, such as the L2 loss (for a Gaussian decoder with a fixed variance), the cross-entropy (for a Bernoulli decoder), etc.

The second term is the Kullback-Leibler divergence between the encoder’s distribution \(q_{\phi}(z \mid x)\) and the standard Gaussian \(p(z)\), where \(p(z)=\mathcal{N}(\mu=0, \sigma^2=1)\). This divergence compares the distribution of the latent vector with a zero-mean, unit-variance Gaussian distribution. It penalizes the VAE if it starts to produce latent vectors that are not from the desired distribution, i.e. if the encoder outputs representations \(z\) that are different from those of a standard normal distribution. For a VAE trained on images of handwritten digits, this regularizer term means ‘keep the representations \(z\) of each digit sufficiently diverse’ (source: Tutorial - What is a variational autoencoder?, by Jaan Altosaar).

Results

The paper tested the VAE on the MNIST and Frey Face datasets and compared the variational lower bound and estimated marginal likelihood, demonstrating improved results over the wake-sleep algorithm.

Detour: why maximize the Expected Value?

In the Bayesian setup, it is common for the loss function to be the sum of the expected values of several terms.

It all comes down to the Law of large numbers. According to the law, the average of the outcomes of a large number of trials should be close to the expected value, and it tends to get closer to the expected value as the number of trials increases.

Imagine a game (Blackjack, Slots, Rock-Paper-Scissors) that can be repeated as many times as desired. Because of the law of large numbers, the value of the wins over \(n\) iterations of the game will approximate \(n \mathbb{E}(X)\), where \(X\) is the value won in one iteration. In practice, the goal is to optimise the problem according to the average outcome of the experiment over many repetitions, i.e. to maximize its expected value, which is not necessarily the most likely outcome.

So in practice, the law of large numbers will often guarantee a better outcome over the long run, and that outcome is maximized by maximizing the expected outcome of the optimisation. In the same way, the expected values in the losses of this post are estimated by averages over the samples of a mini-batch.

Generative Adversarial Network (GAN)

Credit: most content in this section is a summary of the paper Generative Adversarial Networks, published at NeurIPS 2014, by Ian Goodfellow and colleagues at the University of Montreal.

The paper introduces a new generative model composed of two models trained simultaneously:

• a generative model G that captures the data distribution; and
• a discriminative model D that estimates the probability that a sample came from the training data rather than G;

The training procedure for G is to maximize the probability of D making a mistake. This framework corresponds to a minimax two-player game. The name adversarial comes from the fact that the generative model faces the discriminative model, which learns whether a sample is from the model distribution or from the data distribution. Quoting the authors:

“The generative model can be thought of as analogous to a team of counterfeiters, trying to produce fake currency and use it without detection, while the discriminative model is analogous to the police, trying to detect the counterfeit currency. Competition in this game drives both teams to improve their methods until the counterfeits are indistiguishable from the genuine articles.”

Mathematical setup

The goal of the generative model is to generate samples by passing random noise \(z\) through a multilayer perceptron.

• The generator \(G(z; \theta_g)\) is a multilayer perceptron with parameters \(\theta_g\).
• \(p_z(z)\) is the prior on the input noise variables.
• The generator \(G\) implicitly defines a probability distribution \(p_g\) as the distribution of the samples \(G(z)\) obtained when \(z \sim p_z\);
• The goal of the generator is to learn the distribution \(p_g\) over data \(x\), to fake it “well enough” and trick the discriminator;

The discriminative model is also a multilayer perceptron.

• The discriminator \(D(x; \theta_d)\) takes the input \(x\) and, based on its DNN parameters \(\theta_d\), outputs a single scalar representing the probability that \(x\) came from the data rather than \(p_g\);

The overall structure of the generator and discriminator training can be pictured as:

The GAN model. The generator G turns a sample z of the prior p_z into generated data, and the discriminator D outputs the probability that its input is real data (1) rather than generated data (0). Adapted from: Benjamin Striner, lecture notes on GANs, CMU 11-785 Introduction to Deep Learning

Loss and minimax challenge

The discriminator \(D\) is trained to best guess the source of the data, i.e. to maximize the probability of assigning the correct label to both the training examples and the samples from \(G\).

The generator \(G\) is simultaneously trained to minimize \(\log(1-D(G(z)))\), i.e. to be able to fake the data well enough. In practice, \(D\) and \(G\) are playing a two-player minimax game with value function \(V(D, G)\):

\[\min_G \max_D V(D, G) = \mathbb{E}_{\boldsymbol{x} \sim p_{\text{data}}(\boldsymbol{x})} [\log D(\boldsymbol{x})] + \mathbb{E}_{\boldsymbol{z} \sim p_{\boldsymbol{z}}(\boldsymbol{z})} [\log (1 - D(G(\boldsymbol{z})))].\]

The loss function is a sum of two penalization terms:

• in the first term, the discriminator is optimised such that, in the long run, all real inputs (drawn from the data, i.e. \(x \sim p_{data}\)) are expected to be classified as correctly as possible;
• in the second term, the discriminator and the generator are optimised with opposite goals on all fake inputs (generated by the generator from \(z \sim p_z\)): the discriminator maximizes this term, to classify fake inputs as fake, and the generator minimizes it, so that the discriminator classifies fake inputs as real;

Finally, notice there is a \(\log\) term added to both expected value terms. This is because the objective function is the product of multiple probabilities (of \(D\) assigning the correct label to each input), in this case \(D(x) * (1-D(G(z)))\). Adding the \(\log\) makes the optimisation return the same optima (as \(\log\) is monotonic), while making it faster and more numerically stable. The expectation term is also added to the loss because the goal is to minimize the loss that is expected after several samples (law of large numbers). Thus:

\[\mathbb{E} \left[ \log(D(x) * (1-D(G(z)))) \right] = \mathbb{E} \left[ \log(D(x)) \right] + \mathbb{E} \left[ \log(1-D(G(z))) \right]\]

In section 4 (theoretical results), the authors show that, given enough capacity for the discriminator and the generator, it is indeed possible to recover the data-generating distribution. In practice, they show that this minimax game has a global optimum for \(p_g = p_{data}\), and therefore the loss function can be optimized. However, their proof that the training algorithm converges to this optimum assumes that \(D\) reaches its optimum given \(G\) at every step.

Challenge: different convergence speed in the optimization of Generator and Discriminator

Optimizing \(D\) to completion before each step of \(G\) is “computationally prohibitive, and on finite datasets would result in overfitting”. Instead, training alternates between \(k\) steps of optimizing \(D\) and one step of optimizing \(G\) (the authors use \(k=1\)). This keeps \(D\) near its optimal solution, as long as \(G\) changes slowly enough.

Challenge: saturation of the generator loss

The loss equation may not provide sufficient gradient for \(G\) to learn well. Early in learning, when \(G\) is poor, \(D\) can reject samples with high confidence because they are clearly different from the training data. In this case, \(\log(1 − D(G(z)))\) saturates.

Saturation means that \(D(G(z))\) is close to 0, where this term is flat: the output of \(D\) is the sigmoid \(\sigma(a)\) of a logit \(a\), and the derivative of \(\log(1 - \sigma(a))\) with respect to \(a\) is \(-\sigma(a)\), close to 0. In practice, this means that this term of the loss function will not provide any signal (gradients) for the update of the generator, making it useless. I.e. when the discriminator performs significantly better than the generator, the updates to the generator are either inaccurate or disappear.

This leads to two main problems:

• mathematically, the update of the parameters is zero or very small, so progress through the loss landscape is very slow;
• computationally, arithmetic operations over very small values lead to incorrect (or “always zero”) results due to insufficient floating-point precision in the processor (typically 16 or 32 bits);

The workaround suggested by the authors is to change the algorithm for training the generator: instead of minimizing \(\log(1 − D(G(z)))\), they maximize \(\log D(G(z))\). This keeps the same fixed point of the dynamics of \(G\) and \(D\), while the optimisation landscape “provides much stronger gradients early in learning” (the derivative of \(\log \sigma(a)\) with respect to \(a\) is \(1 - \sigma(a)\), close to 1). The pseudocode at the end of this post shows the original update, which minimizes \(\log(1 − D(G(z)))\).

Challenge: Mode collapse

The original publication mentions this challenge only briefly, as “the Helvetica scenario” in which \(G\) collapses too many values of \(z\) to the same value of \(x\). In a regular scenario, GANs are expected to generate a range of outputs, or ideally, a new random valid output for every random input to the generator. Instead, it may happen that sometimes the output is monotonous, i.e. a range of very similar outputs.

This phenomenon is called Mode Collapse and is detailed in this paper as: “This undesirable event occurs when the model can only fit a few modes of the data distribution, while ignoring the majority of them”. In this article, they propose a workaround using second-order gradient information. However, this is a field of intense research and several solutions have been proposed, to name a few: a different loss function (Wasserstein loss), Unrolled GANs, Conditional GANs, VQ-VAEs, etc. Unrolled GANs incorporate the current and future discriminator’s classifications in the loss function, so that the generator cannot over-optimize for a single discriminator.

For the sake of comparison, a VAE is easier to train and less prone to mode collapse, as it maximizes a lower bound of the likelihood of every datapoint, while a GAN generates sharper samples. A VAE is the better choice when a latent representation of the data or a bound on its likelihood is needed, which the original GAN does not provide, and a GAN when sample quality matters most.

Training algorithm

The final training algorithm is the following:

\[\begin{array}{l} \hline \textbf{Algorithm 1} \text{ Minibatch stochastic gradient descent training of generative} \\ \text{adversarial nets. The number of steps to apply to the discriminator, } k \text{, is a} \\ \text{hyperparameter. We used } k = 1 \text{, the least expensive option, in our experiments.} \\ \hline \quad \textbf{for} \text{ number of training iterations } \textbf{do} \\ \qquad \textbf{for}\ k \text{ steps } \textbf{do} \\ \qquad\quad \bullet \text{ Sample minibatch of } m \text{ noise samples } \{\boldsymbol{z}^{(1)}, \ldots, \boldsymbol{z}^{(m)}\} \text{ from noise prior } p_{\boldsymbol{z}}(\boldsymbol{z}). \\ \qquad\quad \bullet \text{ Sample minibatch of } m \text{ examples } \{\boldsymbol{x}^{(1)}, \ldots, \boldsymbol{x}^{(m)}\} \text{ from data generating distribution} \\ \qquad\quad p_{\text{data}}(\boldsymbol{x}). \\ \qquad\quad \bullet \text{ Update the discriminator by ascending its stochastic gradient:} \\ \hspace{9em} \displaystyle \nabla_{\theta_d} \frac{1}{m} \sum_{i=1}^m \left[ \log D\left(\boldsymbol{x}^{(i)}\right) + \log\left(1 - D\left(G\left(\boldsymbol{z}^{(i)}\right)\right)\right) \right]. \\ \qquad \textbf{end for} \\ \qquad \bullet \text{ Sample minibatch of } m \text{ noise samples } \{\boldsymbol{z}^{(1)}, \ldots, \boldsymbol{z}^{(m)}\} \text{ from noise prior } p_{\boldsymbol{z}}(\boldsymbol{z}). \\ \qquad \bullet \text{ Update the generator by descending its stochastic gradient:} \\ \hspace{12.8em} \displaystyle \nabla_{\theta_g} \frac{1}{m} \sum_{i=1}^m \log\left(1 - D\left(G\left(\boldsymbol{z}^{(i)}\right)\right)\right). \\ \quad \textbf{end for} \\ \text{The gradient-based updates can use any standard gradient-based learning rule.} \\ \text{We used momentum in our experiments.} \\ \hline \end{array}\]

To find more about the topic, chapter 20 of the Deep Learning book by Goodfellow, Bengio and Courville covers deep generative models, including VAEs (section 20.10.3) and GANs (section 20.10.4). The NIPS 2016 tutorial on GANs by Ian Goodfellow compares GANs with other generative models and presents the research frontiers of GANs.