壓縮大型語言模型(LLMs)：縮小10倍、性能保持不變
盡管LLMs的巨大規模使其在廣泛的應用場景中表現卓越,但這也為其在實際問題中的應用帶來了挑戰。本文將探討如何通過壓縮LLMs來應對這些挑戰。我們將介紹關鍵概念,然后通過具體的Python代碼實例進行演示。
2023年人工智能領域的主導思想是"更大即更好",改進語言模型的方程相對簡單:更多數據 + 更多參數 + 更多計算資源 = 更優性能。
雖然這種思路可能仍然適用(GPT-5即將問世?),但使用超過100B參數的模型顯然面臨挑戰。例如一個使用FP16的100B參數模型僅存儲就需要200GB空間!
大多數消費級設備(如智能手機、平板電腦、筆記本電腦)無法處理如此龐大的模型。那么,我們是否可以在不損失性能的前提下縮小這些模型呢?
模型壓縮
模型壓縮旨在在保持性能的同時減小機器學習模型的規模[2]。這種方法對(大型)神經網絡特別有效,因為它們通常存在過度參數化的問題(即包含冗余計算單元)[3]。
模型壓縮的主要優勢在于降低推理成本。這意味著強大的機器學習模型可以更廣泛地應用(例如,在個人筆記本電腦上本地運行LLMs),將人工智能集成到消費產品中的成本降低,以及支持設備端推理,從而增強用戶隱私和安全性[3]。
壓縮模型的三種方法
模型壓縮技術多種多樣。本文將重點介紹三大類方法。
- 量化 — 使用低精度數據類型表示模型
- 剪枝 — 移除模型中不必要的組件
- 知識蒸餾 — 利用大型模型訓練小型模型
這些方法都相互獨立，可以結合多種技術以實現最大化壓縮效果!
1、量化
盡管"量化"這個術語聽起來可能晦澀復雜,但其核心概念相對簡單。它指的是降低模型參數的精度。可以將這個過程類比為將高分辨率圖像轉換為低分辨率圖像,同時保持圖像的主要特征。
量化技術主要分為兩類:訓練后量化(PTQ)和量化感知訓訓練(QAT)。
訓練后量化(PTQ)
對于給定的神經網絡,訓練后量化(PTQ)通過將參數替換為低精度數據類型來壓縮模型(例如,從FP16轉換為INT-8)。這是減少模型計算需求最快速和簡單的方法之一,因為它無需額外的訓練或數據標注。
雖然這是一種相對簡便的降低模型成本的方法,但過度使用這種技術進行量化(例如,從FP16轉換為INT4)通常會導致性能下降,這限制了PTQ的潛在收益。
量化感知訓練(QAT)
在需要更高壓縮率的情況下,可以通過使用低精度數據類型從頭開始訓練模型來克服PTQ的局限性。這就是量化感知訓練(QAT)的核心思想。
盡管這種方法在技術上更具挑戰性,但它可以產生顯著更小且性能良好的模型。例如,BitNet架構使用三元數據類型(即1.58位)就達到了與原始Llama LLM相當的性能!
PTQ和從頭開始的QAT之間存在較大的技術差距。介于兩者之間的一種方法是量化感知微調,它包括在量化后對預訓練模型進行額外的訓練[3]。
2、剪枝
剪枝的目標是移除對模型性能影響較小的組件[7]。這種方法之所以有效,是因為機器學習模型(特別是大型模型)往往會學習冗余和噪聲結構。
這個過程可以類比為修剪樹木中的枯枝。移除這些枯枝可以減小樹的體積而不會損害樹的健康。
剪枝方法可以分為兩類:非結構化剪枝和結構化剪枝。
非結構化剪枝
非結構化剪枝從神經網絡中移除不重要的權重(即將其值設為零)。早期的工作如Optimal Brain Damage和Optimal Brain Surgeon通過估計剪枝對損失函數的影響來計算網絡中每個參數的重要性分數。
最近基于幅度的方法(即移除絕對值最小的權重)因其簡單性和可擴展性而變得更加流行。
雖然非結構化剪枝的細粒度特性可以顯著減少參數數量,但這些收益通常需要專門的硬件才能實現。非結構化剪枝會導致稀疏矩陣運算(即乘以包含大量零的矩陣),而標準硬件在執行這類運算時并不比非稀疏運算更有效。
結構化剪枝
相比之下,結構化剪枝從神經網絡中移除整個結構(例如注意力頭、神經元和層)。這種方法避免了稀疏矩陣運算的問題,因為可以直接從模型中刪除整個矩陣,而不是單個參數。
雖然識別待剪枝結構的方法多種多樣,但其基本原則都是試圖移除對性能影響最小的結構。參考文獻[5]提供了結構化剪枝方法的詳細綜述。
3、知識蒸餾
知識蒸餾是一種將知識從(較大的)教師模型轉移到(較小的)學生模型的技術[5]。一種常見的實現方法是使用教師模型生成預測,然后用這些預測來訓練學生模型。從教師模型的輸出logits(即所有可能的下一個標記的概率)中學習,可以提供比原始訓練數據更豐富的信息,從而提高學生模型的性能[8]。
最新的蒸餾應用完全摒棄了對logits的依賴,轉而從教師模型生成的合成數據中學習。一個典型的例子是斯坦福大學的Alpaca模型,它使用OpenAI的text-davinci-003(即原始ChatGPT模型)生成的合成數據對LLaMa 7B(基礎)模型進行了微調,使其能夠遵循用戶指令[9]。
代碼示例：使用知識蒸餾和量化壓縮文本分類器
在了解了各種壓縮技術的基本原理后,讓我們通過一個Python實例來展示如何實際應用這些技術。在這個例子中,我們將壓縮一個具有100M參數的模型,該模型用于將URL分類為安全或不安全(即釣魚網站)。
我們首先使用知識蒸餾將100M參數模型壓縮為50M參數模型。然后,通過應用4位量化,我們進一步將內存占用減少了3倍,最終得到的模型比原始模型小7倍。
首先,我們導入必要的庫:
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上。
# 使用Nvidia GPU
device = torch.device('cuda')
# 加載教師模型和分詞器
model_path = "shawhin/bert-phishing-classifier_teacher"
tokenizer = AutoTokenizer.from_pretrained(model_path)
teacher_model = AutoModelForSequenceClassification.from_pretrained(model_path)
.to(device)教師模型是Google的bert-base-uncased模型的微調版本,用于對釣魚網站URL進行二元分類。
對于學生模型,我們基于distilbert-base-uncased從頭初始化一個新模型。我們通過移除兩層和減少剩余層中的四個注意力頭來修改架構。
# 加載學生模型
my_config = DistilBertConfig(n_heads=8, n_layers=4) # 每層減少4個頭,總共減少2層
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)
# 前向傳播獲取logits
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):
# 從教師logits計算軟目標
soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)
student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)
# 蒸餾的KL散度損失
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 + 1} 輪訓練完成,損失: {loss.item()}")
# 評估教師模型
teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =
evaluate_model(teacher_model, test_dataloader, device)
print(f"教師模型 (測試集) - 準確率: {teacher_accuracy:.4f},
精確率: {teacher_precision:.4f},
召回率: {teacher_recall:.4f},
F1分數: {teacher_f1:.4f}")
# 評估學生模型
student_accuracy, student_precision, student_recall, student_f1 =
evaluate_model(student_model, test_dataloader, device)
print(f"學生模型 (測試集) - 準確率: {student_accuracy:.4f},
精確率: {student_precision:.4f},
召回率: {student_recall:.4f},
F1分數: {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_accuracy:.4f},
精確率: {teacher_precision:.4f},
召回率: {teacher_recall:.4f},
F1分數: {teacher_f1:.4f}")
# 評估學生模型
student_accuracy, student_precision, student_recall, student_f1 =
evaluate_model(student_model, validation_dataloader, device)
print(f"學生模型 (驗證集) - 準確率: {student_accuracy:.4f},
精確率: {student_precision:.4f},
召回率: {student_recall:.4f},
F1分數: {student_f1:.4f}")我們再次觀察到學生模型的表現超過了教師模型。
到目前為止,我們已經將模型從109M參數(438 MB)壓縮到52.8M參數(211 MB)。我們還可以更進一步,對學生模型進行量化處理。
我們使用QLoRA論文中描述的4位NormalFloat數據類型存儲模型參數,并使用bfloat16進行計算。
from transformers import BitsAndBytesConfig
# 以4位精度加載模型
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("量化后性能")
print(f"準確率: {quantized_accuracy:.4f},
精確率: {quantized_precision:.4f},
召回率: {quantized_recall:.4f},
F1分數: {quantized_f1:.4f}")
量化后學生模型在驗證集上的表現。
再次觀察到壓縮后性能略有提升。這可以從奧卡姆剃刀原理的角度理解,該原理認為在其他條件相同的情況下,更簡單的模型通常更優。
在這個案例中,原始模型可能對這個二元分類任務而言過于復雜。簡化模型反而導致了性能的提升。
總結
盡管現代大型語言模型(LLMs)在各種任務上展現出卓越的性能,但它們的規模在實際部署中帶來了諸多挑戰。
近期模型壓縮技術的創新有助于通過降低LLM解決方案的計算成本來緩解這些挑戰。本文討論了三大類壓縮技術(量化、剪枝和知識蒸餾),并通過Python實例演示了它們的實際應用。
