知识蒸馏(Knowledge Distillation)

July 8, 2026 · View on GitHub

知识蒸馏是一种将教师模型(teacher model)的能力迁移到学生模型(student model)的训练方法。其核心思想是:让学生在每个 token 位置上向教师的输出分布靠拢,从而获得比单纯模仿标注答案更丰富的监督信号——教师不仅告诉学生「哪个 token 是对的」,还告诉学生「其他 token 有多好 / 多差」。

本文档自顶向下地介绍:蒸馏为什么有效(第一节)、蒸馏方法的统一设计框架(第二节),最后落到 swift 中三种具体的蒸馏训练方法 GKD / OPD-RL / OPSD(第三节)。


一、为什么需要蒸馏:从稀疏信号到稠密信号

一个语言模型的能力,通常由一连串训练阶段堆叠而成:

  • 预训练(Pre-training):习得语言、世界知识、基础推理等通用能力。
  • 中训练(Mid-training):注入领域知识,如代码、医学、公司内部文档等。
  • 后训练(Post-training):激发目标行为,如指令遵循、数学推理、对话风格等。

蒸馏主要发生在后训练阶段。要理解它的价值,需要从两个彼此独立的维度来看待后训练方法:

  1. 采样方式(数据从哪来):训练序列是由学生自己生成(on-policy),还是来自外部固定数据(off-policy)。
  2. 反馈密度(每条序列能学到多少):是整条序列只有一个奖励(per-sequence,稀疏),还是每个 token 都有信号(per-token,稠密)。

SFT / 离线蒸馏(off-policy + 稠密):在固定数据上对齐标注或教师分布。信号稠密,但训练时见到的都是教师/标注的状态,和学生推理时自己会进入的状态不一致。学生一旦在推理早期犯了教师不会犯的错,就会进入训练中从未见过的状态,误差不断累积,这被称为 exposure bias(曝光偏差)

RL(on-policy + 稀疏):学生自己采样,按最终结果给奖励。分布与学生推理一致,但奖励通常是序列级标量,一般不指明具体哪个 token 出错。

On-policy 蒸馏(on-policy + 稠密):学生自己采样轨迹,再由教师对轨迹的每一个 token 打分。训练分布与学生推理分布一致,且反馈为 per-token 级别。

SFT 是蒸馏的一个特例

理解蒸馏的一个自然切入点是 SFT 的损失函数。SFT 的 cross-entropy loss,等价于以标注 token 的 one-hot 分布 δy\delta_{y^*} 为「教师」的 KL 散度:

logPS(yt)=KL(δyPS)-\log P_S(y^*_t) = \text{KL}(\delta_{y^*} \,\|\, P_S)

知识蒸馏只是把这个确定性的 one-hot「教师」换成了真实教师模型的软分布 PTP_T,在每个 token 上优化 KL(PTPS)\text{KL}(P_T \,\|\, P_S),从而提供比 one-hot 丰富的监督信号。

方法采样方式反馈密度教师分布
SFToff-policy(固定数据)稠密(per-token)one-hot δy\delta_{y^*}
离线(off-policy)蒸馏off-policy(固定数据或教师生成)稠密(per-token)教师软分布
RLon-policy(学生采样)稀疏(per-sequence)
On-policy 蒸馏on-policy(学生采样)稠密(per-token)教师软分布

二、蒸馏的两个核心选择

不同蒸馏方法的差异,几乎都可以归结为两个问题。理解了这两个维度,后面的方法都只是它们的不同组合。

2.1 教师信号怎么算

在每个 token 位置上,量化教师分布 PTP_T 与学生分布 PSP_S 的差异(我们称之为 Teacher KL),有两个子选择。

(a) 散度方向

散度定义优化时的行为(信息论含义)
Forward KLKL(PTPS)\text{KL}(P_T \,\|\, P_S)Mode-covering:学生需对教师概率较高的区域都赋予足够概率
Reverse KLKL(PSPT)\text{KL}(P_S \,\|\, P_T)Mode-seeking:学生主要拟合教师的众数(高概率)区域
广义 JSD(β\beta)βKL(PTM)+(1β)KL(PSM)\beta\,\text{KL}(P_T\|M) + (1-\beta)\,\text{KL}(P_S\|M),其中 M=βPT+(1β)PSM=\beta P_T+(1-\beta)P_S在两者之间插值

