模型压缩综述
为什么需要模型压缩
现代深度学习模型的参数量和计算量急剧增长。以 GPT-3 为例,其拥有 1750 亿参数,单次推理即需消耗大量算力。在实际部署中,我们往往面临严格的约束:
- 延迟约束:实时推理场景(自动驾驶、语音助手)要求毫秒级响应
- 内存约束:边缘设备(手机、IoT)的内存通常仅几百 MB 至几 GB
- 能耗约束:移动端和嵌入式设备对功耗极为敏感
- 成本约束:云端 GPU 推理的计算成本随规模线性增长
模型压缩的目标是在尽量保持模型精度的前提下,减小模型体积、降低计算量、加速推理。
知识蒸馏 (Knowledge Distillation)
核心思想
知识蒸馏由 Hinton 等人于 2015 年提出,其核心是将大模型(Teacher)的"知识"迁移到小模型(Student)中。Teacher 模型的 soft labels(软标签)包含了类别间的相似性信息,比 one-hot hard labels 更具信息量。
Teacher-Student 框架
给定 Teacher 模型输出的 logits \(z_T\) 和 Student 模型输出的 logits \(z_S\),蒸馏损失定义为:
其中:
- \(T\) 是温度参数(temperature),\(T > 1\) 时 softmax 输出更平滑,暴露更多类间关系
- \(\alpha\) 控制蒸馏损失与真实标签损失的权重
- \(\sigma\) 为 softmax 函数
- \(\mathcal{L}_{\text{CE}}\) 为标准交叉熵损失
典型案例:DistilBERT
DistilBERT 将 BERT-base(110M 参数)压缩为一个 6 层的 Student 模型(66M 参数),保留了 97% 的语言理解能力,同时推理速度提升 60%。其蒸馏策略包括:
- 隐藏层蒸馏:对齐 Teacher 与 Student 的隐藏状态
- 注意力蒸馏:对齐注意力权重矩阵
- MLM 损失:保留掩码语言模型目标
模型剪枝 (Pruning)
非结构化剪枝 (Unstructured Pruning)
非结构化剪枝在 单个权重 粒度上操作,将绝对值较小的权重置零:
其中 \(\theta\) 为剪枝阈值。这种方法可以达到极高的稀疏率(如 90%+),但产生的稀疏矩阵不规则,需要专用硬件或稀疏计算库才能真正加速。
代表方法:Magnitude Pruning, Movement Pruning(用于微调阶段)
结构化剪枝 (Structured Pruning)
结构化剪枝以 整个通道(channel)、注意力头(head)、层(layer) 为单位进行剪枝。剪枝后的模型无需特殊硬件即可直接加速。
常见评判准则:
- \(\ell_1\)-norm:移除 \(\ell_1\) 范数最小的卷积核 / 通道
- Taylor 展开:根据损失函数对参数的一阶或二阶 Taylor 展开评估重要性
- BN 缩放因子:利用 Batch Normalization 层的 \(\gamma\) 参数作为通道重要性指标
彩票假说 (Lottery Ticket Hypothesis)
Frankle & Carlin (2019) 提出:一个随机初始化的密集网络中包含一个稀疏子网络("中奖彩票"),该子网络从相同的初始化出发训练,可以在相当甚至更少的迭代次数内达到与原网络相当的精度。
核心流程:
- 随机初始化网络,记录初始权重 \(w_0\)
- 训练网络至收敛,得到 \(w_f\)
- 按 magnitude 剪枝 \(p\%\) 的权重,得到掩码 \(m\)
- 将剩余权重重置为 \(w_0\),即用 \(m \odot w_0\) 重新训练
- 迭代执行(Iterative Magnitude Pruning, IMP)
这一发现对理解神经网络的冗余性和泛化能力具有深远意义。
量化 (Quantization)
量化通过降低权重和激活值的数值精度来压缩模型。典型做法是将 FP32 转为 INT8 甚至 INT4。
均匀量化基本公式
其中 \(S\) 为缩放因子(scale),\(Z\) 为零点(zero-point),\(b\) 为目标位宽。
两大范式
| 范式 | 全称 | 是否需要训练 | 精度保持 | 典型代表 |
|---|---|---|---|---|
| PTQ | Post-Training Quantization | 否 | 中等 | TensorRT, GPTQ, AWQ |
| QAT | Quantization-Aware Training | 是 | 高 | PyTorch QAT, LSQ |
PTQ 只需少量校准数据即可完成,适合快速部署;QAT 在训练中模拟量化误差,精度保持更好但成本更高。
关于量化的详细介绍(包括 PTQ 校准策略、QAT 的 STE 直通估计器、混合精度量化等),请参见 量化专题。
低秩分解 (Low-Rank Factorization)
低秩分解利用矩阵分解将一个大的权重矩阵近似表示为若干小矩阵的乘积,从而减少参数量和计算量。
对于权重矩阵 \(W \in \mathbb{R}^{m \times n}\),进行秩 \(r\) 近似分解:
参数量变化:从 \(mn\) 降低到 \(r(m+n)\)。当 \(r \ll \frac{mn}{m+n}\) 时,压缩效果显著。
常见方法
- SVD 分解:对预训练权重进行奇异值分解 \(W = U \Sigma V^T\),保留前 \(r\) 个奇异值
- Tucker 分解:用于卷积核的高维张量分解
- LoRA:在微调时冻结原始权重 \(W\),仅学习低秩增量 \(\Delta W = BA\),其中 \(B \in \mathbb{R}^{m \times r}\), \(A \in \mathbb{R}^{r \times n}\)
注意:LoRA 严格来说是一种参数高效微调方法(PEFT),但其底层原理正是低秩分解。
各方法对比
| 方法 | 压缩率 | 精度损失 | 是否需要重训练 | 硬件友好性 | 典型应用场景 |
|---|---|---|---|---|---|
| 知识蒸馏 | 2x-10x | 低 | 是(训练 Student) | 高 | 部署轻量模型 |
| 非结构化剪枝 | 10x-100x | 低-中 | 通常需要微调 | 低(需稀疏硬件) | 研究、专用加速器 |
| 结构化剪枝 | 2x-5x | 低-中 | 通常需要微调 | 高 | 通用推理加速 |
| PTQ 量化 | 2x-4x | 低-中 | 否 | 高 | 快速部署 |
| QAT 量化 | 2x-4x | 低 | 是 | 高 | 精度敏感场景 |
| 低秩分解 | 2x-5x | 中 | 可选 | 高 | 全连接/注意力层压缩 |
组合使用
在实际工程中,多种压缩技术往往联合使用:
- 先用知识蒸馏得到一个较小的 Student 模型
- 对 Student 模型进行结构化剪枝,移除冗余通道
- 最后用 PTQ 量化将模型转为 INT8 部署
这种流水线式的压缩策略可以在保持精度的同时实现 10x-50x 的综合压缩比。
参考文献
- Hinton et al., Distilling the Knowledge in a Neural Network, 2015
- Han et al., Deep Compression, ICLR 2016
- Frankle & Carlin, The Lottery Ticket Hypothesis, ICLR 2019
- Sanh et al., DistilBERT, NeurIPS Workshop 2019
- Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022