深度变分信息瓶颈——Deep Variational Information Bottleneck

2024-11-28 17:223020编辑收藏

　　Deep Variational Information Bottleneck (VIB) 变分信息瓶颈论文阅读笔记。本文利用变分推断将信息瓶颈框架适应到深度学习模型中，可视为一种正则化方法。

1 变分信息瓶颈#

　　假设数据输入输出对为 $(X,Y)$ ，假设判别模型 $f_\theta(\cdot)$ 有关于 $X$ 的中间表示 $Z$ ，本文旨在优化 $\theta$ 以最小化互信息 $I(Z;X)$ ，同时最大化互信息 $I(Z;Y)$ ，即：

$\max\limits_{\theta}I(Z;Y|\theta)-\beta I(Z;X|\theta)$

　　其中 $\beta>0$ 为平衡系数。直觉理解，上式期望 $Z$ 能保留更少 $X$ 信息的同时能较好用于预测 $Y$ 。那么如何构造相应的深度学习模型以及相应的优化方案？下面推导上式的下界，使其下界变大，上式即可变大。为了简化，下面去掉 $\theta$ 进行推导。

1.1 上界1#

　　 $I(Z;X)$ 展开为：

$\displaystyle I(Z;X)=\int \int p(x,z)\log \frac{p(z|x)}{p(z)}dx\,dz$

　　其中 $p(z|x)$ 为是原始模型关于 $x$ 对中间表示 $z$ 的推理分布。对于其中的 $p(z)$ ，作者用另一个变分估计 $r(z)$ 来拟合。由于有

$\begin{align*} &\text{KL}(p(Z),r(Z))\geq 0\\ \implies&\int p(z)\log p(z)dz\geq\int p(z)\log r(z)dz\\ \implies&\int\int p(x,z)\log p(z)dx\,dz\geq\int \int p(x,z)\log r(z)dx\,dz \end{align*}$

　　则有

$\begin{align*} I(Z;X) &= \int \int p(x,z)\log p(z|x) - p(x,z) \log p(z)dx\,dz \\ &\leq \int \int p(x,z)\log p(z|x) - p(x,z) \log r(z)dx\,dz\\ &=\int\int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}dx\,dz \end{align*}$

1.2 下界2#

　　 $I(Z;Y)$ 展开为：

$\displaystyle I(Z; Y) = \int \int p(y, z) \log \frac{p(y|z)}{p(y)} \, dy \, dz$

　　其中 $p(y)$ 是数据的标签分布，已知。未知而需要进行处理的是其中的 $p(y,z)$ 和 $p(y|z)$ ，也就是模型需要拟合的分布。对于 $p(y|z)$ ，可以用一个解码器 $q(y|z)$ 来拟合，即文中所谓的变分估计。利用KL散度的大于零性质，有以下不等式：

$\begin{align*} &\text{KL}(p(Y|Z),q(Y|Z))\geq 0\\ \implies &\int \, p(y|z) \log \frac{p(y|z)}{q(y|z)} dy\geq 0\\ \implies &\int \, \frac{p(y,z)}{p(z)} \log \frac{p(y|z)}{q(y|z)} dy\geq 0\\ \implies &\int \, p(y,z) \log p(y|z) dy\geq \int \, p(y,z) \log q(y|z)dy\\ \end{align*}$

　　注意最后一步去掉 $p(z)$ 是由于它没有在积分中，是常数。则有

$\begin{align*} \displaystyle I(Z; Y) &= \int \int p(y, z) \log p(y|z) - p(y, z) \log p(y) \, dy \, dz\\ &\geq \int \int p(y, z) \log q(y|z) - p(y, z) \log p(y) \, dy \, dz\\ &= \int \int \, p(y, z) \log q(y|z) dy \, dz - \int \, p(y) \log p(y) dy \\ &= \int \int \, p(y, z) \log q(y|z)dy \, dz + H(Y) \end{align*}$

　　对于其中的 $p(y,z)$ ，本文基于马尔科夫假设： $Y\leftrightarrow X\leftrightarrow Z$ 。这个假设表明， $Y$ 和 $Z$ 在 $X$ 的条件下独立（那在优化时呢？ $Z$ 是关于 $X$ 和 $Y$ 的联合分布进行更新的）。有：

$\displaystyle p(y,z)=\int p(x,y,z)dx=\int p(x,y)p(z|x,y)dx= \int p(x,y)p(z|x)dx$

　　此外，由于 $H(Y)$ 已知且固定，可忽略。则有

$\displaystyle I(Z; Y) \geq \int \int \int \, p(x,y)p(z|x) \log q(y|z)dx \, dy \, dz$

