VAE 模型基本原理简单介绍...

新用户3677sdB0 2021-09-17

展开全文

VAE 模型基本原理简单介绍

1. 编写目的
2. 推荐资料
3. 相关背景
4. 理论基础
5. 代码相关
6. 更多疑问
- 6.1 为什么叫变分自编码器？
- 6.2 为什么可以从隐变量中拿到均值和方差
7. 总结

1. 编写目的

当前不少关于异常检测的论文提到了基于VAE的算法，为了更好理解那些论文，这里专门记录一下关于VAE 的相关内容。

以下内容主要内容都是通过阅读论文《Tutorial on Variational Autoencoders》和李宏毅教授的教学视频而做的一些笔记，推荐看一下这篇论文。

内容主要包括三部分：相关背景、理论推导和核心代码。请结合代码理解应该会更加简单。

2. 推荐资料

因为时间问题，我并没有完全、足够仔细地学习VAE，但先把一些学习资源分享一下，希望能够帮到大家：

学习视频：
- 2020 李宏毅 Auto Encoder
- 2017 李宏毅 Variational Auto Encoder
- 可以在B站上找到其他一些
论文阅读，这里只推荐一篇。
- Tutorial on Variational Autoencoders 下载地址
- Auto-Encoding Variational Bayes
其他
- 简单介绍 VAE 与 CVAE 的博客
- VAE 相对详细一些的介绍博客
- VAE 的实现1
- VAE 的实现2
- VAE 的实现3
- VAE 的实现4

3. 相关背景

3.1 生成模型(Generative model)：

生成模型是指能够随机生成观测数据的模型，尤其是在给定某些隐含参数的条件下。它给观测值和标注数据序列指定一个联合概率分布。换句话说，生成模型首先研究的是特征 X X X 与标签 y y y 之间的联合分布，然后再求条件概率 P ( y ∣ X ) P(y|X) P(y∣X) ，预测时应用最大后验概率法得到预测结果。

论文(Tuturial on Variational Autoencoders)中举例如下：

图像是一种流行的数据，我们可以为其创建生成模型。每个“数据点”（图像）都有数千或数百万个维度（像素），生成模型的任务是以某种方式捕捉像素之间的依赖关系，例如，附近的像素具有相似的颜色，并被组织成对象。“捕获”这些依赖项的确切含义取决于我们要对模型做什么。一种简单明了的生成模型允许我们用数值计算 P ( X ) P(X) P(X) 。在图像的情况下，看起来像真实图像的 X X X 值应该得到高概率，而看起来像随机噪声的图像应该得到低概率。然而，像这样的模型并不一定有用：知道一个图像不太可能帮助我们合成一个相似的图像。

关于生成模型的相关概念与例子请自行搜索查询资料。

3.2 隐变量模型(Latent Variable Models)

隐变量： 不可直接观测的综合性变量。
隐变量模型： 隐变量模型是一种概率模型，其中某些变量是不可观测的。（即包隐变量的概率模型）A latent variable model is a probability model for which certain variables are never observed.

借用这篇论文中的例子，在生成手写数字的过程中(0~9)，如果数字的左半部分包含数字 5 的左半部分，则右半部分不能包含 0 的左半部分，否则字符将很明显看起来不像任何真正的数字。简单来说，如果模型在给任何像素赋值之前，先决定生成哪个字符（或者说具有某种特征的字符），这对生成字符这个工作将会很有帮助。这种决定在形式上被称为隐变量(Latent Variable)。

也就是说，在我们的模型生成任何数据之前，它首先从集合 [0,1,…,9] 中随机抽取一个数字值 z z z，然后确保接下来生成的数值中的笔划都与该字符匹配。 z z z 之所以称为 “latent”，是因为我们需要的是模型产生一个字符，但是我们并不清楚模型是使用了哪些隐变量而产生这个字符的。

3.3 自编码器 (AutoEncoder)

3.3.1 概述

自编码器是一类在半监督学习和非监督学习中使用的人工神经网络，其功能是通过将输入信息作为学习目标，对输入信息进行表征学习。

AutoEncoder 包括编码器(Encoder) 和解码器(Decoder) 两部分。Encoder 过程是将原先的数据（常用于图像方向）压缩为低维向量；Decoder 则是把低维向量还原为原来数据。

