[LLM 6/10] 上下文蒸馏：把 system prompt 搬进模型权重里

用户每一次给你的聊天机器人发消息，你都会把那段几百 token 的 system prompt 原封不动地一起送过去—— 系统活多久就付多久，每一次 request 都要重付，永远没有尽头。 这一章我们要把那一坨知识从 prompt 搬进模型权重，用的技术叫 Context Distillation（上下文蒸馏），而且是 on-policy 版本（OPCD）。 其中最漂亮的一点是：老师和学生是完全同一个模型——唯一的区别只是谁看得见 prompt。

1. 问题（Problem statement）​

一个典型的泰语客服助手，system prompt 大致长这样：设定 persona， 强制永远用泰语回答，必须礼貌地以 ครับ/ค่ะ 结尾（泰语礼貌助词，分别由男性/女性说话者使用）， 禁止给出医疗和法律建议。认真写下来大约就是 400 token——而且它跟着每一次 request 一起发出去。

算一笔账：一个每天承接 100,000 次 request 的系统，等于为同一段文字反复付出 每天 4,000 万 token、每月 12 亿 token——而这段内容一个字符都没变过。 而这还没算上账单里看不见的另外两项代价：

• Latency —— 模型每次都得先 prefill 完 400 token，才能开始想第一个字
• Context budget —— persona 占掉的每一个 token，都是从对话历史和附件文档那里抢走的

放到本系列的框架里看，知识有三个可以存放的地方，而每个地方的"付款方式"都不一样：

一段已经稳定下来、却仍然每次都要附带的 system prompt，就是被放错地方的知识—— 它应该从这张表的第一行，搬到最后一行去。这一章讲的就是怎么搬。

2. 我们要做什么（Solution）​

Context distillation 就是训练一个看不见 context $c$ 的学生， 让它的行为等同于一个看得见 $c$ 的老师——换句话说，把 $c$ 的作用从 prompt 搬进权重。 最早的 offline 思路可以追溯到 Askell 等人（2021）的工作， 而本章用的版本是 OPCD（On-Policy Context Distillation）， 出自 Ye、Dong、Wu、Huang 和 Wei（2026，arXiv:2602.12275）， 它加了两味关键配料，第 3 节会一味一味拆开看：

1. 学生采样自己的回答（on-policy），并且看不见 $c$
2. 在这些回答上，最小化相对于"看得见 $c$ 的老师"的 reverse KL

有一句话我想现在就先钉在这里，因为第 7 章还会出现另一个常被混淆的 "distillation"：

Context distillation 改变的是"模型不用被告知就已经知道的东西"——model distillation 改变的是"模型的大小"

这一章里模型不会变小哪怕一个参数，它只是不再需要那段 prompt 了。 而第 7 章讲的是把大模型压缩成小模型——完全是另一根轴。

3. 公式（Equation）​

3.1 OPCD 的 objective​

其中每个 token 位置上的 KL，是在整个 vocabulary $\mathcal{V}$ 上求和：

• $c$ = 想搬进 weights 的 context（persona + 安全政策，约 400 token）
• $x$ = 用户的问题，$y$ = 学生自己采样出来的、没看见 $c$ 的回答
• $\pi_\theta$ = 学生（只看得见 $x$ 去预测），$\pi_{\text{teacher}}$ = 老师（同一份权重，但还看得见 $c$）
• $\frac{1}{|y|}$ = 按 token 取平均，防止长回答拿到过高的权重（是不是有点眼熟——第 4 章的 length bias）

注意这里并不是跟任何"标准答案"算 cross-entropy——目标是老师在每一个 token 位置上的 整行概率分布。学生学的不是"下一个词是什么"，而是 "如果眼前摆着 $c$，vocab 里每一个词的概率会长成什么样"。

这个式子里有两个决定，它们各自扛着整套方法的一半重量，我们分开看。

3.2 第一个决定——KL 必须是 reverse（$\pi_\theta$ 放在前面）​

KL 不对称，而它的顺序就是在选择行为：

• Forward KL $\mathbb{D}_{\text{KL}}(\pi_{\text{teacher}} \| \pi_\theta)$ 会在老师有质量而学生没有的地方爆掉 → 学生被迫去"覆盖"老师的每一个 mode（mode-covering） 如果容量不够，它就会把质量摊平铺开去盖住一切，包括老师从来没去过的、两个 mode 之间的山谷—— 用 LLM 的话说，那就是"两种风格混在一起混出怪东西"的回答，或者 hallucination
• Reverse KL $\mathbb{D}_{\text{KL}}(\pi_\theta \| \pi_{\text{teacher}})$ 会在学生有质量而老师没有的地方爆掉 → 学生被迫不去做老师不做的事，然后挑老师的某一个 mode 死死咬住（mode-seeking，模式寻找）

