关注TechLead，复旦AI博士，分享AI领域全维度知识与研究。拥有10+年AI领域研究经验、复旦机器人智能实验室成员，国家级大学生赛事评审专家，发表多篇SCI核心期刊学术论文，上亿营收AI产品研发负责人。
如何在不牺牲性能的情况下将大型语言模型缩小十倍
虽然LLM的巨大规模赋予了它们在各种用例中的出色性能，但这也在其应用于现实世界问题时带来了挑战。在本文中，我将讨论如何通过压缩LLM来克服这些挑战。我将从概述关键概念开始，接着通过Python代码展示一个具体的示例。
2023年AI领域的口号是"越大越好"，提升语言模型的公式非常简单：更多的数据 + 更多的参数 + 更多的计算资源 = 更好的性能。
虽然这仍然是目前的趋势，但处理1000亿以上参数的模型显然存在挑战。例如，一个具有1000亿参数的模型仅在FP16格式下存储就需要200GB的空间！
不用说，大多数消费设备（如手机、平板电脑、笔记本电脑）无法处理如此庞大的模型。但……如果我们可以让模型变小呢？
模型压缩
模型压缩旨在在不牺牲性能的前提下减少机器学习模型的大小。对于（大型）神经网络，这可行，因为它们通常是过参数化的（即由冗余的计算单元组成）。
模型压缩的主要好处是降低推理成本。这意味着功能强大的ML模型可以更广泛地被使用（例如在本地笔记本电脑上运行LLM），并且能够以更低的成本将AI集成到消费产品中，还支持设备上的推理，从而保护用户隐私和安全。
三种压缩模型的方法
模型压缩有多种技术。这里我将重点介绍三种广泛使用的类别。
- 量化——使用更低精度的数据类型表示模型
- 剪枝——从模型中删除不必要的组件
- 知识蒸馏——通过较大的模型训练较小的模型
注意：这些方法是相互独立的。因此，可以组合来自多个类别的技术，以实现最大的压缩效果！
量化
虽然"量化"听起来像是一个复杂的术语，但它实际上是一个简单的概念。它包括降低模型参数的精度。你可以把它想象成将高分辨率图像转换为低分辨率图像，同时保持图片的核心属性。
两类常见的量化技术是后训练量化（PTQ）和量化感知训练（QAT）。
后训练量化（PTQ）
对于一个神经网络，后训练量化（PTQ）通过将参数替换为低精度的数据类型（例如从FP16转换为INT-8）来压缩模型。这是减少模型计算需求最快且最简单的方法之一，因为它不需要额外的训练或数据标注。
虽然这是降低模型成本的相对简单的方法，但以这种方式过度量化（例如从FP16到INT4）通常会导致性能下降，这限制了PTQ的潜在收益。
量化感知训练（QAT）
对于需要更高压缩的情况，可以通过从零开始用低精度数据类型训练模型来克服PTQ的局限性。这就是量化感知训练（QAT）的思想。
虽然这种方法在技术上更具挑战性，但它可以生成一个显著更小且表现良好的模型。例如，BitNet架构使用了一种三进制数据类型（即1.58位），匹配了原始Llama模型的性能！
当然，PTQ和从零开始的QAT之间存在较大的技术鸿沟。一种介于两者之间的方法是量化感知微调，它包括在量化后对预训练模型进行额外的训练。
剪枝
剪枝的目的是删除对模型性能影响较小的组件。这很有效，因为ML模型（尤其是大型模型）往往会学习到冗余和噪声结构。
一个比喻是修剪树木中的枯枝。去除它们可以减少树的大小而不会伤害树。
剪枝方法可以分为两类：非结构化剪枝和结构化剪枝。
非结构化剪枝
非结构化剪枝是从神经网络中移除不重要的权重（即将它们设置为零）。例如，早期的工作如Optimal Brain Damage和Optimal Brain Surgeon通过估计每个参数对损失函数的影响来计算其重要性分数。
最近，基于幅值的方法（即移除绝对值最小的权重）变得更受欢迎，因为它们简单且易于扩展。
虽然非结构化剪枝可以显著减少参数数量，但这些收益需要特殊硬件才能实现。非结构化剪枝会导致稀疏矩阵操作（即乘以大量零的矩阵），而标准硬件无法比非稀疏操作更有效地执行这些操作。
结构化剪枝
相对而言，结构化剪枝是从神经网络中删除整个结构（例如注意力头、神经元和层）。这样可以避免稀疏矩阵操作的问题，因为整个矩阵可以从模型中删除，而不是个别参数。
虽然有多种方法可以确定要剪枝的结构，但原则上，它们都试图删除对性能影响最小的结构。
知识蒸馏
知识蒸馏是将知识从一个（较大的）教师模型传递到一个（较小的）学生模型。一种方法是通过教师模型生成预测，并使用这些预测来训练学生模型。通过学习教师模型的输出logits（即所有可能的下一个标记的概率），学生模型获得了比原始训练数据更丰富的信息，从而提高了性能。
最近的蒸馏应用完全摒弃了logits的需求，而是通过教师模型生成的合成数据进行学习。一个著名的例子是斯坦福的Alpaca模型，它使用来自OpenAI的text-davinci-003（即原始ChatGPT模型）的合成数据微调了LLaMa 7B（基础）模型，使其能够遵循用户指令。
示例代码：通过知识蒸馏和量化压缩文本分类器
在基本了解了各种压缩技术后，让我们看一个如何在Python中进行压缩的实际示例。这里，我们将压缩一个有1亿参数的模型，该模型用于分类URL是否安全（即钓鱼网站）。
我们首先使用知识蒸馏将1亿参数的模型压缩到5000万参数。然后，使用4位量化进一步将内存占用减少3倍，得到的最终模型比原始模型小7倍。
示例代码可以在 github.com/ShawhinT/Yo… 上找到。教师模型 huggingface.co/shawhin/ber… 、学生模型huggingface.co/shawhin/ber… 4位学生模型 huggingface.co/shawhin/ber… 、数据集huggingface.co/datasets/sh… 可以在Hugging Face Hub上免费获取。
首先，我们导入一些有用的库。
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import DistilBertForSequenceClassification, DistilBertConfig
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
然后，我们从Hugging Face Hub加载数据集。包括训练集（2100行）、测试集（450行）和验证集（450行）。
data = load_dataset("shawhin/phishing-site-classification")
接下来，我们加载教师模型。为了加速训练，我将模型加载到Google Colab上提供的免费T4 GPU上。
device = torch.device('cuda')
model_path = "shawhin/bert-phishing-classifier_teacher"
tokenizer = AutoTokenizer.from_pretrained(model_path)
teacher_model = AutoModelForSequenceClassification.from_pre
trained(model_path).to(device)
教师模型是Google的 bert-base-uncased huggingface.co/google-bert… 的微调版本，执行对钓鱼网站URL的二分类。训练教师模型的代码可在 GitHub github.com/ShawhinT/Yo… 上找到。
对于学生模型，我们从 distilbert-base-uncased huggingface.co/distilbert/… 初始化一个新模型。我们通过移除两层和剩余层的四个注意力头修改了架构。
my_config = DistilBertConfig(n_heads=8, n_layers=4)
student_model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", config=my_config).to(device)
在训练学生模型之前，我们需要对数据集进行标记。这是必要的，因为模型期望输入文本以特定的方式表示。
在这里，我根据每个批次的最长示例填充样本。这使批次能够表示为PyTorch张量。
def preprocess_function(examples):
return tokenizer(examples["text"], padding='max_length', truncation=True)
tokenized_data = data.map(preprocess_function, batched=True)
tokenized_data.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
训练前的另一个重要步骤是定义一个评估策略，用于在训练期间评估我们的模型。在下面，我定义了一个函数，该函数在给定模型和数据集的情况下计算准确率、精确率、召回率和F1得分。
def evaluate_model(model, dataloader, device):
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
outputs = model(input_ids, attention_mask=attention_mask)
logits = outputs.logits
preds = torch.argmax(logits, dim=1).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.cpu().numpy())
accuracy = accuracy_score(all_labels, all_preds)
precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='binary')
return accuracy, precision, recall, f1
现在我们可以开始训练了。为了让学生模型同时学习训练集中的真实标签（即硬目标）和教师模型的logits（即软目标），我们需要构建一个特殊的损失函数，该函数考虑到两种目标。
这可以通过结合学生和教师输出概率分布的KL散度与学生logits与真实标签的交叉熵损失来实现。
def distillation_loss(student_logits, teacher_logits, true_labels, temperature, alpha):
soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)
student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)
distill_loss = nn.functional.kl_div(student_soft, soft_targets, reduction='batchmean') * (temperature ** 2)
hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)
loss = alpha * distill_loss + (1.0 - alpha) * hard_loss
return loss
接下来，我们定义超参数、优化器和训练/测试数据集。
batch_size = 32
lr = 1e-4
num_epochs = 5
temperature = 2.0
alpha = 0.5
optimizer = optim.Adam(student_model.parameters(), lr=lr)
dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)
test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)
最后，我们使用PyTorch训练学生模型。
student_model.train()
for epoch in range(num_epochs):
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
with torch.no_grad():
teacher_outputs = teacher_model(input_ids, attention_mask=attention_mask)
teacher_logits = teacher_outputs.logits
student_outputs = student_model(input_ids, attention_mask=attention_mask)
student_logits = student_outputs.logits
loss = distillation_loss(student_logits, teacher_logits, labels, temperature, alpha)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch + 1} completed with loss: {loss.item()}")
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 = evaluate_model(teacher_model, test_dataloader, device)
print(f"Teacher (test) - Accuracy: {teacher_accuracy:.4f}, Precision: {teacher_precision:.4f}, Recall: {teacher_recall:.4f}, F1 Score: {teacher_f1:.4f}")
student_accuracy, student_precision, student_recall, student_f1 = evaluate_model(student_model, test_dataloader, device)
print(f"Student (test) - Accuracy: {student_accuracy:.4f}, Precision: {student_precision:.4f}, Recall: {student_recall:.4f}, F1 Score: {student_f1:.4f}")
print("\n")
student_model.train()
训练结果如下图所示。令人惊讶的是，训练结束时，学生模型在所有评估指标上都超过了教师模型！
接下来，我们可以在独立的验证集上评估模型，即未用于训练模型参数或调整超参数的数据。
validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 = evaluate_model(teacher_model, validation_dataloader, device)
print(f"Teacher (validation) - Accuracy: {teacher_accuracy:.4f}, Precision: {teacher_precision:.4f}, Recall: {teacher_recall:.4f}, F1 Score: {teacher_f1:.4f}")
student_accuracy, student_precision, student_recall, student_f1 = evaluate_model(student_model, validation_dataloader, device)
print(f"Student (validation) - Accuracy: {student_accuracy:.4f}, Precision: {student_precision:.4f}, Recall: {student_recall:.4f}, F1 Score: {student_f1:.4f}")
在这里，我们再次看到学生模型优于教师模型。
到目前为止，我们已经将模型从1.09亿参数（438 MB）缩小到5280万参数（211 MB）。然而，我们可以更进一步，对学生模型进行量化。
首先，我们将模型推送到 Hugging Face Hub huggingface.co/shawhin/ber… 。
student_model.push_to_hub("shawhin/bert-phishing-classifier_student")
然后，我们可以使用4位量化加载模型。为此，我们可以使用transformers库中的BitsAndBytes集成。
我们设置配置以使用 QLoRA medium.com/towards-dat… 论文中描述的4位NormalFloat数据类型存储模型参数，并使用bfloat16进行计算。
from transformers import BitsAndBytesConfig
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True
)
model_nf4 = AutoModelForSequenceClassification.from_pretrained(model_id, device_map=device, quantization_config=nf4_config)
接下来，我们可以在验证集上评估量化模型。
quantized_accuracy, quantized_precision, quantized_recall, quantized_f1 = evaluate_model(model_nf4, validation_dataloader, device)
print("Post-quantization Performance")
print(f"Accuracy: {quantized_accuracy:.4f}, Precision: {quantized_precision:.4f}, Recall: {quantized_recall:.4f}, F1 Score: {quantized_f1:.4f}")
我们再次看到压缩后的性能略有提升（缩小到62.7MB）。一个直观的解释是奥卡姆剃刀原理，即越简单的模型越好。
在这种情况下，模型可能对于这个二分类任务来说过度参数化了。因此，简化模型带来了更好的性能。
