Notes about Variational Information Bottleneck(VIB)

写在前面

变分信息瓶颈（Variational Information Bottleneck，VIB），简单来说，可视为一种正则化方法，用于引导模型只保留必要信息，压缩冗余，预测准确。本文参考了经典论文 《Deep Variational Information Bottleneck》，勉强也算是论文阅读笔记吧，不过只摘录了推导部分以及结合了个人的理解，没有放具体的实验效果。

一、预备知识

信息量：用来消除随机不确定性，信息发生的概率越小，即不确定性越大，信息量越小。设某事件 $x$ 发生的概率为 $p(x)$，则它携带的信息量定义为：$I(x) = -\log p(x)$ ；

信息单位：用 $2$ 为底则单位是 比特（bit），可以简单理解为用多少个比特可以表示这个变量；用 $e$ 为底则单位是 nat 。

信息熵 Entropy：用来表示信息的平均不确定度，若一个离散随机变量 $X$ 的取值空间为 $\{x_1, x_2, \dots, x_n\}$，其概率分布为 $p(x)$，则熵定义为：$H(X) = - \sum_{i} p(x_i) \log p(x_i)$ 。

条件熵 Conditional Entropy：用来表示知道另一个变量后剩下的不确定性。设 $X$ 和 $Y$ 是两个随机变量，则 $H(Y|X) = \mathbb{E}_{x \sim p(x)} [H(Y|X = x)]$ ，表示在知道 $X$ 的条件下，$Y$ 还剩下多少不确定性。

互信息 Mutual Information：用来表示两个变量间共享的信息。它衡量的是已知随机变量 $X$ ，对另一个随机变量 $Y$ 的不确定性的减少程度，即知道 $X$ 对知道 $Y$ 有多少帮助。听着有亿点抽象，但看下图就能基本理解了：

其中，$I(X; Y) = H(Y) - H(Y|X) = H(X) - H(X|Y)$ ；

或者也可以写为 $I(X; Y) = H(X) + H(Y) - H(X, Y)$ ；

用概率代入可得：$I(X; Y) = \sum_{x,y} p(x, y) \log \frac{p(x, y)}{p(x)p(y)}$ 。

若两个变量相互独立，显然互信息为 0 。

相对熵（KL 散度，Kullback-Leibler Divergence）：用来表示两个概率分布的差异。如果对于随机变量 $X$ 有两个分布 $P(X)$ 和 $Q(X)$ ，则可以用 KL 散度度量两个分布的“距离”，定义为：$D_{KL}(P || Q) = - \sum_x P(x) \ln \frac{Q(x)}{P(x)} = \sum_x P(x) \ln \frac{P(x)}{Q(x)}$ ；

写成上面那种形式可能并不太方便理解和记忆，$D_{KL}(P || Q) = \sum_xP(x)\ln{P(X)} - P(X)\ln{Q(X)}$ ，第一部分就是 $P$ 的信息熵，第二部分是 $Q$ 相对于 $P$ 的期望值，即交叉熵，所以 KL 散度可以简单理解为，站在 $P$ 的视角，$Q$ 和 $P$ 的差异性。

一些特性：

• $D_{KL}(P || Q) \ge 0$ ，当且仅当 $P(X) = Q(X)$ 时取等；
• $D_{KL}(P||Q) \neq D_{KL}(Q||P)$ ，即不具有对称性，所以并不是一个真正的度量或距离；

交叉熵（Cross Entropy）：主要用于衡量两个概率分布之间的差异，定义为：$H(p, q) = - \sum_x p(x) \log q(x)$ ，即交叉熵 = 信息熵 + KL 散度，那么，KL 散度其实就是交叉熵和真实熵之间的“差距”；

既然交叉熵和 KL 散度都是用来衡量两个概率分布的差异性，那么他们有什么区别呢？

交叉熵 vs. KL 散度：在监督学习中，或是训练神经网络时：

