信息瓶颈与在图中的应用

Information Bottleneck又称信息瓶颈，是一个基于信息论的算法。

Information Bottleneck

最早来自2000年Naftali Tishby等人的的《The information bottleneck method》

具体来说，它试图优化以下两个目标：

压缩性 (Compression)：压缩表示T应该包含尽可能少的关于原始变量X的信息。这通过最小化互信息 I(X;T) 来实现。
相关性 (Relevance)：压缩表示T应该尽可能多地保留关于相关变量Y的信息。这通过最大化互信息 I(T;Y) 来实现。

故信息瓶颈可以表示为一个带权衡参数 β 的优化问题，最大化：
$$
L[p(t∣x)]=I(T;Y)-βI(X;T)
$$

Variational Information Bottleneck

难点在于互信息很难计算。所以，我们需要使用变分推理来近似它们。

在VIB中，通常会构建一个编码器 (Encoder) $p_\theta(z|x)$和一个解码器/预测器 (Decoder/Predictor)$q_\theta(y|z)$，其中：

编码器$p_\theta(z|x)$：由神经网络参数化（参数为 θ），将输入X映射到一个潜在的随机表示Z（即瓶颈T）。这个潜在表示通常被设计为遵循一个先验分布，比如标准正态分布 r(z)=N(0,I)。
解码器/预测器$q_\theta(y|z)$：也由神经网络参数化（参数为 $\phi$），根据潜在表示Z来预测目标变量Y。

即最大化:

$$
\mathcal{L} _ {IB}[p(z|x)] = I(Z;Y) - \beta I(X;Z)
$$

其中：

$I(Z;Y)$ 是Z和Y之间的互信息，越大越好。
$I(X;Z)$ 是X和Z之间的互信息，越小越好（通过$\beta$来权衡）。
$\beta > 0$ 是拉格朗日乘子，控制压缩程度和保留相关信息之间的平衡。

但是直接计算和优化这两个互信息项非常困难，主要原因如下：

$I(Z;Y)$ 的计算：
$$
I(Z;Y) = \mathbb{E} _ {p(z,y)} \left[ \log \frac{p(z,y)}{p(z)p(y)} \right] = H(Y) - H(Y|Z)
$$
其中 $H(Y)$ 是Y的熵（通常是常数，不影响优化），而 $H(Y|Z) = - \mathbb{E} _ {p(z)} \left[ \mathbb{E} _ {p(y|z)} [\log p(y|z)] \right]$。

要计算 $p(y|z)$，我们需要知道从Z到Y的真实条件概率，这在复杂模型中是未知的。而且，即使知道了 $p(y|z)$，计算期望也可能涉及难以处理的积分。边缘分布 $p(z) = \int p(z|x)p(x)dx$ 的计算也可能非常复杂。

$I(X;Z)$ 的计算：
$$
I(X;Z) = \mathbb{E} _ {p(x,z)} \left[ \log \frac{p(x,z)}{p(x)p(z)} \right] = H(Z) - H(Z|X)
$$

这里同样涉及到难以计算的边缘分布 $p(z)$ 和条件熵 $H(Z|X) = - \mathbb{E} _ {p(x)} \left[ \mathbb{E} _ {p(z|x)} [\log p(z|x)] \right]$。对于连续高维的Z，或者当 $p(z|x)$ 是一个复杂的神经网络时，$p(z)$ 的计算和期望的计算都非常棘手。

变分推断的引入

变分推断的核心思想是，当直接处理一个复杂的概率分布或量（如互信息）很困难时，我们引入一个更简单、更易于处理的变分分布（通常带有可学习的参数），并优化这个变分分布来逼近我们感兴趣的真实分布或量。

在VIB中，我们主要对 $I(Z;Y)$ 和 $I(X;Z)$ 进行变分近似。

1. 近似 $I(Z;Y)$ (保留相关信息项)

我们希望最大化 $I(Z;Y) = H(Y) - H(Y|Z)$。由于 $H(Y)$ 相对于编码器 $p(z|x)$ 是常数，所以最大化 $I(Z;Y)$ 等价于最小化条件熵 $H(Y|Z)$。

$$
H(Y|Z) = - \mathbb{E} _ {p(z,y)}[\log p(y|z)] = - \mathbb{E} _ {p(x)p(y|x)p(z|x)}[\log p(y|z)]
$$

