模型里有很多权重,其中一部分删掉以后,结果可能几乎不变。模型剪枝做的就是这件事。它先判断哪些连接影响较小,再把它们置零或整组移除,随后用原任务重新训练和评测。论文里的压缩率常常很漂亮,部署时却还要回答一个更土的问题,机器会不会真的少算。
我把早期的 Deep Compression、彩票假设和 NVIDIA 的稀疏计算说明放在一起看。三份材料说的是同一条过程里的不同位置。算法决定删什么,训练决定能否保住能力,运行库和芯片决定那些零值能不能换成速度。
权重变成零,只完成了前半段
2015 年的 Deep Compression 先剪掉连接,再量化权重并做编码。论文在 AlexNet 和 VGG-16 上报告,单独剪枝把连接数降到原来的九分之一到十三分之一。三步合起来,AlexNet 从 240MB 降到 6.9MB,VGG-16 从 552MB 降到 11.3MB。作者还在当时的 CPU、GPU 与移动 GPU 上测到分层计算加速和能耗下降。
这些数字属于两种卷积网络和那套压缩实现。今天给语言模型剪枝,不能照抄比例。PyTorch 的通用剪枝工具会给权重套上掩码,前向计算看到许多零,但张量形状可能原封不动。零值还在那里。若底层仍按稠密矩阵计算,文件没有变小,乘法也未必少做。
结构化剪枝更照顾机器。它可以整列、整通道或按固定小组删除权重,剩下的形状比较规整。NVIDIA Ampere 的稀疏 Tensor Core 支持 2 比 4 模式,每四个连续数值里要有两个为零。符合这种模式并走到对应算子后,硬件才有机会跳过零值计算。厂商给出的理论计算吞吐可到稠密计算的两倍,端到端请求还会受到访存、调度与其他算子影响。
剪多少要由任务来定
彩票假设研究发现,稠密网络里能找到很小的子网络,在保留合适初始化时,可以用相近训练轮数达到原网络的测试表现。研究中的子网络常剩原模型 10% 到 20% 权重。换一组随机初始化,效果会明显变差。这个结果提醒人,留下哪些连接和留下什么数值同样重要。
面向预训练语言模型的 Movement Pruning 又给出一个差别。只按权重当前大小剪,在迁移学习阶段未必合适。它改看训练中权重变化的方向,在高稀疏度实验里表现更好。论文把这种剪枝与蒸馏组合后,在特定任务上用约 3% 参数仍保持较小精度损失。这里依然有任务和方法边界,不能扩成所有语言模型都能删掉 97%。
实际验收最好保留四组数。第一组是剪枝前后的任务指标,尤其看少数类和长输入。第二组是落盘文件与运行时内存。第三组是在目标机器上的首字延迟、整段延迟和吞吐。第四组是回退条件,达到哪一档精度损失就停止继续剪。
如果权重零了一半,延迟一毫秒没少,剪枝在这台机器上只得到一张稀疏度报表。先确认运行库支持哪种稀疏结构,再决定怎样剪,能省掉很多无效实验。