模型剪枝与知识蒸馏：压缩大模型的两种路径与工程取舍

一、模型压缩的必要性：精度与效率的永恒博弈

大模型的推理成本与参数量成正比。一个 7B 参数的模型在 FP16 下需要 14GB 显存存储权重，推理时还需要额外的 KV Cache 和激活值空间。在边缘设备或低成本服务器上部署时，模型必须被压缩。两种主流压缩路径是剪枝（Pruning）和知识蒸馏（Knowledge Distillation）。

剪枝直接移除模型中的冗余参数，减少计算量和内存占用。知识蒸馏训练一个小模型（Student）来模仿大模型（Teacher）的输出分布，不改变模型结构但减少参数量。两者的适用场景不同：剪枝适合保留原模型架构的场景，蒸馏适合可以接受更小模型架构的场景。

二、剪枝与蒸馏的机制对比：结构压缩 vs 知识迁移

剪枝分为非结构化剪枝（将个别权重置零）和结构化剪枝（移除整个通道或注意力头）。非结构化剪枝的压缩率高但需要稀疏计算硬件支持，结构化剪枝的压缩率低但可以直接在标准 GPU 上加速。知识蒸馏的核心是让 Student 学习 Teacher 的软标签（Soft Labels）——Teacher 输出的概率分布比硬标签包含更多信息。

flowchart TB
    A[大模型 Teacher] --> B{压缩路径}

    B --> C[剪枝]
    B --> D[知识蒸馏]

    C --> C1[非结构化剪枝<br/>权重级稀疏]
    C --> C2[结构化剪枝<br/>通道/头级移除]

    C1 --> C3[优点: 压缩率高 90%+<br/>缺点: 需稀疏硬件]
    C2 --> C4[优点: 通用 GPU 加速<br/>缺点: 压缩率有限 50-70%]

    D --> D1[软标签蒸馏<br/>学习概率分布]
    D --> D2[特征蒸馏<br/>学习中间表示]
    D --> D3[关系蒸馏<br/>学习样本间关系]

    D1 --> D5[优点: 灵活选择 Student<br/>缺点: 需要训练资源]
    D2 --> D5
    D3 --> D5

    C3 --> E[部署: 稀疏推理引擎]
    C4 --> F[部署: 标准 ONNX Runtime]
    D5 --> G[部署: 标准推理引擎]

关键差异：剪枝是"做减法"，保留原模型结构但移除部分参数；蒸馏是"做迁移"，用小模型继承大模型的知识。两者可以组合——先剪枝再蒸馏，或先蒸馏再剪枝，但组合的收益不一定叠加。

三、生产级代码实现：结构化剪枝与知识蒸馏

3.1 幅度剪枝：基于权重绝对值的结构化剪枝

import torch
import torch.nn as nn
import numpy as np

