变分信息瓶颈

互信息（Mutual Information）

两个随机变量 $X$ 和 $Y$ 的互信息定义为

$$I(X,Y)=\iint p(x,y)\log \frac{p(x,y)}{p(x)p(y)}\mathrm{d}x\mathrm{d}y$$

这是用来衡量两个随机变量相关性的一个量，显而易见若 $X$ 和 $Y$ 独立分布，则 $p(x,y)=p(x)p(y)$ ，它们的互信息是0。

假设深度神经网络的输入是 $X$ ， 输出是 $Y$ ，中间的某些层的特征表示我们把它当成是一种随机编码 $Y$ ，如此，这个网络结构可以看成是

$$X\to Z\to Y$$

我们希望一个好的特征表示 $Z$ 应该满足：

1. $Z$ 尽可能是对 $X$ 的压缩。
2. $Z$ 对于预测 $Y$ 应该具有最大的信息量。

把这两个条件写成数学公式就是:

$$\max I(Z,Y;\theta )s.t.\min I(X,Z;\theta )$$

其中 $\theta$ 是网络的参数。这个目标的约束条件不直观，引入信息约束 $I_c$ ，可以改写成：

$$\max I(Z,Y;\theta )s.t.,I(X,Z;\theta )\le I_c$$

引入拉格朗日乘子，目标转变为最大化：

$$R_{IB}=I(Z,Y;\theta )-\beta I(X,Z;\theta )$$

这个 $R_{IB}$ 就是信息瓶颈[1],它的意义就是要学习到一种编码，能够对于输出的预测具有最大的表达能力，同时对于输入信息具备最大的压缩能力。

两个互信息分别可以展开成，

$$\begin{array}{rl}I(Z,Y)& =\int p(y,z)\log \frac{p(y,z)}{p(y)p(z)}\mathrm{d}y\mathrm{d}z\\ & =\int p(y,z)\log \frac{p(y|z)}{p(y)}\mathrm{d}y\mathrm{d}z\end{array}$$

$$\begin{array}{rl}I(X,Z)& =\int p(x,z)\log \frac{p(x,z)}{p(x)p(z)}\mathrm{d}x\mathrm{d}z\\ & =\int p(x,z)\log \frac{p(z|x)}{p(z)}\mathrm{d}x\mathrm{d}z\end{array}$$

变分信息瓶颈(Variational Information Bottleneck)

基于此再假设 $X$ ， $Y$ ， $Z$ 之间的关系满足如下的马尔可夫链（即 $Z$ 不能直接由 $Y$ 决定）：

$$Y\leftrightarrow X\leftrightarrow Z$$

联合分布可以分解成：

$$p(X,Y,Z)=p(X)p(Y|X)p(Z|X)$$

这个马尔可夫链上的分布可以完全由我们的编码器 $p(Z|X)$ 和马尔可夫链本身的约束得到，利用马尔可夫链，还可以得到，

$$p(y|z)=\int p(x,y|z)\mathrm{d}x=\int p(y|x)p(x|z)\mathrm{d}x=\int \frac{p(y|x)p(z|x)p(x)}{p(z)}\mathrm{d}x$$

其中 $p(z|x)$ 是编码器, 因此 $p(y|z)$ 完全由我们的编码器和这个马尔可夫链本身所决定。

对于 $I(Z,Y)$

因为 $p(y|z)$ 无法直接计算，假设 $q(y|z)$ 是 $p(y|z)$ 的变分近似（我们用模型来计算它，就是我们的解码器模块），利用KL散度非负的特性：

$$\mathrm{KL}\left[p(y|z),q(y|z)\right]\ge 0\Rightarrow \int p(y|z)\log p(y|z)\mathrm{d}y\ge \int p(y|z)\log q(y|z)\mathrm{d}y$$

因此有，

$$\begin{array}{rl}I(Z,Y)& \ge \int p(y,z)\log \frac{q(y|z)}{p(y)}\mathrm{d}y\mathrm{d}z\\ & =\int p(y,z)\log q(y|z)\mathrm{d}y\mathrm{d}z-\int p(y)\log p(y)\mathrm{d}y\\ & =\int p(y,z)\log q(y|z)\mathrm{d}y\mathrm{d}z+H(Y)\end{array}$$

