VAE 理论推导及代码实现

VAE 理论推导及代码实现

熵、交叉熵、KL 散度的概念

熵（Entropy)

假设 p (x）是一个分布函数，满足在 x 上的积分为 1，那么 $p(x)$的熵定义为 $H(p(x))$，这里我们简写为 $H(p)$
$$H(p)=\int p(x)\log \frac{1}{p(x)}dx$$
直观上，越分散的分布函数熵越大。越集中的分布函数熵越小。熵的最小值为 0.

从信息论的角度来说，熵又叫信息熵，它的大小表示信息量的多少，分散的分布函数可能性多、拿到 p (x）后对于 x 的推断不确定性大，即信息量大，而对于 p © =1 这种情况拿到分布函数直接就拿到了结果，因此信息量为 0

交叉熵（Cross-Entropy)

假设$p(x)$、$q(x)$是两个分布函数，交叉熵的小大评价了这两个分布函数的相似与否。$p$ 和 $q$ 的交叉熵记为$H(p,q)$
$$H(p,q)=\int p(x)\log \frac{1}{q(x)}dx$$

交叉熵小一分布相似；交叉熵大一分布不相似。交叉熵最大为无穷大，最小为 $p$ 的熵 $H(p)$

KL 散度

假设 $p(x)$、$q(x)$是两个分布函数，KL 散度的小大评价了这两个分布函数的相似与否，同时考虑了$KL(x)$这个分布的信息量。记为 $KL(p,q)$。注意：$KL(p,q)$也不一定等于 $KL(q,p)$。
$$KL(p,q)=H(p,q)-H(p)$$
$$\begin{array}{rl}& \int p(x)\log \frac{1}{q(x)}dx-\int p(x)\log \frac{1}{p(x)}dx\\ & =\int p(x)\log \frac{p(x)}{q(x)}dx\end{array}$$

KL散度小—分布相似 & $[p(x)$ 分散 | $p(x)$ 信息量大]。
$\mathrm{KL}$ 散度大–分布不相似 & $[p(x)$ 集中 $\mid p(x)$ 信息量小]。
$\mathrm{KL}$ 散度最小值为 $0:p(x)$ 和 $q(x)$ 完全相同时。

概率知识

将p(x)其改写为包含了传入参数的形式
$$p(x)=\sum _{z}p(x\mid z)p(z)$$

连续分布时，该式就变成了
$$p(x)={\int }_z^{}p(x\mid z)p(z)dz$$

$p(z)$可以是任意分布，在VAE中我们常常假设p(z)服从标准正态分布。

变分方法

Intractability:
$$\begin{array}{rl}{p}_{\mathbf{\theta }}(\mathbf{z}\mid \mathbf{x})& =p_{\mathbf{\theta }}(\mathbf{x}\mid \mathbf{z})p_{\mathbf{\theta }}(\mathbf{z})/p_{\mathbf{\theta }}(\mathbf{x})\\ p_{\mathbf{\theta }}(\mathbf{x})& =\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\dot{\mathbf{x}}\mid \mathbf{z})d\mathbf{z}\end{array}$$
$$\begin{array}{rl}p\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)& =\frac{p\left(\mathbf{z},{\mathbf{x}}^{(i)}\right)}{p\left({\mathbf{x}}^{(i)}\right)}\\ & =\frac{p\left(\mathbf{x}={\mathbf{x}}^{(i)}\mid \mathbf{z}={\mathbf{z}}^{(i)}\right)p\left(\mathbf{z}={\mathbf{z}}^{(i)}\right)}{{\int }_{{\mathbf{z}}^{(i)}}p\left(\mathbf{x}={\mathbf{x}}^{(i)}\mid \mathbf{z}={\mathbf{z}}^{(i)}\right)p\left(\mathbf{z}={\mathbf{z}}^{(i)}\right)d{\mathbf{z}}^{(i)}}\\ & =\frac{p\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)p(\mathbf{z})}{{\int }_{\mathbf{z}}p\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)p(\mathbf{z})d\mathbf{z}}\end{array}$$
参考：https://zhuanlan.zhihu.com/p/519448634