因为自编码常用于图像方向，因为一张图每一个像素点都被看作是一个特征，所以随便一张现实中的图都属于高维向量。所以很多教程直接把这里提到的数据用 图片 来代替，Decoder 的过程也就是转换为原来的图像。

这里借用的是百度百科的图片：
在这里插入图片描述

3.3.2 分类

按学习范式分类

收缩自编码器（undercomplete autoencoder）
正则自编码器（regularized autoencoder）
变分自编码器（Variational AutoEncoder, VAE）

其中前两者是判别模型、后者是生成模型。

按构筑类型分类

前馈结构的神经网络
递归结构的神经网络

按损失函数的约束条件分类

稀疏自编码器
去噪自编码器
卷积自编码器

等等。

3.3.3 小结

自编码器具有以下几个特点：

非监督学习 (Unsupervised Learning)
是一种前馈神经网络，没有任何反馈
是一种生成模型
具有较好的特征提取能力
它的降维可以是非线性的，而 PCA 是线性的
常用于特征提取、文档检索、分类和异常检测

在 TensorFlow 和 PyTorch 中都实现了 AutoEncoder，也可以很容易查到相关实验，这里不介绍代码层次的内容。

以上内容只能非常简单地介绍什么是 AutoEncoder，想详细了解学习请自行查阅相关资料。

参考与推荐阅读
百度百科自编码器
 博客 – 深入理解自编码器
实体书籍《TensorFlow 深度学习实战》安东尼奥·古利、阿米塔·卡普尔著机械工业出版社
学习视频李宏毅视频教程
 李宏毅视频教程笔记

3.4 变分自编码 (Variational AutoEncoder)

3.4.1 概述

定义

VAE 模型是一种包含隐变量的生成模型，它利用神经网络训练得到两个函数（也称为推断网络和生成网络），进而生成输入数据中不包含的数据。

作用

VAE 模型在生成多种复杂数据方面已经显示出了巨大的潜力，包括手写数字图像、人脸图像、门牌号图像、CIFAR 图像、场景物理模型、分割图像以及从静态图像进行预测等。

摘自论文 Tutorial on Variational Autoencoders

在这里插入图片描述
一个例子

例子来源于李宏毅老师的视频

在这里插入图片描述
图片左边部分是 AutoEncoder 的简单例子：我们把一张满月的图片 Encoder 后得到 code，这个code被decoder 后又转换为满月图，弦月图也是如此。注意它们直接的一对一关系。

图片右边部分是 VAE 的简单例子，在 code 中添加一些 noise，这样可以让在满月对应 noise 范围内的code 都可以转换为满月，弦月对应的noise 范围内的code也能转换成弦月。
但当我们在code中进行采样时，在不是满月和弦月对应的noise的code中采样时，decoder出来的图片可能是介于满月和弦月之间的图。

也就是说，VAE 产生了输入数据中不包含的数据，（可以认为产生了含有某种特定信息的新的数据），而 AE 只能产生尽可能接近或者就是以前的数据（当数据简单时，编码解码损耗少时）。

图片左边那个问号的意思是当对 AE 中的code进行随机采样时，它介于满月与弦月之间的数据，decoder后可能会输出什么？
可能会输出满月，可能会输出弦月，但是最有可能输出的是奇奇怪怪的图片。

3.4.2 隐变量(Latent Variable)

上面已经讲过隐变量的基本概念，这里介绍隐变量在 VAE模型中的作用及特点。

隐变量 z z z 是可以认为是隐藏层数据，它是不限定数目的符合 高斯分布 特征的数据。（根据实际情况确定数目）
z z z 由输入数据 X X X 的采样以及参数生成，它既包含 X X X 的信息（这个于 AutoEncoder 的隐藏层类似），同时也满足 高斯分布，方便接下来进行梯度下降或者其他优化技术（By having a Gaussian distribution, we can use gradient descent (or any other optimization technique) to increase P ( X ) P(X) P(X) by making f ( z ; θ ) f(z; \theta) f(z;θ) approach X X X for some z z z ,i.e., gradually making the training data more likely under the generative model.）。
隐变量的作用除了让生成网络尽可能还原原来的数据 X X X，同时也能生成原来数据中不存在的数据。

