Bayesian Compression for Deep Learning 阅读笔记

论文： https://papers.nips.cc/paper/6921-bayesian-compression-for-deep-learning.pdf
代码： https://github.com/KarenUllrich/Tutorial_BayesianCompressionForDL

一、论文摘要

深度学习中的压缩和计算效率已经成为一个非常重要的问题。在这项工作中，我们认为解决这个问题最有原则和最有效的方法是采用贝叶斯观点，通过先验分布引入稀疏性，进而对网络进行剪枝。
本文介绍了两个创新点:

• 使用的层次先验最后剪枝的是节点（结构），而非单个权重（参数）；
• 使用后验不确定性确定最优的固定点准确率来编码权重。

Key Points：为网络权重参数设置0均值的高斯先验，讨论了高斯先验的scale的两种先验分布；使用变分贝叶斯估计，全贝叶斯学习，目的是让网络的各层的有效权重减小。

二、模型介绍

2.1 概述

深度学习的复杂的网络在实际应用中意味着计算、带宽的巨大消耗，以及延时。
Alexnet的卷积层占4%的参数，91%的计算量。

现有方法大多数是减少网络结构和有效权重，在控制精度的同时减少无效节点的个数，减掉不必要的连接，再用一个大网络来训练一个小网络。

从贝叶斯的角度看，减少网络结构和权重个数与实现高进度本身是可以同时存在的。因为贝叶斯方法本身就探索最佳模型结构（先验分布引入稀疏性，进而对网络进行剪枝），并通过奖励参数后验分布中的不确定性（即信息量，可以通过bits back argument的方法）移除无关紧要的信息。

本文使用变分贝叶斯近似来进行贝叶斯推理，这也已经在模型压缩方面得到了明确的解释。通过对隐变量(而不是单个权重)采用先验引入稀疏性，我们可以修剪神经元，包括它们所有的输入和输出权重。这避免了修剪或矢量量化单个权重所需的更复杂和低效的编码方案。作为贝叶斯的额外奖励，可以使用变分后验不确定性来评估哪些比特是重要的，并去除那些在近似后验采样下波动太大的比特，由此推导出每层的最佳定点精度。

2.2 变分贝叶斯和最小描述长度

• 变分贝叶斯:
  变分推断（Variational Inference）
• 最小描述长度(MDL)：
  最小描述长度准则

2.2.1 变分推断

问题描述

变分推断就是近似不可观测变量的后验概率，以便通过这些变量作出统计推断，即近似求$P(Z|D)$，D为数据，Z为隐变量。

以上可以总结为这个问题：有一组观测数据D，并且已知模型的形式，求参数与潜变量（或不可观测变量） Z = { Z 1 , . . . , Z n } Z=\{Z1,...,Zn\} Z={Z1,...,Zn} 的后验分布: $P(Z|D)$。

正如上文所描述的后验概率的形式通常是很复杂(Intractable)的，对于一种算法如果不能在多项式时间内求解，往往不是我们所考虑的。因而我们想能不能在误差允许的范围内，用更简单、容易理解(tractable)的数学形式$Q(Z)$来近似$P(Z|D)$,即 $P(Z|D)\approx Q(Z)$，而这里的$Q(Z)$就是变分。

现在等式两边同时对q(z)做期望，即，

$$\int q(z)\log P(x)dz=\int q(z)\log P(x,z)dz-\int q(z)\log q(z)dz+\int q(z)\log \frac{q(z)}{P(z|x)}dz$$
由于$q(z)$与$P(x)$无关，所以$\int q(z)\log P(x)dz=\log P(x)$，原式最终变为：
$$\log P(x)=\underset{L(q(z)),ELOB(EvidenceLowerBound)}{\underbrace{\int q(z)\log P(x,z)dz-\int q(z)\log q(z)dz}}+\underset{KL(q(z)||P(z|x))}{\underbrace{\int q(z)\log \frac{q(z)}{P(z|x)}dz}}$$

我们可以通过如上的式子分解得到以下的结论，KL散度是用来衡量两个分布之间的距离，当距离为0时，表示这两个分布完全一致。$P(x)$不变，那么想让$KL(q(z)||P(z|x))$越小，即让ELOB越大（因为$KL\ge 0$，所以$logP(x)\ge ELOB$）。所以我们想让$P(Z|D)\approx Q(Z)$即让$KL\approx 0$，即最大化ELOB即可。$$\log P(D)=D_{KL}(Q||P)-\sum _{Z}Q(Z)\log \frac{Q(Z)}{P(Z,D)}=D_{KL}(Q||P)+L(Q).$$求解过程(极大化L(q))