如果假设参数 $\theta$ 已知, 那么先验分布 $p_{\theta }(\mathbf{z})$ 和条件似然函数 $p_{\theta }\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)$ 就都是已知的。理论上 来说, 只要把分母里的积分项 ${\int }_{\mathbf{z}}{p}_{\theta }\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)p(\mathbf{z})d\mathbf{z}$ 计算出来, 那整个后验分布 $p\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 就 可以求了, 后验推断问题也就解决了。但是, 现实很骨感, 在没有对 $p_{\theta }(\mathbf{z})$ 和 $p_{\theta }\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)$ 作任何 简化假设的前提下, 这个积分基本上是没有解析解的。你想硬着头皮解, 那么基本意味着你要穷举 隐变量 $\mathbf{z}$ 的所有可能取值, 假设 $\mathbf{z}$ 有 $k$ 个维度, 每个维度采样 $n$ 个取值, 那么这个穷举过程的复 杂度就是 $O\left(n^k\right)$ 。

当然也有人用MCMC来做积分项的估计，虽然这个方案做采样估计很精准，但是费时费力，很难适用于大数据场景。所以一般更常见的方案是采用变分方法（variational method），它可以绕过对积分项的求解，通过把统计推断问题转化成参数优化问题来实现“降维打击”。

首先变分方法会设置一个新的参数化分布 $q_{\varphi }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$, 它的参数是 $\varphi$, 我们把它称作"识别模型" (原文记作recognition model) 。变分方法的核心思想是：直接让“识别模型”去拟合后验分布 $p_{\theta }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$, 只要近似到位, 那么采用 $q_{\varphi }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 作为后验推断的结果就行了。如何做近似呢? 很简单, 直接最小化 $q_{\varphi }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 和 $p_{\theta }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 两者间的KL散度即可。

就这样，变分方法把原来的统计推断问题转化成了优化问题：

Approximation $p_{\theta }(z\mid x)\cong q_{\varphi }(z\mid x)$
D K L ( q ϕ ( z ∣ x ) ∥ p θ ( z ∣ x ) ) = − ∑ decoder  q ϕ ( z ∣ x ) log ⁡ ( p θ ( z ∣ x ) q ϕ ( z ∣ x ) ) = − ∑ z q ϕ ( z ∣ x ) log ⁡ ( p θ ( x , z ) p θ ( x ) q ϕ ( z ∣ x ) ) = − ∑ z q ϕ ( z ∣ x ) [ log ⁡ ( p θ ( x , z ) q ϕ ( z ∣ x ) ) − log ⁡ ( p θ ( x ) ) ‾ ]  non-negative  log ⁡ ( p θ ( x ) ) = K L ( q ϕ ( z ∣ x ) ∥ p θ ( z ∣ x ) ) + ∑ z q ϕ ( z ∣ x ) log ⁡ ( p θ ( x , z ) q ϕ ( z ∣ x ) ) = D K L ( q ϕ ( z ∣ x ) ∣ ∣ p θ ( z ∣ x ) ) + L ( θ , ϕ ; x )  Variational lower bound  \begin{aligned} & D_{K L}\left(q_\phi(z \mid x) \| p_\theta(z \mid x)\right)=-\sum_{\text {decoder }} q_\phi(z \mid x) \log \left(\frac{p_\theta(z \mid x)}{q_\phi(z \mid x)}\right)=-\sum_z q_\phi(z \mid x) \log \left(\frac{\frac{p_\theta(x, z)}{p_\theta(x)}}{q_\phi(z \mid x)}\right) \\ & =-\sum_z q_\phi(z \mid x)\left[\log \left(\frac{p_\theta(x, z)}{q_\phi(z \mid x)}\right)-\underline{\log \left(p_\theta(x)\right)}\right] \\ & \begin{array}{c} \text { non-negative } \\ \log \left(p_\theta(x)\right) \end{array}=K L\left(q_\phi(z \mid x) \| p_\theta(z \mid x)\right)+\sum_z q_\phi(z \mid x) \log \left(\frac{p_\theta(x, z)}{q_\phi(z \mid x)}\right) \\ & =D_{K L}\left(q_\phi(z \mid x)|| p_\theta(z \mid x)\right)+\frac{L(\theta, \phi ; x)}{\text { Variational lower bound }} \\ & \end{aligned} ​DKL​(qϕ​(z∣x)∥pθ​(z∣x))=−decoder ∑​qϕ​(z∣x)log(qϕ​(z∣x)pθ​(z∣x)​)=−z∑​qϕ​(z∣x)log ​qϕ​(z∣x)pθ​(x)pθ​(x,z)​​ ​=−z∑​qϕ​(z∣x)[log(qϕ​(z∣x)pθ​(x,z)​)−log(pθ​(x))​] non-negative log(pθ​(x))​=KL(qϕ​(z∣x)∥pθ​(z∣x))+z∑​qϕ​(z∣x)log(qϕ​(z∣x)pθ​(x,z)​)=DKL​(qϕ​(z∣x)∣∣pθ​(z∣x))+ Variational lower bound L(θ,ϕ;x)​​