　　其中， $p(x,y)$ 是真实数据分布， $p(z|x)$ 是原始模型关于 $x$ 对中间表示 $z$ 的推理分布。

1.3 总体下界和优化#

　　结合下界1和上界2，有：

$\begin{align*} &I(Z; Y) - \beta I(Z; X) \\ \geq&\int \int \int \, p(x,y)p(z|x) \log q(y|z)dx \, dy \, dz - \beta \int\int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}dx\,dz = L \end{align*}$

　　针对上式，用经验分布来代替真实分布。即用 $\frac{1}{N}\sum_{n=1}^N\delta_{x_n}(x)$ 代替 $p(x)$ ，用 $\frac{1}{N}\sum_{n=1}^N\delta_{y_n}(y)$ 代替 $p(y)$ ，用 $\frac{1}{N}\sum_{n=1}^N\delta_{(x_n,y_n)}(x,y)$ 代替 $p(x,y)$ 。其中 $\delta_{x_n}(x)$ 表示狄拉克函数，其空间内积分为1，且仅在 $x_n$ 上非零。假设经验分布有 $N$ 各样本 $\{(x_n,y_n)\}_{n=1}^N$ 。实际上直接把概率积分改成离散样本的求和取平均即可。则上式可被估计为：

$\displaystyle L\approx \frac{1}{N}\sum\limits_{n=1}^N\int p(z|x_n)\log q(y_n|z)-\beta\, p(z|x_n)\log\frac{p(z|x_n)}{r(z)}\, dz$

　　文中将 $z$ 视为隐变量，利用VAE的重参数技巧将 $p(z|x_n)$ 实现为一个关于 $x_n$ 的正态分布 $\mathcal{N}(f_e^\mu(x_n),f_e^{\Sigma}(x_n))$ ，其中 $f_e^\mu(x_n),f_e^\Sigma(x_n)$ 分别为基于 $x_n$ 生成的均值和协方差矩阵。将 $z$ 抽样表示为 $f(x_n,\epsilon)= f_e^{\Sigma}(x_n))\epsilon + f_e^\mu(x_n)$ ，其中 $\epsilon\sim \mathcal{N}(0,1)$ 。则最大化 $L$ 可表示为最小化：

$\displaystyle J_{IB} = \frac{1}{N}\sum\limits_{n=1}^N \mathbb{E}_{\epsilon\sim \mathcal{N}(0,1)}\left[-\log q(y_n|f(x_n,\epsilon))\right] + \beta\,\text{KL}\left[p(Z|x_n);r(z)\right]$

　　其中 $r(z)$ 利用某一特定分布实现，文中使用标准正态分布实现。

2 直觉理解#

　　直觉上理解：模型要把每个 $x_n$ 分别映射到特定的分布，这些分布既不能偏离标准正态分布太远，又需要让模型后续能根据这些分布的抽样来预测 $x_n$ 的标签。那么这种做法为什么能从 $x_n$ 中抽取对预测 $y_n$ 有效的关键信息而忽略无关信息呢（即信息瓶颈）？我的理解是，模型被惩罚以使不同 $x_n$ 得到的 $z_n$ 分布靠近同一分布，但为了有效预测 $y_n$ ，又必须产生一定的不一致。不同 $x_n$ 对应的 $z$ 分布越一致，通过 $z$ 而流向 $y$ 的差异性信息将越少，导致 $q$ 更难利用采样的 $z$ 预测 $y$ ，从而促使模型忽略 $x$ 中的冗余信息而保留预测 $y$ 所需的关键信息。 $\beta$ 则用于控制 $z$ 保留 $x$ 信息的程度，越大保留信息越少。

　　相较于一般的判别模型：当不把 $z$ 视为隐变量，而变成关于 $x$ 唯一确定的中间表示时，就是一般的判别模型。这种方式隐式地假定了表示的连续性，然而无法确保所有 $z$ 都不是被离散地分散在表示空间中。最坏的过拟合情况下，每个 $(x_n,y_n)$ 都孤立地确定了一个中间表示 $z_n$ 来实现一一映射，导致无泛化。而对于使用了信息瓶颈 $z$ 的判别模型，由于 $x$ 仅仅确定 $z$ 的生成分布，不同的 $x_i,x_j$ 依然可能抽样出同一个 $z$ ，这种模式强制这个抽样出的 $z$ 必须共享这两个样本的相似信息并忽略不同的信息，从而表示语义的相似性被强制由线性距离控制，实现表示语义的连续性，从而显式地确定了模型的泛化。