Notes about Variational Information Bottleneck(VIB)

写在前面

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

一、预备知识

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

信息单位:用 22 为底则单位是 比特(bit),可以简单理解为用多少个比特可以表示这个变量;用 ee 为底则单位是 nat

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

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

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

其中,I(X;Y)=H(Y)H(YX)=H(X)H(XY)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) = H(X) + H(Y) - H(X, Y)

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

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

相对熵(KL 散度,Kullback-Leibler Divergence):用来表示两个概率分布的差异。如果对于随机变量 XX 有两个分布 P(X)P(X)Q(X)Q(X) ,则可以用 KL 散度度量两个分布的“距离”,定义为:DKL(PQ)=xP(x)lnQ(x)P(x)=xP(x)lnP(x)Q(x)D_{KL}(P || Q) = - \sum_x P(x) \ln \frac{Q(x)}{P(x)} = \sum_x P(x) \ln \frac{P(x)}{Q(x)}

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

一些特性:

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

交叉熵(Cross Entropy):主要用于衡量两个概率分布之间的差异,定义为:H(p,q)=xp(x)logq(x)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)H(p,q) \Leftrightarrow Minimize DKLD_{KL} ,因为 H(p)H(p) 是固定的常数,即实际上这相当于是在最小化 KL 散度,不过因为 KL 散度需要额外计算信息熵,所以常用交叉熵。

二、关于信息瓶颈理论

2.1 信息瓶颈 IB

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

  • XX:输入变量;
  • YY:输出变量;
  • ZZ:我们希望从 XX 中提取出的中间表示,即把 XX 编码为 ZZ

神经网络结构如下:

XZYX \rightarrow Z \rightarrow Y

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

  • ZZ 尽可能是对 XX 进行压缩;
  • ZZ 对于预测 YY 应该有尽可能大的信息量;

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

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

其中:

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

所以,训练的目标就是要最大化 ZZYY 间的互信息,最小化 ZZXX 之间的互信息

2.2 变分信息瓶颈 VIB

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

I(Z;Y)=p(y,z)logp(y,z)p(y)p(z)dydz=p(y,z)logp(yz)p(y)dydz\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}

I(X;Z)=p(z,x)logp(z,x)p(z)p(x)dzdx=p(z,x)logp(zx)p(z)dzdx\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 建模的核心假设:

YXZY \leftrightarrow X \leftrightarrow Z

马尔可夫链假设,其直观意义是:YYZZXX 的条件下独立,ZZ 需要通过 XXYY 的联合分布来更新。其数学表示就是:

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

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

  • 模型结构中,ZZ 是通过编码器从 XX 中采样出来的
  • 我们没有直接从 YY 生成 ZZ,而是由 XX 决定 ZZ

基于该假设,联合分布可以分解为:

p(X,Y,Z)=p(X)p(YX)p(ZX,Y)(链式法则)=p(X)p(YX)p(ZX)(马尔可夫链性质)\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(X) 表示输入数据分布;
  • p(YX)p(Y|X) 由标签决定输入(数据标签生成过程);
  • p(ZX)p(Z|X) 由编码器生成的 latent 表示。

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

p(yz)=p(x,yz)dx=p(yx)p(xz)dx=p(zx)p(x,y)p(z)dx=p(yx)p(zx)p(x)p(z)dx\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(yz)p(y|z) 完全由编码器决定,即间接地通过设计编码器的 p(zx)p(z|x),来影响模型对 p(yz)p(y|z) 的学习效果。

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

后验概率 p(yz)p(y|z) 是无法直接计算的,需要用解码器(decoder) qϕ(yz)q_\phi(y|z) 变分近似真实后验 p(yz)p(y|z)

利用 KL 散度的非负性质:

DKL(p(yz)q(yz))0p(yz)logp(yz)p(yz)logq(yz)\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}

则有:

I(Z;Y)=p(y,z)logp(yz)p(y)dydzp(y,z)logq(yz)dydzp(y,z)logp(y)dydz=p(y,z)logq(yz)dydz+H(Y)\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)H(Y) 是标签 yy 的概率分布的熵,和优化过程无关,优化时可忽略,则:

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

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

p(y,z)=p(x,y,z)dx=p(x)p(yx)p(zx)dx\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)p(x)p(yx)p(zx)logq(yz)dxdydzI(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)I(X,Z)