这里的 $p(y|z)$ 是真实的从Z到Y的后验概率，我们通常不知道它。因此，我们引入一个变分近似分布 $q _ {\phi}(y|z)$，这是一个由参数 $\phi$（例如神经网络的权重）控制的、试图从Z预测Y的模型。

我们可以推导出 $I(Z;Y)$ 的一个变分下界 (Variational Lower Bound)：

$$
I(Z;Y) = H(Y) - H(Y|Z)
$$

$$
H(Y|Z) = - \mathbb{E} _ {p(z)} [ \sum_y p(y|z) \log p(y|z) ]
$$
我们知道 $KL(p(y|z) || q _ {\phi}(y|z)) \ge 0$，即 $\mathbb{E} _ {p(y|z)}[\log \frac{p(y|z)}{q _ {\phi}(y|z)}] \ge 0$
所以
$$
\mathbb{E} _ {p(y|z)}[\log p(y|z)] \ge \mathbb{E} _ {p(y|z)}[\log q _ {\phi}(y|z)]
$$

因此，
$$
-H(Y|Z) \ge \mathbb{E} _ {p(z)} [ \mathbb{E} _ {p(y|z)}[\log q _ {\phi}(y|z)] ]
$$

或者说，
$$
H(Y|Z) \leq - \mathbb{E} _ {p(z,y)}[\log q _ {\phi}(y|z)]
$$

所以，
$$
I(Z;Y) \ge H(Y) + \mathbb{E} _ {p(z,y)}[\log q _ {\phi}(y|z)]
$$
由于 $H(Y)$ 是常数，最大化 $I(Z;Y)$ 的下界就等价于最大化 $\mathbb{E} _ {p(z,y)}[\log q _ {\phi}(y|z)]$，或者说最小化负对数似然 $-\mathbb{E} _ {p(z,y)}[\log q _ {\phi}(y|z)]$。这个期望可以写作：
$$
\mathbb{E} _ {p(x)p(y|x)} [ \mathbb{E} _ {p _ {\theta}(z|x)} [\log q _ {\phi}(y|z)] ]
$$

这一项实际上就是训练一个从Z预测Y的模型（解码器/预测器 $q _ {\phi}(y|z)$）时的标准监督学习损失（例如，对于分类问题是交叉熵，对于回归问题是均方误差的负数）。我们希望通过学习编码器 $p _ {\theta}(z|x)$（参数为 $\theta$）和解码器 $q _ {\phi}(y|z)$ 来最大化这个下界。

2. 近似 $I(X;Z)$ (压缩信息项)

我们希望最小化 $I(X;Z)$。
$$
\begin{align}
I(X;Z) = H(Z) - H(Z|X)
\\
H(Z|X) = - \mathbb{E} _ {p(x)} \left[ \int p _ {\theta}(z|x) \log p _ {\theta}(z|x) dz \right]
\\
H(Z) = - \int p(z) \log p(z) dz
\end{align}
$$
其中 $p(z) = \int p _ {\theta}(z|x) p(x) dx$。

直接计算 $p(z)$ 非常困难。VIB在这里引入另一个固定的、简单的先验分布 $r(z)$ (例如，标准正态分布 $\mathcal{N}(0,I)$)。然后，目标是将 $p(z)$ 推向 $r(z)$。
我们可以利用KL散度的性质：

$$
\begin{align}
I(X;Z)&=\mathbb{E} _ {p_(x,z)} [ \log \frac{p _ {\theta}(z|x)p(x)}{p(x)p(z)}]
\\&= \mathbb{E} _ {p(x)} [ \mathbb{E} _ {p _ {\theta}(z|x)} [ \log \frac{p _ {\theta}(z|x)}{p(z)} ] ]
\\&= \mathbb{E} _ {p(x)} [KL(p _ {\theta}(z|x) || p(z))]
\\&=\mathbb{E} _ {p(x)}[\int p _ {\theta}(z|x) \log \frac{p _ {\theta}(z|x)}{p(z)} dz]
\\&=\mathbb{E} _ {p(x)}[ \int p _ {\theta}(z|x) \log \left( \frac{p _ {\theta}(z|x)}{r(z)} \cdot \frac{r(z)}{p(z)} \right) dz]
\\&=\mathbb{E} _ {p(x)}[ \int p _ {\theta}(z|x) \left( \log \frac{p _ {\theta}(z|x)}{r(z)} + \log \frac{r(z)}{p(z)} \right) dz]
\\&=\mathbb{E} _ {p(x)}[\int p _ {\theta}(z|x) \log \frac{p _ {\theta}(z|x)}{r(z)} dz+\int p _ {\theta}(z|x) \log \frac{r(z)}{p(z)} dz]
\\&=\mathbb{E} _ {p(x)}[KL(p{\theta}(z|x) || r(z))+\mathbb{E}{p _ {\theta}(z|x)} \left[\log \frac{r(z)}{p(z)}\right]]
\end{align}
$$

