变分Bayes Dropout 模型压缩

bayesian 神经网络

NIPS 17 bayesian compression for deep learning

---

概率分布基础

【log-uniform distribution】
假设$z\in[c,d]$都是正数, $y=\log z \in [\log c,\log d]$是uniform分布，即$p_{y}=\frac{1}{\log d-\log c}$，
则$F_y(m)={\rm Prob}(y\leq m\leq d )=\frac{m-\log c}{\log d-\log c}$
则$F_z(n)={\rm Prob}(z\leq n)={\rm Prob}(y\leq \log n)=\frac{\log n-\log c}{\log d-\log c}$
所以z的pdf为 $p_z(z)=\frac{d}{dz}F_z(z)=\frac{1}{z(\log d-\log c)} \propto \frac{1}{z}$.

【条件分布 换元】
若$q_{\tilde w}=N({\tilde w};\mu_{\tilde w},\sigma_{\tilde w}^2)$，$w= {\tilde w z,}$，则$q_{w|z}=N(w;\mu_{\tilde w}z,\sigma_{\tilde w}^2 z^2 )$
若$p_{\tilde w}=N({\tilde w};0,1)$，$w= {\tilde w}z$，则$p_{w|z}=N(w;0, z^2 )$
若$p_{w|z}=N(w;0, z^2 )$，$w= {\tilde w} z$，则$p_{\tilde w}=N({\tilde w};0,1)$

Proof:
$F_{\tilde w}(m)={\rm Prob}(\tilde w \leq m)={\rm Prob}( w \leq m z)=F_{w|z}(mz)\\ =\int^{mz}\frac{1}{\sqrt{2\pi z^2}}\exp\{-\frac{w^2}{2 z^2}\} dw\\ =\frac{1}{\sqrt{2\pi z^2}}\frac{\sqrt{\pi}}{\sqrt{\frac{4}{2z^2}}}erf(\sqrt{\frac{1}{2z^2}} mz)- \frac{1}{\sqrt{2\pi z^2}}\frac{\sqrt{\pi}}{\sqrt{\frac{4}{2z^2}}}erf( -\infty)\\ =\frac{1}{2} erf(\frac{m}{\sqrt{2}})+\frac{1}{2}$
$p_{\tilde w}(\tilde w)=\frac{d}{d \tilde w}F_{\tilde w}(\tilde w)= \frac{d}{d \tilde w} \frac{1}{2} erf(\frac{\tilde w }{\sqrt{2}}) =\frac{1}{\sqrt{\pi}}\exp( -{\frac{ \tilde w^2}{2}} )\frac{1}{\sqrt{2}} \sim N(\tilde w;0,1)$

【两个高斯分布的KL散度】
${\rm \mathbb D_{KL}}\left[ N(\mu_1,\sigma_1^2)|N(\mu_2,\sigma_2^2)\right] = 。。。$

${\rm \mathbb D_{KL}}\left[ N(\mu_wz,\sigma_w^2z)|N(0,z^2)\right] = {\rm \mathbb D_{KL}}\left[ N(\mu_w,\sigma_w^2)|N(0,1)\right]$

【log-uniform分布与高斯分布的KL散度】
${\rm \mathbb D_{KL}}\left[N(\mu_z,\sigma_z^2)| \frac{1}{|z|} \right] \approx k_1 \sigma\left(k_2+k_3\log\alpha\right)-0.5\ln(1+\alpha^{-1})-k_1$
其中$\alpha=\frac{\sigma_z^2}{\mu_z^2+10^{-8}}，k_1=,k_2=,k_3=$。

ReLU激活函数的平滑近似SoftPlus函数$f(x)=\frac{1}{\beta}\ln(1+e^{\beta x})$.

---

模型定义

全连接网络，第$l$层
$a^{(l)}=o^{(l-1)}w^{(l)}+b^{(l)}\\ o^{(l)}=f(a^{(l)})\\ $
输入 $x=o^{(0)}$
输出 $p_i=\frac{\exp a_i^{(L)}}{\sum_j \exp a_j^{(L)}}$
正确类标签 $y_c=1$，其它为0。
似然 $p(y|a^{(L)})=p_c=\frac{\exp a_i^{(L)}}{\sum_j \exp a_c^{(L)}}$
交叉熵 ${\rm CrossEntropy}(y,p)=-\sum_j y_j \log p_j=-\log p_c=-\log\frac{\exp a_c^{(L)}}{\sum_j \exp a_j^{(L)}} =-\log p(y|a^{(L)})$

