Variational Autoencoder (VAE)¶
Preface¶
Connection of Generative Models to PPCA
Super Simplified Explanation¶
A variational autoencoder doesn't output a single latent vector \(h\) (like how an autoencoder does), it ouputs the parameters - the mean (\(\mu\)) and standard deviation (\(\sigma\)) - of a Gaussian Distribution \(q_{\phi}(h | x)\).
The latent vector \(h\) used for the decoder is sampled from this learned \(q_{\phi}\), represented as:
This probabilistic approach ensures the latent space is smooth and continuous, which is important for generating new data.
To train a VAE, we maximize the variational lower bound (ELBO), which broken down into two components in the loss function.
- Image reconstruction: how well \(x'\) matches \(x\)
- KL-divergence: a regularization term
The KL Divergence component measures the distance between:
- The encoders distribution: \(q_{\phi}(h|x)\), which is a gaussian who's parameters the encoder learns
- The prior distribution: \(p(h)\), a fixed standard gaussian \(\mathcal{N}(0, 1)\).
KL forces the learned distribution \(q\) to stay near a well-behaved standard gaussian, which means that when we sample \(h \sim \mathcal{N}(0, I)\), the decoder can produce a meaningful image \(x'\), allowing for image generation.
A Basic Explanation¶
We want to approximate the underlying true distribution that created our dataset \(\mathcal{D} = \{ x_{i} \}_{i=1}^n\). However, this underlying distribution may be incredibly complex, so instead we want to have a latent vector, \(\mathbf{h}\), that holds important information about our actual \(\mathbf{x}_{i}\)'s .
During training, we train our encoder - a neural network that gets the parameters for the distribution \(q(\mathbf{h| x})\), which needs to closely approximate the true distribution \(p(\mathbf{h | x, W})\). We use KL divergence to achieve a minimization problem.
After training and getting the parameters for \(q(\mathbf{h|x})\) that closely approximates the true distribution, we can sample \(\mathbf{h}\)'s to feed into the decoder.
The decoder is responsible for finding the parameters to approximate the true distribution that generated our \(\mathbf{x}_{i}\)'s. It takes \(\mathbf{h}\) as the input, and outputs the parameters.
We then run backpropagation, which is only possible through the Reparametrization Trick in VAE's (Machine Learning) to update both the weights of the encoder and decoder neural networks.
After we finish training the VAE, we discard the encoder and sample \(\mathbf{h} \sim \mathcal{N}(0, 1)\) as the input into the decoder to generate new samples/outputs.
Definition¶
VAE's are a type of simple generative model that extends PPCA using neural networks. They obtain stochasticity with simple normal random variables and use complex mappings (such as Neural Networks) to transform them into complex distributions (by learning their parameters). ie. consider \(f\) a higher order function, then \(\mathbf{x} = f(\mathbf{h})\) maps each normally distribution \(\mathbf{h}\) to a different vector and the resulting random variable is no longer Gaussian, and can be a highly complex distribution.
For VAE's assume
where \(f_{\mathbf{W}}(\mathbf{h})\) is a multi-layer neural network composed of all of the weights \(\mathbf{W}\), with input \(\mathbf{h}\) and outputting a \(d\)-dimensional vector, the same size as \(\mathbf{x}\). The primary difficulty is in learning \(\mathbf{W}\).
The procedure for VAE is:
- For \(h_{1}, h_{2}, \dots h_{p}\), sample \(h_{i} \sim \mathcal{N}(0,1)\)
- Obtain latent vector \(\mathbf{h} = [h_{1}, h_{2}, \dots h_{p}]\)
- Sample \(\mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu} = f_{\mathbf{W}}(\mathbf{h}), \sigma^2 \mathbf{I})\)
Goal¶
The goal is to find \(\mathbf{W}\) that maximizes the likelihood of an observed dataset \(\mathcal{D}= \{ \mathbf{x}_{i} \}_{i=1}^n\). However, we run into the issue that the latent variable will make estimation more complex. To combat this, we reason over the probabilities of the latent variable conditioned on the inputs \(\mathbf{x}\).
Similar to probabilistic PCA, our density is:
This integral/expectation doesn't evaluate to a Gaussian (like how PCA does), because \(p(\mathbf{x |h,W})\) is a more complex distribution.
Our goal is to maximize the likelihood of the data:
We can attempt to sample many \(\mathbf{h}\)'s and get an average to estimate \(p(\mathbf{x}_{i} | \mathbf{W})\): \(\frac{1}{m} \sum_{k=1}^m p(\mathbf{x}_{i} | \mathbf{h}_{k}, \mathbf{W})\), but there is an issue. Too many \(\mathbf{h}\) play little to no role in the expectation, namely, \(p(\mathbf{x}_{i}| \mathbf{h}, \mathbf{W})\) is likely near zero for most \(\mathbf{h}\). We would have to sample many many \(\mathbf{h}\) to get an accurate estimate of the expectation, and this just isn't feasible.
To combat this issue, we direct the sampling, which means to say that we choose certain \(\mathbf{h}\) that do play a role in \(p(\mathbf{x}_{i}| \mathbf{h, W})\).
To do this, we must find a distribution \(q(\mathbf{h|x})\) (also known as the variational distribution, so that we can sample \(\mathbf{h}\) with a high likelihood for a given \(\mathbf{x}\).
We do not have \(q(\mathbf{h|x})\), but we can learn it by attempting to match it with the true distribution:
By using the Kullback-Leibler Divergence (KL), we are able to measure the differences between \(q\) and \(p\), and so we can create another minimization problem alongside our original one.
Through some derivation (see Derivation for ELBO from KL Divergence in VAE's), we end with the formula for ELBO:
The good thing is that the log is inside of the expectation, so it makes taking the gradient significantly easier (and more numerically stable), than if the log was outside of the expectation.
Optimizing ELBO¶
In this case, since this is a minimization problem, we are optimizing the negative ELBO.
In order to optimize the negative ELBO, we need to define \(q(\mathbf{h|x})\), into a variational distribution with learnable parameters.
Here \(\boldsymbol{\theta}\) are the parameters of a neural network, trained to find the mean and standard deviation of \(q\).
\(f_{\mu, \theta}(\mathbf{x})\) and \(f_{\sigma, \theta}(\mathbf{x})\) are different neural networks, but share mostly the same paremeters, except at the last layer of the network. \(\boldsymbol{\theta}\) hold all of these parameters (shared and private).
It is crucial to note that \(f_{\sigma, \theta}(\mathbf{x})\) uses \(\boldsymbol{\theta}\) to output the diagonal covariance matrix estimate \(\boldsymbol{\Sigma} \in \mathbb{R}^{d\times d}\). Specifically, it outputs: \(\sigma^2_{1}, \sigma^2_{2}, \dots, \sigma^2_{d}\).
Given the last layer, which we call \(\phi(\mathbf{x})\):
- The mean is a linear function of \(\phi(\mathbf{x})\)
- And the variances \(\sigma^2_{i}\) uses a linear function on \(\phi(\mathbf{x})\) followed by a transform to ensure the output is positive. (We need a positive output to ensure numerical stability when calculating the standard deviation).
Note that it is common in VAE architecture to output the log of \(\sigma^2_{i}(\mathbf{x})\), and use a linear function of \(\phi(\mathbf{x})\) to produce this log variance.
Computing the Gradients of the NN¶
The problem is that sampling from the variational distribution \(q(\mathbf{h|x}, \boldsymbol{\theta})\) is a stochastic process. Differentiating through stochastic nodes is very difficult.
To combat this, we use the Reparametrization Trick in VAE's (Machine Learning).
Let \(\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\), then we can rewrite
Now we can put the gradient inside of the expectation. Then for each SGD update, we sample a point \(\mathbf{x}_{i}\) and sample an error vector \(\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0, I})\) and easily compute the gradient \(\ln p(\mathbf{x} | \mathbf{h} = \boldsymbol{\mu}(\mathbf{x})+ \boldsymbol{\sigma}(\mathbf{x})\boldsymbol{\epsilon})\) with respect to both \(\mathbf{W}\) and \(\boldsymbol{\theta}\). (Don't forget that \(\boldsymbol{\theta}\) parametrizes \(\mu_{j}(\mathbf{x}_{i})\) and \(\sigma^2_{j}(\mathbf{x}_{i})\) for all \(j\).
We then use these gradients to update \(\mathbf{W}\) and \(\boldsymbol{\theta}\).
Notice that the negative ELBO has two terms:
In which we have only optimized (2). Now we see to optimize (1). The KL divergence term is like a regularizer that stops \(q\) from becoming too deterministic, and forces it to have some stochasticity that follows the standard Gaussian.
It is easy to compute the gradient of equation (1) because the KL divergence between two Gaussian's has an analytic expression. And this gradient can be simply added alongside the gradient for equation (2).
VAE Issues¶
- Blurry reconstructions
- VAE's average over many possible versions of an image
- This produces a smooth, washed out result instead of sharp details
- Posterior Collapse
- To minimize the KL the quickest, it will output a standard gaussian instead of using the input image \(x\).
- All latent codes become nearly identical, so the model ignores the latent space entirely
- Reconstruction vs Regularization Trade-off
- Balancing image quality and latent structure is difficult
- Too much regularization = poor reconstructions
- Too little regularization = unstable or meaningless latent space
- Loss of Fine Details
- The decoder models pixels independently, ignoring local texture patterns
- Loss of crisp edges, fine textures, and high-frequency detail
A fix to VAE's issues is VQ-VAE
References¶
Conditional Generative Models Connection of Generative Models to PPCA Neural Networks Standard Deviation Autoencoders 1