Note -「Variational Auto-Encoder」VAE 学习与 MNIST 实战

\(1\) 数学基础
- \(1.1\) KL 散度
- \(1.2\) Evidence Lower BOund (ELBO)
\(2\) 模型结构
\(3\) MNIST 实战
\(4\) 参考资料

\[\mathscr{Lorain~y~w~la~Lora~blea.}

\newcommand{\DS}[0]{\displaystyle}

% operators alias
\newcommand{\opn}[1]{\operatorname{#1}}
\newcommand{\card}[0]{\opn{card}}
\newcommand{\lcm}[0]{\opn{lcm}}
\newcommand{\char}[0]{\opn{char}}
\newcommand{\Char}[0]{\opn{Char}}
\newcommand{\Min}[0]{\opn{Min}}
\newcommand{\rank}[0]{\opn{rank}}
\newcommand{\Hom}[0]{\opn{Hom}}
\newcommand{\End}[0]{\opn{End}}
\newcommand{\im}[0]{\opn{im}}
\newcommand{\tr}[0]{\opn{tr}}
\newcommand{\diag}[0]{\opn{diag}}
\newcommand{\coker}[0]{\opn{coker}}
\newcommand{\id}[0]{\opn{id}}
\newcommand{\sgn}[0]{\opn{sgn}}
\newcommand{\Res}[0]{\opn{Res}}
\newcommand{\Ad}[0]{\opn{Ad}}
\newcommand{\ord}[0]{\opn{ord}}
\newcommand{\Stab}[0]{\opn{Stab}}
\newcommand{\conjeq}[0]{\sim_{\u{conj}}}
\newcommand{\cent}[0]{\u{\degree C}}
\newcommand{\Sym}[0]{\opn{Sym}}
\newcommand{\wg}[0]{\wedge}
\newcommand{\Wg}[0]{\bigwedge}

% symbols alias
\newcommand{\E}[0]{\exist}
\newcommand{\A}[0]{\forall}
\newcommand{\l}[0]{\left}
\newcommand{\r}[0]{\right}
\newcommand{\ox}[0]{\otimes}
\newcommand{\lra}[0]{\leftrightarrow}
\newcommand{\llra}[0]{\longleftrightarrow}
\newcommand{\iso}[1]{\overset{\sim}{#1}}
\newcommand{\eps}[0]{\varepsilon}
\newcommand{\Ra}[0]{\Rightarrow}
\newcommand{\Eq}[0]{\Leftrightarrow}
\newcommand{\d}[0]{\mathrm{d}}
\newcommand{\e}[0]{\mathrm{e}}
\newcommand{\i}[0]{\mathrm{i}}
\newcommand{\j}[0]{\mathrm{j}}
\newcommand{\k}[0]{\mathrm{k}}
\newcommand{\Ex}[0]{\mathbb{E}}
\newcommand{\D}[0]{\mathbb{D}}
\newcommand{\oo}[0]{\infty}
\newcommand{\tto}[0]{\rightrightarrows}
\newcommand{\mmap}[0]{\hookrightarrow}
\newcommand{\emap}[0]{\twoheadrightarrow}
\newcommand{\actl}[0]{\curvearrowright}
\newcommand{\actr}[0]{\curvearrowleft}
\newcommand{\nsubg}[0]{\triangleleft}
\newcommand{\nsupg}[0]{\triangleright}
\newcommand{\lin}[0]{\lim_{n\to\oo}}
\newcommand{\linf}[0]{\liminf_{n\to\oo}}
\newcommand{\lsup}[0]{\limsup_{n\to\oo}}
\newcommand{\ser}[0]{\sum_{n=1}^\oo}
\newcommand{\serz}[0]{\sum_{n=0}^\oo}
\newcommand{\isoto}[0]{\overset\sim\to}
\newcommand{\F}[0]{\mathbb F}
\newcommand{\x}[0]{\times}
\newcommand{\M}[0]{\mathbf{M}}
\newcommand{\T}[0]{\intercal}

% symbols with parameters
\newcommand{\der}[1]{\frac{\d}{\d #1}}
\newcommand{\ul}[1]{\underline{#1}}
\newcommand{\ol}[1]{\overline{#1}}
\newcommand{\wt}[1]{\widetilde{#1}}
\newcommand{\br}[1]{\l(#1\r)}
\newcommand{\bk}[1]{\l[#1\r]}
\newcommand{\ev}[1]{\l.#1\r|}
\newcommand{\abs}[1]{\l|#1\r|}
\newcommand{\bs}[1]{\boldsymbol{#1}}
\newcommand{\env}[2]{\begin{#1}#2\end{#1}} % why not?
\newcommand{\ALI}[1]{\env{aligned}{#1}}
\newcommand{\CAS}[1]{\env{cases}{#1}}
\newcommand{\pmat}[1]{\env{pmatrix}{#1}}
\newcommand{\dary}[2]{\l|\begin{array}{#1}#2\end{array}\r|}
\newcommand{\pary}[2]{\l(\begin{array}{#1}#2\end{array}\r)}
\newcommand{\pblk}[4]{\l(\begin{array}{c|c}{#1}&{#2}\\\hline{#3}&{#4}\end{array}\r)}
\newcommand{\u}[1]{\mathrm{#1}}
\newcommand{\lix}[1]{\lim_{x\to #1}}
\newcommand{\ops}[1]{#1\cdots #1}
\newcommand{\seq}[3]{{#1}_{#2}\ops,{#1}_{#3}}
\newcommand{\dedu}[2]{\u{(#1)}\Ra\u{(#2)}}

% SPECIAL
\newcommand{\dat}[1]{\bs{\mathrm{#1}}} % font for data point / data set
\]

限于笔者水平, 本文或仅适合 AEVB 及 VAE 的基础学习. 如果希望更深入地了解 VAE, 推荐阅读参考资料 \([1]\) 及相关文献.

对于数学水平要求, 本文仅假设读者掌握朴素概率论和入门的分析学.

\(1\) 数学基础

\(1.1\) KL 散度

The Kullback–Leibler divergence (also called relative entropy and I-divergence), denoted \(D_{\u{KL}}(P\parallel Q)\), is a type of statistical distance: a measure of how much a model probability distribution \(Q\) is different from a true probability distribution \(P\).

定量地, 离散条件下的 KL 散度定义为

\[\ALI{
D_{\u{KL}}(P\parallel Q) &:= \sum_{\dat x}P(\dat x)\log\frac{P(\dat x)}{Q(\dat x)}\\
&= -\sum_{\dat x}P(\dat x)\log Q(\dat x)+\sum_{\dat x}P(\dat x)\log P(\dat x).
}
\]

从信息熵 (也即 "相对熵" 这个名字) 的角度容易理解. 我们尝试用 \(Q\) 的最优编码方式 (即事件 \(\dat x\) 使用 \(-\log Q(\dat x)\) 个 bit 的编码) 来编码 \(P\), \(D_{\u{KL}}(P\parallel Q)\) 给出的就是这种编码所用 bit 数与直接最优编码 \(P\) 本身的 bit 数 (即 \(P\) 本身的熵) 的差值, 这一差值反应了把编码从 \(Q\) 直接迁移到 \(P\) 的 "某种代价". 在这样的直观理解下, 如果二者是同分布的, 这一差值显然是 \(0\); 而对于一般的 \(P\) 和 \(Q\), 也不难看出 \(D_{\u{KL}}(P\parallel Q)\ge 0\).

\(1.2\) Evidence Lower BOund (ELBO)

这里我们着重研究形如 \(D_{\u{KL}}(Q(\dat z)\parallel P(\dat z\mid\dat x))\) 的 KL 散度, 其中 \(\dat x\) 是某一特定事件, \(P(\dat z\mid\dat x)\) 给出此时 \(\dat z\) 的条件分布. 推导:

\[\ALI{
D_{\u{KL}}(Q\parallel P) &= \sum_{\dat z}Q(\dat z)\log\frac{Q(\dat z)P(\dat x)}{P(\dat x\dat z)} \\
&= \sum_{\dat z}Q(\dat z)\br{\log\frac{Q(\dat z)}{P(\dat x\dat z)}+\log P(\dat x)}\\
&= \sum_{\dat z}Q(\dat z)(\log Q(\dat z)-\log P(\dat x\dat z))+\underbrace{\sum_{\dat z}Q(\dat z)}_{=1}\log P(\dat x) \\
&= \sum_{\dat z}Q(\dat z)(\log Q(\dat z)-\log P(\dat x\dat z))+\log P(\dat x).
}
\]

对分布 \(Q(\dat z)\), 记 \(\Ex_Qf(\dat z):=\sum_{\dat z}Q(\dat z)f(\dat z)\), 则

\[D_{\u{KL}}(Q(\dat z)\parallel P(\dat z\mid\dat x))=\Ex_Q(\log Q(\dat z)-\log P(\dat x\dat z))+\log P(\dat x).
\]

\[\ALI{
\implies \log P(\dat x) &= D_{\u{KL}}(Q(\dat z)\parallel P(\dat z\mid\dat x))-\Ex_Q(\log Q(\dat z)-\log P(\dat x\dat z))\\
&=: D_{\u{KL}}(Q(\dat z)\parallel P(\dat z\mid\dat x))+\mathcal L(Q).
}\tag 1
\]

由于 \(D_{\u{KL}}(Q\parallel P)\ge 0\), 有

\[\log P(\dat X)\ge\mathcal L(Q).\tag 2
\]

即 \(\mathcal L(Q)\) 可以作为 \(\log P(\dat x)\) 的下界估计.

\(2\) 模型结构

\(2.1\) 基本假设

设数据集 \(\dat X=\{\dat x^{(i)}\}_{i=1}^N\) 由 \(N\) 个独立同分布的数据点构成. 我们假设它由以下过程采样而来:

从某个先验分布 \(p_{\dat\theta^*}(\dat z)\) 采样 \(\dat z^{(i)}\);
从某个条件分布 \(p_{\dat\theta^*}(\dat x\mid\dat z=\dat z^{(i)})\) 采样 \(\dat x^{(i)}\).

其中 \(p_{\dat\theta^*}(\dat z)\) 和 \(p_{\dat\theta^*}(\dat x\mid \dat z)\) 来自一族参数化分布 \(p_{\dat\theta}(\dat z)\) 和 \(p_{\dat\theta}(\dat x\mid\dat z)\), 且它们的概率密度函数对 \(\dat\theta\) 和 \(\dat z\) 几乎处处可微.

现在, 数据集 \(\dat X\) 是已知的, 但我们不知道隐变量 \(\dat z^{(i)}\) 和具体的分布参数 \(\dat\theta^*\). 因此, 我们尝试引入一个识别模型 \(q_{\dat\phi}(\dat z\mid \dat x)\) 用来估计真实的后验分布 \(p_{\dat\theta}(\dat z\mid\dat x)\), 并尝试一起学习 \(\dat\phi\) 和 \(\dat\theta\).

我们将在后验分布 \(p_{\dat\theta}(\dat z\mid\dat x)\) (\(q_{\dat\phi}(\dat z\mid\dat x)\)) 上采样 \(\dat z\) 的行为视作对数据 \(\dat x\) 的编码, 在条件分布 \(p_{\dat\theta}(\dat x\mid\dat z)\) 上采样 \(\dat x\) 的行为视作对编码 \(\dat z\) 的解码, 这就是所谓的 encode 和 decode 过程.

\(2.2\) Marginal Likelyhood

为了学到最优的 \(\dat\theta^*\), 我们势必需要引入一个评估分布参数优劣的值. 模仿最大似然的手法, 我们仍然研究数据集 \(\dat X\) 被模型生成的概率. 则对某个数据点 \(\dat x\) 和待评估的参数 \(\dat\theta\), 有

\[p_{\dat\theta}(\dat x)=\int p_{\dat\theta}(\dat x\mid\dat z)p_{\dat\theta}(z)\d z.
\]

(这里忽略了超参数 \(\alpha\). 为了让式子更完整, 可以在所有概率中 condition on \(\alpha\).) 而

\[\log p_{\dat\theta}(\dat X)=\sum_{i=1}^N\log p_{\dat\theta}(\dat x^{(i)}).
\]

利用识别模型 \(q_{\dat\phi}\) 估计后验分布, 套用 \((1)\), 我们知道

\[\log p_{\dat\theta}(\dat x^{(i)})=D_{\u{KL}}(q_{\dat\phi}(\dat z\mid\dat x^{(i)})\parallel p_{\dat\theta}(\dat z\mid\dat x^{(i)}))+\mathcal L(\dat\theta,\dat\phi;\dat x^{(i)}).
\]

同时由 \((2)\),

\[\ALI{
\log_{\dat\theta}(\dat x^{(i)}) &\ge \mathcal L(\dat\theta,\dat\phi;\dat x^{(i)})\\
&= \Ex_{q_{\dat\phi}(\dat z\mid \dat x^{(i)})}(-\log q_{\dat\phi}(\dat z\mid\dat x^{(i)})+\log p_{\dat\theta}(\dat x^{(i)}\dat z))&(3)\\
&= \Ex_{q_{\dat\phi}(\dat z\mid \dat x^{(i)})}(-\log q_{\dat\phi}(\dat z\mid\dat x^{(i)})+\log p_{\dat\theta}(\dat z)+\log p_{\dat\theta}(\dat x^{(i)}\mid\dat z))\\
&= -D_{\u{KL}}(q_{\dat\phi}(\dat z\mid \dat x^{(i)})\parallel p_{\dat \theta}(\dat z))+\Ex_{q_{\dat\phi}(\dat z\mid \dat x^{(i)})}(\log p_{\dat\theta}(\dat x^{(i)}\mid \dat z)).&(4)
}
\\
\]

我们希望通过对 \(\mathcal L(\dat\theta,\dat\phi;\dat x^{(i)})\) 梯度下降来学出优秀的 \(\dat\theta\) 和 \(\dat\phi\).

\(2.3\) 重参数化 (reparameterization) 与 AEVB 算法

然而 \([1]\) 中指出, \(\mathcal L(\dat\theta,\dat\phi;\dat x^{(i)})\) 对 \(\dat\phi\) 的梯度的方差很大, 不适用于数值计算. (不过对此论断, \([2]\) 的评论区中有不同的分析, 可自行了解.) 这里, 我们采用重参数化技巧: 对 \(\dat z\sim q_{\dat\phi}(\dat z\mid\dat x)\), 假定 \(\dat z=g_{\dat\phi}(\dat\epsilon,\dat x)\) 可微, \(\dat\phi\) 是参数, \(\dat\epsilon\sim p(\dat\epsilon)\) 是噪声. 以此为条件, 根据概率密度的定义:

\[q_{\dat\phi}(\dat z\mid\dat x)\d\dat z=p(\dat\epsilon)\d\dat\epsilon.
\]

进而

\[\ALI{
\Ex_{q_{\dat\phi}(\dat z\mid\dat x^{(i)})}f(\dat z) &= \int q_{\dat\phi}(\dat z\mid \dat x^{(i)})f(\dat z)\d\dat z\\
&= \int p(\dat\epsilon)f(g_{\dat\phi}(\dat\epsilon,\dat x^{(i)}))\d\dat\epsilon\\
&\approx \frac{1}{L}\sum_{\ell=1}^L f(\underbrace{g_{\dat\phi}(\dat\epsilon^{(\ell)},\dat x^{(i)})}_{=:\dat z^{(i,\ell)}}),\quad \dat\epsilon^{(\ell)}\sim p(\dat\epsilon).
}
\]

以此估计 \((3)\), 给出

\[\ALI{
\mathcal L(\dat\theta,\dat\phi;\dat x^{(i)}) &\approx \wt{\mathcal L}^A(\dat\theta,\dat\phi;\dat x^{(i)})\\
&:= \frac{1}{L}\sum_{\ell=1}^L(-\log q_{\dat\phi}(\dat z^{(i,\ell)}\mid\dat x^{(i)})+\log p_{\dat\theta}(\dat x^{(i)}\dat z^{(i,\ell)})).
}
\]

或者, 以此估计 \((4)\), 给出

\[\ALI{
\mathcal L(\dat\theta,\dat\phi;\dat x^{(i)}) &\approx \wt{\mathcal L}^B(\dat\theta,\dat\phi;\dat x^{(i)})\\
&:= -D_{\u{KL}}(q_{\dat\phi}(\dat z\mid\dat x^{(i)})\parallel p_{\dat\theta}(\dat z))+\frac{1}{L}\sum_{\ell=1}^L\log p_{\dat\theta}(\dat x^{(i)}\mid\dat z^{(i,\ell)}).
}
\]

前一项散度据 \([1]\) 称通常可以解析地求出.

接着, 在数据集 \(\dat X\) 上采样一个大小为 \(M\) 的 minibatch 来估计给定参数的 marginal likelyhood, 有

\[\ALI{
\mathcal L(\dat\theta,\dat\phi;\dat X) &\approx \wt{\mathcal L}^M(\dat\theta,\dat\phi;\dat X)\\
&:= \frac{N}{M}\sum_{i=1}^M\wt{L}(\dat\theta,\dat\phi;\dat x^{(i)}).
}
\]

(这里的 \(M\) 和单个数据点的采样数量 \(L\) 间可以 trade-off. \([1]\) 指出当 \(M=100\) 时 \(L=1\) 的表现已经出色.)

最终, 嵌套地使用 \(\wt{\mathcal L}^M\) 和 [\(\wt{\mathcal L}^A\) 或 \(\wt{\mathcal L}^B\)] 两次估计, 我们就能对 marginal likelyhood 的下界 ELBO 进行调优了. 这朴素地推导出 Auto-Encoding VB (AEVB) 算法:

\[\begin{array}{r|l}
& \text{Minibatch version of the Auto-Encoding VB algorithm}\\ \hline
0 & M,L\gets 100,1\\
1 & p(\dat\epsilon),p_{\dat\theta}(\dat x\mid\dat z),q_{\dat\phi}(\dat z\mid \dat x),p_{\dat\theta}(\dat z) \gets \text{chosen distri. forms}\\
2 & \dat\theta,\dat\varphi \gets \text{initial parameters}\\
3 & \textbf{repeat}\\
4 & \qquad \dat X^M \gets \text{minibatch sampled from }\dat X\\
5 & \qquad \dat \epsilon \gets \text{noise sampled from }p(\dat\epsilon)\\
6 & \qquad \dat g \gets \nabla_{\dat\theta,\dat\phi}\wt{\mathcal L}^M(\dat\theta,\dat\phi;\dat X^M,\dat\epsilon)\\
7 & \qquad \dat\theta,\dat\phi \gets \text{parameters optimized by }\dat g\\
8 & \textbf{until}~\text{convergence of }(\dat\theta,\dat\phi)\\
9 & \textbf{return}~\dat\theta,\dat\phi
\end{array}
\]

\(2.4\) 实例: VAE 算法

在 AEVB 的框架下, 不平凡的工作是指定分布 \(p(\dat\epsilon),p_{\dat\theta}(\dat x\mid\dat z),q_{\dat\phi}(\dat z\mid \dat x),p_{\dat\theta}(\dat z)\) 的形式. 在 Variational Auto-Encoder (VAE) 中, 我们取

\[\ALI{
p(\dat\epsilon) &= \mathcal N(\dat\epsilon;\bs 0,\bs 1),\\
q_{\dat\phi}(\dat z\mid\dat x^{(i)}) &= \mathcal N(\dat z;\dat\mu^{(i)},(\dat\sigma^2)^{(i)}\bs 1),\\
p_{\dat\theta}(\dat z) &= \mathcal N(\dat z;\bs 0,\bs 1),\\
g_{\dat\phi}(\dat\epsilon^{(\ell)},\dat x^{(i)}) &= \dat\mu^{(i)}+\dat\sigma^{(i)}\odot\dat\epsilon^{(\ell)}.
}
\]

其中 \(\bs 1\) 是适合尺寸的单位矩阵. \((\dat\sigma^2)^{(i)}\bs 1\) 给出的是对角协方差阵, 即每个 \(z_j\sim\mathcal N(\mu^{(i)}_j,(\sigma_j^{(i)})^2)\), 互相独立. (但个人感觉这个记号本身有些奇怪.)

而对于 \(p_{\dat\theta}(\dat x\mid\dat z)\), 可以根据数据类型选择:

对于二元数据, \(p_{\dat\theta}(x_i\mid\dat z)=\mathcal B(x_i;1,y_i)\), 其中 \(\dat y\) 由模型给出;
对于实值数据, \(p_{\dat\theta}(x_i\mid\dat z)=\mathcal N(x_i;\mu'_i,\sigma_i'^2)\), 其中 \(\dat\mu'\) 和 \(\dat\sigma'\) 由模型给出.

这里给出实值数据下 VAE 一次 encode-decode 的示意. 其中 \(\dat x\in\R^5\), \(\dat z\in\R^3\), 蓝色点云表示概率密度:

接下来还需要验证 \(\wt{\mathcal L}\) 的形式. 这里采用 \(\wt{\mathcal L}^B\) 的估计, 需要计算 \(-D_{\u{KL}}(q_{\dat\phi}(\dat z\mid\dat x^{(i)})\parallel p_{\dat\theta}(\dat z))+\frac{1}{L}\sum_{\ell=1}^L\log p_{\dat\theta}(\dat x^{(i)}\mid\dat z^{(i,\ell)})\). 对于前一项, \(q_{\dat\phi}(\dat z\mid\dat x^{(i)})\) 简记作 \(q_{\dat\phi}(\dat z)\), 设向量维度为 \(J\), 根据定义 (这里就是把离散情况的求和对应地变为分布函数上的 Lebesgue 积分, 我们在上文已经假设了这些分布良好的分析性质):

\[\ALI{
-D_{\u{KL}}(q_{\dat\phi}(\dat z)\parallel p_{\dat\theta}(\dat z)) &= \int q_{\dat\phi}(\dat z)\log p_{\dat\theta}(\dat z)\d\dat z-\int q_{\dat\phi}(\dat z)\log q_{\dat\phi}(\dat z)\d\dat z\\
&= \int\mathcal N(\dat z;\dat\mu,\dat\sigma^2)\log \mathcal N(\dat z;\bs 0,\bs 1)\d\dat z-\int\mathcal N(\dat z;\dat\mu,\dat\sigma^2)\log\mathcal N(\dat z;\dat\mu,\dat\sigma^2)\d\dat z\\
&=: I_1-I_2.
}
\]

容易计算:

\[\ALI{
I_1 &= \int\br{\prod_{j=1}^J\mathcal N(z_j;\mu_j,\sigma_j^2)}\sum_{j=1}^J\log \mathcal N(z_j;0,1)\d\dat z\\
&= \sum_{i=1}^J\int\mathcal N(z_i;\mu_i,\sigma_i^2)\log\mathcal N(z_i;0,1)\cdot\prod_{j\neq i}\mathcal N(z_j;\mu_j,\sigma_j^2)\d\dat z\\
&= \sum_{i=1}^J\int\mathcal N(z_i;\mu_i,\sigma_i^2)\log\mathcal N(z_i;0,1)\d z_i\cdot\underbrace{\prod_{j\neq i}\int\mathcal N(z_j;\mu_j,\sigma_j^2)\d z_j}_{=1}\\
&= -\frac{1}{2}\sum_{i=1}^J\int\frac{1}{\sqrt{2\pi}\sigma_i}\e^{-\frac{(z_i-\mu_i)^2}{2\sigma_i^2}}\br{\log(2\pi)+z_i^2}\d z_i\\
&= -\frac{1}{2}\sum_{i=1}^J\br{\log(2\pi)+\frac{1}{\sqrt{2\pi}\sigma_i}\int_{-\oo}^{+\oo}\e^{-\frac{x^2}{2\sigma_i^2}}(x^2+2\mu_ix+\mu_i^2)\d x}\\
&= -\frac{1}{2}\sum_{i=1}^J\br{\log(2\pi)+\mu_i^2+\frac{1}{\sqrt{2\pi}\sigma_i}\int_{-\oo}^{+\oo}x^2\e^{-\frac{x^2}{2\sigma_i^2}}\d x}
}
\]

回忆 Gauss 积分 \(\DS\int_{-\oo}^{+\oo}x^2\e^{-ax^2}\d x=\frac{1}{2}\sqrt{\frac{\pi}{a^3}}\), 代入化简得

\[\ALI{
I_1 &= -\frac{J}{2}\log(2\pi)-\frac{1}{2}\sum_{i=1}^J\br{\mu_i^2+\frac{1}{\sqrt{2\pi}\sigma_i}\cdot\frac{1}{2}\sqrt{8\sigma_i^6\pi}}\\
&= -\frac{J}{2}\log(2\pi)-\frac{1}{2}\sum_{i=1}^J(\mu_i^2+\sigma_i^2).
}
\]

同理

\[I_2=-\frac{J}{2}\log(2\pi)-\frac{1}{2}\sum_{i=1}^J(1+\log\sigma_j^2).
\]

所以

\[-D_{\u{KL}}(q_{\dat\phi}(\dat z)\parallel p_{\dat\theta}(\dat z))=\frac{1}{2}\sum_{i=1}^J(1+\log\sigma_i^2-\mu_i^2-\sigma_i^2).
\]

最终

\[\ALI{
\mathcal L(\dat\theta,\dat\phi;\dat x^{(i)}) &\approx \wt{\mathcal L}^B(\dat\theta,\dat\phi;\dat x^{(i)})\\
&= \frac{1}{2}\sum_{i=1}^J(1+\log\sigma_i^2-\mu_i^2-\sigma_i^2)+\frac{1}{L}\sum_{\ell=1}^L\log p_{\dat\theta}(\dat x^{(i)}\mid \dat z^{(i,\ell)}).
}
\]

这样的良好形式已然可以启动训练了. 在这一表达式中, 前一项即 (负) KL 散度, 后一项一般称为重构损失 (reconstruction loss).

\(3\) MNIST 实战

由于 VAE 和最常见的 "将 batch 输入模型 - 比对模型输出与 ground truth 计算 loss - 反向传播" 的训练方式有些差异, 实现起来可能有些难度. 所以这里以 MNIST 为例实现完整的 VAE, 并通过一些数据实验加深对 VAE 的理解.

(注: 文末提供了本节的完整代码.)

\(3.1\) 数据准备

无需多言. (Tips: MNIST 单图的初始形态为 \((1,28,28)\); ToTensor() 后灰度值在 \([0,1]\) 中.)

import torch

import torch.nn as nn

import torch.nn.functional as F

import torchvision

import matplotlib

matplotlib.use("Agg") # 笔者使用的 WSL

import matplotlib.pyplot as plt

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

train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True)

train_dataset.transform = torchvision.transforms.ToTensor()

# 注意这里的 100 对应了训练量时 M 的值

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=100, shuffle=True)

\(3.2\) 分布选取与框架代码

实践上, 在 decode 时直接采用独立 Bernoulli 分布是一个高质且高效的选择. 这时

\[\log p_{\dat\theta}(\dat x^{(i)}\mid\dat z^{(i,\ell)})=\sum_{j}\br{x^{(i)}_j\log\mu'_j+(1-x^{(i)}_j)\log(1-\mu'_j)},
\]

其中 \(\dat \mu'=\dat \mu'(\dat z^{(i,\ell)})\) 即 decode 样本点 \(\dat z^{(i,\ell)}\) 的模型输出 (不必再如上文图中输出一个 \(\dat\sigma'\)).

Q1: 灰度值是一个实值量, 为什么不如上文所说地使用正态分布来 decode?

A1: 用正态分布的最大问题是范围不匹配. 正态分布会给出 \(\R\) 上的采样, 如果不在训练过程中强制截断, 会导致重构损失非常巨大 (实测 \(10^9\) 倍于 KL 散度) 而难以训练; 而强制截断则会导致边界概率密度的不合理分配.

Q2: 图像灰度值分布的 ground truth 总该是 \([0,1]\) 上的连续分布, 我们用离散的 Bernoulli 分布去拟合合理吗?

A2: 的确, Bernoulli 分布无法建模中间灰度, 理论上有偏差. 如果希望更精确地拟合, 可以采用独立 Beta 分布等分布模型. Bernoulli 分布的优势在于其模型简单, 训练高效且稳定.

给出框架代码:

class Encoder(nn.Module):

    def __init__(self, LATENT_DIM):

        super(Encoder, self).__init__()

        self.W_h = nn.Linear(784, 256)

        self.b_h = nn.Parameter(torch.zeros(256))

        self.W_mu = nn.Linear(256, LATENT_DIM)

        self.b_mu = nn.Parameter(torch.zeros(LATENT_DIM))

        self.W_sgm = nn.Linear(256, LATENT_DIM)

        self.b_sgm = nn.Parameter(torch.zeros(LATENT_DIM))

    def forward(self, x):

        x = x.view((-1, 784))

        h = F.relu(self.W_h(x) + self.b_h) # 也可以用 tanh 等激活

        mu = self.W_mu(h) + self.b_mu

        sgm = self.W_sgm(h) + self.b_sgm # sigma 可能 <0, 其行为和 >0 一致

        return mu, sgm

class Decoder(nn.Module):

    def __init__(self, LATENT_DIM):

        super(Decoder, self).__init__()

        self.W_h = nn.Linear(LATENT_DIM, 256)

        self.b_h = nn.Parameter(torch.zeros(256))

        self.W_mu = nn.Linear(256, 784)

        self.b_mu = nn.Parameter(torch.zeros(784))

    def forward(self, z):

        h = F.relu(self.W_h(z) + self.b_h)

        mu_re = F.sigmoid(self.W_mu(h) + self.b_mu)

        return mu_re # 使用 Bernoulli 分布, 只输出 mu'

class VAE(nn.Module):

    def __init__(self, LATENT_DIM):

        super(VAE, self).__init__()

        self.LATENT_DIM = LATENT_DIM

        self.encoder = Encoder(LATENT_DIM)

        self.decoder = Decoder(LATENT_DIM)

    def generate(self, num=1): # [用于测试] 在隐空间随机采样重构

        imgs = None

        with torch.no_grad():

            z = torch.randn((num, self.LATENT_DIM)).to(device)

            mu_re = self.decoder(z)

            imgs = mu_re.view(-1, 1, 28, 28)

        return imgs.cpu()

    def reconstruct(self, X): # [用于测试] 模拟 encode-decode (如上文图过程)

        mu, sgm = self.encoder(X)

        eps = torch.randn_like(sgm).to(device)

        z = mu + sgm * eps

        mu_re = self.decoder(z)

        return mu_re.view(-1, 1, 28, 28).cpu()

    # 没有必要实现 forward 方法

class ELBO_Estimator(nn.Module):

    def __init__(self):

        super(ELBO_Estimator, self).__init__()

        self.L = 1 # 估算积分时的采样次数

        self.FIX_EPS = 1e-8 # /0, log0 修正

    def forward(self, X_M):

        mu, sgm = model.encoder(X_M)

        kl_div = -0.5 * torch.sum(1 + torch.log(sgm**2 + self.FIX_EPS) - mu**2 - sgm**2)

        re_loss = 0

        for _ in range(self.L):

            e_l = torch.randn_like(sgm).to(device) # 批量采样 epsilon

            z_l = mu + sgm * e_l

            mu_re = model.decoder(z_l)

            re_loss += torch.sum(X_M * torch.log(mu_re + self.FIX_EPS))

            re_loss += torch.sum((1 - X_M) * torch.log(1 - mu_re + self.FIX_EPS))

        re_loss /= self.L

        elbo = -(re_loss - kl_div)  # 负的 ELBO (调优时最小化之), 忽略了常数因子

        return elbo, kl_div, re_loss # 后两项用于输出时观察

\(3.3\) 训练

无需多言.

model = VAE(2).to(device) # 这里 2 是隐空间维度, 可以自由调节

criterion = ELBO_Estimator().to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # 随手写的学习率

def train_vae(model, train_loader, optimizer, epochs=10):

    model.train()

    for epoch in range(epochs):

        total_loss = 0

        for batch_idx, (data, _) in enumerate(train_loader):

            data = data.view((-1, 784)).float().to(device)

            optimizer.zero_grad()

            loss, kl_div, re_loss = criterion(data) # 直接算 criterion, 不必 model.forward

            loss.backward()

            optimizer.step()

            total_loss += loss.item()

            if batch_idx % 100 == 0:

                print(f'batch {batch_idx + 1}/{len(train_loader)} | loss: {loss.item():.2f}',

                      f'| kl_div: {kl_div.item():.2f} | re_loss: {re_loss.item():.2f}')

        print(f'---epoch {epoch + 1}/{epochs} | loss: {total_loss / len(train_loader):.2f}---\n')

# 启动训练

train_vae(model, train_loader, optimizer, epochs=10)

torch.save(model.state_dict(), f'vae.pth')

\(3.4\) 实验

先来观察直接在整个隐空间采样 \(\dat z\) 并重构的效果.

def generate_grid(model):

    model.eval()

    with torch.no_grad():

        imgs = model.generate(16)

        grid = torchvision.utils.make_grid(imgs, nrow=4, padding=2)

        plt.figure(figsize=(10, 10))

        plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap='gray')

        plt.axis('off')

        plt.savefig('grid.png', format='png')

        plt.close()

generate_grid(model)

结果 (LATENT_DIM=10):

每个 "数字" 看上去是若干个标准数字的模糊叠加. 直接这样生成数字虽然勉强能看, 但的确不够理想.

接着再来对比 encode-decode 过程下的原数据 \(\dat x\) 和还原数据 \(\dat x'\).

def reconstruct_compare(model, valid_loader):

    model.eval()

    with torch.no_grad():

        for data, _ in valid_loader:

            data = data.view((-1, 784)).float().to(device)

            recons = model.reconstruct(data)

            data = data.view(-1, 1, 28, 28).cpu()

            # 制作 data 和 recons 的对比网格图

            grid = torchvision.utils.make_grid(torch.cat((data.view(-1, 1, 28, 28),

                                                          recons), dim=0), nrow=8, padding=2)

            plt.figure(figsize=(10, 10))

            plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap='gray')

            plt.axis('off')

            plt.savefig('compare.png', format='png')

            plt.close()

            break

valid_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True)

valid_dataset.transform = torchvision.transforms.ToTensor()

valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=16, shuffle=False)

reconstruct_compare(model, valid_loader)

结果 (LATENT_DIM=10):

效果不错. 像 \(1,0,7\) 这几个不太容易混淆的数字, 还原的数字看上去甚至更圆润美观一些. 但左起第一列的 \(5\), 倒数第二列的 \(4\) 和最后一列的 \(9\) 的还原得效果较差, 这可能是因为原数据就不太容易分辨.

最后, 我们取 LATENT_DIM=2 并观察隐空间形态. 这里我们取验证集全体进行 encode, 并描出每个点的正态中心:

def show_2d_latent_space(model, valid_loader, no_offset=False):

    model.eval()

    assert model.LATENT_DIM == 2, "Latent dimension must be 2 for visualization"

    with torch.no_grad():

        all_z = []

        all_labels = []

        for data, labels in valid_loader:

            data = data.view((-1, 784)).float().to(device)

            mu, sgm = model.encoder(data)

            if no_offset:

                z = mu

            else:

                eps = torch.randn_like(sgm).to(device)

                z = mu + sgm * eps

            all_z.append(z.cpu())

            all_labels.append(labels.cpu())

        all_z = torch.cat(all_z, dim=0)

        all_labels = torch.cat(all_labels, dim=0)

        plt.figure(figsize=(12, 12))

        scatter = plt.scatter(all_z[:, 0], all_z[:, 1], c=all_labels, cmap='tab10', alpha=0.5)

        plt.colorbar(scatter)

        plt.title('2D Latent Space')

        plt.xlabel('Latent Dimension 1')

        plt.ylabel('Latent Dimension 2')

        plt.savefig('latent-space.png', format='png')

        plt.close()

show_2d_latent_space(model, valid_loader, no_offset=True)

结果:

我们难以解释隐空间坐标轴的意义. 从散点来观察, 十个数字大致存在各自聚类的趋势. \(1,0,7\) 与其他数字的距离较远, 这和刚刚的还原效果以及我们区分数字的直观感受相合. 图上看最难区分的事 \(4\) 和 \(9\), 从形态上看可以理解, 且依照笔者在 MNIST 上测试的经验, 很多分辨 \(4\) 和 \(9\) 的任务的确是强人 (指人类) 所难, 所以也模型在此的模糊性也值得原谅.

另外, 在重复试验时, 空间一般会发生一些典范的变化: 例如上下左右翻转, 坐标轴交换等. 但散点的总体形态却总是类似.

\(4\) 参考资料

\([1]\) Kingma, Diederik P., and Max Welling. "Auto-encoding variational bayes." 20 Dec. 2013;

\([2]\) 知乎专栏: 变分自编码器 (VAEs), Gapeng, 2017-11-07 00:28;

\([3]\) 维基百科: Marginal likelihood, 21 February 2025, at 00:14 (UTC);

\([4]\) 维基百科: Kullback–Leibler divergence, 5 July 2025, at 21:27 (UTC).

附完整代码

import torch

import torch.nn as nn

import torch.nn.functional as F

import torchvision

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt

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

train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True)

train_dataset.transform = torchvision.transforms.ToTensor()

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=100, shuffle=True)

class Encoder(nn.Module):

    def __init__(self, LATENT_DIM):

        super(Encoder, self).__init__()

        self.W_h = nn.Linear(784, 256)

        self.b_h = nn.Parameter(torch.zeros(256))

        self.W_mu = nn.Linear(256, LATENT_DIM)

        self.b_mu = nn.Parameter(torch.zeros(LATENT_DIM))

        self.W_sgm = nn.Linear(256, LATENT_DIM)

        self.b_sgm = nn.Parameter(torch.zeros(LATENT_DIM))

    def forward(self, x):

        x = x.view((-1, 784))

        h = F.relu(self.W_h(x) + self.b_h)

        mu = self.W_mu(h) + self.b_mu

        sgm = self.W_sgm(h) + self.b_sgm

        return mu, sgm

class Decoder(nn.Module):

    def __init__(self, LATENT_DIM):

        super(Decoder, self).__init__()

        self.W_h = nn.Linear(LATENT_DIM, 256)

        self.b_h = nn.Parameter(torch.zeros(256))

        self.W_mu = nn.Linear(256, 784)

        self.b_mu = nn.Parameter(torch.zeros(784))

    def forward(self, z):

        h = F.relu(self.W_h(z) + self.b_h)

        mu_re = F.sigmoid(self.W_mu(h) + self.b_mu)

        return mu_re

class VAE(nn.Module):

    def __init__(self, LATENT_DIM):

        super(VAE, self).__init__()

        self.LATENT_DIM = LATENT_DIM

        self.encoder = Encoder(LATENT_DIM)

        self.decoder = Decoder(LATENT_DIM)

    def generate(self, num=1):

        imgs = None

        with torch.no_grad():

            z = torch.randn((num, self.LATENT_DIM)).to(device)

            mu_re = self.decoder(z)

            imgs = mu_re.view(-1, 1, 28, 28)

        return imgs.cpu()

    def reconstruct(self, X):

        mu, sgm = self.encoder(X)

        eps = torch.randn_like(sgm).to(device)

        z = mu + sgm * eps

        mu_re = self.decoder(z)

        return mu_re.view(-1, 1, 28, 28).cpu()

class ELBO_Estimator(nn.Module):

    def __init__(self):

        super(ELBO_Estimator, self).__init__()

        self.L = 1

        self.FIX_EPS = 1e-8

    def forward(self, X_M):

        mu, sgm = model.encoder(X_M)

        kl_div = -0.5 * torch.sum(1 + torch.log(sgm**2 + self.FIX_EPS) - mu**2 - sgm**2)

        re_loss = 0

        for _ in range(self.L):  # sampling integral ranges

            e_l = torch.randn_like(sgm).to(device)

            z_l = mu + sgm * e_l

            mu_re = model.decoder(z_l)

            re_loss += torch.sum(X_M * torch.log(mu_re + self.FIX_EPS))

            re_loss += torch.sum((1 - X_M) * torch.log(1 - mu_re + self.FIX_EPS))

        re_loss /= self.L

        elbo = -(re_loss - kl_div)  # negated ELBO, constant factors ignored

        return elbo, kl_div, re_loss

model = VAE(2).to(device)

criterion = ELBO_Estimator().to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

def train_vae(model, train_loader, optimizer, epochs=10):

    model.train()

    for epoch in range(epochs):

        total_loss = 0

        for batch_idx, (data, _) in enumerate(train_loader):

            data = data.view((-1, 784)).float().to(device)

            optimizer.zero_grad()

            loss, kl_div, re_loss = criterion(data)

            loss.backward()

            optimizer.step()

            total_loss += loss.item()

            if batch_idx % 100 == 0:

                print(f'batch {batch_idx + 1}/{len(train_loader)} | loss: {loss.item():.2f}',

                      f'| kl_div: {kl_div.item():.2f} | re_loss: {re_loss.item():.2f}')

        print(f'---epoch {epoch + 1}/{epochs} | loss: {total_loss / len(train_loader):.2f}---\n')

def generate_grid(model):

    model.eval()

    with torch.no_grad():

        imgs = model.generate(16)

        grid = torchvision.utils.make_grid(imgs, nrow=4, padding=2)

        plt.figure(figsize=(10, 10))

        plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap='gray')

        plt.axis('off')

        plt.savefig('grid.png', format='png')

        plt.close()

def reconstruct_compare(model, valid_loader):

    model.eval()

    with torch.no_grad():

        for data, _ in valid_loader:

            data = data.view((-1, 784)).float().to(device)

            recons = model.reconstruct(data)

            data = data.view(-1, 1, 28, 28).cpu()

            # 制作 data 和 recons 的对比网格图

            grid = torchvision.utils.make_grid(torch.cat((data.view(-1, 1, 28, 28),

                                                          recons), dim=0), nrow=8, padding=2)

            plt.figure(figsize=(10, 10))

            plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap='gray')

            plt.axis('off')

            plt.savefig('compare.png', format='png')

            plt.close()

            break

def show_2d_latent_space(model, valid_loader, no_offset=False):

    model.eval()

    assert model.LATENT_DIM == 2, "Latent dimension must be 2 for visualization"

    with torch.no_grad():

        all_z = []

        all_labels = []

        for data, labels in valid_loader:

            data = data.view((-1, 784)).float().to(device)

            mu, sgm = model.encoder(data)

            if no_offset:

                z = mu

            else:

                eps = torch.randn_like(sgm).to(device)

                z = mu + sgm * eps

            all_z.append(z.cpu())

            all_labels.append(labels.cpu())

        all_z = torch.cat(all_z, dim=0)

        all_labels = torch.cat(all_labels, dim=0)

        plt.figure(figsize=(12, 12))

        scatter = plt.scatter(all_z[:, 0], all_z[:, 1], c=all_labels, cmap='tab10', alpha=0.5)

        plt.colorbar(scatter)

        plt.title('2D Latent Space')

        plt.xlabel('Latent Dimension 1')

        plt.ylabel('Latent Dimension 2')

        plt.savefig('latent-space.png', format='png')

        plt.close()

# train_vae(model, train_loader, optimizer, epochs=10)

# torch.save(model.state_dict(), f'vae.pth')

valid_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True)

valid_dataset.transform = torchvision.transforms.ToTensor()

valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=16, shuffle=False)

model.load_state_dict(torch.load(f'vae.pth'))

# generate_grid(model)

# reconstruct_compare(model, valid_loader)

show_2d_latent_space(model, valid_loader, no_offset=True)