3.5 VAE 与 AE

区别

VAE 中隐藏层服从高斯分布，AE 中的隐藏层无分布要求
训练时，AE 训练得到 Encoder 和 Decoder 模型，而 VAE 除了得到这两个模型，还获得了隐藏层的分布模型（即高斯分布的均值与方差）
AE 只能重构输入数据X，而 VAE 可以生成含有输入数据某些特征与参数的新数据。

联系

VAE 与 AE 完全不同，但是从结构上看都含有 Decoder 和 Encoder 过程。

4. 理论基础

4.1 `Step 1` 引入隐变量 z

首先当我们引入隐变量 Z Z Z 以后，可以用Z来表示 P ( x ) P(x) P(x) ：

P ( X ) = ∫ z P ( X ∣ z ) P ( z ) d z ( 1 ) P(X) = \int_z{P(X|z)P(z)} \,{\rm d}z \ \ \ \ \ \ \ \ (1) P(X)=∫zP(X∣z)P(z)dz (1)

其中，用 P ( X ∣ z ) P(X|z) P(X∣z) 替代了 f ( z ) f(z) f(z) ，这样可以用概率公式明确表示 X X X 对 z z z 的依赖性；

在这里插入图片描述
如图所示，标准的VAE模型是一个图模型，注意明显缺乏任何结构，甚至没有“编码器”路径：可以在不输入的情况下从模型中采样。这里，矩形是“盘子表示法（plate notation）（图模型中表示重复变量的方法）”，这意味着我们可以在模型参数 θ \theta θ 保持不变的情况下，从 z z z 和 X X X 中采样N次。

4.2 `Step 2` Decoder 过程

上面也曾讲过，VAE 模型中同样有 Encoder 与 Decoder 过程，也就是说模型对于相同的输入时，也应该有尽可能相同的输出。所以这里再次遇到 Maximum likelihood（极大似然）。

在公式(1) 中，将 P ( X ) P(X) P(X) 用 Z Z Z 来表示，所以对任何输入数据，应该尽量保证最后由隐变量而转换回输出数据与输入数据尽可能相等。

L = ∑ x l o g P ( X ) ( 2 ) L = \sum_x{log P(X)} \ \ \ \ \ \ \ \ \ (2) L=x∑logP(X) (2)

为了让公式2的输出极大似然 X X X，神经网络登场，它的作用就是调参，来达到极大似然的目的。（注意这里虽然介绍在前，但其实在训练的时候与后面的 NN 是同时开始训练）

本步介绍的是如何实现与 AE 类似的功能，即保证输出数据极大似然与输入数据。

这里用到的网络称为 VAE 中的生成网络，即根据隐变量 Z Z Z 生成输出数据的网络。

4.3 `Step 3` Encoder 过程

这个步骤需要确定隐变量 Z Z Z ，也就是 Encoder 过程。

这里需要用到另外一个分布， q ( z ∣ x ) q(z|x) q(z∣x) ，其中 q ( z ∣ x ) q(z|x) q(z∣x) ~ N ( μ ′ ( x ) , σ ′ ( x ) ) N(\mu'(x), \sigma'(x) ) N(μ′(x),σ′(x))，注意这里的均值和方差都是对 X X X 进行的。

同样地，这个任务也交给 NN 去完成。

这里的网络称为推断网络。

在这里插入图片描述

4.4 `Step 4` ELBO

这里对 l o g P ( X ) log\ P(X) log P(X) 进行变形转换。
log ⁡ P ( X ) = ∫ z q ( z ∣ x ) log ⁡ P ( x ) d z ( 3 ) \log P(X) = \int_z{q(z|x)\log P(x)} \,{\rm d}z \ \ \ \ \ \ \ \ (3) logP(X)=∫zq(z∣x)logP(x)dz (3)

在公式3中， q ( z ∣ x ) q(z|x) q(z∣x) 可以是任意分布。

为了让 Encode 过程也参与进来，这里引入 q ( z ∣ x ) q(z|x) q(z∣x)，推导步骤如下：