Maximize the lower bound
$$\begin{array}{rl}& L(\theta ,\varphi ;x)=\sum _z{q}_{\varphi }(z\mid x)\log \left(\frac{p_{\theta }(x,z)}{q_{\varphi }(z\mid x)}\right)=\sum _z{q}_{\varphi }(z\mid x)\log \left(\frac{p_{\theta }(x\mid z)p_{\theta }(z)}{q_{\varphi }(z\mid x)}\right)\\ & =\sum _z{q}_{\varphi }(z\mid x)\left[\log \left(p_{\theta }(x\mid z)\right)+\log \left(\frac{p_{\theta }(z)}{q_{\varphi }(z\mid x)}\right)\right]\\ & =\frac{E_{q_{\varphi }(z\mid x)}\left[\log \left(p_{\theta }(x\mid z)\right)\right]}{\text{ Reconstruction Loss }}-\frac{D_{KL}\left(q_{\varphi }(z\mid x)\Vert p_{\theta }(z)\right)}{\text{ Regularization Loss }}\end{array}$$

$$\begin{array}{c}L(\theta ,\varphi ;x)=\sum _z{q}_{\varphi }(z\mid x)\log \left(\frac{p_{\theta }(x,z)}{q_{\varphi }(z\mid x)}\right)=\sum _z{q}_{\varphi }(z\mid x)\log \left(\frac{p_{\theta }(x\mid z)p_{\theta }(z)}{q_{\varphi }(z\mid x)}\right)\\ =\sum _z{q}_{\varphi }(z\mid x)\left[\log \left(p_{\theta }(x\mid z)\right)+\log \left(\frac{p_{\theta }(z)}{q_{\varphi }(z\mid x)}\right)\right]\end{array}$$

Regularization Loss（ 重参数化）

而在实践中, 一般不对 $q_{\varphi }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 直接作采样, 采用 reparameterization trick 来简化操作, 我们 设 ${\mathbf{z}}^{(i,l)}=g_{\varphi }\left({ϵ}^{(i;l)};{\mathbf{x}}^{(i)}\right)$, 其中 $g_{\varphi }$ 是一个拟合函数 (e.g. 神经网络) , 而噪声 ${ϵ}^{(i;l)}$ 可以通过 采样得到, 一般直接采样自简单的标准正态分布。

$$\int q_{\theta }(z\mid x)\log p(z)dz=\int N\left(z;\mu ,{\sigma }^2\right)\log N(z;0,I)dz$$
$$f(x)=\frac{1}{\sigma \sqrt{2\pi }}{e}^{-\frac{1}{2}{\left(\frac{x-\mu }{\sigma }\right)}^2}$$

