{ "id": "cd59f7ee3fdb", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/", "source_excel_row": 1600, "source_url": "https://lilianweng.github.io/posts/2018-08-12-vae/", "is_paid_flag": "", "images": [ { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/autoencoder-architecture.png", "bytes": 137292, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-db59b90d279b037a97d1.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/autoencoder-architecture.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/denoising-autoencoder-architecture.png", "bytes": 256981, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-ab7006570b99c84db1b9.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/denoising-autoencoder-architecture.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/kl-metric-sparse-autoencoder.png", "bytes": 26225, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-1c83c8f4d8be7a5b4c49.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/kl-metric-sparse-autoencoder.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/k-sparse-autoencoder.png", "bytes": 1878475, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-2e93bdc2c46162899b3f.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/k-sparse-autoencoder.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/VAE-graphical-model.png", "bytes": 151943, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-13d9a58c515022883c19.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/VAE-graphical-model.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/forward_vs_reversed_KL.png", "bytes": 482588, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-ee2f9d9fbcd3aed80f31.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/forward_vs_reversed_KL.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/reparameterization-trick.png", "bytes": 188866, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-2351ecad009ac2895562.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/reparameterization-trick.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/vae-gaussian.png", "bytes": 274889, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-fae191e6a32c13b2540d.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/vae-gaussian.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/VQ-VAE.png", "bytes": 806890, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-9610ae1d4a5d7cbb7454.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/VQ-VAE.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/VQ-VAE-2.png", "bytes": 943633, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-e326acd6fdf8195c78b9.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/VQ-VAE-2.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/VQ-VAE-2-algo.png", "bytes": 594737, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-91ba8add21dff9815090.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/VQ-VAE-2-algo.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/TD-VAE-state-space.png", "bytes": 207374, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-7020e37618b19a5457cc.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/TD-VAE-state-space.png", "alt": "", "downloaded": true }, { "status": "ok", "url": "https://lilianweng.github.io/posts/2018-08-12-vae/TD-VAE.png", "bytes": 946267, "reused": true, "local_path": "attachments/cd59f7ee3fdb/images/body-da329904d7433579a338.png", "src": "https://lilianweng.github.io/posts/2018-08-12-vae/TD-VAE.png", "alt": "", "downloaded": true } ], "videos": [ { "url": "https://www.youtube.com/watch?v=9zKuYvjFFS8", "src": "https://www.youtube.com/watch?v=9zKuYvjFFS8", "status": "ok", "kind": "single_video", "local_path": "attachments/cd59f7ee3fdb/videos/repair-9zKuYvjFFS8.mp4", "bytes": 119566312, "valid": true, "height": 1080, "width": 1920, "duration": "905.001000", "reused": false } ], "body_markdown": "[Updated on 2019-07-18: add a section on\nVQ-VAE & VQ-VAE-2\n.]\n[Updated on 2019-07-26: add a section on\nTD-VAE\n.]\nAutocoder is invented to reconstruct high-dimensional data using a neural network model with a narrow bottleneck layer in the middle (oops, this is probably not true for\nVariational Autoencoder\n, and we will investigate it in details in later sections). A nice byproduct is dimension reduction: the bottleneck layer captures a compressed latent encoding. Such a low-dimensional representation can be used as en embedding vector in various applications (i.e. search), help data compression, or reveal the underlying data generative factors.\nNotation\n#\nSymbol\nMean\n$\\mathcal{D}$\nThe dataset, $\\mathcal{D} = \\{ \\mathbf{x}^{(1)}, \\mathbf{x}^{(2)}, \\dots, \\mathbf{x}^{(n)} \\}$, contains $n$ data samples; $\\vert\\mathcal{D}\\vert =n $.\n$\\mathbf{x}^{(i)}$\nEach data point is a vector of $d$ dimensions, $\\mathbf{x}^{(i)} = [x^{(i)}_1, x^{(i)}_2, \\dots, x^{(i)}_d]$.\n$\\mathbf{x}$\nOne data sample from the dataset, $\\mathbf{x} \\in \\mathcal{D}$.\n$\\mathbf{x}’$\nThe reconstructed version of $\\mathbf{x}$.\n$\\tilde{\\mathbf{x}}$\nThe corrupted version of $\\mathbf{x}$.\n$\\mathbf{z}$\nThe compressed code learned in the bottleneck layer.\n$a_j^{(l)}$\nThe activation function for the $j$-th neuron in the $l$-th hidden layer.\n$g_{\\phi}(.)$\nThe\nencoding\nfunction parameterized by $\\phi$.\n$f_{\\theta}(.)$\nThe\ndecoding\nfunction parameterized by $\\theta$.\n$q_{\\phi}(\\mathbf{z}\\vert\\mathbf{x})$\nEstimated posterior probability function, also known as\nprobabilistic encoder\n.\n$p_{\\theta}(\\mathbf{x}\\vert\\mathbf{z})$\nLikelihood of generating true data sample given the latent code, also known as\nprobabilistic decoder\n.\nAutoencoder\n#\nAutoencoder\nis a neural network designed to learn an identity function in an unsupervised way to reconstruct the original input while compressing the data in the process so as to discover a more efficient and compressed representation. The idea was originated in\nthe 1980s\n, and later promoted by the seminal paper by\nHinton & Salakhutdinov, 2006\n.\nIt consists of two networks:\nEncoder\nnetwork: It translates the original high-dimension input into the latent low-dimensional code. The input size is larger than the output size.\nDecoder\nnetwork: The decoder network recovers the data from the code, likely with larger and larger output layers.\nIllustration of autoencoder model architecture.\nThe encoder network essentially accomplishes the\ndimensionality reduction\n, just like how we would use Principal Component Analysis (PCA) or Matrix Factorization (MF) for. In addition, the autoencoder is explicitly optimized for the data reconstruction from the code. A good intermediate representation not only can capture latent variables, but also benefits a full\ndecompression\nprocess.\nThe model contains an encoder function $g(.)$ parameterized by $\\phi$ and a decoder function $f(.)$ parameterized by $\\theta$. The low-dimensional code learned for input $\\mathbf{x}$ in the bottleneck layer is $\\mathbf{z} = g_\\phi(\\mathbf{x})$ and the reconstructed input is $\\mathbf{x}’ = f_\\theta(g_\\phi(\\mathbf{x}))$.\nThe parameters $(\\theta, \\phi)$ are learned together to output a reconstructed data sample same as the original input, $\\mathbf{x} \\approx f_\\theta(g_\\phi(\\mathbf{x}))$, or in other words, to learn an identity function. There are various metrics to quantify the difference between two vectors, such as cross entropy when the activation function is sigmoid, or as simple as MSE loss:\n$$\nL_\\text{AE}(\\theta, \\phi) = \\frac{1}{n}\\sum_{i=1}^n (\\mathbf{x}^{(i)} - f_\\theta(g_\\phi(\\mathbf{x}^{(i)})))^2\n$$\nDenoising Autoencoder\n#\nSince the autoencoder learns the identity function, we are facing the risk of “overfitting” when there are more network parameters than the number of data points.\nTo avoid overfitting and improve the robustness,\nDenoising Autoencoder\n(Vincent et al. 2008) proposed a modification to the basic autoencoder. The input is partially corrupted by adding noises to or masking some values of the input vector in a stochastic manner, $\\tilde{\\mathbf{x}} \\sim \\mathcal{M}_\\mathcal{D}(\\tilde{\\mathbf{x}} \\vert \\mathbf{x})$. Then the model is trained to recover the original input (note: not the corrupt one).\n$$\n\\begin{aligned}\n\\tilde{\\mathbf{x}}^{(i)} &\\sim \\mathcal{M}_\\mathcal{D}(\\tilde{\\mathbf{x}}^{(i)} \\vert \\mathbf{x}^{(i)})\\\\\nL_\\text{DAE}(\\theta, \\phi) &= \\frac{1}{n} \\sum_{i=1}^n (\\mathbf{x}^{(i)} - f_\\theta(g_\\phi(\\tilde{\\mathbf{x}}^{(i)})))^2\n\\end{aligned}\n$$\nwhere $\\mathcal{M}_\\mathcal{D}$ defines the mapping from the true data samples to the noisy or corrupted ones.\nIllustration of denoising autoencoder model architecture.\nThis design is motivated by the fact that humans can easily recognize an object or a scene even the view is partially occluded or corrupted. To “repair” the partially destroyed input, the denoising autoencoder has to discover and capture relationship between dimensions of input in order to infer missing pieces.\nFor high dimensional input with high redundancy, like images, the model is likely to depend on evidence gathered from a combination of many input dimensions to recover the denoised version rather than to overfit one dimension. This builds up a good foundation for learning\nrobust\nlatent representation.\nThe noise is controlled by a stochastic mapping $\\mathcal{M}_\\mathcal{D}(\\tilde{\\mathbf{x}} \\vert \\mathbf{x})$, and it is not specific to a particular type of corruption process (i.e. masking noise, Gaussian noise, salt-and-pepper noise, etc.). Naturally the corruption process can be equipped with prior knowledge\nIn the experiment of the original DAE paper, the noise is applied in this way: a fixed proportion of input dimensions are selected at random and their values are forced to 0. Sounds a lot like dropout, right? Well, the denoising autoencoder was proposed in 2008, 4 years before the dropout paper (\nHinton, et al. 2012\n) ;)\nSparse Autoencoder\n#\nSparse Autoencoder\napplies a “sparse” constraint on the hidden unit activation to avoid overfitting and improve robustness. It forces the model to only have a small number of hidden units being activated at the same time, or in other words, one hidden neuron should be inactivate most of time.\nRecall that common\nactivation functions\ninclude sigmoid, tanh, relu, leaky relu, etc. A neuron is activated when the value is close to 1 and inactivate with a value close to 0.\nLet’s say there are $s_l$ neurons in the $l$-th hidden layer and the activation function for the $j$-th neuron in this layer is labelled as $a^{(l)}_j(.)$, $j=1, \\dots, s_l$. The fraction of activation of this neuron $\\hat{\\rho}_j$ is expected to be a small number $\\rho$, known as\nsparsity parameter\n; a common config is $\\rho = 0.05$.\n$$\n\\hat{\\rho}_j^{(l)} = \\frac{1}{n} \\sum_{i=1}^n [a_j^{(l)}(\\mathbf{x}^{(i)})] \\approx \\rho\n$$\nThis constraint is achieved by adding a penalty term into the loss function. The KL-divergence $D_\\text{KL}$ measures the difference between two Bernoulli distributions, one with mean $\\rho$ and the other with mean $\\hat{\\rho}_j^{(l)}$. The hyperparameter $\\beta$ controls how strong the penalty we want to apply on the sparsity loss.\n$$\n\\begin{aligned}\nL_\\text{SAE}(\\theta) \n&= L(\\theta) + \\beta \\sum_{l=1}^L \\sum_{j=1}^{s_l} D_\\text{KL}(\\rho \\| \\hat{\\rho}_j^{(l)}) \\\\\n&= L(\\theta) + \\beta \\sum_{l=1}^L \\sum_{j=1}^{s_l} \\rho\\log\\frac{\\rho}{\\hat{\\rho}_j^{(l)}} + (1-\\rho)\\log\\frac{1-\\rho}{1-\\hat{\\rho}_j^{(l)}}\n\\end{aligned}\n$$\nThe KL divergence between a Bernoulli distribution with mean $\\rho=0.25$ and a Bernoulli distribution with mean $0 \\leq \\hat{\\rho} \\leq 1$.\n$k$-Sparse Autoencoder\nIn $k$-Sparse Autoencoder (\nMakhzani and Frey, 2013\n), the sparsity is enforced by only keeping the top k highest activations in the bottleneck layer with linear activation function.\nFirst we run feedforward through the encoder network to get the compressed code: $\\mathbf{z} = g(\\mathbf{x})$.\nSort the values in the code vector $\\mathbf{z}$. Only the k largest values are kept while other neurons are set to 0. This can be done in a ReLU layer with an adjustable threshold too. Now we have a sparsified code: $\\mathbf{z}’ = \\text{Sparsify}(\\mathbf{z})$.\nCompute the output and the loss from the sparsified code, $L = |\\mathbf{x} - f(\\mathbf{z}’) |_2^2$.\nAnd, the back-propagation only goes through the top k activated hidden units!\nFilters of the k-sparse autoencoder for different sparsity levels k, learnt from MNIST with 1000 hidden units.. (Image source:\nMakhzani and Frey, 2013\n)\nContractive Autoencoder\n#\nSimilar to sparse autoencoder,\nContractive Autoencoder\n(\nRifai, et al, 2011\n) encourages the learned representation to stay in a contractive space for better robustness.\nIt adds a term in the loss function to penalize the representation being too sensitive to the input, and thus improve the robustness to small perturbations around the training data points. The sensitivity is measured by the Frobenius norm of the Jacobian matrix of the encoder activations with respect to the input:\n$$\n\\|J_f(\\mathbf{x})\\|_F^2 = \\sum_{ij} \\Big( \\frac{\\partial h_j(\\mathbf{x})}{\\partial x_i} \\Big)^2\n$$\nwhere $h_j$ is one unit output in the compressed code $\\mathbf{z} = f(x)$.\nThis penalty term is the sum of squares of all partial derivatives of the learned encoding with respect to input dimensions. The authors claimed that empirically this penalty was found to carve a representation that corresponds to a lower-dimensional non-linear manifold, while staying more invariant to majority directions orthogonal to the manifold.\nVAE: Variational Autoencoder\n#\nThe idea of\nVariational Autoencoder\n(\nKingma & Welling, 2014\n), short for\nVAE\n, is actually less similar to all the autoencoder models above, but deeply rooted in the methods of variational bayesian and graphical model.\nInstead of mapping the input into a\nfixed\nvector, we want to map it into a distribution. Let’s label this distribution as $p_\\theta$, parameterized by $\\theta$. The relationship between the data input $\\mathbf{x}$ and the latent encoding vector $\\mathbf{z}$ can be fully defined by:\nPrior $p_\\theta(\\mathbf{z})$\nLikelihood $p_\\theta(\\mathbf{x}\\vert\\mathbf{z})$\nPosterior $p_\\theta(\\mathbf{z}\\vert\\mathbf{x})$\nAssuming that we know the real parameter $\\theta^{*}$ for this distribution. In order to generate a sample that looks like a real data point $\\mathbf{x}^{(i)}$, we follow these steps:\nFirst, sample a $\\mathbf{z}^{(i)}$ from a prior distribution $p_{\\theta^*}(\\mathbf{z})$.\nThen a value $\\mathbf{x}^{(i)}$ is generated from a conditional distribution $p_{\\theta^*}(\\mathbf{x} \\vert \\mathbf{z} = \\mathbf{z}^{(i)})$.\nThe optimal parameter $\\theta^{*}$ is the one that maximizes the probability of generating real data samples:\n$$\n\\theta^{*} = \\arg\\max_\\theta \\prod_{i=1}^n p_\\theta(\\mathbf{x}^{(i)})\n$$\nCommonly we use the log probabilities to convert the product on RHS to a sum:\n$$\n\\theta^{*} = \\arg\\max_\\theta \\sum_{i=1}^n \\log p_\\theta(\\mathbf{x}^{(i)})\n$$\nNow let’s update the equation to better demonstrate the data generation process so as to involve the encoding vector:\n$$\np_\\theta(\\mathbf{x}^{(i)}) = \\int p_\\theta(\\mathbf{x}^{(i)}\\vert\\mathbf{z}) p_\\theta(\\mathbf{z}) d\\mathbf{z} \n$$\nUnfortunately it is not easy to compute $p_\\theta(\\mathbf{x}^{(i)})$ in this way, as it is very expensive to check all the possible values of $\\mathbf{z}$ and sum them up. To narrow down the value space to facilitate faster search, we would like to introduce a new approximation function to output what is a likely code given an input $\\mathbf{x}$, $q_\\phi(\\mathbf{z}\\vert\\mathbf{x})$, parameterized by $\\phi$.\nThe graphical model involved in Variational Autoencoder. Solid lines denote the generative distribution $p\\_\\theta(.)$ and dashed lines denote the distribution $q\\_\\phi (\\mathbf{z}\\vert\\mathbf{x})$ to approximate the intractable posterior $p\\_\\theta (\\mathbf{z}\\vert\\mathbf{x})$.\nNow the structure looks a lot like an autoencoder:\nThe conditional probability $p_\\theta(\\mathbf{x} \\vert \\mathbf{z})$ defines a generative model, similar to the decoder $f_\\theta(\\mathbf{x} \\vert \\mathbf{z})$ introduced above. $p_\\theta(\\mathbf{x} \\vert \\mathbf{z})$ is also known as\nprobabilistic decoder\n.\nThe approximation function $q_\\phi(\\mathbf{z} \\vert \\mathbf{x})$ is the\nprobabilistic encoder\n, playing a similar role as $g_\\phi(\\mathbf{z} \\vert \\mathbf{x})$ above.\nLoss Function: ELBO\n#\nThe estimated posterior $q_\\phi(\\mathbf{z}\\vert\\mathbf{x})$ should be very close to the real one $p_\\theta(\\mathbf{z}\\vert\\mathbf{x})$. We can use\nKullback-Leibler divergence\nto quantify the distance between these two distributions. KL divergence $D_\\text{KL}(X|Y)$ measures how much information is lost if the distribution Y is used to represent X.\nIn our case we want to minimize $D_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) | p_\\theta(\\mathbf{z}\\vert\\mathbf{x}) )$ with respect to $\\phi$.\nBut why use $D_\\text{KL}(q_\\phi | p_\\theta)$ (reversed KL) instead of $D_\\text{KL}(p_\\theta | q_\\phi)$ (forward KL)? Eric Jang has a great explanation in his\npost\non Bayesian Variational methods. As a quick recap:\nForward and reversed KL divergence have different demands on how to match two distributions. (Image source:\nblog.evjang.com/2016/08/variational-bayes.html\n)\nForward KL divergence: $D_\\text{KL}(P|Q) = \\mathbb{E}_{z\\sim P(z)} \\log\\frac{P(z)}{Q(z)}$; we have to ensure that Q(z)>0 wherever P(z)>0. The optimized variational distribution $q(z)$ has to cover over the entire $p(z)$.\nReversed KL divergence: $D_\\text{KL}(Q|P) = \\mathbb{E}_{z\\sim Q(z)} \\log\\frac{Q(z)}{P(z)}$; minimizing the reversed KL divergence squeezes the $Q(z)$ under $P(z)$.\nLet’s now expand the equation:\n$$\n\\begin{aligned}\n& D_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}\\vert\\mathbf{x}) ) & \\\\\n&=\\int q_\\phi(\\mathbf{z} \\vert \\mathbf{x}) \\log\\frac{q_\\phi(\\mathbf{z} \\vert \\mathbf{x})}{p_\\theta(\\mathbf{z} \\vert \\mathbf{x})} d\\mathbf{z} & \\\\\n&=\\int q_\\phi(\\mathbf{z} \\vert \\mathbf{x}) \\log\\frac{q_\\phi(\\mathbf{z} \\vert \\mathbf{x})p_\\theta(\\mathbf{x})}{p_\\theta(\\mathbf{z}, \\mathbf{x})} d\\mathbf{z} & \\scriptstyle{\\text{; Because }p(z \\vert x) = p(z, x) / p(x)} \\\\\n&=\\int q_\\phi(\\mathbf{z} \\vert \\mathbf{x}) \\big( \\log p_\\theta(\\mathbf{x}) + \\log\\frac{q_\\phi(\\mathbf{z} \\vert \\mathbf{x})}{p_\\theta(\\mathbf{z}, \\mathbf{x})} \\big) d\\mathbf{z} & \\\\\n&=\\log p_\\theta(\\mathbf{x}) + \\int q_\\phi(\\mathbf{z} \\vert \\mathbf{x})\\log\\frac{q_\\phi(\\mathbf{z} \\vert \\mathbf{x})}{p_\\theta(\\mathbf{z}, \\mathbf{x})} d\\mathbf{z} & \\scriptstyle{\\text{; Because }\\int q(z \\vert x) dz = 1}\\\\\n&=\\log p_\\theta(\\mathbf{x}) + \\int q_\\phi(\\mathbf{z} \\vert \\mathbf{x})\\log\\frac{q_\\phi(\\mathbf{z} \\vert \\mathbf{x})}{p_\\theta(\\mathbf{x}\\vert\\mathbf{z})p_\\theta(\\mathbf{z})} d\\mathbf{z} & \\scriptstyle{\\text{; Because }p(z, x) = p(x \\vert z) p(z)} \\\\\n&=\\log p_\\theta(\\mathbf{x}) + \\mathbb{E}_{\\mathbf{z}\\sim q_\\phi(\\mathbf{z} \\vert \\mathbf{x})}[\\log \\frac{q_\\phi(\\mathbf{z} \\vert \\mathbf{x})}{p_\\theta(\\mathbf{z})} - \\log p_\\theta(\\mathbf{x} \\vert \\mathbf{z})] &\\\\\n&=\\log p_\\theta(\\mathbf{x}) + D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z})) - \\mathbb{E}_{\\mathbf{z}\\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})}\\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) &\n\\end{aligned}\n$$\nSo we have:\n$$\nD_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}\\vert\\mathbf{x}) ) =\\log p_\\theta(\\mathbf{x}) + D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z})) - \\mathbb{E}_{\\mathbf{z}\\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})}\\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z})\n$$\nOnce rearrange the left and right hand side of the equation,\n$$\n\\log p_\\theta(\\mathbf{x}) - D_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}\\vert\\mathbf{x}) ) = \\mathbb{E}_{\\mathbf{z}\\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})}\\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) - D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}))\n$$\nThe LHS of the equation is exactly what we want to maximize when learning the true distributions: we want to maximize the (log-)likelihood of generating real data (that is $\\log p_\\theta(\\mathbf{x})$) and also minimize the difference between the real and estimated posterior distributions (the term $D_\\text{KL}$ works like a regularizer). Note that $p_\\theta(\\mathbf{x})$ is fixed with respect to $q_\\phi$.\nThe negation of the above defines our loss function:\n$$\n\\begin{aligned}\nL_\\text{VAE}(\\theta, \\phi) \n&= -\\log p_\\theta(\\mathbf{x}) + D_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}\\vert\\mathbf{x}) )\\\\\n&= - \\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})} \\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) + D_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}) ) \\\\\n\\theta^{*}, \\phi^{*} &= \\arg\\min_{\\theta, \\phi} L_\\text{VAE}\n\\end{aligned}\n$$\nIn Variational Bayesian methods, this loss function is known as the\nvariational lower bound\n, or\nevidence lower bound\n. The “lower bound” part in the name comes from the fact that KL divergence is always non-negative and thus $-L_\\text{VAE}$ is the lower bound of $\\log p_\\theta (\\mathbf{x})$.\n$$\n-L_\\text{VAE} = \\log p_\\theta(\\mathbf{x}) - D_\\text{KL}( q_\\phi(\\mathbf{z}\\vert\\mathbf{x}) \\| p_\\theta(\\mathbf{z}\\vert\\mathbf{x}) ) \\leq \\log p_\\theta(\\mathbf{x})\n$$\nTherefore by minimizing the loss, we are maximizing the lower bound of the probability of generating real data samples.\nReparameterization Trick\n#\nThe expectation term in the loss function invokes generating samples from $\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})$. Sampling is a stochastic process and therefore we cannot backpropagate the gradient. To make it trainable, the reparameterization trick is introduced: It is often possible to express the random variable $\\mathbf{z}$ as a deterministic variable $\\mathbf{z} = \\mathcal{T}_\\phi(\\mathbf{x}, \\boldsymbol{\\epsilon})$, where $\\boldsymbol{\\epsilon}$ is an auxiliary independent random variable, and the transformation function $\\mathcal{T}_\\phi$ parameterized by $\\phi$ converts $\\boldsymbol{\\epsilon}$ to $\\mathbf{z}$.\nFor example, a common choice of the form of $q_\\phi(\\mathbf{z}\\vert\\mathbf{x})$ is a multivariate Gaussian with a diagonal covariance structure:\n$$\n\\begin{aligned}\n\\mathbf{z} &\\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x}^{(i)}) = \\mathcal{N}(\\mathbf{z}; \\boldsymbol{\\mu}^{(i)}, \\boldsymbol{\\sigma}^{2(i)}\\boldsymbol{I}) & \\\\\n\\mathbf{z} &= \\boldsymbol{\\mu} + \\boldsymbol{\\sigma} \\odot \\boldsymbol{\\epsilon} \\text{, where } \\boldsymbol{\\epsilon} \\sim \\mathcal{N}(0, \\boldsymbol{I}) & \\scriptstyle{\\text{; Reparameterization trick.}}\n\\end{aligned}\n$$\nwhere $\\odot$ refers to element-wise product.\nIllustration of how the reparameterization trick makes the $\\mathbf{z}$ sampling process trainable.(Image source: Slide 12 in Kingma’s NIPS 2015 workshop\ntalk\n)\nThe reparameterization trick works for other types of distributions too, not only Gaussian.\nIn the multivariate Gaussian case, we make the model trainable by learning the mean and variance of the distribution, $\\mu$ and $\\sigma$, explicitly using the reparameterization trick, while the stochasticity remains in the random variable $\\boldsymbol{\\epsilon} \\sim \\mathcal{N}(0, \\boldsymbol{I})$.\nIllustration of variational autoencoder model with the multivariate Gaussian assumption.\nBeta-VAE\n#\nIf each variable in the inferred latent representation $\\mathbf{z}$ is only sensitive to one single generative factor and relatively invariant to other factors, we will say this representation is disentangled or factorized. One benefit that often comes with disentangled representation is\ngood interpretability\nand easy generalization to a variety of tasks.\nFor example, a model trained on photos of human faces might capture the gentle, skin color, hair color, hair length, emotion, whether wearing a pair of glasses and many other relatively independent factors in separate dimensions. Such a disentangled representation is very beneficial to facial image generation.\nβ-VAE (\nHiggins et al., 2017\n) is a modification of Variational Autoencoder with a special emphasis to discover disentangled latent factors. Following the same incentive in VAE, we want to maximize the probability of generating real data, while keeping the distance between the real and estimated posterior distributions small (say, under a small constant $\\delta$):\n$$\n\\begin{aligned}\n&\\max_{\\phi, \\theta} \\mathbb{E}_{\\mathbf{x}\\sim\\mathcal{D}}[\\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})} \\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z})]\\\\\n&\\text{subject to } D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x})\\|p_\\theta(\\mathbf{z})) < \\delta\n\\end{aligned}\n$$\nWe can rewrite it as a Lagrangian with a Lagrangian multiplier $\\beta$ under the\nKKT condition\n. The above optimization problem with only one inequality constraint is equivalent to maximizing the following equation $\\mathcal{F}(\\theta, \\phi, \\beta)$:\n$$\n\\begin{aligned}\n\\mathcal{F}(\\theta, \\phi, \\beta) &= \\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})} \\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) - \\beta(D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x})\\|p_\\theta(\\mathbf{z})) - \\delta) & \\\\\n& = \\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})} \\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) - \\beta D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x})\\|p_\\theta(\\mathbf{z})) + \\beta \\delta & \\\\\n& \\geq \\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})} \\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) - \\beta D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x})\\|p_\\theta(\\mathbf{z})) & \\scriptstyle{\\text{; Because }\\beta,\\delta\\geq 0}\n\\end{aligned}\n$$\nThe loss function of $\\beta$-VAE is defined as:\n$$\nL_\\text{BETA}(\\phi, \\beta) = - \\mathbb{E}_{\\mathbf{z} \\sim q_\\phi(\\mathbf{z}\\vert\\mathbf{x})} \\log p_\\theta(\\mathbf{x}\\vert\\mathbf{z}) + \\beta D_\\text{KL}(q_\\phi(\\mathbf{z}\\vert\\mathbf{x})\\|p_\\theta(\\mathbf{z}))\n$$\nwhere the Lagrangian multiplier $\\beta$ is considered as a hyperparameter.\nSince the negation of $L_\\text{BETA}(\\phi, \\beta)$ is the lower bound of the Lagrangian $\\mathcal{F}(\\theta, \\phi, \\beta)$. Minimizing the loss is equivalent to maximizing the Lagrangian and thus works for our initial optimization problem.\nWhen $\\beta=1$, it is same as VAE. When $\\beta > 1$, it applies a stronger constraint on the latent bottleneck and limits the representation capacity of $\\mathbf{z}$. For some conditionally independent generative factors, keeping them disentangled is the most efficient representation. Therefore a higher $\\beta$ encourages more efficient latent encoding and further encourages the disentanglement. Meanwhile, a higher $\\beta$ may create a trade-off between reconstruction quality and the extent of disentanglement.\nBurgess, et al. (2017)\ndiscussed the distentangling in $\\beta$-VAE in depth with an inspiration by the\ninformation bottleneck theory\nand further proposed a modification to $\\beta$-VAE to better control the encoding representation capacity.\nVQ-VAE and VQ-VAE-2\n#\nThe\nVQ-VAE\n(“Vector Quantised-Variational AutoEncoder”;\nvan den Oord, et al. 2017\n) model learns a discrete latent variable by the encoder, since discrete representations may be a more natural fit for problems like language, speech, reasoning, etc.\nVector quantisation (VQ) is a method to map $K$-dimensional vectors into a finite set of “code” vectors. The process is very much similar to\nKNN\nalgorithm. The optimal centroid code vector that a sample should be mapped to is the one with minimum euclidean distance.\nLet $\\mathbf{e} \\in \\mathbb{R}^{K \\times D}, i=1, \\dots, K$ be the latent embedding space (also known as “codebook”) in VQ-VAE, where $K$ is the number of latent variable categories and $D$ is the embedding size. An individual embedding vector is $\\mathbf{e}_i \\in \\mathbb{R}^{D}, i=1, \\dots, K$.\nThe encoder output $E(\\mathbf{x}) = \\mathbf{z}_e$ goes through a nearest-neighbor lookup to match to one of $K$ embedding vectors and then this matched code vector becomes the input for the decoder $D(.)$:\n$$\n\\mathbf{z}_q(\\mathbf{x}) = \\text{Quantize}(E(\\mathbf{x})) = \\mathbf{e}_k \\text{ where } k = \\arg\\min_i \\|E(\\mathbf{x}) - \\mathbf{e}_i \\|_2\n$$\nNote that the discrete latent variables can have different shapes in differnet applications; for example, 1D for speech, 2D for image and 3D for video.\nThe architecture of VQ-VAE (Image source:\nvan den Oord, et al. 2017\n)\nBecause argmin() is non-differentiable on a discrete space, the gradients $\\nabla_z L$ from decoder input $\\mathbf{z}_q$ is copied to the encoder output $\\mathbf{z}_e$. Other than reconstruction loss, VQ-VAE also optimizes:\nVQ loss\n: The L2 error between the embedding space and the encoder outputs.\nCommitment loss\n: A measure to encourage the encoder output to stay close to the embedding space and to prevent it from fluctuating too frequently from one code vector to another.\n$$\nL = \\underbrace{\\|\\mathbf{x} - D(\\mathbf{e}_k)\\|_2^2}_{\\textrm{reconstruction loss}} + \n\\underbrace{\\|\\text{sg}[E(\\mathbf{x})] - \\mathbf{e}_k\\|_2^2}_{\\textrm{VQ loss}} + \n\\underbrace{\\beta \\|E(\\mathbf{x}) - \\text{sg}[\\mathbf{e}_k]\\|_2^2}_{\\textrm{commitment loss}}\n$$\nwhere $\\text{sg}[.]$ is the\nstop_gradient\noperator.\nThe embedding vectors in the codebook is updated through EMA (exponential moving average). Given a code vector $\\mathbf{e}_i$, say we have $n_i$ encoder output vectors, $\\{\\mathbf{z}_{i,j}\\}_{j=1}^{n_i}$, that are quantized to $\\mathbf{e}_i$:\n$$\nN_i^{(t)} = \\gamma N_i^{(t-1)} + (1-\\gamma)n_i^{(t)}\\;\\;\\;\n\\mathbf{m}_i^{(t)} = \\gamma \\mathbf{m}_i^{(t-1)} + (1-\\gamma)\\sum_{j=1}^{n_i^{(t)}}\\mathbf{z}_{i,j}^{(t)}\\;\\;\\;\n\\mathbf{e}_i^{(t)} = \\mathbf{m}_i^{(t)} / N_i^{(t)}\n$$\nwhere $(t)$ refers to batch sequence in time. $N_i$ and $\\mathbf{m}_i$ are accumulated vector count and volume, respectively.\nVQ-VAE-2 (\nAli Razavi, et al. 2019\n) is a two-level hierarchical VQ-VAE combined with self-attention autoregressive model.\nStage 1 is to\ntrain a hierarchical VQ-VAE\n: The design of hierarchical latent variables intends to separate local patterns (i.e., texture) from global information (i.e., object shapes). The training of the larger bottom level codebook is conditioned on the smaller top level code too, so that it does not have to learn everything from scratch.\nStage 2 is to\nlearn a prior over the latent discrete codebook\nso that we sample from it and generate images. In this way, the decoder can receive input vectors sampled from a similar distribution as the one in training. A powerful autoregressive model enhanced with multi-headed self-attention layers is used to capture the prior distribution (like\nPixelSNAIL; Chen et al 2017\n).\nConsidering that VQ-VAE-2 depends on discrete latent variables configured in a simple hierarchical setting, the quality of its generated images are pretty amazing.\nArchitecture of hierarchical VQ-VAE and multi-stage image generation. (Image source:\nAli Razavi, et al. 2019\n)\nThe VQ-VAE-2 algorithm. (Image source:\nAli Razavi, et al. 2019\n)\nTD-VAE\n#\nTD-VAE\n(“Temporal Difference VAE”;\nGregor et al., 2019\n) works with sequential data. It relies on three main ideas, described below.\nState-space model as a Markov Chain model.\n1. State-Space Models\nIn (latent) state-space models, a sequence of unobserved hidden states $\\mathbf{z} = (z_1, \\dots, z_T)$ determine the observation states $\\mathbf{x} = (x_1, \\dots, x_T)$. Each time step in the Markov chain model in Fig. 13 can be trained in a similar manner as in Fig. 6, where the intractable posterior $p(z \\vert x)$ is approximated by a function $q(z \\vert x)$.\n2. Belief State\nAn agent should learn to encode all the past states to reason about the future, named as\nbelief state\n, $b_t = belief(x_1, \\dots, x_t) = belief(b_{t-1}, x_t)$. Given this, the distribution of future states conditioned on the past can be written as $p(x_{t+1}, \\dots, x_T \\vert x_1, \\dots, x_t) \\approx p(x_{t+1}, \\dots, x_T \\vert b_t)$. The hidden states in a recurrent policy are used as the agent’s belief state in TD-VAE. Thus we have $b_t = \\text{RNN}(b_{t-1}, x_t)$.\n3. Jumpy Prediction\nFurther, an agent is expected to imagine distant futures based on all the information gathered so far, suggesting the capability of making jumpy predictions, that is, predicting states several steps further into the future.\nRecall what we have learned from the variance lower bound\nabove\n:\n$$\n\\begin{aligned}\n\\log p(x) \n&\\geq \\log p(x) - D_\\text{KL}(q(z|x)\\|p(z|x)) \\\\\n&= \\mathbb{E}_{z\\sim q} \\log p(x|z) - D_\\text{KL}(q(z|x)\\|p(z)) \\\\\n&= \\mathbb{E}_{z \\sim q} \\log p(x|z) - \\mathbb{E}_{z \\sim q} \\log \\frac{q(z|x)}{p(z)} \\\\\n&= \\mathbb{E}_{z \\sim q}[\\log p(x|z) -\\log q(z|x) + \\log p(z)] \\\\\n&= \\mathbb{E}_{z \\sim q}[\\log p(x, z) -\\log q(z|x)] \\\\\n\\log p(x) \n&\\geq \\mathbb{E}_{z \\sim q}[\\log p(x, z) -\\log q(z|x)]\n\\end{aligned}\n$$\nNow let’s model the distribution of the state $x_t$ as a probability function conditioned on all the past states $x_{