对本章这件事来说——persona 和安全政策——我们想要的显然是后者，根本不用犹豫： 一个"能稳稳当当地按老师的某一种方式做事"的学生， 远比一个"给老师的每一条路都留概率，连老师明令禁止的路也留"的学生有价值。

如果你觉得在哪儿见过——没错，第 3–4 章 RLHF 公式里的 KL 也是把 $\pi$ 放在前面的， 理由完全一样：我们要管住的是正在训练的那一个的行为，而不是参考模型的。

3.3 第二个决定——rollout 必须是学生自己的（on-policy）​

注意 3.1 式里的 $y\sim\pi_\theta(\cdot|x)$：用来训练的回答是从学生那里采样的，不是从老师那里。

更省事的做法是让老师（看得见 $c$）先写好一批回答，然后让学生照着做 SFT—— 但那条路有一个结构性的毛病，叫 exposure bias：学生只在老师写出来的文本路径上被教过， 可真正上线时，它必须从自己写出来的 prefix 往下接。 错一个 token 就会掉进一个从没被教过的状态，然后误差一路复利滚下去。

on-policy 采样从结构上消掉了这个问题：学生在训练中遇到的状态， 和它在 inference 时会遇到的状态是同一类，因为两边都是它自己造出来的。 老师只有一个职责，就是"站在学生的路径上做批改"——告诉它，在你刚刚走到的这个点上， 如果有 $c$，接下来该往哪边走（这和第 5 章必须采样自己的回答、而不是继续用 DPO 的理由是同一个）。

一行诚实的备注：算 gradient 的时候，我们把采样出来的 $y$ 当作常量， 不让 gradient 反向穿过采样这一步——这是 on-policy distillation 的标准做法。

3.4 必须打赢的 baseline：offline context distillation​

大多数博客说的 "context distillation"，指的是 offline 版本：

直白地读出来就是：让看得见 $c$ 的老师写出回答，然后拿这些回答去 SFT 看不见 $c$ 的学生 ——就是在老师文本上的普通 cross-entropy，没有整行 KL，也没有 on-policy。

这可不是个稻草人，它是一个真正结实而且更便宜的 baseline（训练方式跟第 2 章一模一样）。 第 9 节会让 OPCD 在同一批数据上跟它公平较量。 如果 OPCD 那两味配料（reverse KL + on-policy）真的有价值，它就应该赢在理论说它会赢的地方： 向没见过的 prompt 类型上泛化。

4. 把公式画出来（Visualize）​

整套方法就在这一张图里​

看这张图的时候，请留意它有多节省：没有第二个模型，没有 reward model，没有标准答案数据集。 只有同一份权重的两种 forward pass——一种看得见 $c$，一种看不见—— 再加上一个 LoRA adapter，负责把这两者之间的"差值"收进权重里。

为什么 KL 的方向就决定了行为​

左边这一面板，就是 distillation 语境下 hallucination 这个词的解释： 按 forward KL 算出来的最优 q，会实实在在地把质量放在p 几乎为零的地方， 因为它宁可付这个代价，也不愿漏掉任何一个 mode。 右边这一面板才是我们对一个安全导向的学生的要求：选一条老师认可的路，然后死死咬住。

OPCD 给我们买到了什么​

本章的目标用几何语言写出来就是：把那个蓝点往左推 400 个 token，同时让它掉下去的高度尽可能小。

在继续之前，我们先放大到 token 级别看看：下面这句话正是 persona 所规定的行为。 看看学生（看不见 $c$ 的那个）逐 token 的 log-prob 在训练之后有什么变化—— 训练前，这样的概率需要靠 $c$ 顶上去；训练后，它变成了模型自己的默认值：

5. 准备环境（Environment）​

打开 Colab，选择 Runtime → Change runtime type → T4 GPU（免费额度就够用）。

torch_dtype=torch.float16      # 不是 bfloat16
attn_implementation="sdpa"     # 不是 flash_attention_2

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True（把 emulation 也算上了！）

一个连一个字节 VRAM 都不多占的老师​

OPCD 同时需要老师和学生，听上去像是要加载两个模型——不用。 因为两者就是同一份权重，区别只在 adapter 和 prompt：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", padding_side="left")
base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
).cuda()

policy = get_peft_model(base, LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
))

• 学生 = policy（base + LoRA），forward 时没有 $c$
• 老师 = 同一个模型，在 policy.disable_adapter() 之下 forward，前面带着 $c$

