1. 知識蒸餾簡介

什么是知識蒸餾？

知識蒸餾（Knowledge Distillation）是一種模型壓縮技術，目標是讓一個較小的模型（學生模型，Student Model）學習一個較大、性能更優的模型（教師模型，Teacher Model）的知識。這樣，我們可以在保持較高準確率的同時，大幅減少計算和存儲成本。

為什么需要知識蒸餾？

• 降低計算成本：大模型（如 DeepSeek、GPT-4）通常計算量巨大，不適合部署到移動設備或邊緣設備上。
• 加速推理：較小的模型可以更快地推理，減少延遲。
• 減少內存占用：適用于資源受限的環境，如嵌入式設備或低功耗服務器。

知識蒸餾的核心思想是：學生模型不僅僅學習教師模型的硬標簽（one-hot labels），更重要的是學習教師模型輸出的概率分布，從而獲得更豐富的表示能力。

2. KL 散度的數學原理

2.1 KL 散度公式

在知識蒸餾過程中，我們通常使用Kullback-Leibler 散度（KL Divergence） 來衡量兩個概率分布（教師模型和學生模型）之間的差異。

2.2 直觀理解

KL 散度可以理解為如果用分布 Q 來近似分布 P，會損失多少信息。

• 當 KL 散度為 0，表示兩個分布完全相同。
• KL 散度不是對稱的，即

3. DeepSeek 中的 KL 散度應用

DeepSeek 作為一個強大的開源大語言模型（LLM），在模型蒸餾時廣泛使用了 KL 散度。例如，在訓練較小版本的 DeepSeek 時，研究人員采用了溫度標度（Temperature Scaling） 來調整教師模型的輸出，使其更適合學生模型學習。

教師模型的 softmax 輸出使用溫度參數 TT 進行調整：

當 T 增大時，softmax 輸出的概率分布變得更平滑，從而讓學生模型更容易學習教師模型的知識。

在 DeepSeek 的蒸餾過程中，常見的損失函數是加權組合：

其中：

• 第一項是 KL 散度損失，使得學生模型的輸出接近教師模型。
• 第二項是交叉熵損失，確保學生模型仍然學習真實標簽。
• λ是一個超參數，控制兩者的平衡。

4. 代碼示例：用 Keras 進行知識蒸餾

下面我們用 TensorFlow/Keras 訓練一個簡單的學生模型，讓它學習一個教師模型的知識。

4.1 定義教師模型

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers# 構建一個簡單的教師模型
teacher_model = keras.Sequential([layers.Dense(128, activation="relu", input_shape=(784,)),layers.Dense(10, activation="softmax")
])

4.2 訓練教師模型

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train, x_test = x_train.reshape(-1, 784) / 255.0, x_test.reshape(-1, 784) / 255.0
y_train, y_test = keras.utils.to_categorical(y_train, 10), keras.utils.to_categorical(y_test, 10)teacher_model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
teacher_model.fit(x_train, y_train, epochs=5, batch_size=32, validation_data=(x_test, y_test))

4.3 讓教師模型生成 soft labels

temperature = 5.0
def soft_targets(logits):return tf.nn.softmax(logits / temperature)y_teacher = soft_targets(teacher_model.predict(x_train))

4.4 訓練學生模型

student_model = keras.Sequential([layers.Dense(64, activation="relu", input_shape=(784,)),layers.Dense(10, activation="softmax")
])student_model.compile(optimizer="adam",loss=tf.keras.losses.KLDivergence(),  # 使用 KL 散度metrics=["accuracy"]
)student_model.fit(x_train, y_teacher, epochs=5, batch_size=32, validation_data=(x_test, y_test))

5. 真實應用場景

5.1 輕量級大模型

• DistilBERT：使用 BERT 作為教師模型進行蒸餾，訓練更小的 Transformer。
• TinyBERT：針對任務優化蒸餾，提高學生模型的表現。
• DeepSeek-Chat 小模型：使用 KL 散度訓練高效版本，提高推理速度。

5.2 知識蒸餾的優勢

• 可以訓練更小的模型，適用于移動端、嵌入式設備。
• 學生模型比直接訓練的模型泛化性更強，能更好地模仿教師模型。
• 結合 KL 散度 + 交叉熵 可以提升訓練效果。

結論

KL 散度損失是知識蒸餾的核心，它讓學生模型學習教師模型的概率分布，從而獲得更好的表現。DeepSeek 這樣的 LLM 在蒸餾過程中廣泛使用 KL 散度，使得較小模型也能高效推理。希望本文能幫助你理解 KL 散度在知識蒸餾中的應用！

其它

代碼示例一，

假設我們有兩個概率分布 p（真實分布）和 q（預測分布），我們使用 KLDivergence 計算它們之間的 KL 散度損失。

import tensorflow as tf
import numpy as np# 定義 KLDivergence 損失函數
kl_loss = tf.keras.losses.KLDivergence()# 真實分布 p (標簽)
p = np.array([0.1, 0.4, 0.5], dtype=np.float32)# 預測分布 q
q = np.array([0.2, 0.3, 0.5], dtype=np.float32)# 計算 KL 散度損失
loss_value = kl_loss(p, q)print(f'KL Divergence Loss: {loss_value.numpy()}')

代碼示例二，

一個完整的 Keras 代碼示例，展示了如何在分類任務中使用 KLDivLoss 作為損失函數。這個示例使用一個簡單的神經網絡對 手寫數字 MNIST 數據集 進行分類，并使用 KLDivLoss 計算真實分布和模型預測分布之間的散度。

import tensorflow as tf
from tensorflow import keras
from keras import layers
import numpy as np# 加載 MNIST 數據集
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()# 歸一化數據到 [0,1] 之間
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0# 將標簽轉換為概率分布 (one-hot 編碼)
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)# 構建一個簡單的神經網絡模型
model = keras.Sequential([layers.Flatten(input_shape=(28, 28)),layers.Dense(128, activation="relu"),layers.Dense(10, activation="softmax")  # 輸出層用 softmax 歸一化
])# 編譯模型，使用 KLDivLoss 作為損失函數
model.compile(optimizer="adam",loss=tf.keras.losses.KLDivergence(),metrics=["accuracy"])# 訓練模型
model.fit(x_train, y_train, epochs=5, batch_size=32, validation_data=(x_test, y_test))# 評估模型
test_loss, test_acc = model.evaluate(x_test, y_test)
print(f"Test Accuracy: {test_acc:.4f}")