詳解模型蒸餾，破解DeepSeek性能謎題
不少關注 DeepSeek 最新動態的朋友，想必都遇到過 “Distillation”（蒸餾）這一術語。但它究竟指代何種技術，又為何在 AI 領域占據重要地位呢？
本文為你深度剖析模型蒸餾技術的原理，同時借助 TensorFlow 框架中的實例進行詳細演示。相信通過閱讀本文，你會對模型蒸餾有全新的認識，輕松解鎖深度學習優化的新視角。
1 模型蒸餾原理
在深度學習領域，模型蒸餾是優化模型的關鍵技術。它讓小的學生模型不再單純依賴原始標簽，而是基于大的教師模型軟化后的概率輸出進行訓練。
以圖像分類為例，普通模型只是簡單判斷圖像內容，而運用模型蒸餾技術的學生模型，能從教師模型的置信度分數（如80%是狗，15%是貓，5%是狐貍）中獲取更豐富信息，從而保留更細致知識。
這樣一來，學生模型能用更少參數實現與教師模型相近的性能，在保持高精度的同時，減小模型規模、降低計算需求，為深度學習模型優化開辟了新路徑。
讓我們通過一個例子來看看具體是如何操作的。以使用MNIST數據集訓練卷積神經網絡（CNN）為例。
MNIST （Modified National Institute of Standards and Technology database）數據集在機器學習和計算機視覺里常用，有 70,000 張 28x28 像素的手寫數字（0 - 9）灰度圖，60,000 張訓練圖、10,000 張測試圖。
模型蒸餾要先建教師模型，是用 MNIST 數據集訓練的 CNN，參數多、結構復雜。
再建個更簡單、規模更小的學生模型。
目的是讓學生模型模仿教師模型性能，還能減少計算量和訓練時間。
訓練時，兩個模型都用 MNIST 數據集預測，接著算它們輸出的 Kullback-Leibler（KL）散度。這個值能確定梯度，指導調整學生模型。
一番操作后，學生模型就能達到和教師模型差不多的準確率，成功 “出師”。
2 用 TensorFlow 和 MNIST 構建模型
接下來，我們借助 TensorFlow 和 MNIST 數據集，搭建一個模型蒸餾示例項目。
先訓練一個教師模型，再通過模型蒸餾技術，訓練出一個更小的學生模型。這個學生模型能模仿教師模型的性能，而且所需資源更少。
2.1 使用MNIST數據集
確保你已經安裝了TensorFlow：
!pip install tensorflow然后加載MNIST數據集：
from tensorflow import keras
import matplotlib.pyplot as plt
# 加載數據集（MNIST）
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
fig = plt.figure()
# 可視化部分數字
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數據集中的前9個樣本數字及其標簽：
還需要對圖像數據進行歸一化處理，并擴展數據集的維度，為訓練做準備：
import tensorflow as tf
import numpy as np
# 歸一化圖像
x_train, x_test = x_train / 255.0, x_test / 255.0
# 為卷積神經網絡擴展維度
x_train = np.expand_dims(x_train, axis=-1)
x_test = np.expand_dims(x_test, axis=-1)
# 將標簽轉換為分類（獨熱編碼）
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)2.2 定義教師模型
在基于模型蒸餾的示例項目構建中，定義并訓練教師模型是關鍵的環節。這里，我們構建一個多層卷積神經網絡（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用于蒸餾
])需要注意的是，模型最后一層設置了 10 個單元，對應 0 - 9 這 10 個數字，但未采用 softmax 激活函數，而是輸出原始的 logits。這一設計對于模型蒸餾很重要，因為在后續的蒸餾過程里，會借助 softmax 函數來計算教師模型與學生模型之間的 Kullback-Leibler（KL）散度，以此衡量二者差異，為學生模型的優化提供方向。
完成模型定義后，要使用compile()方法對其進行配置，設置優化器、損失函數以及評估指標：
teacher_model.compile(
optimizer = 'adam',
loss = tf.keras.losses.CategoricalCrossentropy(from_logits = True),
metrics = ['accuracy']
)配置完成，就可以使用fit()方法啟動模型訓練：
# 訓練教師模型
teacher_model.fit(x_train, y_train,
epochs = 5,
batch_size = 64,
validation_data = (x_test, y_test))本次訓練設定了 5 個訓練周期，訓練過程中的詳細信息如下：
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從這些數據中，我們可以直觀地看到模型在訓練過程中的準確率和損失變化，了解模型的學習效果，為后續的模型優化和評估提供依據。
2.3 定義學生模型
教師模型訓練完成后，就該定義學生模型了。與教師模型相比，學生模型的架構更簡單，層數更少：
# 學生模型
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 定義蒸餾損失函數
在模型蒸餾的實現過程中，distillation_loss() 函數發揮著核心作用，它借助 Kullback-Leibler（KL）散度來精準計算蒸餾損失，從而推動學生模型向教師模型 “看齊”。下面，我們就來詳細解讀這個函數的代碼實現。
蒸餾損失函數的計算，依賴于教師模型和學生模型的預測結果，具體步驟如下：
- 使用教師模型為輸入批次生成軟目標（概率）。
- 使用學生模型的預測計算其軟概率。
- 計算教師模型和學生模型軟概率之間的Kullback-Leibler（KL）散度。
- 返回蒸餾損失。
軟概率和常見的硬標簽不同。硬標簽是明確分類，像判斷郵件是否為垃圾郵件，結果只有 “是”（1）或 “否”（0）。而軟概率會給出多種結果的概率，比如某郵件是垃圾郵件的概率為 0.85，不是的概率為 0.15，能更全面反映模型判斷。
計算軟概率要用到 softmax 函數，且受溫度參數影響。在知識蒸餾里，教師模型的軟概率包含類間豐富信息，學生模型學習后，能提升泛化能力和性能，更好地模仿教師模型。
以下是distillation_loss()函數的定義：
def distillation_loss(y_true, y_pred, x_batch, teacher_model, temperature=5):
"""
使用KL散度計算蒸餾損失。
"""
# 計算當前批次的教師模型logits
teacher_logits = teacher_model(x_batch, training=False)
# 將logits轉換為軟概率
teacher_probs = tf.nn.softmax(teacher_logits / temperature)
student_probs = tf.nn.softmax(y_pred / temperature)
# KL散度損失（教師模型和學生模型分布之間的差異）
return tf.reduce_mean(tf.keras.losses.KLDivergence()(teacher_probs, student_probs))Kullback-Leibler（KL）散度，也稱為相對熵，用于衡量一個概率分布與另一個參考概率分布之間的差異。
2.5 使用知識蒸餾訓練學生模型
現在你已經準備好使用知識蒸餾來訓練學生模型了。首先，定義train_step()函數：
optimizer = tf.keras.optimizers.Adam()
@tf.function
def train_step(x_batch, y_batch, student_model, teacher_model):
with tf.GradientTape() as tape:
# 獲取學生模型的預測
student_preds = student_model(x_batch, training=True)
# 計算蒸餾損失（顯式傳入教師模型）
loss = distillation_loss(y_batch, student_preds, x_batch, teacher_model, temperature=5)
# 計算梯度
gradients = tape.gradient(loss, student_model.trainable_variables)
# 應用梯度 - 訓練學生模型
optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))
return loss這個函數執行單個訓練步驟：
- 計算學生模型的預測。
- 使用教師模型的預測計算蒸餾損失。
- 計算梯度并更新學生模型的權重。
為了訓練學生模型，需要創建一個訓練循環，遍歷數據集，在每一步更新學生模型的權重，并在每個訓練周期結束時打印損失，以監控訓練進度：
# 訓練循環
epochs = 5
batch_size = 32
# 準備數據集批次
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!")訓練后，應該會看到類似以下的輸出：
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 評估學生模型
學生模型已經訓練完成，可以使用測試集（x_test和y_test）對其進行評估，觀察其性能：
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}")如預期，學生模型達到了相當不錯的準確率：
Student Model Accuracy: 0.98632.7 使用教師模型和學生模型進行預測
現在可以使用教師模型和學生模型進行一些預測，看看是否都能準確預測MNIST測試數據集中的數字：
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()
# 現在可以進行預測
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))以下是前兩個結果：
如果測試更多的數字，你會發現學生模型的表現與教師模型一樣好。
本文轉載自???AI科技論談???，作者：AI科技論談