6. 准备数据（Data）​

要准备两样东西：要搬进 weights 的 context，以及给学生练习采样回答用的问题。

Context：persona + 安全政策（约 400 token）​

PERSONA = """คุณคือ "น้องใจดี" ผู้ช่วยฝ่ายบริการลูกค้าของร้านค้าออนไลน์
กฎที่ต้องปฏิบัติตามทุกข้อ ไม่มีข้อยกเว้น:
1. ตอบเป็นภาษาไทยเท่านั้น ห้ามสลับเป็นภาษาอังกฤษกลางประโยค
   แม้ผู้ใช้จะถามมาเป็นภาษาอังกฤษก็ตอบกลับเป็นภาษาไทยอย่างสุภาพ
2. ใช้ภาษาสุภาพเสมอ ลงท้ายประโยคด้วย "ครับ/ค่ะ" อย่างสม่ำเสมอ
3. ห้ามให้คำแนะนำทางการแพทย์ การวินิจฉัยโรค หรือการใช้ยา
   ให้ปฏิเสธอย่างสุภาพ แล้วแนะนำให้ปรึกษาแพทย์หรือเภสัชกรโดยตรง
4. ห้ามให้คำแนะนำทางกฎหมาย ให้ปฏิเสธอย่างสุภาพ
   แล้วแนะนำให้ปรึกษาทนายความหรือหน่วยงานที่เกี่ยวข้อง
5. ..."""                      # 完整版约 400 token，在 notebook 里

print("context length:", len(tok(PERSONA).input_ids), "tokens")

这段 PERSONA 保持泰语原文不动——它是真正喂给模型的那份操作性文本， 第 8 节的检查器正是照着里面的字面串（比如句尾助词 "ครับ/ค่ะ"）来匹配的， 译成中文会让这一章的实验和它的验收标准对不上。它的中文大意是：

你是网店的客服助手"น้องใจดี（贴心小妹）"。以下规则必须条条遵守，没有任何例外：

1. 只能用泰语回答，不许在句中切换成英语；即使用户用英语提问，也要礼貌地用泰语回答。
2. 始终使用礼貌用语，句尾一律稳定地加上 "ครับ/ค่ะ"。
3. 不得给出医疗建议、疾病诊断或用药指导；要礼貌地拒绝，然后建议直接咨询医生或药师。
4. 不得给出法律建议；要礼貌地拒绝，然后建议咨询律师或相关机构。

这就是我们的 $c$——注意它纯粹是行为，里面没有任何需要背诵的事实 （这个观察会在本章末尾的局限性框里变成一件大事）。

训练问题：来自泰语数据集的 300 道​

from datasets import load_dataset

ds = load_dataset("airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k",
                  split="train")
prompts = [r["instruction"] for r in ds.shuffle(seed=42).select(range(300))]

数据集里的答案列我们一行都不用——OPCD 不需要标准答案， 它只需要足够多样的问题，让学生在各种情境下试着回答，再交给老师批改。

评测集单独留出，训练期间绝不碰，而且刻意包含了训练集里没有的 prompt 类型：

• 40 道：与训练集同一路数的普通问题（in-distribution）
• 20 道：医疗/法律类问题——检验"拒绝"这条政策是不是真的进了权重
• 20 道：英语问题——在最容易破功的情境下检验"永远用泰语回答"这条规则

后两组就是第 9 节里的 OOD compliance 那一列——它是区分 "记住了例子"和"吸收了政策"的分水岭。

7. 核心代码（Main code）​

OPCD 的循环有三拍：学生采样 → 老师批改 → 按 reverse KL 调整权重。

7.1 学生采样自己的 rollout（看不见 $c$）​

@torch.no_grad()
def rollout(policy, x_texts, G=4):
    """on-policy 这个词的核心：回答来自学生，不是来自老师"""
    batch = tok(x_texts, return_tensors="pt", padding=True).to("cuda")
    out = policy.generate(**batch,
                          do_sample=True, temperature=1.0, top_p=1.0,
                          max_new_tokens=192, num_return_sequences=G)
    return out              # [len(x_texts) * G, |x| + |y|]

temperature=1.0 不是随手挑的——第 9 节的第 4 个坑会讲，为什么调得比这更低就危险了。

7.2 相对老师的 reverse KL——全章的核心代码​

import torch.nn.functional as F

K = 128         # 只保留老师的 top-K —— 理由在下面的算术框里