log ⁡ P ( X ) = ∫ z q ( z ∣ x ) log ⁡ P ( x ) d z = ∫ z q ( z ∣ x ) log ⁡ ( P ( z , x ) P ( z ∣ x ) ) d z = ∫ z q ( z ∣ x ) log ⁡ ( P ( z , x ) q ( z ∣ x ) q ( z ∣ x ) P ( z ∣ x ) ) d z = ∫ z q ( z ∣ x ) log ⁡ ( P ( z , x ) q ( z ∣ x ) ) d z + ∫ z q ( z ∣ x ) log ⁡ ( q ( z ∣ x ) P ( z ∣ x ) ) d z ( 4 ) \log P(X) = \int_z{q(z|x)\log P(x)} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \int_z{q(z|x)\log ({P(z,x)\over P(z|x)})} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \int_z{q(z|x)\log ({P(z,x)\over q(z|x)}{q(z|x)\over P(z|x)})} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \int_z{q(z|x)\log ({P(z,x)\over q(z|x)})} \,{\rm d}z + \int_z{q(z|x)\log ({q(z|x)\over P(z|x)})} \,{\rm d}z \ \ \ \ \ \ \ \ \ (4) logP(X)=∫zq(z∣x)logP(x)dz =∫zq(z∣x)log(P(z∣x)P(z,x))dz =∫zq(z∣x)log(q(z∣x)P(z,x)P(z∣x)q(z∣x))dz =∫zq(z∣x)log(q(z∣x)P(z,x))dz+∫zq(z∣x)log(P(z∣x)q(z∣x))dz (4)

然后，因为公式 ∫ z q ( z ∣ x ) log ⁡ ( q ( z ∣ x ) P ( z ∣ x ) ) d z \int_z{q(z|x)\log ({q(z|x)\over P(z|x)})} \,{\rm d}z ∫zq(z∣x)log(P(z∣x)q(z∣x))dz 即计算 KL 散度( K L ( q ( z ∣ x ) ∣ ∣ P ( z ∣ x ) ) KL(q(z|x)||P(z|x)) KL(q(z∣x)∣∣P(z∣x)))，所以这个式子运算值一定大于等于0。

所以公式4 一定大于等于 ∫ z q ( z ∣ x ) log ⁡ ( P ( z , x ) q ( z ∣ x ) ) d z \int_z{q(z|x)\log ({P(z,x)\over q(z|x)})} \,{\rm d}z ∫zq(z∣x)log(q(z∣x)P(z,x))dz，（下界 lower bound）

为了方便，我们把这个公式记作 L b L_b Lb，(ELBO)
记 L b = ∫ z q ( z ∣ x ) log ⁡ ( P ( z , x ) q ( z ∣ x ) ) d z ( 5 ) 记 L_b = \int_z{q(z|x)\log ({P(z,x)\over q(z|x)})} \,{\rm d}z \ \ \ \ \ \ \ \ \ \ (5) 记Lb=∫zq(z∣x)log(q(z∣x)P(z,x))dz (5)
在这里插入图片描述

4.5 `Step 5` 消除 KL 项

根据公式4和公式5，可以把 log ⁡ P ( X ) \log P(X) logP(X) 简单记为
log ⁡ P ( X ) = L b + K L ( q ( z ∣ x ) ∣ ∣ P ( z ∣ x ) ) ( 6 ) \log P(X) = L_b + KL(q(z|x)\ ||\ P(z|x)) \ \ \ \ \ \ \ \ \ \ \ \ (6) logP(X)=Lb+KL(q(z∣x) ∣∣ P(z∣x)) (6)

这里首先需要提出一个重要结论：

log P(X) 值的大小与 q(z|x) 无关。

所以不管怎么只调 q ( z ∣ x ) q(z|x) q(z∣x)，都不会让 log ⁡ P ( X ) \log P(X) logP(X) 增大。

所以可以通过调 q(z|x) 让 K L ( q ( z ∣ x ) ∣ ∣ P ( z ∣ x ) ) KL(q(z|x)\ ||\ P(z|x)) KL(q(z∣x) ∣∣ P(z∣x)) 尽可能小（调整为0），再通过调 E L B O ELBO ELBO 来实现最大化 P(X)。