其中 β=0\beta=0 退化为 Forward KL,β=1\beta=1 退化为 Reverse KL。SFT 等价于 Forward KL(教师为 one-hot)。

在 swift 中:

  • GKD 默认 β=0.5\beta=0.5(JSD),可通过 --beta 在 Forward / JSD / Reverse 之间选择;
  • OPD-RL 的实现固定使用 Reverse KL 的 k1 估计量 logπteacher(yt)logπstudent(yt)\log\pi_{\text{teacher}}(y_t)-\log\pi_{\text{student}}(y_t) 作为 per-token advantage。

(b) 计算粒度

计算粒度需要的教师信息说明
全词表教师完整的 next-token 分布可计算散度的精确值;显存开销大
Top-K教师概率最高的 K 个 token在 top-K 子集上重新归一化后的近似;适合外部 API(受 max_logprobs 限制)
采样 token教师在学生实际采样 token 上的单个 logpReverse KL 的单样本蒙特卡洛估计;通信开销最低

精度 vs 开销:全词表需要物化完整 logits;采样 token 只需教师在已采样 token 上的 logp(可走远程 API)。DeepSeek-V4技术报告指出,仅用采样 token 的 log-ratio 作 advantage 时梯度估计方差较大,因此其全词表 OPD 采用完整 logit 蒸馏

2.2 信号怎么传给学生