I(X;Z)=p(z,x)logp(zx)p(z)dzdx=p(z,x)logp(zx)dzdxp(z,x)logp(z)dzdx\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) 可以变换为 p(zx)p(x)p(z|x)\cdot p(x) ,但 p(z)p(z) 的边际分布 p(z)=p(zx)p(x)dxp(z)=\int p(z|x)p(x)\,dx 并不是很好计算,这时候有需要变分近似出场了,用 r(z)r(z) 作为 p(z)p(z) 的变分近似分布(一般采用一些比较特殊的分布,比如标准正态分布),同样由 KL 散度的非负性,可得:

DKL(p(z)r(z))0p(z)logp(z)p(z)logr(z)\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}

则有:

I(X;Z)=p(z,x)logp(zx)dzdxp(z,x)logp(z)dzdxp(z,x)logp(zx)dzdxp(z,x)logr(z)dzdx=p(zx)p(x)logp(zx)r(z)dzdx\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 变分信息瓶颈下界

结合以上推导的最终结果,可以得到:

RIB=I(Z;Y)βI(Z;X)p(x)p(yx)p(zx)logq(yz)dxdydzβp(zx)p(x)logp(zx)r(z)dzdx\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}

将下界记为 L\mathcal{L}

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

设我们有 NN 个样本对:{(xn,yn)}n=1N\{(x_n, y_n)\}_{n=1}^N

  • 对于 xxp^(x)=1Nn=1Nδxn(x)\widehat{p}(x) = \frac{1}{N} \sum_{n=1}^N \delta_{x_n}(x)
  • 对于 yyp^(y)=1Nn=1Nδyn(y)\widehat{p}(y) = \frac{1}{N} \sum_{n=1}^N \delta_{y_n}(y)
  • 对于联合分布 (x,y)(x, y)p^(x,y)=1Nn=1Nδ(xn,yn)(x,y)\widehat{p}(x, y) = \frac{1}{N} \sum_{n=1}^N \delta_{(x_n, y_n)}(x, y)

其中 δxn(x)\delta_{x_n}(x)狄拉克 delta 函数,表示一个“把质量全部放在样本点 xnx_n”的分布:

δxn(x)={+x=xn0xxn,δxn(x)dx=1\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

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

L=1Nn=1N[p(zxn)logq(ynz)βp(zxn)logp(zxn)r(z)dz]\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)是一种生成模型,它学习数据 xx 的潜在表示 zz,并能够从 zz 重新生成数据

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

  • 均值:fμ(xn)f_\mu(x_n)
  • 协方差(或标准差):feΣ(xn)f_e^\Sigma(x_n)
  • fef_e 表示编码器网络

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

z=f(xn,ϵ)=feΣ(xn)ϵ+feμ(xn),ϵN(0,1)z = f(x_n, \epsilon) = f_e^\Sigma(x_n) \cdot \epsilon + f_e^\mu(x_n), \quad \epsilon \sim \mathcal{N}(0, 1)

那么,最大化信息瓶颈等价于最小化:

JIB=1Nn=1NEϵN(0,1)[logq(ynf(xn,ϵ))]+βDKL(p(zx,n)r(z))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)中,我们把输入 xnx_n 编码为一个概率分布 p(zxn)p(z|x_n)。模型有两股对抗的力量:

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

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

JIB=1Nn=1NEϵN(0,1)[logq(ynf(xn,ϵ))]预测损失+βDKL(p(zx,n)r(z))压缩信息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 项的惩罚促使不同输入 xnx_n 的编码分布 p(zxn)p(z|x_n) 靠近一个公共 prior r(z)r(z),比如标准正态分布 N(0,1)\mathcal{N}(0, 1) ,但是,如果过于一致,会导致通过 zz 流向 yy 的差异性信息减少,q(yz)q(y|z) 难以通过采样的 zz 预测 yy ,即第一项损失增大,从而促使模型忽略 xx 中的冗余信息而保留预测 yy 所需的关键信息。

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

那么信息瓶颈做了什么?

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

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

四、实验效果

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

参考资料

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

Notes about Variational Information Bottleneck(VIB)
https://blog.yokumi.cn/2025/07/01/Notes about Variational Information Bottleneck(VIB)/
作者
Yokumi
发布于
2025年7月1日
更新于
2025年7月7日
许可协议