模型压缩技术

Posted on 2022-04-25 In Learning , Project Views: Waline:

参考自李宏毅机器学习课程中网络压缩部分内容，主要涉及：

模型剪枝(Pruning)：移除“冗余”参数/神经单元

知识蒸馏(Distillation)：将大模型知识转移给小模型

模型量化(Quantization)：压缩存储参数所需空间而非参数量

构架设计(Network Design)：参数矩阵低秩分解/1x1卷积、参数共享

剪枝(Pruning)

剪枝即移除网络中“冗余”参数/神经单元，是网络压缩最直接的一种方式，LeCun等人在上世纪就曾提出“最优脑损伤”(1989)的剪枝技术。剪枝的基本思路为：评估参数或神经元重要性，移除其中相对不重要的，微调整个模型；之后循环操作、渐进剪枝，避免单次移除太多参数或单元破坏模型稳定性。剪枝可以针对单个权重连接或是整个神经元(行/列)，前者更精细灵活，但其问题在于最终模型连接非常规则，代码实现和计算加速都相对困难(参数通常是固定为0而非移除)，速度提升有限(甚至降速)，移除单元可避免类似问题。

大模型可视为众多小模型的集成，单独训练小模型很难找到全局最优，而大模型表现好，意味其中存在着“中奖”的小模型，模型剪枝就是要找出这个小模型，但关于其中最关键的因素有不同观点。大乐透假说(Lottery Ticket Hypothesis)作者发现大模型剪枝得到的小模型重新初始化表现并不好，但如果初始化参数与大模型一致，训练效果又快又好，这意味这组“初始化”就是中奖参数。之后又有文章发现初始化的数值其实不重要，正负号才是关键。但另一方面，也有文章指出剪枝后的小模型随机初始化也可获得很好的效果(更多epoch)，因此“网络结构”才更为关键，而大乐透假说只在权重剪枝且学习率较小时成立。不过这些结论都属于经验猜想，并没有严格证明，很难评判绝对的对与错。

蒸馏(Distillation)

更大的模型或者更多模型集成通常能获得更优秀和稳定的表现，而知识蒸馏是希望将大模型的知识转移给小模型，最具代表性的是Hinton于2015年发表的Distilling the Knowledge in NN。其原理猜想为大模型可学习到大量通用特征，而具体的下游任务只需要用到其中一部分，知识蒸馏可将这部分有效特征提取出来。

具体的大模型(以及模型集成)的主要优势是更强的表示能力和泛化能力：如果将知识视为从输入到输出的某种映射，这种映射代表了模型的表示能力；而预测输出中非正确分类的概率值(softmax)则反映了模型泛化的倾向。因此要让小模型获得大模型的泛化能力，可直接以大模型输出的softmax概率值为目标进行学习。更为巧妙的，可在softmax中引入温度参数：softmax对不同分类的区分度随温度升高而降低，因此在教学阶段大小模型可取(相同的)较大温度，从而重点强调泛化能力的学习；在学习完后，则恢复正常温度(t=1)，用于正常输出预测，这一升温再降温过程类似于“蒸馏”操作。

除了将原本的标签数据替换为软标签(soft label)，还可以让大模型去标记无标签数据，生成伪标签(pseudo label)，扩大小模型监督学习的数据量。后一点对于语言任务尤为重要，因为语言任务中标签数据通常很有限，大规模预训练模型通常是基于自监督任务，而得到预训练模型后就可生成大量(软)标签数据用于有监督的训练小模型。

Turc et al.(2019)中就使用预训练(+微调)作为Teacher模型，让Teacher去生成软标签数据用于进一步训练经过预训练的小模型(缩小50~100倍)，最终小模型效果可接近Teacher，远超直接预训练+微调同尺寸模型。DistillBERT中则进一步合并了两步操作，直接在小模型预训练阶段就引入知识蒸馏。

量化(Quantization)

参数量化的目标是减少参数存储所需空间而非参数量，最经典的做法就是通过降低参数数值精度(字节数)来减小所需存储空间，而最极端的情况是权重取值为-1和1二值，即只需要单个二进制位存储，如Binary Connect, Binarized NN, XNOR-net等。