• p(x)：表示真实分布（通常是 one-hot 标签）
• q(x)：是模型输出的概率分布（softmax 后）

通常会用交叉熵作为 loss function，目标就是 Minimize $H(p,q)$ $\Leftrightarrow$ Minimize $D_{KL}$ ，因为 $H(p)$ 是固定的常数，即实际上这相当于是在最小化 KL 散度，不过因为 KL 散度需要额外计算信息熵，所以常用交叉熵。

二、关于信息瓶颈理论

2.1 信息瓶颈 IB

在深度学习中，信息瓶颈理论事实上就是把神经网络理解为一个编码器和解码器，设：

• $X$：输入变量；
• $Y$：输出变量；
• $Z$：我们希望从 $X$ 中提取出的中间表示，即把 $X$ 编码为 $Z$ ；

神经网络结构如下：

$$X \rightarrow Z \rightarrow Y$$

IB 的核心目标就是，在保留与输出 $Y$ 的有关信息的同时，最大可能地压缩输入 $X$ 。或者说它认为一个好的特征 $Z$ ，应该满足：

• $Z$ 尽可能是对 $X$ 进行压缩；
• $Z$ 对于预测 $Y$ 应该有尽可能大的信息量；

可以用数学语言描述如下：

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

其中：

• $\beta$ ：拉格朗日乘子（越大越关注任务表现，越小越关注压缩）；
• $\theta$ ：神经网络的参数，也就是要优化的东西，或者说编码方式；（后文推导过程中省略了）

所以，训练的目标就是要最大化 $Z$ 和 $Y$ 间的互信息，最小化 $Z$ 和 $X$ 之间的互信息。

2.2 变分信息瓶颈 VIB

优化的目标包含两个互信息，先对两个互信息进行展开，

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

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

2.2.1 模型假设

我们先给出 VIB 建模的核心假设：

$$Y \leftrightarrow X \leftrightarrow Z$$

即马尔可夫链假设，其直观意义是：$Y$ 和 $Z$ 在 $X$ 的条件下独立，$Z$ 需要通过 $X$ 和 $Y$ 的联合分布来更新。其数学表示就是：

$$p(z|x,y)=p(z|x)$$

也可以从任务的角度理解：

• 模型结构中，$Z$ 是通过编码器从 $X$ 中采样出来的；
• 我们没有直接从 $Y$ 生成 $Z$，而是由 $X$ 决定 $Z$ ；

基于该假设，联合分布可以分解为：

$$\begin{aligned} p(X, Y, Z) &= p(X) \cdot p(Y|X) \cdot p(Z|X, Y) \quad \text{(链式法则)}\\ &= {p(X) \cdot p(Y|X) \cdot p(Z|X)} \quad \text{(马尔可夫链性质)} \end{aligned}$$

其中：

• $p(X)$ 表示输入数据分布；
• $p(Y|X)$ 由标签决定输入（数据标签生成过程）；
• $p(Z|X)$ 由编码器生成的 latent 表示。

再看后验分布 $p(y|z)$ ：

$$\begin{aligned} p(y|z) &= \int p(x, y | z) \,dx \\ &= \int p(y|x) \cdot p(x|z) \,dx \\ &= \int \frac{p(z|x) p(x, y)}{p(z)} \,dx \\ &= \int p(y|x) \cdot \frac{p(z|x)p(x)}{p(z)} \,dx \end{aligned}$$

可以发现，$p(y|z)$ 完全由编码器决定，即间接地通过设计编码器的 $p(z|x)$，来影响模型对 $p(y|z)$ 的学习效果。

2.2.2 互信息 $I(Z;Y)$

后验概率 $p(y|z)$ 是无法直接计算的，需要用解码器（decoder） $q_\phi(y|z)$ 变分近似真实后验 $p(y|z)$ 。

利用 KL 散度的非负性质：