class MagnitudePruner:
    """幅度剪枝器：移除绝对值最小的权重"""

    def __init__(self, model, pruning_ratio=0.5):
        self.model = model
        self.pruning_ratio = pruning_ratio
        self.masks = {}

    def compute_masks(self):
        """计算剪枝掩码"""
        for name, param in self.model.named_parameters():
            if "weight" not in name:
                continue

            # 计算每个输出通道的 L1 范数
            # 为什么用 L1 范数而非 L2：L1 范数对
            # 小权重更敏感，更适合识别"不重要"的通道；
            # L2 范数会被少数大权重主导
            if param.dim() >= 2:
                # 对卷积/线性层：按输出通道计算重要性
                importance = param.abs().sum(
                    dim=tuple(range(1, param.dim())))
            else:
                importance = param.abs()

            # 确定阈值：保留 top-k 通道
            k = int(len(importance) * (1 - self.pruning_ratio))
            if k <= 0:
                k = 1

            threshold = torch.topk(importance, k).values[-1]

            # 创建掩码
            if param.dim() >= 2:
                channel_mask = (importance >= threshold).float()
                # 扩展掩码到所有维度
                expand_shape = [-1] + [1] * (param.dim() - 1)
                mask = channel_mask.view(*expand_shape).expand_as(
                    param)
            else:
                mask = (importance >= threshold).float()

            self.masks[name] = mask

    def apply_masks(self):
        """应用剪枝掩码"""
        for name, param in self.model.named_parameters():
            if name in self.masks:
                # 用掩码将不重要的权重置零
                param.data.mul_(self.masks[name])

    def fine_tune(self, train_loader, epochs=5, lr=1e-4):
        """剪枝后微调恢复精度"""
        # 为什么剪枝后需要微调：直接剪枝会导致
        # 精度大幅下降，微调让剩余权重重新适应
        # 被移除通道的功能
        optimizer = torch.optim.AdamW(
            self.model.parameters(), lr=lr)

        for epoch in range(epochs):
            for batch in train_loader:
                optimizer.zero_grad()
                output = self.model(batch["input"])
                loss = nn.CrossEntropyLoss()(
                    output, batch["label"])
                loss.backward()
                optimizer.step()

                # 每步训练后重新应用掩码
                # 为什么每步都应用：梯度更新可能
                # 让被剪枝的权重变为非零，
                # 必须持续掩码才能维持稀疏结构
                self.apply_masks()

            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

3.2 知识蒸馏：软标签与温度调节

class KnowledgeDistiller:
    """知识蒸馏训练器"""

    def __init__(self, teacher, student,
                 temperature=4.0, alpha=0.7):
        self.teacher = teacher
        self.student = student
        self.temperature = temperature
        # alpha: 蒸馏损失权重, 1-alpha: 硬标签损失权重
        # 为什么需要两个损失：纯蒸馏损失可能忽略
        # 真实标签的信息，混合损失兼顾两者
        self.alpha = alpha

        # Teacher 冻结参数
        for param in self.teacher.parameters():
            param.requires_grad = False
        self.teacher.eval()

    def distillation_loss(self, student_logits,
                          teacher_logits, labels):
        """计算蒸馏损失"""
        # 软标签损失：KL 散度
        # 为什么用 KL 散度而非 MSE：KL 散度衡量
        # 两个概率分布的差异，与交叉熵等价；
        # MSE 衡量 logits 的数值差异，不保证
        # 概率分布的语义一致性
        soft_targets = nn.functional.softmax(
            teacher_logits / self.temperature, dim=-1)
        soft_student = nn.functional.log_softmax(
            student_logits / self.temperature, dim=-1)

        # KL 散度 × T^2：补偿温度缩放导致的梯度缩小
        # 为什么乘 T^2：温度 T 使概率分布变平滑，
        # 梯度被缩小 T 倍；乘 T^2 恢复梯度量级
        kd_loss = nn.functional.kl_div(
            soft_student, soft_targets,
            reduction="batchmean"
        ) * (self.temperature ** 2)

        # 硬标签损失：标准交叉熵
        ce_loss = nn.functional.cross_entropy(
            student_logits, labels)

        # 加权组合
        total_loss = (
            self.alpha * kd_loss +
            (1 - self.alpha) * ce_loss
        )
        return total_loss

    def train_step(self, batch, optimizer):
        """单步蒸馏训练"""
        inputs = batch["input"]
        labels = batch["label"]

        # Teacher 推理（不计算梯度）
        with torch.no_grad():
            teacher_logits = self.teacher(inputs)

        # Student 推理
        student_logits = self.student(inputs)

        # 计算蒸馏损失
        loss = self.distillation_loss(
            student_logits, teacher_logits, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        return loss.item()

    def train(self, train_loader, val_loader,
              epochs=20, lr=3e-4):
        """完整蒸馏训练流程"""
        optimizer = torch.optim.AdamW(
            self.student.parameters(), lr=lr)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=epochs)

        best_val_acc = 0.0
        for epoch in range(epochs):
            self.student.train()
            total_loss = 0

            for batch in train_loader:
                loss = self.train_step(batch, optimizer)
                total_loss += loss

            scheduler.step()

            # 验证
            val_acc = self._evaluate(val_loader)
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(self.student.state_dict(),
                    "best_student.pt")

            avg_loss = total_loss / len(train_loader)
            print(f"Epoch {epoch}: Loss={avg_loss:.4f}, "
                  f"Val Acc={val_acc:.4f}")

        return best_val_acc

    def _evaluate(self, val_loader):
        self.student.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for batch in val_loader:
                outputs = self.student(batch["input"])
                preds = outputs.argmax(dim=-1)
                correct += (preds == batch["label"]).sum().item()
                total += len(batch["label"])
        return correct / total