概率描述：

先验分布
$p_{w|z}(w)=N(w;0,z^2)$
$p_{z}(z)=\frac{1}{|z|(\log d-\log c)}, \ \ \ \ \ \ \ \ |z|\in [c,d]$，$|z|$服从log-uniform分布:
${\tilde w}=\frac{w}{z}$， $w={\tilde w}z$
$p_{\tilde w}(\tilde w)=N(\tilde w;0,1)$。

$p(w,z)=p(z)p(w|z) =\frac{1}{|z|(\log d-\log c)}\frac{1}{\sqrt{2\pi z^2}}\exp\{-\frac{w^2}{2 z^2}\}\\ p(\tilde w,z)=p(z)p(\tilde w)= \frac{1}{|z|(\log d-\log c)}\frac{1}{\sqrt{2\pi }}\exp\{-\frac{\tilde w^2}{2 }\}\\ $
相差一个$\frac{1}{\sqrt{z^2}}$?!

后验分布
$q_z(z)=N(z;\mu_z,\sigma_z^2) \\ q_{\tilde w}(\tilde w)=N(\tilde w; \mu_{\tilde w},\sigma_{\tilde w}^2)\\ q_{w|z}(w)=N(w;\mu_{\tilde w}z, \sigma_{\tilde w}^2z^2) \\ $

模型训练

目标函数Evidence Lower Bound

$$L =\mathbb E_{q(w,z)}\left[\log p(D|w,z)\right]-{\rm \mathbb D_{KL}}\left[q(w,z)|p(w,z)\right]$$

联合分布的KL划分成两个KL：
${\rm \mathbb D_{KL}}\left[q(w,z)|p(w,z)\right] = \rm E_{q(z)}\left[{\rm \mathbb D_{KL}}\left[q(w|z)|p(w|z)\right]\right] +{\rm \mathbb D_{KL}}\left[q(z)|p(z)\right]$

高斯分布与log-uniform分布的KL散度用近似技巧：
${\rm \mathbb D_{KL}}\left[q(z)|p(z)\right]={\rm \mathbb D_{KL}}\left[N(\mu_z,\sigma_z^2)| \frac{1}{|z|} \right] \approx k_1 \sigma\left(k_2+k_3\log\alpha\right)-0.5\ln(1+\alpha^{-1})-k_1$其中$\alpha=\frac{\sigma_z^2}{\mu_z^2+10^{-8}}$.

高斯分布与高斯分布的KL散度可以利用计算公式：
${\rm \mathbb D_{KL}}\left[q(w|z)|p(w|z)\right]={\rm \mathbb D_{KL}}\left[ N(\mu_{\tilde w}z,\sigma_{\tilde w}^2z)|N(0,z^2)\right]\\ =\frac{1}{2}\left( \frac{\sigma^2_w z^2}{z^2} +\frac{(-\mu_w z)^2}{z^2} -1+\ln\frac{z^2}{\sigma_w^2 z^2} \right) = \frac{1}{2}\left( \sigma_{\tilde w}^2+\mu_{\tilde w}^2-1-\ln\sigma_{\tilde w}^2 \right)$
与z无关，关于q(z)的求期望自动消去。

计算$\mathbb E_{q(w,z)}\left[\log p(D|w,z)\right]$时采用mini-batch蒙特卡洛估计，从后验分布q(w,z)采样w,z。
而且采用reparameterization技巧，将采样表示成一个随机噪声与后验分布的参数的可积分函数的形式，如$N(w;\mu,\sigma^2)$的采样的函数形式为$w=\mu+\sigma \epsilon$,其中随机噪声$\epsilon\sim N(0,1)$。（为啥随机噪声可以直接从高斯分布采样，后验分布不能直接从高斯分布采样呀？...)
z变量从后验分布取$q(z)=N(\mu_z,\sigma_z^2)$,
为mini-batch里的每一个输入$x^i$,随机噪声$\epsilon^i$,得到采样$z^i=\mu_z+\sigma_z \epsilon^i$。

