变分互信息蒸馏（Variational mutual information KD）

原文标题是Variational Information Distillation for Knowledge Transfer，是CVPR2019的录用paper。

VID方法

思路比较简单，就是利用互信息（mutual information，MI）的角度，增加teacher网络与student网络中间层特征的MI，motivation是因为MI可以表示两个变量的依赖程度，MI越大，表明两者的输出越相关。
首先定义输入数据$\mathbf{x}\sim p(\mathbf{x})$,给定一个样本$\mathbf{x}$,得到关于teacher和student输出的$K$个对集合 R = { ( t ( k ) , s ( k ) ) } k = 1 K \mathcal{R}=\{(\bm{t}^{(k)},\bm{s}^{(k)})\}_{k=1}^{K} R={(t(k),s(k))}k=1K​,$K$表示选择的层数。变量对的MI被定义为$$\begin{gathered} I(\mathbf{t};\mathbf{s})=H(\mathbf{t})-H(\mathbf{t}|\mathbf{s}) \\ =-{\mathbb{E}}_{\mathbf{t}}[\log p(\mathbf{t})]+{\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log p(\mathbf{t}|\mathbf{s})] \end{gathered}$$
之后可以设计如下的loss函数来增大teacher和student之间的输出特征的互信息：
$$\mathcal{L}={\mathcal{L}}_{\mathcal{S}}-\sum _{k=1}^K{\lambda }_{k}I({\mathbf{t}}^{(k)},{\mathbf{s}}^{(k)})$$
其中${\mathcal{L}}_{\mathcal{S}}$表示task-specific的误差，${\lambda }_k$是超参数用于平衡误差。因为精确的计算MI是困难的，这里采用了变分下界(variational lower bound)的trick,采用variational的思想使用一个variational分布$q(\mathbf{t}|\mathbf{s})$去近似真实分布$p(\mathbf{t}|\mathbf{s})$。
Note that variational的思想就是针对某个分布很难求解的时候，采用另外一个分布来近似这个分布的做法，并使用变分信息最大化 (论文：The IM algorithm: A variational approach to information maximization) 的方法求解变分下界（variational low bound），这方法也被用在InfoGAN中。
$$\begin{gathered} I(\mathbf{t};\mathbf{s})=H(\mathbf{t})-H(\mathbf{t}|\mathbf{s}) \\ =H(\mathbf{t})+{\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log p(\mathbf{t}|\mathbf{s})] \\ =H(\mathbf{t})+{\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log q(\mathbf{t}|\mathbf{s})]+{\mathbb{E}}_{\mathbf{s}}[D_{KL}(p(\mathbf{t}|\mathbf{s})||q(\mathbf{t}|\mathbf{s}))] \\ \ge H(\mathbf{t})+{\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log q(\mathbf{t}|\mathbf{s})] \end{gathered}$$
${\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log p(\mathbf{t}|\mathbf{s})]={\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log q(\mathbf{t}|\mathbf{s})]+{\mathbb{E}}_{\mathbf{s}}[D_{KL}(p(\mathbf{t}|\mathbf{s})||q(\mathbf{t}|\mathbf{s}))]$这个关系是由变分信息最大化中得到的，真实分布$\log p(\mathbf{t}|\mathbf{s})$的期望等于变分分布${\mathbb{E}}_{\mathbf{t},\mathbf{s}}[\log q(\mathbf{t}|\mathbf{s})]$的期望+两分布的KL散度期望。因为KL散度的值是恒大于0的，所以得到变分下界。进一步可以得到如下的误差函数：
$$\tilde{\mathcal{L}}={\mathcal{L}}_{\mathcal{S}}-\sum _{k=1}^K{\lambda }_k{\mathbb{E}}_{{\mathbf{t}}^{(\mathbf{k})},{\mathbf{s}}^{(\mathbf{k})}}[\log q({\mathbf{t}}^{(\mathbf{k})}|{\mathbf{s}}^{(\mathbf{k})})]$$
$H(\mathbf{t})$由于和待优化的student参数无关，所以是常数。联合的训练学生网络利用target task和最大化条件似然去拟合teacher激活值。

作者采用高斯分布来实例化变分分布，这里的采用heteroscedastic的均值$\mathbf{\mu }(\cdot )$,即$\mathbf{\mu }(\cdot )$是关于student输出的函数；同时采用homoscedastic的方差$\mathbf{\sigma }$,即不是关于student输出的函数，作者尝试采用heteroscedastic的均值$\mathbf{\sigma }(\cdot )$，但是容易训练不稳定且提升不大。$\mathbf{\mu }(\cdot )$其实就是相当于在feature KD时teacher与student之间的回归器，包含卷积等操作。
$$\begin{gathered} -\log q(\mathbf{t}|\mathbf{s})=-\sum _{c=1}^C\sum _{h=1}^H\sum _{w=1}^W\log q(t_{c,h,w}|\mathbf{s}) \\ =\sum _{c=1}^C\sum _{h=1}^H\sum _{w=1}^W\log {\sigma }_c+\frac{(t_{c,h,w}-{\mu }_{c,h,w}(\mathbf{s}){)}^2}{2{\sigma }_c^2}+\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{t} \end{gathered}$$
由${\sigma }_c=\log (1+exp({\alpha }_c))$，${\alpha }_c$是一个可学习的参数。
对于logit层，$$\begin{gathered} -\log q(\mathbf{t}|\mathbf{s})=-\sum _{n=1}^N\log q(t_n|\mathbf{s}) \\ =\sum _{n=1}^N\log {\sigma }_n+\frac{(t_n-{\mu }_n(\mathbf{s}){)}^2}{2{\sigma }_n^2}+\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{t} \end{gathered}$$
这里$\mathbf{\mu }(\cdot )$是一个线性的变换矩阵。

与MSE的区别

作者认为当前基于MSE的方法是该方法在方差相同时的特例，即为：
$$-\log q(\mathbf{t}|\mathbf{s})=\sum _{n=1}^N\frac{(t_n-{\mu }_n(\mathbf{s}){)}^2}{2}+\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{t}$$
VID比MSE的好处为建模了不同维度的方差，使得更加灵活的方式来避免一些model capacity用来到一些无用的信息。MSE采用一样的方差会高度限制student，如果teacher的无用信息也同样的地位拟合，会造成过拟合问题，浪费掉了student的网络capacity。