3.3 剪枝效果评估

def evaluate_pruning(original_model, pruned_model,
                     test_loader):
    """评估剪枝效果"""
    # 精度对比
    orig_acc = evaluate_accuracy(original_model, test_loader)
    pruned_acc = evaluate_accuracy(pruned_model, test_loader)

    # 参数量对比
    orig_params = sum(p.numel()
        for p in original_model.parameters())
    pruned_params = sum(p.numel()
        for p in pruned_model.parameters())
    # 非零参数量
    nonzero_params = sum(
        (p != 0).sum().item()
        for p in pruned_model.parameters())

    # 推理速度对比
    orig_latency = benchmark_latency(original_model)
    pruned_latency = benchmark_latency(pruned_model)

    print(f"原始模型: 参数={orig_params/1e6:.1f}M, "
          f"精度={orig_acc:.4f}, 延迟={orig_latency:.2f}ms")
    print(f"剪枝模型: 参数={nonzero_params/1e6:.1f}M, "
          f"精度={pruned_acc:.4f}, 延迟={pruned_latency:.2f}ms")
    print(f"压缩率: {1 - nonzero_params/orig_params:.2%}")
    print(f"精度损失: {orig_acc - pruned_acc:.4f}")
    print(f"加速比: {orig_latency/pruned_latency:.2f}x")

四、模型压缩的架构权衡：精度、加速比与部署复杂度

剪枝的精度恢复瓶颈：50% 剪枝率下，微调通常能恢复大部分精度；70% 以上剪枝率时，微调的精度恢复越来越困难，因为被移除的参数中包含了不可替代的信息。建议从 30% 剪枝率开始，逐步增加直到精度下降不可接受。

蒸馏的 Student 架构选择：Student 太小（如 2 层 Transformer）无法学习 Teacher 的复杂表示，精度损失大；Student 太大（如与 Teacher 同构）则压缩效果有限。经验法则：Student 参数量约为 Teacher 的 1/4 到 1/2，层数约为 Teacher 的 1/2。

剪枝与蒸馏的组合顺序：先剪枝再蒸馏，Teacher 是原始大模型，Student 是剪枝后的模型——蒸馏帮助恢复剪枝损失的精度。先蒸馏再剪枝，先得到一个中等大小的 Student，再对 Student 剪枝——最终模型更小但精度损失可能更大。建议优先尝试"先剪枝再蒸馏"的路径。

部署端的稀疏推理支持：非结构化剪枝的加速依赖稀疏矩阵运算，但大多数推理引擎（ONNX Runtime、TensorRT）对稀疏运算的优化有限。实际加速比可能远低于理论压缩率。结构化剪枝虽然压缩率低，但加速比更可预测。

五、总结

模型压缩的两种路径各有适用场景。剪枝适合需要保留原模型架构的场景，结构化剪枝的加速比更可预测；蒸馏适合可以接受更小模型架构的场景，灵活性更高。落地时建议先尝试知识蒸馏（实现更简单、风险更低），如果压缩比不够再叠加结构化剪枝。压缩后的模型必须在实际业务数据上验证精度，不能只看公开数据集的结果。温度参数和 alpha 权重是蒸馏效果的关键超参数，需要网格搜索确定。