知識蒸餾技術的核心目標是將大型教師模型的知識遷移至小型學生模型，以降低運算成本和資源消耗。這種方法特別適用於資源受限的場景，例如移動裝置或嵌入式系統。透過模仿教師模型的輸出機率分佈，學生模型可以學習到更豐富的資訊，從而提升效能。在實務應用中，我們通常會選擇與教師模型結構相似的學生模型，以最大程度地保留知識遷移的有效性。例如，使用 DistilBERT 作為 BERT 的學生模型，可以有效地壓縮模型大小，同時保持良好的效能表現。
提升模型效率：知識蒸餾技術詳解
在前面的章節中，我們已經完成了 PerformanceBenchmark 類別的實作，並且對 BERT 基線模型進行了基準測試，獲得了模型的效能指標。現在，我們將探討第一種模型壓縮技術：知識蒸餾（Knowledge Distillation）。
知識蒸餾的基本原理
知識蒸餾是一種通用方法，用於訓練一個較小的學生模型來模仿一個較慢、較大但效能更好的教師模型的行為。這種方法最初於 2006 年在整合模型中被提出，後來在 2015 年的一篇著名論文中被推廣到深度神經網路，並應用於影像分類別和自動語音識別。
知識蒸餾在微調中的應用
在監督任務（如微調）中，知識蒸餾的主要思想是將地面真實標籤（ground truth labels）與教師模型的「軟機率」（soft probabilities）分佈相結合，為學生模型提供額外的學習資訊。例如，如果 BERT-base 分類別器為多個意圖分配高機率，那麼這可能是這些意圖在特徵空間中彼此接近的跡象。透過訓練學生模型來模仿這些機率，目標是將教師模型學習到的「暗知識」（dark knowledge）——即無法從標籤中獲得的知識——傳遞給學生模型。
知識蒸餾的數學原理
假設我們將輸入序列 x 輸入教師模型，生成一個 logits 向量 𝐀 x = [z1 x, ..., zN x]。我們可以透過應用 softmax 函式將這些 logits 轉換為機率：
import numpy as np
def softmax(logits):
exp_logits = np.exp(logits)
return exp_logits / np.sum(exp_logits)
# 示例 logits
logits = np.array([1.0, 2.0, 3.0])
probabilities = softmax(logits)
print(probabilities)
內容解密：
此程式碼定義了一個 softmax 函式，用於將輸入的 logits 向量轉換為機率分佈。首先，計算 logits 的指數值，然後將每個指數值除以所有指數值的總和，得到歸一化的機率。
然而，在許多情況下，教師模型會為一個類別分配高機率，而其他類別的機率接近零。這時，教師模型並沒有提供太多超出地面真實標籤的資訊。因此，我們透過在應用 softmax 之前使用溫度超引數 T 對 logits 進行縮放來「軟化」機率：
def softened_softmax(logits, temperature):
exp_logits = np.exp(logits / temperature)
return exp_logits / np.sum(exp_logits)
# 示例 logits 和溫度
logits = np.array([1.0, 2.0, 3.0])
temperature = 2.0
softened_probabilities = softened_softmax(logits, temperature)
print(softened_probabilities)
內容解密：
此程式碼定義了一個 softened_softmax 函式，用於計算軟化的機率分佈。透過將 logits 除以溫度 T，然後應用 softmax 函式，可以獲得更平滑的機率分佈。較高的溫度值會產生更軟化的機率分佈，揭示了教師模型學習到的決策邊界的更多資訊。
Kullback-Leibler 發散度
由於學生模型也會產生自己的軟化機率 qi(x)，我們可以使用 Kullback-Leibler（KL）發散度來衡量兩個機率分佈之間的差異：
import numpy as np
def kl_divergence(p, q):
return np.sum(p * np.log(p / q))
# 示例機率分佈
p = np.array([0.2, 0.3, 0.5])
q = np.array([0.1, 0.4, 0.5])
kl_div = kl_divergence(p, q)
print(kl_div)
內容解密：
此程式碼定義了一個 kl_divergence 函式，用於計算兩個機率分佈 p 和 q 之間的 KL 發散度。KL 發散度衡量了兩個分佈之間的差異，用於訓練學生模型以模仿教師模型的行為。
透過最小化 KL 發散度，學生模型可以學習到教師模型的「暗知識」，從而提高自己的效能。知識蒸餾技術在壓縮大型預訓練語言模型方面具有廣泛的應用前景。
知識蒸餾技術在模型壓縮中的應用
知識蒸餾是一種有效的模型壓縮技術，透過將大型教師模型的知識轉移至小型學生模型，從而實作模型的輕量化。在本章中，我們將探討知識蒸餾的基本原理、實作方法以及在實際應用中的案例。
知識蒸餾的基本原理
知識蒸餾的核心思想是利用教師模型的輸出機率分佈來指導學生模型的訓練。教師模型的輸出包含了豐富的知識和資訊，透過將這些知識轉移至學生模型，可以提高學生模型的效能。
知識蒸餾損失函式
知識蒸餾損失函式（$L_{KD}$）用於衡量學生模型與教師模型之間的差異。其定義如下：
$L_{KD} = T^2 \cdot D_{KL}$
其中，$T$ 是溫度引數，用於控制輸出機率分佈的平滑程度；$D_{KL}$ 是 KL 散度，用於衡量兩個機率分佈之間的差異。
學生模型損失函式
在分類別任務中，學生模型的損失函式是交叉熵損失函式（$L_{CE}$）與知識蒸餾損失函式（$L_{KD}$）的加權平均：
$L_{student} = \alpha \cdot L_{CE} + (1 - \alpha) \cdot L_{KD}$
其中，$\alpha$ 是超引數，用於控制兩個損失函式之間的權重。
知識蒸餾在預訓練中的應用
知識蒸餾不僅可以用於微調階段，也可以用於預訓練階段，以建立一個通用型的學生模型。DistilBERT 是一個典型的例子，它透過知識蒸餾技術將 BERT 模型的知識轉移至一個更小的模型。
DistilBERT 損失函式
DistilBERT 的損失函式結合了掩碼語言建模損失函式（$L_{mlm}$）、知識蒸餾損失函式（$L_{KD}$）和餘弦嵌入損失函式（$L_{cos}$）：
$L_{DistilBERT} = \alpha \cdot L_{mlm} + \beta \cdot L_{KD} + \gamma \cdot L_{cos}$
實作知識蒸餾訓練器
要實作知識蒸餾，我們需要建立一個自定義的訓練器。以下是使用 PyTorch 實作的示例程式碼：
from transformers import TrainingArguments, Trainer
import torch.nn as nn
import torch.nn.functional as F
class DistillationTrainingArguments(TrainingArguments):
def __init__(self, *args, alpha=0.5, temperature=2.0, **kwargs):
super().__init__(*args, **kwargs)
self.alpha = alpha
self.temperature = temperature
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):
outputs_stu = model(**inputs)
loss_ce = outputs_stu.loss
logits_stu = outputs_stu.logits
with torch.no_grad():
outputs_tea = self.teacher_model(**inputs)
logits_tea = outputs_tea.logits
loss_fct = nn.KLDivLoss(reduction="batchmean")
loss_kd = self.args.temperature ** 2 * loss_fct(
F.log_softmax(logits_stu / self.args.temperature, dim=-1),
F.softmax(logits_tea / self.args.temperature, dim=-1))
loss = self.args.alpha * loss_ce + (1. - self.args.alpha) * loss_kd
return (loss, outputs_stu) if return_outputs else loss
內容解密：
DistillationTrainingArguments類別繼承自TrainingArguments，新增了alpha和temperature兩個超引數，用於控制知識蒸餾的權重和溫度。DistillationTrainer類別繼承自Trainer，新增了teacher_model屬性，用於儲存教師模型。- 在
compute_loss方法中，首先計算學生模型的輸出和交叉熵損失函式。 - 然後，使用教師模型計算輸出，並使用 KL 散度計算知識蒸餾損失函式。
- 最後，將交叉熵損失函式和知識蒸餾損失函式加權平均，得到最終的損失函式。
知識蒸餾最佳化模型效能
選擇適當的學生模型初始化
在知識蒸餾的過程中，選擇一個合適的學生模型至關重要。一般來說，我們會選擇比教師模型更小的模型，以降低延遲和記憶體佔用。文獻中的一個經驗法則是，當教師和學生模型屬於同一型別時，知識蒸餾的效果最佳。
為什麼選擇 DistilBERT？
由於我們的教師模型是 BERT，因此 DistilBERT 成為初始化學生模型的自然選擇。DistilBERT 的引數比 BERT 少 40%，並且在下游任務中表現出色。首先，我們需要對查詢進行標記化和編碼。讓我們例項化 DistilBERT 的標記器，並建立一個簡單的 tokenize_text() 函式來處理預處理：
from transformers import AutoTokenizer
student_ckpt = "distilbert-base-uncased"
student_tokenizer = AutoTokenizer.from_pretrained(student_ckpt)
def tokenize_text(batch):
return student_tokenizer(batch["text"], truncation=True)
clinc_enc = clinc.map(tokenize_text, batched=True, remove_columns=["text"])
clinc_enc = clinc_enc.rename_column("intent", "labels")
處理文字資料
在這裡，我們刪除了文字列，因為我們不再需要它，並且將 intent 列重新命名為 labels，以便訓練器可以自動檢測它。
定義超引數和計算指標
接下來，我們需要定義超引數和 compute_metrics() 函式，用於我們的 DistillationTrainer。我們還將把所有模型推播到 Hugging Face Hub，因此讓我們先登入到我們的帳戶：
from huggingface_hub import notebook_login
notebook_login()
定義計算指標函式
我們將使用準確率作為主要指標，這意味著我們可以重用 accuracy_score() 函式，在 compute_metrics() 函式中：
def compute_metrics(pred):
predictions, labels = pred
predictions = np.argmax(predictions, axis=1)
return accuracy_score.compute(predictions=predictions, references=labels)
設定訓練引數
batch_size = 48
finetuned_ckpt = "distilbert-base-uncased-finetuned-clinc"
student_training_args = DistillationTrainingArguments(
output_dir=finetuned_ckpt,
evaluation_strategy="epoch",
num_train_epochs=5,
learning_rate=2e-5,
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
alpha=1,
weight_decay=0.01,
push_to_hub=True
)
初始化學生模型
為了初始化學生模型，我們建立了一個 student_init() 函式，該函式將在每次呼叫 train() 方法時初始化一個新的模型。我們還需要提供學生模型與每個 intent 和 label ID 之間的對映。這些對映可以從我們下載的 BERT-base 模型中獲得：
id2label = pipe.model.config.id2label
label2id = pipe.model.config.label2id
from transformers import AutoConfig
num_labels = intents.num_classes
student_config = AutoConfig.from_pretrained(
student_ckpt,
num_labels=num_labels,
id2label=id2label,
label2id=label2id
)
建立學生模型
import torch
from transformers import AutoModelForSequenceClassification
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def student_init():
return AutoModelForSequenceClassification.from_pretrained(
student_ckpt,
config=student_config
).to(device)
執行知識蒸餾訓練
現在我們已經準備好了所有必要的元件，讓我們載入教師模型並進行微調：
teacher_ckpt = "transformersbook/bert-base-uncased-finetuned-clinc"
teacher_model = AutoModelForSequenceClassification.from_pretrained(
teacher_ckpt,
num_labels=num_labels
).to(device)
distilbert_trainer = DistillationTrainer(
model_init=student_init,
teacher_model=teacher_model,
args=student_training_args,
train_dataset=clinc_enc['train'],
eval_dataset=clinc_enc['validation'],
compute_metrics=compute_metrics,
tokenizer=student_tokenizer
)
distilbert_trainer.train()
訓練結果
| Epoch | Training Loss | Validation Loss | Accuracy |
|---|---|---|---|
| 1 | 4.2923 | 3.289337 | 0.742258 |
| 2 | 2.6307 | 1.883680 | 0.828065 |
| 3 | 1.5483 | 1.158315 | 0.896774 |
| 4 | 1.0153 | 0.861815 | 0.909355 |
| 5 | 0.7958 | 0.777289 | 0.917419 |
驗證集上的 92% 準確率與 BERT-base 教師模型的 94% 相比已經相當不錯。現在我們已經微調了 DistilBERT，讓我們將模型推播到 Hub，以便稍後重用：
結果分析
透過知識蒸餾，我們成功地將一個較大的 BERT 模型的知識轉移到一個較小的 DistilBERT 模型上，並在 CLINC 資料集上取得了良好的效能。這種方法不僅減小了模型的規模，還保留了大部分的效能，使其更適合佈署在資源受限的環境中。
程式碼重點解析：
tokenize_text函式：使用AutoTokenizer對文字進行標記化處理，確保輸入資料格式正確。compute_metrics函式：計算模型的準確率，作為評估指標。DistillationTrainingArguments：定義了訓練過程中的超引數，如學習率、訓練輪數等。student_init函式：動態初始化學生模型，確保每次訓練都使用新的模型例項。DistillationTrainer：整合了教師模型、學生模型和訓練引數，負責執行知識蒸餾的訓練過程。
知識蒸餾的優勢：
- 模型壓縮：透過將大型模型的知識轉移到小型模型上，實作了模型的壓縮。
- 效能保持：儘管模型規模減小，但在目標任務上的效能仍然保持在較高水平。
- 佈署友好：較小的模型更適合在資源有限的環境中佈署，如移動裝置或嵌入式系統。
未來改進方向：
- 進一步最佳化超引數：透過調整訓練引數，如學習率、批次大小等，進一步提升模型的效能。
- 嘗試不同的學生模型：探索其他小型模型結構，以找到更優的效能與效率平衡點。
- 擴充套件到其他任務：將知識蒸餾的方法應用到其他 NLP 任務中，驗證其通用性。