$$\begin{array}{rl}& =\int N\left(z;\mu ,{\sigma }^2\right)\left(-\frac{1}{2}{z}^2-\frac{1}{2}\log (2\pi )\right)dz=-\frac{1}{2}\int N\left(z;\mu ,{\sigma }^2\right)z^{2}dz-\frac{J}{2}\log (2\pi )\\ & =-\frac{J}{2}\log (2\pi )-\frac{1}{2}{E}_{z\sim N\left(z;\mu ,{\sigma }^2\right)}\left[Z^2\right]\\ & =-\frac{J}{2}\log (2\pi )-\frac{1}{2}\left(E_{z\sim N\left(z;\mu ,{\sigma }^2\right)}[Z{]}^2+\mathrm{Var}(Z)\right)\\ & =-\frac{J}{2}\log (2\pi )-\frac{1}{2}\sum _{j=1}^J\left({\mu }_j^2+{\sigma }_j^2\right)\text{ Let }J\text{ be the dimensionality of }z\end{array}$$

L1用于最小化$KL(q(z|x)||p(z))$，VAE假设$q(z|x)$的分布为正态分布，而$p(z)$为标准正态分布。计算两个正态分布之间的KL散度的公式如下：
$$KL\left(N\left({\mu }_1,{\sigma }_1^2\right),N\left({\mu }_2,{\sigma }_2^2\right)\right)=\log \frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+{\left({\mu }_1-{\mu }_2\right)}^2}{2{\sigma }_2^2}-\frac{1}{2}$$

由于此处p(z)为标准正态分布，因此其μ为0，σ为1，那么我们带入后可得
$$L_1=-\frac{1}{2}\left(\log {\sigma }^2-{\sigma }^2-{\mu }^2+1\right)$$

采用reparameterization trick有两大好处：

• 由于分布 $q_{\varphi }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 可能是一个比较复杂的函数, 直接采样操作费时费力, 而且采样方差可能很 大, 不利于收玫, 通过reparameterization可以简化操作, 提高效率, 提高数值上的稳定性;
• 假设我们不考虑采样难度, 直接对 $q_{\varphi }\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)$ 采样, 那么梯度反向传播的时候, 损失函数中的 $\frac{1}{L}{\sum }_{l=1}^L\left[\log p_{\theta }\left({\mathbf{x}}^{(i)}\mid {\mathbf{z}}^{(i,l)}\right)\right]$ 是没法对 $\varphi$ 求导的, 这样损失函数 $\mathcal{L}\left(\varphi ,\theta ,{\mathbf{x}}^{(\mathbf{i})}\right)$ 只能通过KL散度 的梯度对 $\varphi$ 做优化, 这和我们做联合参数优化的意图是违背的。所以使用reparameterization trick 让 ${\mathbf{z}}^{(i,l)}=g_{\varphi }\left({ϵ}^{(i;l)};{\mathbf{x}}^{(i)}\right)$, 实际上是让参数 $\theta ,\varphi$ 可以同时得到期望项和KL散度项的反传 梯度进行优化, 让模型学得更好。

Reconstruction Loss

$$\begin{array}{rl}& \mathcal{L}\left(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)}\right)=\underline{-D_{KL}\left(q_{\mathbf{\varphi }}\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)\Vert p_{\mathbf{\theta }}(\mathbf{z})\right)}+\underline{{\mathbb{E}}_{q_{\mathbf{\varphi }}\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)}\left[\log p_{\mathbf{\theta }}\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)\right]}\\ & -D_{KL}(\left(q_{\mathbf{\varphi }}(\mathbf{z})\Vert p_{\mathbf{\theta }}(\mathbf{z})\right)=\int q_{\mathbf{\theta }}(\mathbf{z})\left(\log p_{\mathbf{\theta }}(\mathbf{z})-\log q_{\mathbf{\theta }}(\mathbf{z})\right)d\mathbf{z}\\ & =\frac{1}{2}\sum _{j=1}^J\left(1+\log \left({\left({\sigma }_j\right)}^2\right)-{\left({\mu }_j\right)}^2-{\left({\sigma }_j\right)}^2\right)\\ & f^{\ast }=\underset{f\in F}{\mathrm{argmax}}{\mathbb{E}}_{z\sim q_x^{\ast }}(\log p(x\mid z))\\ & =\underset{f\in F}{\mathrm{argmax}}{\mathbb{E}}_{z\sim q_x^{\ast }}\left(-\frac{\Vert x-f(z){\Vert }^2}{2c}\right)\\ & \end{array}$$

$$\mathcal{L}\left(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)}\right)=-D_{KL}\left(q_{\mathbf{\varphi }}\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)\Vert p_{\mathbf{\theta }}(\mathbf{z})\right)+{\mathbb{E}}_{q_{\mathbf{\varphi }}\left(\mathbf{z}\mid {\mathbf{x}}^{(i)}\right)}\left[\log p_{\mathbf{\theta }}\left({\mathbf{x}}^{(i)}\mid \mathbf{z}\right)\right]$$

