变分自编码器（四）：一步到位的聚类方案

原文链接 | 2018-09-17 | 苏剑林

由于VAE中既有编码器又有解码器（生成器），同时隐变量分布又被近似编码为标准正态分布，因此VAE既是一个生成模型，又是一个特征提取器。在图像领域中，由于VAE生成的图片偏模糊，因此大家通常更关心VAE作为图像特征提取器的作用。提取特征都是为了下一步的任务准备的，而下一步的任务可能有很多，比如分类、聚类等。本文来关心“聚类”这个任务。

一般来说，用AE或者VAE做聚类都是分步来进行的，即先训练一个普通的VAE，然后得到原始数据的隐变量，接着对隐变量做一个K-Means或GMM之类的。但是这样的思路的整体感显然不够，而且聚类方法的选择也让我们纠结。本文介绍基于VAE的一个“一步到位”的聚类思路，它同时允许我们完成无监督地完成聚类和条件生成。

理论#

一般框架#

回顾VAE的loss（如果没印象请参考《变分自编码器（二）：从贝叶斯观点出发》）：

\[KL\Big(p(x,z)\Big\Vert q(x,z)\Big) = \iint p(z|x)\tilde{p}(x)\ln \frac{p(z|x)\tilde{p}(x)}{q(x|z)q(z)} dzdx\tag{1}\]

通常来说，我们会假设\(q(z)\)是标准正态分布，\(p(z|x),q(x|z)\)是条件正态分布，然后代入计算，就得到了普通的VAE的loss。

然而，也没有谁规定隐变量一定是连续变量吧？这里我们就将隐变量定为\((z, y)\)，其中\(z\)是一个连续变量，代表编码向量；\(y\)是离散的变量，代表类别。直接把\((1)\)中的\(z\)替换为\((z,y)\)，就得到

\[KL\Big(p(x,z,y)\Big\Vert q(x,z,y)\Big) = \sum_y \iint p(z,y|x)\tilde{p}(x)\ln \frac{p(z,y|x)\tilde{p}(x)}{q(x|z,y)q(z,y)} dzdx\tag{2}\]

这就是用来做聚类的VAE的loss了。

分步假设#

啥？就完事了？呃，是的，如果只考虑一般化的框架，\((2)\)确实就完事了。

不过落实到实践中，\((2)\)可以有很多不同的实践方案，这里介绍比较简单的一种。首先我们要明确，在\((2)\)中，我们只知道\(\tilde{p}(x)\)（通过一批数据给出的经验分布），其他都是没有明确下来的。于是为了求解\((2)\)，我们需要设定一些形式。一种选取方案为

\[p(z,y|x)=p(y|z)p(z|x),\quad q(x|z,y)=q(x|z),\quad q(z,y)=q(z|y)q(y)\tag{3}\]

代入\((2)\)得到

\[KL\Big(p(x,z,y)\Big\Vert q(x,z,y)\Big) = \sum_y \iint p(y|z)p(z|x)\tilde{p}(x)\ln \frac{p(y|z)p(z|x)\tilde{p}(x)}{q(x|z)q(z|y)q(y)} dzdx\tag{4}\]

其实\((4)\)式还是相当直观的，它分布描述了编码和生成过程：

1、从原始数据中采样到\(x\)，然后通过\(p(z|x)\)可以得到编码特征\(z\)，然后通过分类器\(p(y|z)\)对编码特征进行分类，从而得到类别；2、从分布\(q(y)\)中选取一个类别\(y\)，然后从分布\(q(z|y)\)中选取一个随机隐变量\(z\)，然后通过生成器\(q(x|z)\)解码为原始样本。

具体模型#

\((4)\)式其实已经很具体了，我们只需要沿用以往VAE的做法：\(p(z|x)\)一般假设为均值为\(\mu(x)\)方差为\(\sigma^2(x)\)的正态分布，\(q(x|z)\)一般假设为均值为\(G(z)\)方差为常数的正态分布（等价于用MSE作为loss），\(q(z|y)\)可以假设为均值为\(\mu_y\)方差为1的正态分布，至于剩下的\(q(y),p(y|z)\)，\(q(y)\)可以假设为均匀分布（它就是个常数），也就是希望每个类大致均衡，而\(p(y|z)\)是对隐变量的分类器，随便用个softmax的网络就可以拟合了。

