詳解模型蒸餾，破解DeepSeek性能謎題
不少關(guān)注 DeepSeek 最新動(dòng)態(tài)的朋友，想必都遇到過(guò) “Distillation”（蒸餾）這一術(shù)語(yǔ)。但它究竟指代何種技術(shù)，又為何在 AI 領(lǐng)域占據(jù)重要地位呢？
本文為你深度剖析模型蒸餾技術(shù)的原理，同時(shí)借助 TensorFlow 框架中的實(shí)例進(jìn)行詳細(xì)演示。相信通過(guò)閱讀本文，你會(huì)對(duì)模型蒸餾有全新的認(rèn)識(shí)，輕松解鎖深度學(xué)習(xí)優(yōu)化的新視角。
1 模型蒸餾原理
在深度學(xué)習(xí)領(lǐng)域，模型蒸餾是優(yōu)化模型的關(guān)鍵技術(shù)。它讓小的學(xué)生模型不再單純依賴原始標(biāo)簽，而是基于大的教師模型軟化后的概率輸出進(jìn)行訓(xùn)練。
以圖像分類為例，普通模型只是簡(jiǎn)單判斷圖像內(nèi)容，而運(yùn)用模型蒸餾技術(shù)的學(xué)生模型，能從教師模型的置信度分?jǐn)?shù)（如80%是狗，15%是貓，5%是狐貍）中獲取更豐富信息，從而保留更細(xì)致知識(shí)。
這樣一來(lái)，學(xué)生模型能用更少參數(shù)實(shí)現(xiàn)與教師模型相近的性能，在保持高精度的同時(shí)，減小模型規(guī)模、降低計(jì)算需求，為深度學(xué)習(xí)模型優(yōu)化開(kāi)辟了新路徑。
讓我們通過(guò)一個(gè)例子來(lái)看看具體是如何操作的。以使用MNIST數(shù)據(jù)集訓(xùn)練卷積神經(jīng)網(wǎng)絡(luò)（CNN）為例。
MNIST （Modified National Institute of Standards and Technology database）數(shù)據(jù)集在機(jī)器學(xué)習(xí)和計(jì)算機(jī)視覺(jué)里常用，有 70,000 張 28x28 像素的手寫數(shù)字（0 - 9）灰度圖，60,000 張訓(xùn)練圖、10,000 張測(cè)試圖。
模型蒸餾要先建教師模型，是用 MNIST 數(shù)據(jù)集訓(xùn)練的 CNN，參數(shù)多、結(jié)構(gòu)復(fù)雜。
再建個(gè)更簡(jiǎn)單、規(guī)模更小的學(xué)生模型。
目的是讓學(xué)生模型模仿教師模型性能，還能減少計(jì)算量和訓(xùn)練時(shí)間。
訓(xùn)練時(shí)，兩個(gè)模型都用 MNIST 數(shù)據(jù)集預(yù)測(cè)，接著算它們輸出的 Kullback-Leibler（KL）散度。這個(gè)值能確定梯度，指導(dǎo)調(diào)整學(xué)生模型。
一番操作后，學(xué)生模型就能達(dá)到和教師模型差不多的準(zhǔn)確率，成功 “出師”。
2 用 TensorFlow 和 MNIST 構(gòu)建模型
接下來(lái)，我們借助 TensorFlow 和 MNIST 數(shù)據(jù)集，搭建一個(gè)模型蒸餾示例項(xiàng)目。
先訓(xùn)練一個(gè)教師模型，再通過(guò)模型蒸餾技術(shù)，訓(xùn)練出一個(gè)更小的學(xué)生模型。這個(gè)學(xué)生模型能模仿教師模型的性能，而且所需資源更少。
2.1 使用MNIST數(shù)據(jù)集
確保你已經(jīng)安裝了TensorFlow：
!pip install tensorflow然后加載MNIST數(shù)據(jù)集：
from tensorflow import keras
import matplotlib.pyplot as plt
# 加載數(shù)據(jù)集（MNIST）
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
fig = plt.figure()
# 可視化部分?jǐn)?shù)字
for i in range(9):
plt.subplot(3,3,i+1)
plt.tight_layout()
plt.imshow(x_train[i], interpolatinotallow='none')
plt.title("Digit: {}".format(y_train[i]))
# 不顯示x軸和y軸刻度
plt.xticks([])
plt.yticks([])以下是MNIST數(shù)據(jù)集中的前9個(gè)樣本數(shù)字及其標(biāo)簽：
還需要對(duì)圖像數(shù)據(jù)進(jìn)行歸一化處理，并擴(kuò)展數(shù)據(jù)集的維度，為訓(xùn)練做準(zhǔn)備：
import tensorflow as tf
import numpy as np
# 歸一化圖像
x_train, x_test = x_train / 255.0, x_test / 255.0
# 為卷積神經(jīng)網(wǎng)絡(luò)擴(kuò)展維度
x_train = np.expand_dims(x_train, axis=-1)
x_test = np.expand_dims(x_test, axis=-1)
# 將標(biāo)簽轉(zhuǎn)換為分類（獨(dú)熱編碼）
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)2.2 定義教師模型
在基于模型蒸餾的示例項(xiàng)目構(gòu)建中，定義并訓(xùn)練教師模型是關(guān)鍵的環(huán)節(jié)。這里，我們構(gòu)建一個(gè)多層卷積神經(jīng)網(wǎng)絡(luò)（CNN）作為教師模型。
代碼如下：
# 教師模型
teacher_model = keras.Sequential([
keras.layers.Conv2D(32, (3, 3), activatinotallow='relu', input_shape=(28, 28, 1)),
keras.layers.MaxPooling2D((2, 2)),
keras.layers.Conv2D(64, (3, 3), activatinotallow='relu'),
keras.layers.MaxPooling2D((2, 2)),
keras.layers.Flatten(),
keras.layers.Dense(128, activatinotallow='relu'),
keras.layers.Dense(10) # 不使用softmax，輸出原始logits用于蒸餾
])需要注意的是，模型最后一層設(shè)置了 10 個(gè)單元，對(duì)應(yīng) 0 - 9 這 10 個(gè)數(shù)字，但未采用 softmax 激活函數(shù)，而是輸出原始的 logits。這一設(shè)計(jì)對(duì)于模型蒸餾很重要，因?yàn)樵诤罄m(xù)的蒸餾過(guò)程里，會(huì)借助 softmax 函數(shù)來(lái)計(jì)算教師模型與學(xué)生模型之間的 Kullback-Leibler（KL）散度，以此衡量二者差異，為學(xué)生模型的優(yōu)化提供方向。
完成模型定義后，要使用compile()方法對(duì)其進(jìn)行配置，設(shè)置優(yōu)化器、損失函數(shù)以及評(píng)估指標(biāo)：
teacher_model.compile(
optimizer = 'adam',
loss = tf.keras.losses.CategoricalCrossentropy(from_logits = True),
metrics = ['accuracy']
)配置完成，就可以使用fit()方法啟動(dòng)模型訓(xùn)練：
# 訓(xùn)練教師模型
teacher_model.fit(x_train, y_train,
epochs = 5,
batch_size = 64,
validation_data = (x_test, y_test))本次訓(xùn)練設(shè)定了 5 個(gè)訓(xùn)練周期，訓(xùn)練過(guò)程中的詳細(xì)信息如下：
Epoch 1/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 8s 8ms/step - accuracy: 0.8849 - loss: 0.3798 - val_accuracy: 0.9844 - val_loss: 0.0504
Epoch 2/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 9ms/step - accuracy: 0.9847 - loss: 0.0494 - val_accuracy: 0.9878 - val_loss: 0.0361
Epoch 3/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 10ms/step - accuracy: 0.9907 - loss: 0.0302 - val_accuracy: 0.9898 - val_loss: 0.0316
Epoch 4/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 10s 10ms/step - accuracy: 0.9928 - loss: 0.0223 - val_accuracy: 0.9895 - val_loss: 0.0303
Epoch 5/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 10s 11ms/step - accuracy: 0.9935 - loss: 0.0197 - val_accuracy: 0.9919 - val_loss: 0.0230從這些數(shù)據(jù)中，我們可以直觀地看到模型在訓(xùn)練過(guò)程中的準(zhǔn)確率和損失變化，了解模型的學(xué)習(xí)效果，為后續(xù)的模型優(yōu)化和評(píng)估提供依據(jù)。
2.3 定義學(xué)生模型
教師模型訓(xùn)練完成后，就該定義學(xué)生模型了。與教師模型相比，學(xué)生模型的架構(gòu)更簡(jiǎn)單，層數(shù)更少：
# 學(xué)生模型
student_model = keras.Sequential([
keras.layers.Conv2D(16, (3, 3), activatinotallow='relu', input_shape=(28, 28, 1)),
keras.layers.MaxPooling2D((2, 2)),
keras.layers.Flatten(),
keras.layers.Dense(64, activatinotallow='relu'),
keras.layers.Dense(10) # 不使用softmax，輸出原始logits用于蒸餾
])2.4 定義蒸餾損失函數(shù)
在模型蒸餾的實(shí)現(xiàn)過(guò)程中，distillation_loss() 函數(shù)發(fā)揮著核心作用，它借助 Kullback-Leibler（KL）散度來(lái)精準(zhǔn)計(jì)算蒸餾損失，從而推動(dòng)學(xué)生模型向教師模型 “看齊”。下面，我們就來(lái)詳細(xì)解讀這個(gè)函數(shù)的代碼實(shí)現(xiàn)。
蒸餾損失函數(shù)的計(jì)算，依賴于教師模型和學(xué)生模型的預(yù)測(cè)結(jié)果，具體步驟如下：
- 使用教師模型為輸入批次生成軟目標(biāo)（概率）。
- 使用學(xué)生模型的預(yù)測(cè)計(jì)算其軟概率。
- 計(jì)算教師模型和學(xué)生模型軟概率之間的Kullback-Leibler（KL）散度。
- 返回蒸餾損失。
軟概率和常見(jiàn)的硬標(biāo)簽不同。硬標(biāo)簽是明確分類，像判斷郵件是否為垃圾郵件，結(jié)果只有 “是”（1）或 “否”（0）。而軟概率會(huì)給出多種結(jié)果的概率，比如某郵件是垃圾郵件的概率為 0.85，不是的概率為 0.15，能更全面反映模型判斷。
計(jì)算軟概率要用到 softmax 函數(shù)，且受溫度參數(shù)影響。在知識(shí)蒸餾里，教師模型的軟概率包含類間豐富信息，學(xué)生模型學(xué)習(xí)后，能提升泛化能力和性能，更好地模仿教師模型。
以下是distillation_loss()函數(shù)的定義：
def distillation_loss(y_true, y_pred, x_batch, teacher_model, temperature=5):
"""
使用KL散度計(jì)算蒸餾損失。
"""
# 計(jì)算當(dāng)前批次的教師模型logits
teacher_logits = teacher_model(x_batch, training=False)
# 將logits轉(zhuǎn)換為軟概率
teacher_probs = tf.nn.softmax(teacher_logits / temperature)
student_probs = tf.nn.softmax(y_pred / temperature)
# KL散度損失（教師模型和學(xué)生模型分布之間的差異）
return tf.reduce_mean(tf.keras.losses.KLDivergence()(teacher_probs, student_probs))Kullback-Leibler（KL）散度，也稱為相對(duì)熵，用于衡量一個(gè)概率分布與另一個(gè)參考概率分布之間的差異。
2.5 使用知識(shí)蒸餾訓(xùn)練學(xué)生模型
現(xiàn)在你已經(jīng)準(zhǔn)備好使用知識(shí)蒸餾來(lái)訓(xùn)練學(xué)生模型了。首先，定義train_step()函數(shù)：
optimizer = tf.keras.optimizers.Adam()
@tf.function
def train_step(x_batch, y_batch, student_model, teacher_model):
with tf.GradientTape() as tape:
# 獲取學(xué)生模型的預(yù)測(cè)
student_preds = student_model(x_batch, training=True)
# 計(jì)算蒸餾損失（顯式傳入教師模型）
loss = distillation_loss(y_batch, student_preds, x_batch, teacher_model, temperature=5)
# 計(jì)算梯度
gradients = tape.gradient(loss, student_model.trainable_variables)
# 應(yīng)用梯度 - 訓(xùn)練學(xué)生模型
optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))
return loss這個(gè)函數(shù)執(zhí)行單個(gè)訓(xùn)練步驟：
- 計(jì)算學(xué)生模型的預(yù)測(cè)。
- 使用教師模型的預(yù)測(cè)計(jì)算蒸餾損失。
- 計(jì)算梯度并更新學(xué)生模型的權(quán)重。
為了訓(xùn)練學(xué)生模型，需要?jiǎng)?chuàng)建一個(gè)訓(xùn)練循環(huán)，遍歷數(shù)據(jù)集，在每一步更新學(xué)生模型的權(quán)重，并在每個(gè)訓(xùn)練周期結(jié)束時(shí)打印損失，以監(jiān)控訓(xùn)練進(jìn)度：
# 訓(xùn)練循環(huán)
epochs = 5
batch_size = 32
# 準(zhǔn)備數(shù)據(jù)集批次
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
for epoch in range(epochs):
total_loss = 0
num_batches = 0
for x_batch, y_batch in train_dataset:
loss = train_step(x_batch, y_batch, student_model, teacher_model)
total_loss += loss.numpy()
num_batches += 1
avg_loss = total_loss / num_batches
print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")
print("Student Model Training Complete!")訓(xùn)練后，應(yīng)該會(huì)看到類似以下的輸出：
Epoch 1, Loss: 0.1991
Epoch 2, Loss: 0.0588
Epoch 3, Loss: 0.0391
Epoch 4, Loss: 0.0274
Epoch 5, Loss: 0.0236
Student Model Training Complete!2.6 評(píng)估學(xué)生模型
學(xué)生模型已經(jīng)訓(xùn)練完成，可以使用測(cè)試集（x_test和y_test）對(duì)其進(jìn)行評(píng)估，觀察其性能：
student_model.compile(
optimizer='adam',
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
metrics=['accuracy']
)
student_acc = student_model.evaluate(x_test, y_test, verbose=0)[1]
print(f"Student Model Accuracy: {student_acc:.4f}")如預(yù)期，學(xué)生模型達(dá)到了相當(dāng)不錯(cuò)的準(zhǔn)確率：
Student Model Accuracy: 0.98632.7 使用教師模型和學(xué)生模型進(jìn)行預(yù)測(cè)
現(xiàn)在可以使用教師模型和學(xué)生模型進(jìn)行一些預(yù)測(cè)，看看是否都能準(zhǔn)確預(yù)測(cè)MNIST測(cè)試數(shù)據(jù)集中的數(shù)字：
import numpy as np
_, (x_test, y_test) = keras.datasets.mnist.load_data()
for index in range(5):
plt.figure(figsize=(2, 2))
plt.imshow(x_test[index], interpolatinotallow='none')
plt.title("Digit: {}".format(y_test[index]))
# 不顯示x軸和y軸刻度
plt.xticks([])
plt.yticks([])
plt.show()
# 現(xiàn)在可以進(jìn)行預(yù)測(cè)
x = x_test[index].reshape(1,28,28,1)
predictions = teacher_model.predict(x)
print(predictions)
print("Predicted value by teacher model: ", np.argmax(predictions, axis=-1))
predictions = student_model.predict(x)
print(predictions)
print("Predicted value by student model: ", np.argmax(predictions, axis=-1))以下是前兩個(gè)結(jié)果：
如果測(cè)試更多的數(shù)字，你會(huì)發(fā)現(xiàn)學(xué)生模型的表現(xiàn)與教師模型一樣好。
本文轉(zhuǎn)載自???AI科技論談???，作者：AI科技論談