调整的最终结果是使得 q ( z ∣ x ) q(z|x) q(z∣x) 尽可能接近 p ( z ∣ x ) p(z|x) p(z∣x)，换句话说，最终的 K L ( q ( z ∣ x ) ∣ ∣ P ( z ∣ x ) ) ≈ 0 KL(q(z|x)\ ||\ P(z|x)) \approx 0 KL(q(z∣x) ∣∣ P(z∣x))≈0 。所以这个步骤我们消除了这个 KL 项，剩下 ELBO 等待解决。
在这里插入图片描述

4.6 `Step 6` 解决 ELBO (Evidence Lower Bound)

上面已经明确了目标：找到 P ( x ∣ z ) P(x|z) P(x∣z) 和 q ( z ∣ x ) q(z|x) q(z∣x) 使得 L b L_b Lb 尽可能大。

公式推导如下：
L b = ∫ z q ( z ∣ x ) log ⁡ ( P ( z , x ) q ( z ∣ x ) ) d z = ∫ z q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) P ( z ) q ( z ∣ x ) ) d z = ∫ z q ( z ∣ x ) log ⁡ ( P ( z ) q ( z ∣ x ) ) d z + ∫ z q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) ) d z = − K L ( q ( z ∣ x ) ∣ ∣ P ( z ) ) + ∫ z q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) ) d z ( 7 ) L_b = \int_z{q(z|x)\log ({P(z,x)\over q(z|x)})} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ \ \ \ = \int_z{q(z|x)\log ({P(x|z)P(z)\over q(z|x)})} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = \int_z{q(z|x)\log ({P(z)\over q(z|x)})} \,{\rm d}z + \int_z{q(z|x)\log ({P(x|z)})} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ = -KL(q(z|x)\ || \ P(z)) + \int_z{q(z|x)\log ({P(x|z)})} \,{\rm d}z \ \ \ \ \ \ \ (7) Lb=∫zq(z∣x)log(q(z∣x)P(z,x))dz =∫zq(z∣x)log(q(z∣x)P(x∣z)P(z))dz =∫zq(z∣x)log(q(z∣x)P(z))dz+∫zq(z∣x)log(P(x∣z))dz =−KL(q(z∣x) ∣∣ P(z))+∫zq(z∣x)log(P(x∣z))dz (7)

公式中第2步到第3步的过程用到的是对数函数的性质，即
log ⁡ A B = log ⁡ A + log ⁡ B \log AB = \log A + \log B \\ logAB=logA+logB

接下来需要 Minimize K L ( q ( z ∣ x ) ∣ ∣ P ( z ) ) KL(q(z|x)\ || \ P(z)) KL(q(z∣x) ∣∣ P(z)) ，也就是调整 q ( z ∣ x ) q(z|x) q(z∣x) 来最小化，这样的苦差事交给 NN 吧。

并且需要最大化另外一项，即 ∫ z q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) ) d z \int_z{q(z|x)\log ({P(x|z)})} \,{\rm d}z ∫zq(z∣x)log(P(x∣z))dz，同样这份苦差事也交给 NN 去完成。
在这里插入图片描述

在这里插入图片描述

4.7 推导结果

综合上面的推导，简单概括一下可以得到：
log ⁡ P ( X ) = L b + K L ( q ( z ∣ x ) ∣ ∣ P ( z ∣ x ) ) \log P(X) = L_b + KL(q(z|x)\ ||\ P(z|x)) \ \ \ \ \ \ \ \ \ \ \ \ logP(X)=Lb+KL(q(z∣x) ∣∣ P(z∣x))
最终经过NN调参后， K L ( q ( z ∣ x ) ∣ ∣ P ( z ∣ x ) ) = 0 KL(q(z|x)\ ||\ P(z|x))=0 KL(q(z∣x) ∣∣ P(z∣x))=0， L b L_b Lb 在这里即 ELBO.

E L B O = ∫ z q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) P ( z ) q ( z ∣ x ) ) d z ( 8 ) ELBO = \int_z{q(z|x)\log ({P(x|z)P(z)\over q(z|x)})} \,{\rm d}z \ \ \ \ \ \ \ \ (8) ELBO=∫zq(z∣x)log(q(z∣x)P(x∣z)P(z))dz (8)