$H(Y)$ 是 标签 $y$ 的概率分布的熵，这个和我们的优化过程无关，可以被忽略掉，

$$I(Z,Y)\ge \int p(y,z)\log q(y|z)\mathrm{d}y\mathrm{d}z$$

将 $p(y,z)$ 写成 $p(y,z)=p(x)p(y|x)p(z|x)$ ，可以得到新的下界：

$$I(Z,Y)\ge \int p(x)p(y|x)p(z|x)\log q(y|z)\mathrm{d}x\mathrm{d}y\mathrm{d}z$$

对于 $I(Z,X)$

对于 $X$ 和 $Z$ 之间的互信息，

$$I(Z,X)=\int p(x,z)\log \frac{p(z|x)}{p(z)}\mathrm{d}x\mathrm{d}z=\int p(x,z)\log p(z|x)\mathrm{d}x\mathrm{d}z-\int p(z)\log p(z)\mathrm{d}z$$

计算 $Z$ 的边际分布 $p(z)=\int p(z|x)p(x)\mathrm{d}x$ 不是一件容易的事情，因此，让 $r(z)$ 作为这个边际分布的一个变分近似，利用 $\mathrm{KL}\left[p(z),r(z)\right]\ge 0$ ，可以得到：

$$\int p(z)\log p(z)\mathrm{d}z\ge \int p(z)\log r(z)\mathrm{d}z$$

因此，

$$\begin{array}{r}I(Z,X)\le \int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}\mathrm{d}x\mathrm{d}z\end{array}$$

Variational Information Bottleneck ($R_{VIB}$)

结合 $I(Z,Y)$ 的下界和 $I(Z,X)$ 的上界，可以得到，

$$R_{IB}=I(Z,Y)-\beta I(Z,X)\ge \int p(x)p(y|x)p(z|x)\log q(y|z)\mathrm{d}x\mathrm{d}y\mathrm{d}z-\int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}\mathrm{d}x\mathrm{d}z=R_{VIB}$$

在实际计算中，将 $p(x,y)$ 用经验分布 $p(x,y)=\frac{1}{N}{\sum }_{i=1}^N{\delta }_{x_n}(x){\delta }_{y_n}(y)$ 来替代，可以得到

$$R_{VIB}=\frac{1}{N}\sum _{i=1}^N\left[\int p(z|x_n)\log q(y_n|z)-\beta p(z|x_n)\log \frac{p(z|x_n)}{r(z)}\mathrm{d}z\right]$$

假设编码器是类似VAE的结构 $p(z|x)=\mathcal{N}(z|f_e^{\mu }(x),f_e^{\Sigma }(x))$ , $f_e$ 是编码器网络，可以利用重参数化技巧，得到

$$p(z|x)\mathrm{d}z=p(ϵ)\mathrm{d}ϵ$$

其中 $ϵ$ 是高斯随机变量。

假设 $p(z|x)$ 和 $r(z)$ 的选择使得我们能够计算KL散度，于是，最大化变分信息瓶颈在实际计算中可以转变为最小化如下目标：

$$J_{IB}=\frac{1}{N}\sum _{n=1}^N{\mathbb{E}}_{ϵ\sim p(ϵ)}[-\log q(y_n|f(x_n,ϵ))]+\beta \mathrm{KL}[p(z|x_n)|r(z)]$$

参考文献

1. Tishby, Naftali, Fernando C. Pereira, and William Bialek. “The information bottleneck method.” arXiv preprint physics/0004057 (2000).
2. Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin and Katja Hofmann. “Generalization in Reinforcement Learning with Selective Noise Injection and Information Bottleneck.” arXiv preprint cs.LG/1910.12911 (2019).
3. Alemi, Alexander A., et al. “Deep variational information bottleneck.” arXiv preprint arXiv:1612.00410 (2016).