5.1 模型蒸馏与压缩

大模型虽然性能强大，但其庞大的参数量和计算开销严重制约了在生产环境中的部署。一个拥有数十亿参数的 Transformer 模型需要数 GB 的显存，推理延迟也难以满足实时性要求。模型蒸馏与压缩技术旨在将大模型的能力迁移到小模型中，在保持性能的同时大幅降低部署成本。

5.1.1 知识蒸馏

Teacher-Student 框架

知识蒸馏（Knowledge Distillation）由 Geoffrey Hinton 于 2015 年提出，其核心思想是：一个大模型（Teacher）经过充分训练后，它的输出中蕴含的"知识"不仅包含最终预测结果，还包含类别之间的相似度关系。通过让学生模型（Student）学习教师模型的软化输出，学生可以在不显著降低性能的前提下大幅减小模型规模。

软标签与温度参数

标准分类器输出的是经过 Softmax 的概率分布，通常分布很尖锐（接近 one-hot）。知识蒸馏的关键在于引入温度参数 T 来软化概率分布：

softmax(z_i / T) = exp(z_i / T) / Σ exp(z_j / T)
​

当 T=1 时，输出的是正常的 Softmax 概率。当 T 增大时，分布变得更平滑，暴露出类别之间的相似度信息。

例如，一只猫的图片，标准 Softmax 可能输出 [猫: 0.95, 狗: 0.03, 老虎: 0.01, ...]，而 T=5 的软化输出可能是 [猫: 0.6, 狗: 0.15, 老虎: 0.1, ...]，清晰反映了猫和狗、猫和老虎之间的语义关系。

Transformer 蒸馏实现

HuggingFace Transformers 内置了蒸馏训练支持：

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
)

# Teacher 模型（大模型，已训练好，冻结参数）
teacher_model = AutoModelForSequenceClassification.from_pretrained(
    "bert-large-uncased",  # 340M 参数
    num_labels=2,
)
teacher_model.eval()  # 冻结，不参与训练
for param in teacher_model.parameters():
    param.requires_grad = False

# Student 模型（小模型，待训练）
student_model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased",  # 66M 参数（约为 Teacher 的 1/5）
    num_labels=2,
)

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
​

自定义蒸馏损失函数：

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

class DistillationTrainer(Trainer):
    def __init__(self, *args, teacher_model=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model
    
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        # Student 前向传播
        outputs = model(**inputs)
        student_logits = outputs.logits
        labels = inputs.pop("labels")
        
        # 标准交叉熵损失（与真实标签对比）
        ce_loss = F.cross_entropy(student_logits, labels)
        
        # 蒸馏损失（与 Teacher 的软化输出对比）
        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs)
            teacher_logits = teacher_outputs.logits
        
        temperature = 4.0  # 温度参数
        alpha = 0.5       # 蒸馏损失权重
        
        # KL 散度作为蒸馏损失
        kd_loss = F.kl_div(
            F.log_softmax(student_logits / temperature, dim=-1),
            F.softmax(teacher_logits / temperature, dim=-1),
            reduction="batchmean",
        ) * (temperature ** 2)
        
        # 总损失 = α × 蒸馏损失 + (1-α) × 交叉熵损失
        total_loss = alpha * kd_loss + (1 - alpha) * ce_loss
        
        return (total_loss, outputs) if return_outputs else total_loss
​

中间层蒸馏

除了输出层的软标签蒸馏，还可以让 Student 学习 Teacher 中间层的表示：

# 在蒸馏损失中增加中间层匹配损失
# hidden_loss = MSE(student_hidden_states[layer], teacher_hidden_states[layer])
# total_loss = alpha * kd_loss + beta * hidden_loss + (1-alpha-beta) * ce_loss
​

DistilBERT 就是使用这种方法训练的——通过同时匹配 BERT 的蒸馏损失、中间层 Cosine Embedding 损失和语言模型损失，将 BERT 压缩了 40% 的参数量，保留了 97% 的性能，推理速度提升 60%。

5.1.2 模型剪枝

结构化剪枝 vs 非结构化剪枝

非结构化剪枝：将权重矩阵中绝对值较小的单个参数设为零。稀疏度可以很高（90%+），但需要专用稀疏计算库才能获得实际加速。

结构化剪枝：移除整个神经元、注意力头或 Transformer 层。不会引入稀疏性，模型仍然保持密集计算，因此推理加速是确定性的。

Transformer 层剪枝

Transformer 模型不同层的贡献差异很大——研究表明，BERT 的中间层对最终性能贡献最大，而靠近输入和输出的几层贡献较小，可以安全移除：

from transformers import BertModel

model = BertModel.from_pretrained("bert-base-uncased")
print(f"原始层数: {len(model.encoder.layer)}")  # 12 层

