Skip to content

模型压缩技术

模型压缩是在保持性能的前提下减少模型大小和计算成本的关键技术。

为什么需要模型压缩?

  • 部署成本:大模型需要昂贵 GPU(H100/A100)
  • 推理延迟:压缩后响应更快
  • 内存限制:移动设备内存有限
  • 能耗:边缘设备功耗敏感

主要压缩方法

1. 模型剪枝(Pruning)

移除模型中不重要的参数或结构。

剪枝类型

类型描述压缩率精度损失
结构化剪枝移除整个神经元/通道高 (50-70%)中等 (~2%)
非结构化剪枝移除单个权重超高 (90%+)小 (<1%)
层剪枝移除整个层较大

示例:非结构化剪枝

python
import torch
import torch.nn.utils.prune as prune

# 1. 对卷积层进行剪枝
module = model.conv1
prune.l1_unstructured(module, name='weight', amount=0.3)  # 剪枝 30%

# 2. 查看稀疏性
print(f"稀疏度: {(module.weight == 0).float().mean():.2%}")

# 3. 移除剪枝掩码(永久删除)
prune.remove(module, 'weight')

剪枝流程

训练模型 → 评估重要性 → 剪枝 → 微调 → 重复直到目标压缩率

重要性标准:

  • 权重幅值:接近 0 的权重不重要
  • 梯度幅值:梯度小的连接不重要
  • 激活值:输出为 0 的神经元可移除

2. 知识蒸馏(Knowledge Distillation)

用大模型(Teacher)训练小模型(Student)。

核心思想

Student 学习 Teacher 的输出分布,而不仅是硬标签。

损失 = α × 软标签交叉熵 + (1-α) × 硬标签交叉熵

温度参数

软标签更平滑, Temperature T > 1:

soft_target = softmax(teacher_logits / T)
student_target = softmax(student_logits / T)

T 越大,分布越平滑,学到的知识越"软"。

代码示例

python
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.7):
        super().__init__()
        self.T = temperature
        self.alpha = alpha

    def forward(self, student_logits, teacher_logits, labels):
        # 软标签损失
        soft_loss = F.kl_div(
            F.log_softmax(student_logits / self.T, dim=-1),
            F.softmax(teacher_logits / self.T, dim=-1),
            reduction='batchmean'
        ) * (self.T ** 2)

        # 硬标签损失
        hard_loss = F.cross_entropy(student_logits, labels)

        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

# 训练循环
criterion = DistillationLoss(temperature=4, alpha=0.7)
for batch in dataloader:
    teacher_out = teacher(batch)  # Teacher 不更新梯度
    student_out = student(batch)
    loss = criterion(student_out, teacher_out, labels)
    loss.backward()
    optimizer.step()

DistilBERT 案例

  • 模型:BERT-base (110M 参数) → DistilBERT (67M 参数)
  • 方法:知识蒸馏 + 层数减半
  • 效果:速度快 60%,性能保留 97%

3. 低秩分解(Low-Rank Factorization)

将大矩阵分解为小矩阵乘积。

SVD 分解

W (m×n) ≈ U (m×r) × V (r×n)

其中 r << min(m,n),参数量从 m×n 降到 r(m+n)。

python
import torch

def low_rank_decomposition(weight, rank):
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
    U_r = U[:, :rank]
    S_r = torch.diag(S[:rank])
    Vh_r = Vh[:rank, :]
    return U_r, S_r, Vh_r

应用

  • 全连接层:2 个大矩阵替代 1 个
  • 卷积层:核张量分解

缺点:需要重新训练或微调恢复精度。

4. 动态推理(Early Exit)

在推理过程中提前输出,减少计算。

python
class EarlyExitModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([...])
        self.exit_heads = nn.ModuleList([
            nn.Linear(d_model, num_classes) for _ in range(N)
        ])

    def forward(self, x, threshold=0.9):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.exit_heads):
                pred = self.exit_heads[i](x)
                confidence = F.softmax(pred, dim=-1).max()
                if confidence > threshold:
                    return pred, f"exit_at_layer_{i}"
        return pred, "exit_at_final"

适用场景:简单样本提前退出,复杂样本走完全流程。

组合策略

实际应用通常组合多种方法:

模型 → 剪枝 → 知识蒸馏 → 量化 → 部署

案例:TinyBERT

  1. 预训练 BERT (110M)
  2. Transformer 层剪枝 (6 层 → 4 层)
  3. 知识蒸馏 (6 层 Teacher → 4 层 Student)
  4. 量化 (FP32 → INT8)
  5. 结果:14.8× 压缩,性能保留 97%

工具与框架

工具功能支持模型
NNCF (Intel)剪枝、量化、蒸馏PyTorch, ONNX
Distiller (Intel)压缩综合框架PyTorch
NNI (Microsoft)自动化压缩多框架
PocketFlow (ByteDance)动态推理CV 模型
Optimum (HuggingFace)量化、剪枝Transformers

评估指标

压缩不仅要看压缩率,更要看实际效果:

  • 参数量 / FLOPs:模型大小和计算量
  • 推理速度:实际吞吐量、P99 延迟
  • 精度:任务指标(准确率、BLEU、PPL)
  • 内存占用:GPU/CPU 峰值内存

最佳实践

  1. 先量化再剪枝:INT8 量化后的模型再剪枝效果更好
  2. 逐步压缩:不要一次压缩太多(单次 <50%)
  3. 微调必不可少:每步压缩后必须微调恢复精度
  4. 任务特定评估:不要只看 PPL,测真实下游任务
  5. 硬件感知:针对目标硬件优化(GPU/CPU/移动端)

典型压缩结果(BERT-base)

方法参数量精度 (SQuAD)速度提升
原始110M88.5 F1
剪枝 30%77M88.2 F11.5×
DistilBERT67M86.9 F1
量化 INT8110M88.3 F12.5×
剪枝+蒸馏+量化35M87.5 F1

模型压缩是 LLM 落地必学技能,掌握这些技术能大幅降低部署成本!