路径 A:GKD(直接损失)路径 B:OPD-RL(RL Advantage)
训练范式--rlhf_type gkd--rlhf_type grpo + 教师
信号传递把信号作为 loss把信号当 advantage,走 policy gradient
梯度流经学生全词表 logits(或 top-k)仅学生采样 tokenlogπ(yt)\nabla\log\pi(y_t)
教师信息需求全词表分布(或 top-k logits)采样 token 上的单个 logp
散度选择Forward / Reverse / JSD(--betaReverse KL(k1 log-ratio)
与任务奖励组合通过 sft_alpha 混合 SFT loss可与 GRPO reward 叠加为 advantage

两者共享同一套教师基础设施(见下文),区别只在如何使用 teacher KL 信号。

蒸馏的常见用法

  1. 能力融合:多个专家模型蒸馏到统一模型。
  2. 强到弱:大模型向小模型传递能力。
  3. 防遗忘:用旧 checkpoint 作教师,在多阶段训练后恢复先前能力。

三、swift 中的蒸馏方法

swift 提供三种蒸馏训练方法,它们共享同一套教师基础设施:

方法信号传递路径启用方式一句话
GKD直接损失(路径 A)--rlhf_type gkd教师散度作为 loss 反向传播;支持全词表 / top-k 散度
OPD-RLRL advantage(路径 B)--rlhf_type grpo + 教师教师 log-ratio 注入 GRPO advantage,可与任务奖励叠加
OPSDA 或 B 均可在上面基础上提供 teacher_prompt单模型自蒸馏:教师输入含特权信息(如参考解答)

教师的三种来源(GKD 与 OPD-RL 通用):

  • --teacher_model:在训练进程中加载一个独立的冻结教师模型。
  • --teacher_model_server:连接一个外部教师服务(swift deploy启动的vllm服务),不在训练卡上加载教师。GKD 使用 API 时需同时设置 --gkd_logits_topk。支持单 URL 与多 teacher JSON 配置。
  • 自蒸馏:教师与学生同源。LoRA 训练且 --teacher_model--model 相同时,自动用 disable_adapter() 以基座为固定教师,无需额外加载;不传 --teacher_model 且设置 teacher_prompt, 则以「学生当前权重」为动态教师。
参数默认值说明
--teacher_modelNone教师模型路径;GKD 下不传则为动态自蒸馏
--teacher_model_serverNone教师 API 地址(与 teacher_model 互斥);见下方格式说明
--teacher_tag_key"dataset"多 teacher 路由时,样本用于匹配 teacher tags 的字段名
--teacher_deepspeedNone教师模型的 DeepSpeed 配置(如 zero3
--offload_teacher_modelFalse非前向阶段将教师卸载到 CPU(仅 teacher_model 生效)

完整参数说明见命令行参数

外部教师 API

通过 swift deploy --model xxx --infer_backend vllm 部署教师服务后,训练进程按 prompt 向 API 请求 logprobs,无需在训练卡上加载教师权重。

  • GKD:需设置 --gkd_logits_topk(API 仅返回 top-k logprobs),用于计算 JSD 散度损失。
  • OPD-RL:取采样 token 的 logp(prompt_logprobs=0),注入 advantage;系数由全局 --teacher_kl_coef 控制(见 3.2)。

Multi-Teacher(多教师路由)

Multi-Teacher 允许同时连接多个外部教师 API,按 tag 将每条样本路由到一个教师,实现领域专家蒸馏。

teacher_model_server 格式

# 单 teacher
--teacher_model_server http://localhost:8000

# 多 teacher(`tags` 与 `--dataset` 或数据列取值对应,见下方「模式一」)
--teacher_model_server '[{"url":"http://t1:8000","tags":["data/math.jsonl"]},{"url":"http://t2:8001","tags":["data/code.jsonl"]}]'

每个 teacher 配置项:

字段说明
url教师 API 地址
tags该教师负责的数据来源;单 teacher 可省略;多 teacher 时各条目不重叠,且与下方路由方式中的标识一致

路由方式

模式一:按数据集路由(默认)

传入多个 --dataset 时,样本按来源数据集匹配 teacher。默认 --teacher_tag_key"dataset"tags--dataset 中对应项保持一致:

# Hub ID
--dataset AI-ModelScope/alpaca-gpt4-data-en AI-ModelScope/alpaca-cleaned \
--teacher_model_server '[{"url":"http://t1:8000","tags":["AI-ModelScope/alpaca-gpt4-data-en"]},{"url":"http://t2:8001","tags":["AI-ModelScope/alpaca-cleaned"]}]'

# 本地路径
--dataset data/math.jsonl data/code.jsonl \
--teacher_model_server '[{"url":"http://t1:8000","tags":["data/math.jsonl"]},{"url":"http://t2:8001","tags":["data/code.jsonl"]}]'

模式二:按样本路由

在数据文件中为每条样本增加列(如 teacher_tag),设置 --teacher_tag_key teacher_tagtags 与该列取值对应。适用于只传一个 --dataset 但仍需多个 teacher 的场景。

示例:GKD + 多 teacher

# 部署两个教师服务(GKD 需 max_logprobs >= gkd_logits_topk)
CUDA_VISIBLE_DEVICES=1 swift deploy --model Qwen/Qwen3.5-4B --port 8000 --max_logprobs 64
CUDA_VISIBLE_DEVICES=2 swift deploy --model Qwen/Qwen3.5-1.7B --port 8001 --max_logprobs 64

CUDA_VISIBLE_DEVICES=0 swift rlhf \
    --rlhf_type gkd \
    --model Qwen/Qwen3.5-0.6B \
    --teacher_model_server '[{"url":"http://localhost:8000","tags":["data/math.jsonl"]},{"url":"http://localhost:8001","tags":["data/code.jsonl"]}]' \
    --gkd_logits_topk 64 \
    --dataset data/math.jsonl data/code.jsonl \
    ...

示例:OPD-RL + 多 teacher

CUDA_VISIBLE_DEVICES=1 swift deploy --model Qwen/Qwen3.5-4B --port 8000 --max_logprobs 1
CUDA_VISIBLE_DEVICES=2 swift deploy --model Qwen/Qwen3.5-1.7B --port 8001 --max_logprobs 1

CUDA_VISIBLE_DEVICES=0 swift rlhf \
    --rlhf_type grpo \
    --model Qwen/Qwen3.5-0.6B \
    --teacher_model_server '[{"url":"http://localhost:8000","tags":["data/math.jsonl"]},{"url":"http://localhost:8001","tags":["data/code.jsonl"]}]' \
    --teacher_kl_coef 1.0 \
    --dataset data/math.jsonl data/code.jsonl \
    --use_vllm true --vllm_mode colocate \
    ...

3.1 GKD:散度作为直接损失

GKD(Generalized Knowledge Distillation)直接把教师-学生间的散度作为损失函数反向传播。

损失函数

LGKD(x,y)=t=1yDJSD(β)(Pteacher(x,y<t),Pstudent(x,y<t))\mathcal{L}_{\text{GKD}}(x, y) = \sum_{t=1}^{|y|} D_{\text{JSD}(\beta)}\big(P_{\text{teacher}}(\cdot|x,y_{<t}),\, P_{\text{student}}(\cdot|x,y_{<t})\big)

其中散度 DD--beta 选择(见 2.1):β=0\beta=0 为 Forward KL,β=1\beta=1 为 Reverse KL,$0<\beta<1 为广义 JSD(默认 \0.5$)。

On-Policy vs Off-Policy:lmbda

GKD 通过 lmbda 控制每个 batch 用学生在线采样的概率:

if random() <= lmbda:
    y = student.generate(x)   # on-policy:学生自己采样
else:
    y = y_ground_truth        # off-policy:使用数据集标注
loss = D(P_teacher(·|x, y), P_student(·|x, y))
  • lmbda=0:纯离线(传统 SFT 蒸馏)。
  • lmbda=1:纯在线(学生从自身错误中学习,即 on-policy 蒸馏)。
  • 0<lmbda<1:混合。

若希望走teacher生成数据路径,需先用教师离线生成响应写入数据集,再设 lmbda=0 训练

GKD参数

参数类型默认值说明
--betafloat0.5散度插值:0=Forward KL,0.5=JSD,1=Reverse KL
--lmbdafloat0.5在线采样概率:0=离线,1=纯在线
--sft_alphafloat0混合 SFT loss 比例,最终 loss = gkd_loss + sft_alpha * sft_loss(仅对非学生生成的数据生效)
--gkd_logits_topkintNone仅用教师 top-K logits 计算 KL;使用 teacher_model_server 时为必填

Top-K 蒸馏(省显存)

默认用完整词表计算 KL,词表很大时容易 OOM,可使用--gkd_logits_topk参数。

外部教师 API

设置 --teacher_model_server 时需同时设置 --gkd_logits_topk(API 仅返回 top-k logprobs)。示例如下:

# 步骤 1:部署教师模型(max_logprobs 需 >= gkd_logits_topk)
CUDA_VISIBLE_DEVICES=0 swift deploy \
    --model Qwen/Qwen3.5-9B \
    --infer_backend vllm \
    --port 8000 \
    --max_logprobs 64

# 步骤 2:启动 GKD 训练
CUDA_VISIBLE_DEVICES=1,2,3,4 \
NPROC_PER_NODE=4 \
swift rlhf \
    --rlhf_type gkd \
    --model Qwen/Qwen3.5-2B \
    --teacher_model_server http://localhost:8000 \
    --gkd_logits_topk 64 \
    --lmbda 1.0 \
    --beta 1.0 \
    --dataset xxx

在线采样加速

lmbda > 0 时学生需在线生成序列,建议用 vLLM 加速采样(colocate / server 两种模式,与 GRPO 一致),参考 GRPO 文档

多轮 GKD

GKD 多轮训练,与 GRPO 共享同一套 MultiTurnScheduler 基础设施。

完整接口和自定义方式请参考 GRPO 多轮训练文档

参考脚本


3.2 OPD-RL:KL 作为 RL Advantage

OPD(On-Policy Distillation)RL 把教师 KL 信号注入 GRPO 的 per-token advantage,通过 policy gradient 更新学生。

原理

标准 GRPO 的 advantage 来自组内归一化的任务奖励(per-sequence 标量)。OPD-RL 在 advantage 归一化之后逐 token 注入教师信号:

At=Atbase+α(logπteacher(ytx,y<t)logπstudent(ytx,y<t))A_t = A_t^{\text{base}} + \alpha \cdot \big(\log \pi_{\text{teacher}}(y_t|x,y_{<t}) - \log \pi_{\text{student}}(y_t|x,y_{<t})\big)
  • AtbaseA_t^{\text{base}}:GRPO 归一化后的任务奖励 advantage(无奖励函数时为 0)。
  • α\alpha--teacher_kl_coef,教师信号强度。
  • logπteacher(yt)logπstudent(yt)\log\pi_{\text{teacher}}(y_t) - \log\pi_{\text{student}}(y_t):教师 log-ratio,即 Reverse KL 梯度中 θlogπθ(yt)\nabla_\theta\log\pi_\theta(y_t) 系数对应的 k1 估计量(见 compute_teacher_logratio)。

纯蒸馏模式:不设 --reward_funcs 时 base advantage 为 0,教师信号是唯一驱动力,此时 At=α(logπteacher(yt)logπstudent(yt))A_t = \alpha\cdot(\log\pi_{\text{teacher}}(y_t)-\log\pi_{\text{student}}(y_t))

监控指标:日志中的 teacher_kl 是 k3 估计量 edd1e^{d}-d-1d=logπteacherlogπstudentd=\log\pi_{\text{teacher}}-\log\pi_{\text{student}}),衡量学生与教师的距离。

启用方式:在 --rlhf_type grpo 下设置 --teacher_model--teacher_model_server 即自动启用 OPD-RL,无需额外开关。教师相关参数见上文共享参数表。

OPD-RL 特有参数

参数默认值说明
--teacher_kl_coef1.0教师 log-ratio 注入 advantage 的系数 α\alpha

参考脚本


3.3 OPSD:On-Policy Self-Distillation

OPSD(On-Policy Self-Distillation)是一种单模型自蒸馏方法:同一模型分别构造学生输入与教师输入,教师侧额外接收特权信息(如参考解答),再对齐两者在学生采样响应上的输出分布。

核心机制

  • 学生:仅看到问题,正常推理。
  • 教师:看到问题 + 参考解答(通过 teacher_prompt 列提供特权信息)。
  • 训练目标:用散度(JSD / KL)对齐学生与教师在同一份学生采样响应上的输出分布。

OPSD 既可走 GKD 路径,也可走 OPD-RL 路径:

  • GKD + OPSD--rlhf_type gkd,教师 KL 作为直接损失。
  • OPD-RL + OPSD--rlhf_type grpo;动态模式不传 --teacher_model,固定模式设 --teacher_model--model 相同。

两种自蒸馏权重模式

模式参数配置教师权重说明
Dynamic(动态)不传 --teacher_model学生当前权重教师随训练同步更新
Fixed(固定)--teacher_model 设为与 --model 相同初始教师权重教师权重固定

数据格式

OPSD 数据集需包含 teacher_prompt 列,可通过 --external_plugins 加载数据处理插件来构建。以数学推理数据集 open-r1/OpenThoughts-114k-math 为例:

from swift.dataset import DatasetMeta, RowPreprocessor, register_dataset

class OpenThoughtsOPSDPreprocessor(RowPreprocessor):
    def preprocess(self, row):
        if not row.get('correct', True):
            return None
        problem = row.get('problem', '')
        solution = row.get('solution', '')
        teacher_prompt = f'{problem}\n\nReference solution:\n{solution}\n\nNow articulate your own reasoning.'
        messages = [
            {'role': 'system', 'content': 'Please reason step by step, and put your final answer within \\boxed{}.'},
            {'role': 'user', 'content': problem},
        ]
        return {'messages': messages, 'teacher_prompt': teacher_prompt}

register_dataset(DatasetMeta(
    ms_dataset_id='open-r1/OpenThoughts-114k-math',
    preprocess_func=OpenThoughtsOPSDPreprocessor(),
    tags=['math', 'opsd'],
))

参考脚本


Reference