Auto-encoding Variational Bayes 阅读笔记

Notation

• $p_{\theta}(z|x)$: intractable posterior
• $p_{\theta}(x|z)$: probabilistic decoder
• $q_{\phi}(z|x)$: recognition model, variational approximation to $p_{\theta}(z|x)$, also regarded as a probabilistic encoder
• $p_{\theta}(z)p_{\theta}(x|z)$: generative model
• $\phi$: variational parameters
• $\theta$: generative parameters

Abbreviation

• SGVB: Stochastic Gradient Variational Bayes
• AEVB: auto-encoding VB
• ML: maximum likelihood
• MAP: maximum a posteriori

---

Motivation

• Problem

  • How to perform efficient inference and learning in directed probabilistic models, in the presence of continuous latent variables with intractable posterior distribution $p_{\theta}(z|x)$ and large datasets?

• Existing Solution and Difficulty

  • VB: involves the optimization of an approximation to the intractable posterior
  • mean-field: requires analytical solutions of expectations w.r.t. the approximate posterior, which are also intractable in the general case

---

Contribution of this paper

• (1) SGVB estimator: an estimator of the variational lower bound
  
  • yielded by a reparameterization of the variational lower bound
  • simple & differentiable & unbiased
  • straightforwad to optimize using standard SG ascent techniques
• (2) AEVB algorithm
  
  • using SGVB to optimize a recognition model that allows us to perform very efficient approximate posterior inference using simple ancestral sampling, which in turn allows us to efficiently learn the model parameters, without the need of expensive iterative inference schemes (such as MCMC) per datapoint.
  • condition: i.i.d. datasets $X=\{x^{(i)}\}_{i=1}^N$ & continuous latent variable $z$ per datapoint

---

Methodology

assumption

• directed graphical models with continuous latent variables
• i.i.d. dataset with latent variables per datapoint
  
  • where we like to perform
    
    • ML or MAP inference on the (global) paramters $\theta$
    • variational inference on the latent variable $z$
• $p_{\theta}(z)$ and $p_{\theta}(x|z)$: both PDFs are differentiable almost everywhere w.r.t. both $\theta$ and $z$

target case

• intractability
  
  • $p_{\theta}(x) = \int p_{\theta}(z)p_{\theta}(x|z)dz$: so we cannot evaluate or differentiate it.
  • $p_{\theta}(z|x)=\frac{p_{\theta}(x|z)p_{\theta}(z)}{p_{\theta}(x)}$: so the EM algorithm cannot be used.
  • the required integrals for any reasonable mean-field VB algorithm: so the VB algorithm cannot be used.
  • in cases of moderately complicated likelihood function, e.g. in a neural network with a nonlinear hidden layer
• a large dataset
  
  • batch optimization is too costly => minibatch or single datapoints
  • sampling bases solutions are too slow, e.g. Monte Carlo EM, since it involves a typically expensive sampling loop per datapoint.

solution and application

• efficient approximate ML or MAP estimation for $\theta$ (Full): Appendix F
  
  • allow us to mimic the hidden random process and generate artificial data that resemble the real data
• efficient approximate posterior inference $p_{\theta}(z|x)$ for a choice of $\theta$
  
  • useful for coding or data representation tasks
