Deep Variational Information Bottleneck (VIB) 变分信息瓶颈 论文阅读笔记。本文利用变分推断将信息瓶颈框架适应到深度学习模型中,可视为一种正则化方法。
变分信息瓶颈
假设数据输入输出对为(X,Y),假设判别模型f_\\theta(\\cdot)有关于X的中间表示Z,本文旨在优化\\theta以最小化互信息I(Z;X) ,同时最大化互信息I(Z;Y),即:
\\max\\limits_{\\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)} dy\geq 0\\ \implies &\int \, \frac{p(y,z)}{p(z)} \log \frac{p(y|z)}{q(y|z)} dy\geq 0\\ \implies &\int \, p(y,z) \log p(y|z) dy\geq \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),本文基于马尔科夫假设:Y\\leftrightarrow X\\leftrightarrow 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)dz\geq\int p(z)\log r(z)dz\\ \implies&\int\int p(x,z)\log p(z)dx\,dz\geq\int \int p(x,z)\log r(z)dx\,dz \end{align*}
则有
\begin{align*} I(Z;X)\leq \int\int 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 \int\int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}dx\,dz = L \end{align*}
针对上式,用经验分布来代替真实分布。即用\\frac{1}{N}\\sum_{n=1}\^N\\delta_{x_n}(x)代替p(x),用\\frac{1}{N}\\sum_{n=1}\^N\\delta_{y_n}(y)代替p(y),用\\frac{1}{N}\\sum_{n=1}\^N\\delta_{(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 L\\approx \\frac{1}{N}\\sum\\limits_{n=1}\^N\\int p(z\|x_n)\\log q(y_n\|z)-\\beta\\, p(z\|x_n)\\log\\frac{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) ,其中\epsilon\sim \mathcal{N}(0,1)。则最大化L$可表示为最小化:
\\displaystyle J_{IB} = \\frac{1}{N}\\sum\\limits_{n=1}\^N \\mathbb{E}_{\\epsilon\\sim \\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 $必须共享这两个样本的相似信息并忽略不同的信息,从而表示语义的相似性被强制由线性距离控制,实现表示语义的连续性,从而显式地确定了模型的泛化。
实验
表1:信息瓶颈加成的模型和各种正则化后模型的对比。
图1:不同\\beta、z维度K下VIB模型在MNIST上的错误率,以及两个互信息的平衡。
图2:z维度K=2时,1000张图片的z分布的可视化。
后续是一些对抗鲁棒的实验,不记录