$$\begin{aligned} & D_{KL}(p(y|z)||q(y|z))\ge 0 \\ \Rightarrow & \int p(y|z)\log{p(y|z)} \ge \int p(y|z)\log{q(y|z)}\\ \end{aligned}$$

则有：

$$\begin{aligned} I(Z;Y) & =\int p(y,z)\log{\frac{p(y|z)}{p(y)}}\,dy\,dz \\ & \ge \int p(y,z)\log{q(y|z)}\,dy\,dz - \int p(y,z)\log{p(y)}\,dy\,dz\\ & = \int p(y,z)\log{q(y|z)}\,dy\,dz + H(Y) \end{aligned}$$

其中 $H(Y)$ 是标签 $y$ 的概率分布的熵，和优化过程无关，优化时可忽略，则：

$$I(Z;Y)\ge \int p(y,z)\log{q(y|z)}\,dy\,dz$$

对于 $p(y,z)$ 可以写为：

$$\begin{aligned} p(y,z)&=\int p(x,y,z)\, dx\\ &=\int p(x) \cdot p(y|x) \cdot p(z|x)\,dx \\ \end{aligned}$$

所以，下界可以写为：

$$I(Z;Y)\ge \int p(x)\cdot p(y|x)\cdot p(z|x)\log{q(y|z)}\,dx\,dy\,dz$$

2.2.3 互信息 $I(X,Z)$

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

$p(z,x)$ 可以变换为 $p(z|x)\cdot p(x)$ ，但 $p(z)$ 的边际分布 $p(z)=\int p(z|x)p(x)\,dx$ 并不是很好计算，这时候有需要变分近似出场了，用 $r(z)$ 作为 $p(z)$ 的变分近似分布（一般采用一些比较特殊的分布，比如标准正态分布），同样由 KL 散度的非负性，可得：

$$\begin{aligned} & D_{KL}(p(z)||r(z))\ge 0 \\ \Rightarrow & \int p(z)\log{p(z)} \ge \int p(z)\log{r(z)}\\ \end{aligned}$$

则有：

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

2.2.4 变分信息瓶颈下界

结合以上推导的最终结果，可以得到：

$$\begin{aligned} R_{IB}&=I(Z;Y)-\beta I(Z;X)\\ &\ge \int p(x)\cdot p(y|x)\cdot p(z|x)\log{q(y|z)}\,dx\,dy\,dz - \beta\int p(z|x)\cdot p(x)\cdot \log{\frac{p(z|x)}{r(z)}}\,dz\,dx \\ \end{aligned}$$

将下界记为 $\mathcal{L}$ 。

我们用经验分布来代替真实分布，即我们在不知道真实分布 $p(x)$, $p(y)$, $p(x,y)$ 的情况下，使用有限个观测样本来近似这些分布：

设我们有 $N$ 个样本对：$\{(x_n, y_n)\}_{n=1}^N$ ；

• 对于 $x$ ，$\widehat{p}(x) = \frac{1}{N} \sum_{n=1}^N \delta_{x_n}(x)$ ；
• 对于 $y$ ，$\widehat{p}(y) = \frac{1}{N} \sum_{n=1}^N \delta_{y_n}(y)$ ；
• 对于联合分布 $(x, y)$ ，$\widehat{p}(x, y) = \frac{1}{N} \sum_{n=1}^N \delta_{(x_n, y_n)}(x, y)$ ；

其中 $\delta_{x_n}(x)$ 是 狄拉克 delta 函数，表示一个“把质量全部放在样本点 $x_n$”的分布：

$$\delta_{x_n}(x) = \begin{cases} +\infty & x = x_n \\ 0 & x \ne x_n \end{cases} ,\quad \int \delta_{x_n}(x) dx = 1$$

如果不太好理解，你也可以当成直接把概率积分改成离散样本的求和取平均，总之有：

$$\mathcal{L}=\frac{1}{N}\sum_{n=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)}}}\,dz\right]$$

如果编码器（Encoder）是类似 VAE 的结构：