另一种做法是将参数矩阵按具体取值分区(聚类)，每个分区参数统一平均，存储时只需要保留分区表和对应参数取值，此外还可通过霍夫曼编码进一步压缩所需存储空间。

构架设计(Design)

深度可分离卷积(Depthwise Separable Convolution)是通过调整网络构架实现参数缩减的经典例子。假设卷积核尺寸为 $k × k$ ，输入输出通道数为 $C_1, C_2$ ，则卷积操作的参数量为 $C_1× k× k × C_2$ 。减少参数量的一个可行做法是将通道内卷积和通道混合分开操作：前者对应depthwise卷积操作(各通道独立卷积)，参数量为 $C_1 × k× k × 1$ ；后者则为1×1卷积，参数量 $C_1× 1× 1 × C_2$ ，从而总参数量为 $C_1× k^2+ C_1× C_2$ ，这就是深度可分离卷积。经过分离操作参数量减少为最初的 $\frac{1}{C_2} + \frac{1}{k^2}$ ，考虑到通道数通常远大于卷积核尺寸 $C_2\gg k^2$ ，因此参数压缩为约 $\frac{1}{k^2}$ 。进一步的，我们还可以将上面逐通道的的 $k×k$ 卷积“拆解”为 $k×1$ 与 $1×k$ 卷积两个操作，继续缩减参数量，也被称为空间可分离卷积(spatially separable convolution)。最终，原本由三维矩阵实现的卷积，被替换为由三个一维向量实现的纵向(depthwise k×1)、横向(depthwise 1×k)以及深度方向(elementwise 1×1)的卷积。

其背后的原理是更为通用的矩阵低秩分解：网络存在冗余参数也就意味着参数矩阵为低秩矩阵，而为减少参数量可将较大的参数矩阵拆解为两个(低秩)小矩阵相乘，由 $M × N$ 的矩阵变为 $M × k$ 和 $k × N$ 的两个矩阵相乘，只要 $k\ll M,N$ 就可减少总的参数量。上述深度可分离卷积与这里矩阵分解的思想是一致的，而最近的MLP-Mixer同样是这一思路的延续。此外，这种技巧在注意力机制的发展中同样有所体现，具体可参考NLP基本概念II中对注意力机制的介绍。

除了参数矩阵的低秩分解，还有其他构架优化策略，参数共享就是另一种思路，其实如果回归本质，CNN, RNN等结构相对于全连接网络共通的优势就在于参数共享，Transformer则是一种新的参数共享构架。而在此之上还可进一步探索其他共享方式或网络结构，比如ALBERT中就将BERT网络中所有Transformer模块之间进行参数共享(同时也使用了矩阵低秩分解技巧)，最终参数量减少一个量级以上。不过虽然ALBERT大大降低参数量，提升了训练速度，但效果有一定退化；而如果将参数量增加到近似量级，效果可实现反超，但训练速度却又显著落后，总之不是很理想。

其他策略(Misc.)

渐进替换：Xu et al. 2020针对BERT模型提出通过渐进替换模型中的Transformer模块实现模型压缩，并类比忒修斯之船将其命名为忒修斯压缩(Theseus Compression)。具体的文章会将BERT中每2层Transformer模块作为一组(predecessor)，预训练时以一定概率随机替换为对应的1层Transformer模块(successor)。类似于老手带新手，让successor与predecessor一起训练获得提升，并在预训练结束后只保留successor实现模型的压缩。
动态计算：动态计算的目标并非直接减少参数量或计算量，而是希望网络能根据计算环境或任务，如运行平台、设备电量、任务难度等，自适应调整运算。具体通常有调整深度和宽度两种策略，而最基本的思想就是在训练时让不同层/宽度网络同时接近训练目标，代表性的工作分别为MSDNet和SlimmableNN。对于具体深度或宽度选择，除了根据外部条件手动设定，还可进一步让模型自适应调整，如BlockDrop, SkipNet, Runtime Neural Pruning。

最后需要注意，不同技术并非互斥的，而多种技术的综合使用也是一个值得更多研究的方向。关于模型压缩以及加速可参考相关的综述，如Efficient Deep Learning。