根据平均长理论，变分分布Q(Z)可以通过参数和潜在变量的划分（partition）因式分解，比如将$Z$划分为$Z_1\dots Z_M$。根据以上假设，我们来最大化下界$L(q)$，因为假设$q_i(Z_i)$分布之间都是独立的，所以我们依次轮流来优化，以$q_j(Z_j)$为例(为了简单起见，缩写为$q_j$)。
L ( q j ) = ∫ ∏ i q j { ln ⁡ p ( X , Z ) − ∑ i ln ⁡ q i } d Z = ∫ q j { ∫ ln ⁡ p ( X , Z ) ∏ i ≠ j q i d Z i } d Z j − ∫ q j ln ⁡ q j d Z j + c o n s t = ∫ q j ln ⁡ p ~ ( X , Z j ) d Z j − ∫ q j ln ⁡ q j d Z j + c o n s t \begin{aligned} L(q_j) & =\int \prod_{i}q_j\{\ln p(\mathbf{X},\mathbf{Z})-\sum_i \ln q_i\} d\mathbf{Z} \\ &= \int q_j\{\int \ln p(\mathbf{X},\mathbf{Z})\prod_{i\neq j}q_i d\mathbf{Z}_i\}d\mathbf{Z}_j-\int q_j\ln q_j d\mathbf{Z}_j+const \\ &=\int q_j \ln \tilde{p}(\mathbf{X},\mathbf{Z}_j)d\mathbf{Z}_j-\int q_j\ln q_j d\mathbf{Z}_j +const \\ \end{aligned} L(qj​)​=∫i∏​qj​{lnp(X,Z)−i∑​lnqi​}dZ=∫qj​{∫lnp(X,Z)i̸​=j∏​qi​dZi​}dZj​−∫qj​lnqj​dZj​+const=∫qj​lnp~​(X,Zj​)dZj​−∫qj​lnqj​dZj​+const​

这里我们定义一个新分布$lnp\hat{}(X,Z_j)$为
$$\ln \tilde{p}(\mathbf{X},{\mathbf{Z}}_j)=\int \ln p(\mathbf{X},\mathbf{Z})\prod _{i̸=j}{q}_{i}d{\mathbf{Z}}_i$$我们发现它刚好是除去与$q_j$分布相关的$z_j$之后原似然的期望值，有$$\ln \tilde{p}(\mathbf{X},{\mathbf{Z}}_j)={\mathbb{E}}_{i̸=j}[\ln p(\mathbf{X},\mathbf{Z})]+const$$我们观察$L(q_j)$的最后一行推导，会发现就是$q_j(Z_j)$和$p\hat{}(X,Z_j)$的KL散度的负值，这里我们固定$qi̸=j$不变，那么最大化$L(q)$就变成了最小化这个KL散度，而KL散度的最小值在$qj(Zj)=lnp\hat{}(X,Z_j)$时取到。所以，最优解$q_j^{\ast }(Z_j)$为$$\ln q_j^{\ast }({\mathbf{Z}}_j)={\mathbb{E}}_{i̸=j}[\ln p(\mathbf{X},\mathbf{Z})]+C$$另加的这个C是为了归一化整个分布，有$C=\int exp({\mathbb{E}}_{i̸=j}[\ln p(\mathbf{X},\mathbf{Z})])d{\mathbf{Z}}_j$ ,然后依次更新其他Zj，最终相互迭代达到稳定。

2.2.2 最小描述长度准则

提出最小描述长度(MDL)的目的是为了根据信息论中的基本概念来解释极大后验假设(MAP)。

信息论中的一个基本定理是最小描述长度原则。它与压缩直接相关，因为它将最佳假设定义为用最小位数传递模型(复杂性成本$L_C$)和数据失配(错误成本$L_E$)的总和的假设。即ELOB可以分解为：

其中$H()$为信息熵。
数据拟合误差Error Cost： $$L^E=E_{q_{\varphi }(w)}\left[\log p(D|w)\right]$$模型复杂度Complexity Cost： $$L^C=E_{q_{\varphi }(w)}\left[\log p(w)\right]+H(q_{\varphi }(w))$$Evidence下界ELBO $L(\varphi )=L^E+L^C$ 表示出了通讯的最小代价。

为权重引入先验得到稀疏性往往可以压缩模型，也可以通过$q_{\varphi }(w)$噪声权重编码，利用熵的bits−back argument，使得熵项无限接近负无限$H(\delta (w))=-\inf$。