• efficient approximate marginal inference of $x$: Appendix D
  
  • allow us to perform all kinds of inference tasks where $p(x)$ is required, such as image denoising, inpainting, and super-resolution.

  1. derivation of the variational bound

  $$\log p_{\theta}(x^{(1)}， \cdots, x^{(N)}) = \sum_{i=1}^N \log p_{\theta}(x^{(i)})$$
  Here we use $x^i$ to represent $x^{(i)}$
  $$\log p(x^i) = \int_zq(z|x^i)\log p(x^i)dz ~~\text{($q$ can be any distribution)}\\ = \int_z q(z|x^i)\log \frac{p(z,x^i)}{p(z|x^i)}dz\\ = \int_z q(z|x^i) \log [ \frac{q(z|x^i)}{p(z|x^i)}\cdot\frac{p(z,x^i)}{q(z|x^i)} ]dz\\ = \int_z q(z|x^i) \log \frac{q(z|x^i)}{p(z|x^i)} dz + \int_z q(z|x^i) \log \frac{p(z,x^i)}{q(z|x^i)} dz\\ = D_{KL}[q_{\phi}(z|x^i)\parallel p_{\theta}(z|x^i)] + L_B(\theta,\phi;x^i)$$

  Becuase $D_{KL}\ge 0$, so $L_B$ is called the (variational) lower bound, then we have
  

  $$\log p_{\theta}(x^i) \ge L_B(\theta, \phi; x^i)$$
  $$L_B(\theta, \phi; x^i) = E_{q_{\phi}(z|x)}[\log p_{\theta}(x,z) - \log q_{\phi}(z|x)]\\ = \int_z q(z|x^i) \log \frac{p(z,x^i)}{q(z|x^i)} dz\\ = \int_z q(z|x^i) \log \frac{p(x^i|z)p(z)}{q(z|x^i)} dz\\ = -D_{KL}[q(z|x^i)\parallel p(z)] + \int_z q(z|x^i)\log p(x^i|z) dz\\ = -D_{KL}[q_{\phi}(z|x^i)\parallel p_{\theta}(z)] + E_{q_{\phi}(z|x^i)}[\log p_{\theta}(x^i|z)]$$
  • $-D_{KL}[q_{\phi}(z|x^i)\parallel p_{\theta}(z)]$: act as a regularizer
  • $E_{q_{\phi}(z|x^i)}[\log p_{\theta}(x^i|z)]$: a an expected negative reconstruction error
  • TARGET: defferentiate and optimize $L_B$ w.r.t. both $\phi$ and $\theta$. However, $\nabla_{\phi}L_B$ is problematic.

  2. Solution-1 Naive Monte Carlo gradient estimator

  • disadvantages: high variance & impractical for this purpose, 原文中出现的公式如下：
    $$\nabla_{\phi}E_{q_{\phi}(z)}[f(z)] = E_{q_{\phi}(z)}[f(z)\nabla_{q_{\phi}(z)}\log q_{\phi}(z)] \simeq \frac{1}{L}\sum_{l=1}^Lf(z)\nabla_{q_{\phi}(z^l)}\log q_{\phi}(z^l)\\ \text{where}\ z^l\sim q_{\phi}(z|x^i)$$

  但由于我还没来得及学Monte Carlo Gradient Estimator的理论，根据VAE这边论文后面的公式，个人觉得…上面的公式应该是（数学功底不够深厚，不确定二者是否等价）：
  

  $$\nabla_{\phi}E_{q_{\phi}(z)}[f(z)] = E_{q_{\phi}(z)}[f(z)\nabla_{q_{\phi}(z)}\log q_{\phi}(z)] \simeq \frac{1}{L}\sum_{l=1}^Lf(z^l)\\ \text{where}\ z^l\sim q_{\phi}(z|x^i)$$

  3. Solution-2 SGVB estimator

  • reparamterize $\tilde{z} \sim q_{\phi}(z|x)$ using a differentiable transformation $g_{\phi}(\epsilon, x)$ of an (auxiliary) noise variable $\epsilon$
    
    • under certain mild conditions for a chosen approximate posterior $q_{\phi}(z|x)$
      $$\tilde{z} = g_{\phi}(\epsilon,x) ~with~\epsilon\sim p(\epsilon)$$
  • form Monte Carlo estimates as follows:
    $$q_{\phi}(z|x)\prod_idz_i = p(\epsilon)\prod_id\epsilon_i\\ \int q_{\phi}(z|x) f(z)dz = \int p(\epsilon)f(z)d\epsilon = \int p(\epsilon)f(g_{\phi}(\epsilon,x))d\epsilon\\ E_{q_{\phi}(z|x^i)}[f(z)] = E_{p(\epsilon)}[f(g_{\phi}(\epsilon,x^i))]\simeq \frac{1}{L}\sum_{l=1}^Lf(g_{\phi}(\epsilon^l,x^i))\\ where~ \epsilon^l \sim p(\epsilon)$$
  • apply the MC estimator technique to $L_B(\theta,\phi;x^i)$, yielding 2 SGVB estimator $\tilde{L}^A$ and $\tilde{L}^B$:
    $$\tilde{L}^A(\theta,\phi;x^i) = \frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^i,z^{i,l}) - \log q_{\phi}(z^{i,l}|x^i)\\ \tilde{L}^B(\theta,\phi;x^i) =-D_{KL}(q_{\phi}(z|x^i)\parallel p_{\theta}(z)) + \frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^i|z^{i,l})\\ \text{where } z^{i,l}=g_{\phi}(\epsilon^{i,l},x^i)~\text{and}~\epsilon^l\sim p(\epsilon)$$
  • minibatch with size $M$
    $$L_B(\theta,\phi;X)\simeq \tilde{L}^M(\theta,\phi;X^M) = \frac{N}{M}\sum{i=1}^M\tilde{L}(\theta,\phi;x^i)$$
    
    • the number of samples L per datapoint can be set to 1 as long as the minibatch size $M$ was large enough, e.g. $M=100$

  3.1 Minibatch version of AEVB algorithm

  • $\theta,\phi \leftarrow$ Initialize parameters
  • repeate
    
    • $X^M \leftarrow$ Random minibatch of $M$ datapoints
    • $\epsilon \leftarrow$ Random samples from noise distribution $p(\epsilon)$
    • $g \leftarrow \nabla_{\theta,\phi}\tilde{L}^M(\theta,\phi;X^M,\epsilon)$
    • $\theta,\phi \leftarrow$ Update parameters using gradients $g$ (e.g. SGD or Adagrad)
  • until convergence of parameters $\theta,\phi$
  • return $\theta,\phi$

  chosen of $q_{\phi}(z|x)$, $p(\epsilon)$, and $g_{\phi}(\epsilon, x)$

  There are three basic approaches:

  • Tractable inverse CDF. Let $\epsilon\sim U(0,I)$, $g_{\phi}(\epsilon,x)$ be the inverse CDF of $q_{\phi}(z|x)$.
    
    • Examples: Exponential, Cauchy, Logistic, Rayleigh, Pareto, Weibull, Reciprocal, Gompertz, Gumbel and Erlang distributions.
  • ”location-scale” family of distributions: choose the standard distribution (with location = 0, scale = 1) as the auxiliary variable $\epsilon$, and let $g(\cdot)=location + scale\cdot \epsilon$
    
    • Examples: Laplace, Elliptical, Student’s t, Logistic, Uniform, Triangular and Gaussian distributions.
    • 下文介绍的VAE即适用于此种情况
  • Composition: It is often possible to express random variables as different transformations of auxiliary variables.
    
    • Examples: Log-Normal (exponentiation of normally distributed variable), Gamma (a sum over exponentially distributed variables), Dirichlet (weighted sum of Gamma variates), Beta, Chi-Squared, and F distributions.

  ---

  Variational Auto-Encoder

  • let $p_{\theta}(z)=N(z;0,I)$
  • $p_{\theta}(z|x)$ is intractable
  • use a neural network for $q_{\phi}(z|x)$
  • $\phi$ and $\theta$ are optimized jointly with the AEVB algorithm
  • params of $p_{\theta}(x|z)$ are computed from $z$ with a MLP (multi-layered perceptrons, a fully-connected neural network with a hidden layer)
    
    • multivariae Gaussian: in case of real-valued data
      $$\log p(x|z) = \log N(x;\mu,\sigma^2I)\\ \text{where}~\mu=W_{\mu}h+b_{\mu}\\ \log \sigma^2 = W_{\sigma}h+b_{\sigma}\\ h = \tanh(W_hz+b_h)$$
    • Bernouli: incase of binary data
      $$\log p(x|z) = \sum_{i=1}^D x_i\log y_i + (1-x_i)\cdot \log(1-y_i)\\ \text{where}~y=f_{\sigma}(W_y\tanh(W_h+b_h)+b_y)\\ f_{\sigma}(\cdot): \text{elementwise sigmoid activation function}$$
      
      

  From what mentioned above, we have:
  

  $$\log q_{\phi}(z|x^i) = \log N(z;\mu^i,\sigma^{2,i}I)$$
  where $\mu^i$ and $\sigma^i$ are outputs of the encoding MLP.

  We sample form $z^{i,l}\sim q_{\phi}(z|x^i)$ using $z^{i,l}=g_{\phi}(x^i,\epsilon^l)=\mu^i + \sigma^i\odot \epsilon^l$, where $\epsilon^l\sim N(0,I)$ and $\odot$ denotes element-wize product.

  Here both $p_{\theta}(z)$ and $q_{\phi}(z|x)$ are Gaussian, so we can use the estimator $\tilde{L}^B$, since the KL divergence item is analytical. Then we have:
  

  $$L(\theta,\phi;x^i) \simeq -D_{KL}(q_{\phi}(z|x^i)\parallel p_{\theta}(z)) + \frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^i|z^{i,l})\\ \simeq \frac{1}{2}\sum_{j=1}^L\left(1+\log(\sigma_j^{i2})-\mu_j^{i2}-\sigma_j^{i2} \right) + \frac{1}{L}\sum_{l=1}^L \log p_{\theta}(x^i|z^{i,l})\\ \text{where}~z^{il}=\mu^i+\sigma^i\odot\epsilon^l ~\text{and}~\epsilon^l\sim N(0,I)$$

  Solution of $-D_{KL}(q_{\phi}(z)\parallel p_{\theta}(z))$, Gaussian case

  Let $J$ be the dimensionality of $z$, then we have:

  $$\int q_{\theta}(z) \log p_{\theta}(z)dz = \int N(z;\mu,\sigma^2)\log N(z;0,I)dz\\ =-\frac{J}{2}\log(2\pi) - \frac{1}{2}\sum_{j=1}^J(\mu_j^2+\sigma_j^2)$$

  And:
  

  $$\int q_{\theta}(z) \log q_{\theta}(z)dz = \int N(z;\mu,\sigma^2)\log N(z;\mu,\sigma^2)dz\\ =-\frac{J}{2}\log(2\pi) - \frac{1}{2}\sum_{j=1}^J(1+\log\sigma_j^2)$$

  Therefore:
  

  $$-D_{KL}(q_{\phi}(z)\parallel p_{\theta}(z)) = \int q_{\theta}(z) \left(\log p_{\theta}(z) - \log q_{\theta}(z) \right) dz\\ =\frac{1}{2}\sum_{j=1}^L\left(1+\log(\sigma_j^{i2})-\mu_j^{i2}-\sigma_j^{i2} \right)$$

  此处附上相关证明：
  

  Visulization

  Since the prior of the latent space is Gaussian, linearly spaced coordinates on the unit square were transformed through the inverse CDF of the Gaussian to produce values of the latent variables .