import argparse
import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image


parser = argparse.ArgumentParser(description='VAE MNIST Example with Different Losses')
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=100, metavar='N',
                    help='number of epochs to train (default: 100)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=100000, metavar='N',
                    help='how many batches to wait before logging training status')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)

device = torch.device("cuda" if args.cuda else "cpu")

kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./datasets', train=True, download=True,
                   transform=transforms.ToTensor()),
    batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./datasets', train=False, transform=transforms.ToTensor()),
    batch_size=args.batch_size, shuffle=True, **kwargs)


class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x.view(-1, 784))
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar


model = VAE().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

### 1.
# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function_original(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')

    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

### 2. 
# using the loss function which only consider reconstruction term.
def loss_function_only_recon(recon_x, x):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
    return BCE

### 3. 
# be careful of the way two losses calculated.
# the only difference of this loss function is that the second term - KLD
# is "mean".
def loss_function_o1(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')

    KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

### 4.
def loss_function_o2(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='mean')

    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

### 5.
def loss_function_kld(recon_x, x, mu, logvar):
    KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    return KLD

### 6.
# apply the l1 loss
def loss_function_l1(recon_x, x, mu, logvar):
    L1 = F.l1_loss(recon_x, x.view(-1, 784), reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return L1 + KLD

### 7.
# apply the MSE loss
def loss_function_l2(recon_x, x, mu, logvar):
    L1 = F.mse_loss(recon_x, x.view(-1, 784), reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return L1 + KLD

def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        data = data.to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        which_loss = 7
        
        if which_loss==1:
            loss = loss_function_original(recon_batch, data, mu, logvar)
        elif which_loss==2:
            loss = loss_function_only_recon(recon_batch, data)
        elif which_loss==3:
            loss = loss_function_o1(recon_batch, data, mu, logvar)
        elif which_loss==4:
            loss = loss_function_o2(recon_batch, data, mu, logvar)
        elif which_loss==5:
            loss = loss_function_kld(recon_batch, data, mu, logvar)
        elif which_loss==6:
            loss = loss_function_l1(recon_batch, data, mu, logvar)
        elif which_loss==7:
            loss = loss_function_l2(recon_batch, data, mu, logvar)
            
        loss.backward()
        train_loss += loss.item()
        optimizer.step()

    print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, train_loss / len(train_loader.dataset)))


def test(epoch):
    model.eval()
    with torch.no_grad():
        for i, (data, _) in enumerate(test_loader):
            data = data.to(device)
            recon_batch, mu, logvar = model(data)
            if (i == 0) and (epoch % 10 == 0):
                n = min(data.size(0), 8)
                comparison = torch.cat([data[:n],
                                      recon_batch.view(args.batch_size, 1, 28, 28)[:n]])
                save_image(comparison.cpu(),
                         'vae_img/7_m_reconstruction_' + str(epoch) + '.png', nrow=n)

if __name__ == "__main__":
    for epoch in range(1, args.epochs + 1):
        train(epoch)
        test(epoch)
        if epoch%10 == 0:
            with torch.no_grad():
                sample = torch.randn(64, 20).to(device)
                sample = model.decode(sample).cpu()
                save_image(sample.view(64, 1, 28, 28),
                        'vae_img/7_m_sample_' + str(epoch) + '.png')

参考

https://zhuanlan.zhihu.com/p/345360992