接着对公式8进一步展开，可以得到
E L B O = − K L ( q ( z ∣ x ) ∣ ∣ P ( z ) ) + ∫ z q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) ) d z = − K L ( q ( z ∣ x ) ∣ ∣ P ( z ) ) + E q ( z ∣ x ) log ⁡ ( P ( x ∣ z ) ) ( 9 ) ELBO = -KL(q(z|x)\ || \ P(z)) + \int_z{q(z|x)\log ({P(x|z)})} \,{\rm d}z \\ \ \ \ \ \ \ \ \ \ = -KL(q(z|x)\ || \ P(z)) + E_{q(z|x)}\log ({P(x|z)}) \ \ \ \ \ (9) ELBO=−KL(q(z∣x) ∣∣ P(z))+∫zq(z∣x)log(P(x∣z))dz =−KL(q(z∣x) ∣∣ P(z))+Eq(z∣x)log(P(x∣z)) (9)

公式9 中的 E E E 是指均值，把一个求积分转换为对 log ⁡ ( P ( x ∣ z ) ) \log({P(x|z)}) log(P(x∣z)) 在 q ( z ∣ x ) q(z|x) q(z∣x) 条件范围内的均值。

在对应的代码实现中，有两种实现方式，一种是用公式 9 ，一种直接用公式8。如果直接用公式8 的话，可以理解为： log ⁡ ( P ( x ∣ z ) ) \log({P(x|z)}) log(P(x∣z)) 在 q ( z ∣ x ) q(z|x) q(z∣x) 条件范围内 log ⁡ ( P ( x ∣ z ) P ( z ) q ( z ∣ x ) ) \log ({P(x|z)P(z)\over q(z|x)}) log(q(z∣x)P(x∣z)P(z))的均值。

5. 代码相关

VAE 模型的实现比较简单，有需要的可以在前面的 2. 推荐资料 中找到相应的实现，但为了让上面的理论部分读起来更加舒服，这里从源码中摘取一些数学理论相关计算代码进行说明。

这里的部分源码摘录于 https://tensorflow.google.cn/tutorials/generative/cvae ，对应的英文版本。

5.1 网络结构

三部分内容：Encoder 、 Deconder 与隐变量，对应的代码也非常简单，这里以对 MNIST 数据集为例，输入数据维度是 (28,28,1) ，经过Encoder 过程得到隐变量，然后再通过Decoder 过程把隐变量转换为和输入数据尽可能相似的数据。

所以如果定义 VAE 类的话，至少包括这三部分内容：

import tensorflow as tf

class VAE(tf.keras.Model):
    def __init__(self, latent_dim):
        super(VAE, self).__init__()
        self.latent_dim = latent_dim
        
        # Encoder NN
        self.encoder = tf.keras.Sequential(
          [
            tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(64, activation='relu'),
            tf.keras.layers.Dense(latent_dim + latent_dim),
          ]
        )

        # Decoder NN
        self.decoder = tf.keras.Sequential(
          [
            tf.keras.layers.InputLayer(input_shape=(latent_dim)),
            tf.keras.layers.Dense(64, activation='relu'),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(784),
            tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
          ]
        )

这里使用的测试数据是 MNIST (28, 28, 1)。
这些只是 Encoder 与 Decoder 结构，并不是 encode 与 decode 过程。
在实例化 VAE 对象的时候需要指定隐变量的维度，即 latent_dim.

5.2 μ \mu μ 和 σ \sigma σ 的生成 (即 Encode 过程)

Encoder 返回的是均值 μ \mu μ 与 log ⁡ σ 2 \log \sigma^2 logσ2 ，因此代码非常简单，就是把隐变量分成两部分即可。

    # encode  
    def encode(self, x):
        mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)
        return mean, logvar

tf.split 函数用于均匀分割。

5.3 Decode 过程

整个步骤非常简单，单纯地把数据传入NN就可以了。

    def decode(self, z):
        logits = self.decoder(z)
        return logits

5.3 重参数化 reparameterize