def opcd_loss(policy, c_ids, x_ids, y_ids, y_mask):
    # 学生：只看得见 x + y（adapter 开着）
    s_in = torch.cat([x_ids, y_ids], dim=1)
    s_logits = policy(s_in).logits[:, x_ids.size(1) - 1 : -1]

    # 老师：同一份底座权重，关掉 adapter 并且"看得见 c" —— 不要梯度
    with torch.no_grad(), policy.disable_adapter():
        t_in = torch.cat([c_ids, x_ids, y_ids], dim=1)
        t_logits = policy(t_in).logits[:, c_ids.size(1) + x_ids.size(1) - 1 : -1]

    # 裁剪到老师给出最高质量的 top-K 这个 support 上，然后两边都重新归一化
    topk = t_logits.topk(K, dim=-1).indices
    t_logp = torch.log_softmax(t_logits.gather(-1, topk).float(), dim=-1)
    s_logp = torch.log_softmax(s_logits.gather(-1, topk).float(), dim=-1)

    # reverse KL：π_θ 放在"前面" —— 每一项的权重来自学生，不是老师
    kl = (s_logp.exp() * (s_logp - t_logp)).sum(-1)          # [B, |y|]
    return (kl * y_mask).sum() / y_mask.sum()                # 按 token 取平均 = 1/|y|

7.3 训练循环​

opt = torch.optim.AdamW(
    [p for p in policy.parameters() if p.requires_grad], lr=1e-5)

for epoch in range(2):
    for x_texts in batches(prompts, batch_size=1):
        seqs = rollout(policy, x_texts, G=4)
        c_ids, x_ids, y_ids, y_mask = split_and_pad(seqs, c_len)   # 见 notebook
        loss = opcd_loss(policy, c_ids, x_ids, y_ids, y_mask)
        loss.backward()
        opt.step(); opt.zero_grad()

• lr=1e-5 —— 比 DPO（5e-6）高，但远低于 SFT（2e-4）： 我们是在把分布往一个近在咫尺的老师那边掰，而不是在教新知识
• 300 个 prompt × 4 个 rollout × 2 个 epoch，在 T4 上大约 16 分钟
• 训练过程中，notebook 每 20 step 会 log 一次输出的平均 entropy 和回答的平均长度—— 这两个就是 mode collapse 的矿井金丝雀（canary），见第 9 节的第 1 个坑

8. 结果（Results）​

notebook 会测 4 项并写入 results.json：

1. Persona-compliance rate —— 用一个 deterministic 的检查器测（代码见下），附 Wilson 95% CI， 在"老师 + 完整 context"和"OPCD 学生、不带 context"之间做对比
2. 每次 request 的 prompt tokens —— 用 tokenizer 实测：每次 request 应当减少约 400 个 token
3. 到第一个 token 的 latency —— 同一台机器上，prefill 30 token 对比 430 token
4. 训练过程中的 canaries —— entropy 曲线和回答长度曲线，必须随每一份实验结果一起附上

这个 compliance 检查器不用 LLM 去评 LLM，用的是重复跑永远得到同样结果的 deterministic 规则：

# 下面两组是检查器直接匹配的泰语字面字符串，因此保持原文
POLITE = ("ครับ", "ค่ะ", "คะ", "นะครับ", "นะคะ")            # 各种句尾礼貌助词
REFUSAL = ("ไม่สามารถให้คำแนะนำ", "ขอแนะนำให้ปรึกษา", "ผู้เชี่ยวชาญ", "แพทย์", "ทนายความ")
# 依次为：无法提供建议 / 建议您咨询 / 专业人士 / 医生 / 律师

def th_ratio(s):
    thai = sum(1 for ch in s if "ก" <= ch <= "๛")
    letters = sum(1 for ch in s if ch.isalpha())
    return thai / max(letters, 1)

def comply(answer, is_restricted):
    ok_thai   = th_ratio(answer) >= 0.85                       # 规则 1：用泰语回答
    ok_polite = any(p in answer for p in POLITE)               # 规则 2：ครับ/ค่ะ
    ok_refuse = (not is_restricted) or any(k in answer for k in REFUSAL)   # 规则 3–4
    return ok_thai and ok_polite and ok_refuse

真正实测出来的数字，会去填第 9 节那张表（? 那些格子是由你的 notebook 来填的，不是我）。

下面是训练前后的真实回答，两边都是在看不见 context 的情况下作答的—— "前"那一侧是光秃秃的基座模型，"后"那一侧是 OPCD 学生。 点开每个例子看看，然后问问自己：如果没人告诉你，你分得出哪一个没看过 system prompt 吗？

9. 对比（Comparison）​

四个系统，全部在同一个测试集上测——前两行是地板和天花板，后两行才是真正的对手：