而

$$
\begin{align}
\mathbb{E} _ {p(x)}\left[\mathbb{E}{p _ {\theta}(z|x)} \left[\log \frac{r(z)}{p(z)}\right]\right]&=\mathbb{E}{p(x,z)} \left[\log \frac{r(z)}{p(z)}\right]
\\&=\int_z \int_x p(x)p_\theta(z|x)log\frac{r(z)}{p(z)}dxdz
\\&=\int_z \left(\int_x p(x)p_\theta(z|x)dx\right)log\frac{r(z)}{p(z)}dz
\\&=\int_z p(z)log\frac{r(z)}{p(z)}dz
\\&=-KL(p(z)||r(z))
\end{align}
$$

故利用非负性即有：

$$
I(X;Z) \le \mathbb{E} _ {p(x)}[KL(p _ {\theta}(z|x) || r(z))]
$$

一个更直接且在VIB论文中被采用的思路是，我们希望 $Z$ 尽可能独立于 $X$，同时又保留预测 $Y$ 所需的信息。如果 $Z$ 对 $X$ 的依赖性很小，那么 $p _ {\theta}(z|x)$ 对于不同的 $x$ 应该看起来都差不多，并且都接近某个简单的先验 $r(z)$。因此，我们直接将 $I(X;Z)$ 替换为一个可以计算的上界或一个相关的量，即 $\mathbb{E} _ {p(x)}[KL(p _ {\theta}(z|x) || r(z))]$。

当 $r(z)$ 是一个固定的先验分布时，最小化 $\mathbb{E} _ {p(x)}[KL(p _ {\theta}(z|x) || r(z))]$ 会迫使编码器 $p _ {\theta}(z|x)$ 的输出（对于所有x）都接近这个先验 $r(z)$。这就限制了Z中能包含的关于X的特定信息量，从而达到了压缩的目的。

组合得到最终的VIB目标函数

将上述两部分的近似代入原始IB目标 $\max [I(Z;Y) - \beta I(X;Z)]$，我们得到VIB的目标（通常是最小化其负值）：

最小化
$$
\mathcal{L} _ {VIB}(\theta, \phi) = \mathbb{E} _ {p(x)p(y|x)} \left[ - \mathbb{E} _ {p _ {\theta}(z|x)}[\log q _ {\phi}(y|z)] + \beta KL(p _ {\theta}(z|x) || r(z)) \right]
$$

这里：

第一项 $-\mathbb{E} _ {p _ {\theta}(z|x)}[\log q _ {\phi}(y|z)]$ 是预测损失（或重构损失，取决于任务）。我们希望它小，对应于最大化 $I(Z;Y)$ 的下界。
第二项 $KL(p _ {\theta}(z|x) || r(z))$ 是压缩损失或正则化项。我们希望它小，对应于最小化（或约束）$I(X;Z)$。
$p _ {\theta}(z|x)$ 是编码器网络，将输入X映射到潜在表示Z的分布。
$q _ {\phi}(y|z)$ 是解码器/预测器网络，从Z预测Y。
$r(z)$ 是一个预先选择的简单先验分布（例如，$\mathcal{N}(0,I)$）。
$\beta$ 是权衡这两个损失项的超参数。

Graph Information Bottleneck

我们希望一个好的图表示是鲁棒的、稳定的，并且拥有迁移的能力。具体来说，我们希望即使数据出现了分布漂移或者受到人为的扰动后，好的图表示仍然可拥有之前所述的性质。我们给出的核心想法是：一个好的图表示需要捕捉最少充分信息量minimal sufficient information。