实践中 噪声权重编码的神经网络的数据拟合误差项是非常复杂的，常采用蒙特卡洛积分。 连续的$q_{\varphi }(w)$可以使用reparametrization技巧，不是从$q_{\varphi }(w)$中采样，而是通过噪声变量$ϵ$和变分参数$\varphi$的确定函数采样。
$$L^E=E_{p(ϵ)}\left[\log p(D|f(\varphi ,ϵ))\right]$$

2.2.3 Reparametrization Trick

说到reparametrization trick，其实挺简单的，但是对于变分推断的可解性上面却有很大的作用。

对于reparameter trick的理解就是，上对于上面求出的ELOB：
很多时候我们很难精确地计算其的值，但我们可以通过随机采样$z$来实现优化目标，但这会导致每次优化都要重新采样，进而得到的解方差很大。所以有人提出了使用reparameter trick的方法：$$z=\tilde{g}(x,ϵ)withϵ\sim p(ϵ)$$ 类似这样的一个固定函数进行表示$z$，通过采样$ϵ$来得到$z$，将$z$的随机部分和确定性部分分开，可以有效减少解的方差。

后面又有人提出了local reparameter trick的方法进一步减少这个方差，其思想是为在mini-batch中的每个数据的权重单独采样。
$$q_{\varphi }(w_{i,j})=N({\mu }_{i,j},{\sigma }_{i,j}^2)对于任意w_{i,j}\in W$$这样就可以确定权重参数之间的协方差为0。

2.3 通过高斯分布的比例混合实现贝叶斯压缩

2.3.1 scale-mixtures of normals分布族

将$z$看做随机变量，得到$w$的边界先验分布（积分掉$z$后得$w$的相对于数据D的先验分布）有长尾，且更聚集在0点。随后w的后验分布就会偏向稀疏。
$$z\sim p(z);w\sim N(w;0,z^2)$$$z$ 是参数$w$的先验分布的scale。

这是一种比较通用的分布族，许多能导致稀疏性的分布都是它的特例，例如：

• spike&slab分布p(w)是伯努利分布。神经网络的dropout正来源于这个分布
• Laplace分布 $p(z^2)=exp(\lambda )$。laplace先验的后验分布是Lasso估计。有被用到稀疏化神经网络。但是容易压榨大信号，且只提供点估计，没有保留不确定性，容易过拟合且压缩性不够好。
• 不含超参数的log-uniform先验，half-cauchy先验得到horseshoe分布，spike$slab先验的连续松弛。

2.3.2 Reparametrizing variational dropout for group sparsity

什么是 variational dropout ？

首先先说一下什么是 variational dropout，对于神经网络中任意一层的dropout实际上可以写成如下形式：

其中$A$和$B$分别是输入和输出， $\Xi$是一个噪声矩阵，而$W$是权重参数。我们最常见的dropout其实就是 ${\xi }_{ij}\sim Bernoulli(1-p)$，使得输入矩阵的权重参数以$p$的概率为0，其可以称为binary dropout。

后面又有人提出了高斯dropout，这里 ${\xi }_{ij}\sim N(1,\alpha =\frac{p}{1-p})$，即均值为1，方差为$\alpha$的高斯分布，引入连续噪声而不是离散噪声很重要，因为向输入中添加高斯噪声相当于将高斯噪声引入到权重中，该过程可用于获得模型权重的后验分布。因此，将高斯噪声加入到权重中，实际上相当于在 $q(w_{ij})=N({\theta }_{ij},\alpha {\theta }_{ij}^2)$ 的分布上采样权重$w_{ij}$，即

这里同样也可以用到上面提过的reparameter trick的方法。

而 variational dropout 实际上就是用$q(w|\theta ,\alpha )$分布作为参数的后验近似。当$\alpha$是固定值的时候，其就是高斯dropout，只不过这里将$\alpha$作为了一个学习的参数而不是超参数。同样这里$\alpha$可以是任意值而不是像高斯dropout一样$\alpha \le 1$。

Additive Noise Reparameterization

但是在加入了这样的噪声以后，当$\alpha$特别大的时候梯度也会很大，因为

所以论文中的方法是增加了一个额外的噪声项Additive Noise $${\sigma }_{ij}\cdot {ϵ}_{ij}$$其中${\sigma }_{ij}^2={\alpha }_{ij}{\theta }_{ij}^2$，梯度就变为了

Approximation of KL Divergence

为了方便计算损失，可以将KL散度写成了一个具体的$\alpha$的函数，其中σ是sigmoid函数。

这样做的好处是当$\alpha$趋近于无穷的时候KL散度会趋近于0，同时还利于比较不同尺度网络的损失。