最后，可以形象地将\((4)\)改写为

\[\mathbb{E}_{x\sim\tilde{p}(x)}\Big[-\log q(x|z) + \sum_y p(y|z) \log \frac{p(z|x)}{q(z|y)} + KL\big(p(y|z)\big\Vert q(y)\big)\Big],\quad z\sim p(z|x) \tag{5}\]

其中\(z\sim p(z|x)\)是重参数操作，而方括号中的三项loss，各有各的含义：

1、\(-\log q(x|z)\)希望重构误差越小越好，也就是\(z\)尽量保留完整的信息；2、\(\sum_y p(y|z) \log \frac{p(z|x)}{q(z|y)}\)希望\(z\)能尽量对齐某个类别的“专属”的正态分布，就是这一步起到聚类的作用；3、\(KL\big(p(y|z)\big\Vert q(y)\big)\)希望每个类的分布尽量均衡，不会发生两个几乎重合的情况（坍缩为一个类）。当然，有时候可能不需要这个先验要求，那就可以去掉这一项。

实验#

实验代码自然是Keras完成的了（^_），在mnist和fashion-mnist上做了实验，表现都还可以。实验环境：Keras 2.2 + tensorflow 1.8 + Python 2.7。

代码实现#

代码位于：https://github.com/bojone/vae/blob/master/vae_keras_cluster.py

其实注释应该比较清楚了，而且相比普通的VAE改动不大。可能稍微有难度的是\(\sum_y p(y|z) \log \frac{p(z|x)}{q(z|y)}\)这个怎么实现。首先我们代入

\[\begin{aligned}p(z|x)&=\frac{1}{\prod\limits_{i=1}^d\sqrt{2\pi\sigma_i^2(x)}}\exp\left\{-\frac{1}{2}\left\Vert\frac{z - \mu(x)}{\sigma(x)}\right\Vert^2\right\}\\ q(z|y)&=\frac{1}{(2\pi)^{d/2}}\exp\left\{-\frac{1}{2}\left\Vert z - \mu_y\right\Vert^2\right\}\end{aligned}\tag{6}\]

得到

\[\log \frac{p(z|x)}{q(z|y)}=-\frac{1}{2}\sum_{i=1}^d \log \sigma_i^2(x)-\frac{1}{2}\left\Vert\frac{z - \mu(x)}{\sigma(x)}\right\Vert^2 + \frac{1}{2}\left\Vert z - \mu_y\right\Vert^2 \tag{7}\]

注意其实第二项是多余的，因为重参数操作告诉我们\(z = \varepsilon\otimes \sigma(x) + \mu(x),\,\varepsilon\sim \mathcal{N}(0,1)\)，所以第二项实际上只是\(-\Vert \varepsilon\Vert^2/2\)，跟参数无关，所以\[\log \frac{p(z|x)}{q(z|y)}\sim -\frac{1}{2}\sum_{i=1}^d \log \sigma_i^2(x) + \frac{1}{2}\left\Vert z - \mu_y\right\Vert^2 \tag{8}\]

然后因为\(y\)是离散的，所以事实上\(\sum_y p(y|z) \log \frac{p(z|x)}{q(z|y)}\)就是一个矩阵乘法（相乘然后对某个公共变量求和，就是矩阵乘法的一般形式），用K.batch_dot实现。

其他的话，读者应该清楚普通的VAE的实现过程，然后才看本文的内容和代码，不然估计是一脸懵的。

mnist#

这里是mnist的实验结果图示，包括类内样本图示和按类采样图示。最后还简单估算了一下，以每一类对应的数目最多的那个真实标签为类标签的话，最终的test准确率大约有83%，对比这篇文章《Unsupervised Deep Embedding for Clustering Analysis》的结果（最高也是84%左右），感觉应该很不错了。