Variational Autoencoder（VAE）是一种生成模型，它学习数据 $x$ 的潜在表示 $z$，并能够从 $z$ 重新生成数据。

神经网络是“可导”的函数，但是采样操作是不可导的，VAE 使用了重参数技巧（Reparameterization Trick），首先，给定样本 $x_n$ ，我们用一个神经网络（encoder）输出其对应的条件正态分布 $p(z|x_n) = \mathcal{N}(z|f_e^\mu(x_n), f_e^\Sigma(x_n))$

• 均值：$f_\mu(x_n)$
• 协方差（或标准差）：$f_e^\Sigma(x_n)$
• $f_e$ 表示编码器网络

然后将随机变量的“采样”过程拆成“可导的函数 + 固定随机数”：

$$z = f(x_n, \epsilon) = f_e^\Sigma(x_n) \cdot \epsilon + f_e^\mu(x_n), \quad \epsilon \sim \mathcal{N}(0, 1)$$

那么，最大化信息瓶颈等价于最小化：

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

三、个人的一些理解

在变分信息瓶颈（VIB）中，我们把输入 $x_n$ 编码为一个概率分布 $p(z|x_n)$。模型有两股对抗的力量：

• 一方面想让所有 $p(z|x_n)$ 彼此接近，以压缩信息（KL 惩罚项）；
• 另一方面又希望保留足够信息使得从 $z \sim p(z|x_n)$ 可以预测 $y_n$（负 log likelihood 项）。

这与我们上面推导出的优化目标是一致的，即：

$$J_{IB}=\frac{1}{N}\sum_{n=1}^N \underbrace{\mathbb{E}_{\epsilon \sim \mathcal{N}(0, 1)}\left[-\log{q(y_n|f(x_n,\epsilon))}\right ]}_{\text{预测损失}} + \underbrace{\beta D_{KL}(p(z|x,n)||r(z))}_{\text{压缩信息}}$$

其中，KL 项的惩罚促使不同输入 $x_n$ 的编码分布 $p(z|x_n)$ 靠近一个公共 prior $r(z)$，比如标准正态分布 $\mathcal{N}(0, 1)$ ，但是，如果过于一致，会导致通过 $z$ 流向 $y$ 的差异性信息减少，$q(y|z)$ 难以通过采样的 $z$ 预测 $y$ ，即第一项损失增大，从而促使模型忽略 $x$ 中的冗余信息而保留预测 $y$ 所需的关键信息。

与一般的判别模型做比较可能会更容易理解，即把 $z$ 作为对于 $x$ 的唯一确定的中间表示，而不是隐变量，这会导致导致的最坏结果就是过拟合，每个 $(x_n,y_n)$ 都有唯一确定的中间表示 $z_n$ 与之构成一一映射的关系，毫无泛化能力可言！

那么信息瓶颈做了什么？

• $x$ 只决定一个分布 $p(z|x)$ ，不是一个固定点；
• 所有 $z \sim p(z|x)$ 都要共享一个 decoder $q(y|z)$ ，这强制了：
  • 表示空间必须连续；
  • 相近的 $z$ 必须有相似的语义；
  • 泛化能力得到了显式编码；

举个例子，如果两个样本 $x_i, x_j$ 的编码 $p(z|x_i), p(z|x_j)$ 有重叠，就说明抽样可能落在相同的 $z$ 上；这要求解码器 $q(y|z)$ 必须给出一个同时对多个样本都合理的输出；模型只能从 $x$ 中提取“多个样本共有的信息”来表示 $y$，冗余的、孤立的特征会被舍弃，这就是上面所说的预测 $y$ 的关键信息。通过强迫语义共享，显式地确定了模型的泛化。

四、实验效果

略了。等读到具体论文可以补充到这。

参考资料

1. Alemi, Alexander A., et al. “Deep variational information bottleneck.” arXiv preprint arXiv:1612.00410 (2016). ↩