于是我们使用信息瓶颈，类似地：
$$
\min _ {\mathbb{P}\left(Z _ {X}^{(L)} \mid \mathcal{D}\right) \in \Omega} \operatorname{GIB} _ {\beta}\left(\mathcal{D}, Y ; Z _ {X}^{(L)}\right) \triangleq\left[-I\left(Y ; Z _ {X}^{(L)}\right)+\beta I\left(\mathcal{D} ; Z _ {X}^{(L)}\right)\right]
$$
同样地，我们要推出上下界：
$$
I\left(Y ; Z _ {X}^{(L)}\right) \geq 1+\mathbb{E}\left[\log \frac{\prod _ {v \in V} \mathbb{Q} _ {1}\left(Y _ {v} \mid Z _ {X, v}^{(L)}\right)}{\mathbb{Q} _ {2}(Y)}\right]+\mathbb{E} _ {\mathbb{P}(Y) \mathbb{P}\left(Z _ {X}^{(L)}\right)}\left[\frac{\prod _ {v \in V} \mathbb{Q} _ {1}\left(Y _ {v} \mid Z _ {X, v}^{(L)}\right)}{\mathbb{Q} _ {2}(Y)}\right]
$$

$$
I\left(\mathcal{D} ; Z _ {X}^{(L)}\right) \leq I\left(\mathcal{D} ;\left{Z _ {X}^{(l)}\right} _ {l \in S _ {X}} \cup\left{Z _ {A}^{(l)}\right} _ {l \in S _ {A}}\right) \leq \sum _ {l \in S _ {A}} \mathrm{AIB}^{(l)}+\sum _ {l \in S _ {X}} \mathrm{XIB}^{(l)},
$$

其中
$$
\mathrm{AIB}^{(l)}=\mathbb{E}\left[\log \frac{\mathbb{P}\left(Z _ {A}^{(l)} \mid A, Z _ {X}^{(l-1)}\right)}{\mathbb{Q}\left(Z _ {A}^{(l)}\right)}\right], \mathrm{XIB}^{(l)}=\mathbb{E}\left[\log \frac{\mathbb{P}\left(Z _ {X}^{(l)} \mid Z _ {X}^{(l-1)}, Z _ {A}^{(l)}\right)}{\mathbb{Q}\left(Z _ {X}^{(l)}\right)}\right]
$$

作者提出了GIB-cat和GIB-Bern分别对应不同的边分布。

GIB-cat假设：
$$
Z_{A, v}=\cup_{t=1}^{\mathcal{T}}\left{u \in V_{v t} \mid u \stackrel{\text { iid }}{\sim} \operatorname{Cat}\left(\frac{1}{\left|V_{v t}\right|}\right)\right}
$$
GIB-Bern假设：
$$
Z_{A, v}=\cup_{t=1}^{\mathcal{T}}\left{u \in V_{v t} \mid u \stackrel{\text { iid }}{\sim}Bernoulli(\alpha)\right}
$$
AIB的经验估计为：
$$
\widehat{\mathrm{AIB}}^{(l)}=\mathbb{E}{\mathbb{P}\left(Z{A}^{(l)} \mid A, Z_{X}^{(l-1)}\right)}\left[\log \frac{\mathbb{P}\left(Z_{A}^{(l)} \mid A, Z_{X}^{(l-1)}\right)}{\mathbb{Q}\left(Z_{A}^{(l)}\right)}\right]
$$
对于XIB，假设$Z_{X, v} \sim \sum_{i=1}^{m} w_{i} \operatorname{Gaussian}\left(\mu_{0, i}, \sigma_{0, i}^{2}\right)$，则
$$
\widehat{\mathrm{XIB}}^{(l)}=\log \frac{\mathbb{P}\left(Z_{X}^{(l)} \mid Z_{X}^{(l-1)}, Z_{A}^{(l)}\right)}{\mathbb{Q}\left(Z_{X}^{(l)}\right)}=\sum_{v \in V}\left[\log \Phi\left(Z_{X, v}^{(l)} ; \mu_{v}, \sigma_{v}^{2}\right)-\log \left(\sum_{i=1}^{m} w_{i} \Phi\left(Z_{X, v}^{(l)} ; \mu_{0, i}, \sigma_{0, i}^{2}\right)\right)\right]
$$