为什么 variational dropout 可以引入稀疏性？

当${\alpha }_{ij}\to +\infty$的时候其实相当于binary dropout中$p=1$的情况。可以试想如果我们给予网络的每个权重都非常大的扰动，网络不管怎么改变都不能减少loss，所以其最好的办法就是所有的权重都为0。

可以从另一个角度来看待这种情况。大的${\alpha }_{ij}$对应于$w_{ij}$中无限大的乘性噪声。这意味着该权重的值将是完全随机的，其大小将是无界的，这将破坏模型预测并降低预期的对数可能性。因此，以将${\alpha }_{i}j{\theta }_{ij}^2$变为零的方式将相应的权重${\theta }_{ij}$置零是有益的，从中引入了稀疏性。这意味着后面这个重量的边缘分布实际上是一个$\delta$函数，以零为中心。

其中$\delta (w_{ij})$为Dirac delta function，以0为中心。即$$\delta \left(x\right)=\{\begin{array}{cc}\infty & \text{if }x=0\\ 0& \text{if }x̸=0\end{array}$$

如何利用variational dropout搭建全连接层和卷积层

为了通过变分差训练稀疏化模型，我们用KL散度的近似值优化变分差目标，KL散度对于dropout rate $\alpha$的所有值都是精确的。

在代码实现的时候为了减少损失函数的方差，同时使用 Local Reparameterization trick 和 Additive Noise Reparameterization。

对于一个全连接层，其输出$b_{ij}$的计算方式如下：
$$b_{ij}\sim N({\mu }_{ij},{\delta }_{ij})$$其中${\mu }_{ij}$和${\delta }_{ij}$的由dropout rate $\alpha$得到，$${\mu }_{ij}=\sum _{k=1}^K{\alpha }_{ik}{\theta }_{kj},{\delta }_{ij}={\alpha }_{ik}\sum _{k=1}^K{\alpha }_{ik}^2{\theta }_{kj}^2$$而dropout rate $\alpha$由变分参数${\theta }_k^{h\times w\times C}$和${\sigma }_k^{h\times w\times C}$得到$${\alpha }_{ik}=exp(log{\sigma }_{ij}^2-log{\theta }_{ij}^2)$$用代码实现起来其实很简单，由于${\mu }_{ij}$和${\delta }_{ij}$是由$\alpha$得到，而$\alpha$又是由两个参数矩阵的计算得到的，所以${\mu }_{ij}$和${\delta }_{ij}$可以直接如下定义：

def reparametrize(mu, logvar, sampling=True):
    if sampling: # 训练的时候使用
        std = logvar.mul(0.5).exp_() # exp(0.5*logvar)
        eps = torch.FloatTensor(std.size()).normal_()
        eps = Variable(eps)
        return mu + eps * std
    else: # 测试的时候使用
        return mu
        
# dropout params
z_mu = Parameter(torch.Tensor(in_features))
z_logvar = Parameter(torch.Tensor(in_features))  # = z_mu^2 * alpha
# weight params
weight_mu = Parameter(torch.Tensor(out_features, in_features))
weight_logvar = Parameter(torch.Tensor(out_features, in_features))
bias_mu = Parameter(torch.Tensor(out_features))
bias_logvar = Parameter(torch.Tensor(out_features))

z = reparametrize(z_mu.repeat(batch_size, 1), z_logvar.repeat(batch_size, 1))
# apply local reparametrisation trick (这一步完成了上面mu和delta的计算）
xz = x * z # x is input
mu_activations = F.linear(xz, weight_mu, bias_mu) # input / weight / bias
var_activations = F.linear(xz.pow(2), weight_logvar.exp(), bias_logvar.exp())

output = reparametrize(mu_activations, var_activations.log())

卷积层的计算方式类似，直接贴图：

其中所有的$({)}^2$、$log()$和$exp()$都是element-wise的，$⨀$代表点乘，$\ast$代表卷积操作，$vec()$代表将矩阵或者张量reshape成向量。

损失函数

至于损失函数即ELOB，等于分类误差加上归一化后的KL散度，具体代码如下：

d_error = cross_entropy(output, target)
N = len(train_loader.dataset)
variational_bound = d_error + kl_divergence / N 

2.4 通过半柯西先验实现贝叶斯压缩

（这块内容未来有空再补，代码方面实现较为简单，和2.3类似）

三、实验结果

在不同的网络和不同的数据集上都在不损失精度的前提下达到了不错的压缩效果。

并且随着训练，每一层的参数矩阵会有越来越多的权重置零。