蒙特卡洛估计$\mathbb E_{q(w,z)}\left[\log p(D|w,z)\right]$，要求variance小。local reparameterization技巧，使得不同样本间的协方差为0，将variance限制在单个样本局部。不直接为每个输入样本从后验分布采样参数$w$, 而是计算各个样本的logit的均值、方差，直接采样logit计算响应。
第$l$层，
$a =o w +b \\ o =f(a )$
先采样$z^i=\mu_z+\sigma_z \epsilon^i$,
再根据 $q_{w|z}(w)=N(w;\mu_{\tilde w}z,\sigma_{\tilde w}^2 z^2)$
计算响应的后验分布$q_{a|o,w,z}=N(\mu_a,\sigma_o^2)$的均值、方差
$\mu_a=o^i \mu_{\tilde w} z^i+\mu_b \\ \sigma_a^2=(o^i)^2 \sigma_{\tilde w}^2 (z^i)^2+\sigma_b^2$
再从响应的后验分布$q_{a|o,w,z}=N(\mu_a,\sigma_o^2)$采样响应
$a^i=\mu_a+\sigma_a \hat \epsilon^i \\ o^i={\rm ReLU}(a^i)$

最后一层不加相应函数logit直接作为输出o计算交叉熵
${\rm CrossEntropy}( c,o)=-\log \frac{ e^{a_c} }{ \sum_j e^{a_j} }$
即为-log-likelihood关于后验分布的期望的蒙特卡洛估计 $-\mathbb E_{q(w,z)}\left[\log p(D|w,z)\right]=-\log p(y|a^{(L)})$
$\log p(y|a^{(L)})=p_c=\frac{\exp a_i^{(L)}}{\sum_j \exp a_c^{(L)}}$。

训练结束后最后的测试中，各层确定，不再采样z，不再采样activation，使用后验均值计算。
$\mu_w=\mathbb E_{q(w,z)}\left[w \right]=\mathbb E_{q(z)q(\tilde w)}\left[\tilde w z\right]=\mu_{\tilde w}\mu_z$
$a=o\mu_w+\mu_b$。

$\alpha=\frac{\sigma^2_z}{\mu_z^2}$和阈值 确定掩码（移除）哪些$\mu_w$。

计算压缩率时候，使用后验方差
$\sigma_w^2={\rm Var}({\tilde w}z)={\rm Var}({\tilde w}){\mathbb E}^2[z]+{\rm Var}(z){\mathbb E}^2[{\tilde w}]+{\rm Var}({\tilde w}){\rm Var}(z)\\ =\mu_z^2\sigma^2_{\tilde w}+\sigma^2_z\mu_{\tilde w}^2+\sigma^2_z\sigma^2_{\tilde w}$
根据$\sigma_w^2$的数值大小 确定$\mu_w$需要的精度（有效数字个数）。

$q(w)=\int q(z)q(w|z){\rm d}z=\int N(z;\mu_z,\sigma_z^2)N(w;\mu_{\tilde w}z,\sigma_{\tilde w}^2z^2)dz=N(\mu_w,\sigma_w^2\ )$ ?
Proof:
$q(z)=N(z;\mu_z,\sigma_z^2)=\frac{1}{\sqrt{2\pi \sigma_z^2}}\exp\{-\frac{(z-\mu_z)^2}{2\sigma_z^2}\}\\ =\frac{1}{\sqrt{2\pi \sigma_z^2}}\exp\{ -\frac{1}{2\sigma_z^2}z^{2}+\frac{\mu_z}{\sigma_z^2}z-\frac{\mu_z^2}{2\sigma_z^2}\}$
$q(w|z)=N(w;\mu_wz,\sigma_w^2z^2)=\frac{1}{\sqrt{2\pi \sigma_w^2z^2}}\exp\{-\frac{(w-\mu_wz)^2}{2\sigma_w^2z^2}\} \\ =\frac{1}{\sqrt{2\pi \sigma_w^2z^2}}\exp\{ -\frac{w^2}{2\sigma_w^2}z^{-2}+\frac{w\mu_w}{\sigma_w^2}z^{-1}-\frac{\mu_w^2}{2\sigma_w^2}\}$
$\int q(z)q(w|z) {\rm d}z=\frac{1}{\sqrt{2\pi \sigma_z^2}}\frac{1}{\sqrt{2\pi \sigma_w^2 }} \exp\{-\frac{\mu_z^2}{2\sigma_z^2} -\frac{\mu_w^2}{2\sigma_w^2} \}\\ \int \frac{1}{|z|} \exp\{-\frac{1}{2\sigma_z^2}z^{2}+\frac{\mu_z}{\sigma_z^2}z+\frac{w\mu_w}{\sigma_w^2}z^{-1}-\frac{w^2}{2\sigma_w^2}z^{-2} \} {\rm d} z \\ =。。。?$

