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

技术分享 1年前 (2024-11-29) 0 999+

关注

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

变分信息瓶颈

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

$maxlimits_{theta}I(Z;Y|theta)-beta I(Z;X|theta)$

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

下界1

　　首先，$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)} dygeq 0\ implies &int , frac{p(y,z)}{p(z)} log frac{p(y|z)}{q(y|z)} dygeq 0\ implies &int , p(y,z) log p(y|z) dygeq int , p(y,z) log q(y|z)dy\ end{align*}

　　则有

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)$，本文基于马尔科夫假设：$Yleftrightarrow Xleftrightarrow 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$的推理分布。

上界2

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

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

　　对于其中的$p(z)$，作者用另一个变分估计$r(z)$来拟合。由于有

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

　　则有

begin{align*} I(Z;X)leq intint p(x)p(z|x)log frac{p(z|x)}{r(z)}dx,dz end{align*}

总体下界和优化

　　结合下界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 intint p(x)p(z|x)log frac{p(z|x)}{r(z)}dx,dz = L end{align*}

　　针对上式，用经验分布来代替真实分布。即用$frac{1}{N}sum_{n=1}^Ndelta_{x_n}(x)$代替$p(x)$，用$frac{1}{N}sum_{n=1}^Ndelta_{y_n}(y)$代替$p(y)$，用$frac{1}{N}sum_{n=1}^Ndelta_{(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 Lapprox frac{1}{N}sumlimits_{n=1}^Nint p(z|x_n)log q(y_n|z)-beta, p(z|x_n)logfrac{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) $，其中$epsilonsim mathcal{N}(0,1)$。则最大化$L$可表示为最小化：

$displaystyle J_{IB} = frac{1}{N}sumlimits_{n=1}^N mathbb{E}_{epsilonsim 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)$利用某一特定分布实现，文中使用标准正态分布实现。

直觉理解

　　直觉上理解：模型要把每个$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 $必须共享这两个样本的相似信息并忽略不同的信息，从而表示语义的相似性被强制由线性距离控制，实现表示语义的连续性，从而显式地确定了模型的泛化。