为了达到梯度下降，需要明确表示隐变量 z z z 与 μ \mu μ、 σ \sigma σ 之间的关系（线性关系方便梯度下降运算）：
z = μ + σ ⊙ ε z=\mu + \sigma \odot \varepsilon z=μ+σ⊙ε

可以认为 ε \varepsilon ε 是用来保持 z z z 随机性的随机噪声。

    # reparameterize
    def reparameterize(self, mean, logvar):
        eps = tf.random.normal(shape=mean.shape)
        return eps * tf.exp(logvar * .5) + mean

5.4 损失函数 compute_loss

上面的 4.7 已经总结了计算 Decoder 输出与原来数据误差的结果，一般有两种计算方法。这里采用的是公式8。

为了更加方便代码实现，需要做简单变形。

E L B O = E q ( z ∣ x ) ( log ⁡ P ( x ∣ z ) + log ⁡ ( P ( z ) − l o g ( q ( z ∣ x ) ) ) ) ( 10 ) ELBO = E_{q(z|x)}(\log P(x|z)+\log(P(z)-log(q(z|x)))) \ \ \ \ \ \ (10) ELBO=Eq(z∣x)(logP(x∣z)+log(P(z)−log(q(z∣x)))) (10)

公式10对应的代码如下：

# 参考正态分布的概率密度公式
def log_normal_pdf(sample, mean, logvar, raxis=1):
  log2pi = tf.math.log(2. * np.pi)
  return tf.reduce_sum(
      -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
      axis=raxis)
      
# 损失函数 
def compute_loss(model, x):
  # 1. 获得均值与方差
  mean, logvar = model.encode(x)
  # 2. 重参数化 z
  z = model.reparameterize(mean, logvar)
  # 3. 获得 decoder 输出（可以认为是 log x’ ）
  x_logit = model.decode(z)
  # 4. 计算交叉熵，即 q(z|x)
  cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
  # 5. 计算 log p(x|z)
  logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
  # 6. 计算 log p(z)
  logpz = log_normal_pdf(z, 0., 0.)
  # 7. 计算 log q(z|x)
  logqz_x = log_normal_pdf(z, mean, logvar)
  # 8. 计算 -(log p(x|z) + log p(z) - log q(z|x)) 的均值
  return -tf.reduce_mean(logpx_z + logpz - logqz_x)

关于 log_normal_pdf 函数可能需要解释一下，正态分布概率密度公式如下：

f ( x ) = 1 2 π σ e x p ( − ( x − μ ) 2 2 σ 2 ) f(x) = {1\over{\sqrt{2\pi }\sigma}}exp{(-{(x-\mu)^2\over{2\sigma ^2}})} f(x)=2π σ1exp(−2σ2(x−μ)2)

那么 log f(x) 等于

log ⁡ f ( x ) = log ⁡ 1 2 π σ + ( − ( x − μ ) 2 2 σ 2 ) log ⁡ f ( x ) = − log ⁡ 2 π σ − ( x − μ ) 2 2 σ 2 log ⁡ f ( x ) = − 1 2 ( log ⁡ 2 π + log ⁡ σ 2 + ( x − μ ) 2 σ 2 ) ( 11 ) \log f(x) = \log {1\over \sqrt {2\pi}\sigma} + (-{(x-\mu)^2\over 2\sigma^2})\\ \log f(x) = -\log {\sqrt {2\pi}\sigma} -{(x-\mu)^2\over 2\sigma^2} \ \ \ \ \\ \log f(x) = -{1\over 2}(\log{2\pi}+\log \sigma^2 +{(x-\mu)^2\over \sigma^2} )\ \ \ (11) logf(x)=log2π σ1+(−2σ2(x−μ)2)logf(x)=−log2π σ−2σ2(x−μ)2 logf(x)=−21(log2π+logσ2+σ2(x−μ)2) (11)

对照公式11，就可能理解源码了，其中logvar 是指 log ⁡ σ 2 \log \sigma^2 logσ2。

5.5 训练过程

训练过程非常简单粗暴，复杂的内容都在前面了。

@tf.function
def train_step(model, x, optimizer):
  """Executes one training step and returns the loss.

  This function computes the loss and gradients, and uses the latter to
  update the model's parameters.
  """
  with tf.GradientTape() as tape:
    loss = compute_loss(model, x)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))