---

Dropout变量 与 权重后验分布参数

• 2012_Hinton_improving neural networks by preventing co-adaptation of feature detectors 提出二元伯努利Dropout。在输入上乘以一个伯努利噪声。
  $a=\frac{1}{1-p}o \xi w+b$, $\xi \sim {\rm Bernoulli}(1-p)$,
  $E[\frac{1}{1-p} \xi]=1$,
  ${\rm Var}(\frac{1}{1-p} \xi )=\frac{p}{1-p}$
  $p$的概率取0删除,$p$ 叫Dropout速率。

• 2014_Srivastava_dropout-a simple way to prevent neural networks from overfiting 提出连续的Dropout，在输入上乘以一个高斯噪声。
  $a=o \xi w+b$, $\xi \sim {\rm N}(1，\alpha=\frac{p}{1-p})$

• 2013_Wang_fasting dropout traininga 指出在输入上乘以一个高斯噪声后响应还是高斯的，均值方差可以直接计算。
  $a=o \xi w$, $\xi \sim {\rm N}(1，\alpha)$，$a \sim N(ow; o^2 \alpha w^2)$
  也相当于输入乘以的权重w，且权重w是从一个高斯分布取得的
  $a=ov$, $v\sim N(w,\alpha w^2)$

• 2015_Kingma_variational dropout and the local reparameterization trick 把2013年的两个工作命名为了 局部重参数化 和 变分dropout。
  提出了KL[norm|log-uniform]的估计表达。

• 2017_Molchanov_variational dropout sparsifies deep neural networks为每个权重后验定义不同的variance进而带来稀疏性。同时将 高斯后验$N(w,\alpha w^2)$ 看成$N(w,\sigma^2)$变成独立的两个变量因而可以 加法重参数化进而让alpha可以无穷大取值。
  $a=ov$, $v\sim N(w,\sigma^2)$, $v=w+\sigma \epsilon$，$\alpha=\frac{\sigma^2}{w^2}$
  2014年权重乘以高斯Dropout变量$a=o \xi w+b$, $\xi \sim {\rm N}(1，\alpha)$的采样版本：
  $\xi=1+\sqrt \alpha \epsilon$, $\xi w=(1+\sqrt \alpha \epsilon)w$
  两者等价：$\xi w=w+\sqrt \alpha w \epsilon=w+\sigma \epsilon$。

• 2017_louizos_bayesian compression for deep learning 为权重变量w的分布参数z设置后验高斯分布。
  先高斯采样$z= \mu_z+\sigma_z \epsilon$，再根据 $q(w|z)=N(\mu_{\tilde w}z,\sigma^2_{\tilde w}z^2)$计算响应的均值，
  相当于在权重后验高斯$N(\mu_{\tilde w},\sigma_{\tilde w}^2)$基础上乘以了高斯的Dropout变量$z\sim N(\mu_z,\sigma_z^2)$。
  $a=ow+b$
  $\mu_a=o\mu_{\tilde w}z+\mu_b$
  $\sigma_a^2=o ^2 \sigma_{\tilde w}^2 z^2+\sigma_b^2$

• 权重后验是$N(\mu_w,\sigma^2_w)$相当于权重$\mu_w$上乘以了一个高斯Dropout变量$\xi\sim N(1,\alpha=\frac{\sigma^2_w}{\mu_w^2})$。

• 权重后验是$N(\mu_z,\sigma_z^2)N(\mu_{\tilde w}z,\sigma_{\tilde w}^2z^2)$, 相当于高斯权重$N(\mu_{\tilde w},\sigma_{\tilde w}^2)$乘以高斯Dropout变量$z \sim N(\mu_z,\sigma_z^2)$采样, 当$\alpha=\frac{\sigma^2_z}{\mu_z^2}$大于某个阈值则移除对应的W。

---

数据处理

MNIST数据集：

• [0,255]的数据转换到[0,1]
• 2*(x-0.5)线性变换

初始化 $N(0,0.01^2)$。
当方差大于1后 clip成1。

5遍迭代
每批128个

---