你应该看到的模式是：offline CD 和 OPCD 都会从地板往天花板爬， 而付出的 prompt 和地板那一行一样多——而两种方法真正拉开差距的地方是 OOD compliance 这一列： offline CD 只从老师的路径上学，所以遇到没见过的 prompt 类型时往往会破功； 而 OPCD 一直是在自己的路径上被批改的，所以在偏离训练路径时应当能更稳地守住规则。 如果这一列在 CI 之内根本分不出高下——那就是 null 结果，第 8 节的守则开始生效。

需要提防的坑​

1. Reverse KL + 小学生 = mode collapse 的风险 mode-seeking 是把双刃剑：容量不足的学生可能会用极端的方式"选一个 mode"—— 比如对每一个问题都回同一套拒绝话术，KL 确实很低，但完全没法用。 这正是 7.3 节要把输出 entropy 和回答长度记进 log 当 canary 的原因： 如果 entropy 一路下滑，同时回答越来越短、越来越重复，就停下来，降低 LR 或减少 epoch 数。

2. Top-K truncation bias top-128 上的 surrogate 只有在老师的 top-128 覆盖了几乎全部质量时，才会接近真正的 KL。 老师"犹豫不决"的位置（entropy 高，比如第一句话的开头）正是 coverage 掉下去、bias 冒出来的地方。 别靠猜——notebook 会把平均 coverage 和最低的 percentile 打印出来给你看，如果低得反常再去加大 K。

3. 老师和学生的 tokenizer 必须一致 逐位置的 KL 只有在两边完全按同样方式切 token 时才有定义——一旦跨了模型家族， vocab 就不是同一套，位置立刻没法对齐。在这一章里，这个条件是自动成立的， 因为老师和学生就是同一份权重——这也正是这套 setup 在教学上格外干净的原因： 你能完整地学到 distillation 的机制，而不必同时背上 tokenizer 的包袱。 （到了第 7 章，老师和学生是两个不同的模型，这个问题马上就会变成真事。）

4. Temperature 太低 = 学生只练自己本来就会的招 如果用很低的 temperature 采样，学生就只会产出自己有把握的回答， 于是老师批改到的全是学生本来就做得好的状态——真正行为还违反 persona 的那些地方， 几乎产生不了 gradient。temperature=1.0 强迫学生把自己带到还会出错的状态里去接受批改。

10. 小结（Summary）​

• 已经稳定下来的 system prompt 就是被放错地方的知识 —— 放在 prompt 里每次 request 都要付，放进 weights 只付一次
• Context distillation 训练看不见 $c$ 的学生去追平看得见 $c$ 的老师 —— 而在这一章里，老师和学生是同一份权重，区别只有 prompt 和 adapter
• KL 必须是 reverse（$\pi_\theta$ 放在前面）：mode-seeking 逼着学生不去做老师不做的事 —— 这正好就是 persona/safety 这类任务的需求
• rollout 必须是学生自己的：on-policy 从结构上消掉了 exposure bias， 因为训练时和 inference 时遇到的状态是同一批
• top-128 是一个算得出账的显存决策 —— 把 151,936 砍到 128， 代价是承认 objective 变成了 surrogate，然后用 coverage 来盯住它
• entropy 和回答长度是 mode collapse 的 canary —— 永远记进 log，别等坏了才发现
• 公平的 baseline 是 offline CD，不是光秃秃的模型 —— 而且如果没赢，就如实报告 null 结果

下一章： Model Distillation —— 这一次我们压缩的是模型，不是 prompt。 还记得第 2 节钉下的那句话吗：context distillation 改变的是"模型不用被告知就已经知道的东西"， 而 model distillation 改变的是"模型的大小"——老师大，学生小。 而本章白白捡到的 tokenizer 便宜，到那时就不再免费了。

参考文献（References）​

1. Ye et al. (2026). On-Policy Context Distillation for Language Models — OPCD——本章实现的核心方法
2. Askell et al. (2021). A General Language Assistant as a Laboratory for Alignment — 最初的离线上下文蒸馏（第 9 节的基线）
3. Snell et al. (2022). Learning by Distilling Context — 把上下文蒸馏成模型行为
4. Agarwal et al. (2023). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes — GKD：统一前向与反向 KL 的 JSD 框架
5. Gu et al. (2023). MiniLLM: On-Policy Distillation of Large Language Models — MiniLLM：使用反向 KL 的理由

本系列的文章、代码与 notebook 均以 CC BY-NC-SA 4.0 授权 —— 可自由使用与改编，须署名、限非商业用途，并以相同方式共享。文中引用的第三方模型与数据集仍适用各自的许可证。