# 移除第 0, 1, 10, 11 层（输入/输出层）
layers_to_keep = [2, 3, 4, 5, 6, 7, 8, 9]  # 保留 8 层
model.encoder.layer = nn.ModuleList([
    model.encoder.layer[i] for i in layers_to_keep
])
model.config.num_hidden_layers = len(layers_to_keep)

print(f"剪枝后层数: {len(model.encoder.layer)}")  # 8 层，减少 33%
​

注意力头剪枝

每个 Transformer 层有多个注意力头，研究发现部分注意力头是冗余的，可以直接移除：

# Michel et al. (2019) 的方法
# 评估每个注意力头的重要性，移除不重要的头
# BERT-base: 12层 × 12头 = 144 个注意力头
# 实验发现可以移除约 20-40% 的注意力头而不显著降低性能
​

迭代剪枝与微调

剪枝后的模型性能通常会下降，需要通过微调来恢复：

训练 → 剪枝 20% → 微调 → 剪枝 40% → 微调 → 剪枝 60% → 微调
​

这种渐进式剪枝（Gradual Pruning）策略比一次性大幅度剪枝效果更好。

5.1.3 模型量化

量化基础

模型量化将浮点数权重映射到低精度表示，是部署加速最直接的手段：

• FP32：标准 32 位浮点数，每个参数 4 字节
• FP16：半精度浮点数，每个参数 2 字节，速度提升约 2 倍
• INT8：8 位整数，每个参数 1 字节，速度提升约 4 倍
• INT4：4 位整数，每个参数 0.5 字节，速度提升约 8 倍，但精度损失较大

训练后量化（PTQ）

训练后量化不需要重新训练，直接将训练好的模型参数转换为目标精度：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")

# 动态量化（运行时计算量化参数）
# 适用于 LSTM、Transformer 等模型
dynamically_quantized = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# 模型大小对比
import os
def get_size(model):
    torch.save(model.state_dict(), "/tmp/tmp.pt")
    size = os.path.getsize("/tmp/tmp.pt") / 1024 / 1024
    os.remove("/tmp/tmp.pt")
    return size

print(f"FP32: {get_size(model):.1f} MB")
print(f"INT8: {get_size(dynamically_quantized):.1f} MB")
# FP32: 1100 MB, INT8: ~280 MB（约 1/4）
​

GPTQ 量化

GPTQ（GPT Quantization）是一种先进的量化方法，专门针对 Transformer 类大语言模型设计，可以在 3-bit 或 4-bit 量化下保持几乎无损的性能：

# 使用 auto-gptq 进行 GPTQ 量化
# pip install auto-gptq

from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# GPTQ 量化配置
gptq_config = GPTQConfig(
    bits=4,           # 4-bit 量化
    group_size=128,    # 分组量化
    desc_act=True,     # 激活感知量化
    dataset="c4",      # 校准数据集
)

# 加载并量化模型
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=gptq_config,
    torch_dtype=torch.float16,
    device_map="auto",
)

print(f"原始模型: ~14 GB (FP16)")
print(f"GPTQ-4bit: ~4 GB（约 1/3.5）")
​

AWQ 量化

AWQ（Activation-aware Weight Quantization）是一种激活感知的权重量化方法，通过分析输入激活的分布来确定哪些权重通道最重要，对重要通道保留更高精度：

# AWQ 量化（通常效果优于 GPTQ）
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config={
        "quant_method": "awq",
        "bits": 4,
        "group_size": 128,
    },
    device_map="auto",
)
​

AWQ 和 GPTQ 在 4-bit 量化下的性能损失通常小于 1%，而模型大小缩减到原来的 1/4，是当前大模型部署的主流方案。

5.1.4 边缘设备部署

ONNX 导出与优化

将 PyTorch 模型导出为 ONNX 格式，利用推理引擎优化：

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import pipeline

# 导出并优化 ONNX 模型
ort_model = ORTModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english",
    export=True,
)

# 推理
classifier = pipeline("text-classification", model=ort_model)
result = classifier("This movie is great!")
​

移动端部署

对于手机、树莓派等边缘设备，可以考虑更激进的压缩方案：

# 使用 TensorFlow Lite 部署到移动端
# 支持的优化：
# 1. INT8 量化（CPU 加速 3-4 倍）
# 2. 16-bit 浮点（GPU 加速 2 倍）
# 3. 模型剪枝（减少计算量）
# 4. Neural Architecture Search（搜索最优子网络）

# HuggingFace 导出为 TFLite 格式
# 使用 Optimum TFLite exporter
​

压缩方案选择指南

模型压缩不是零和博弈，通过合理组合蒸馏、剪枝和量化技术，可以在特定资源约束下找到性能与效率的最优平衡点。