
Transformer 相关的文章里讲 KL 散度的不少但绝大多数都是一句话带过“损失函数里加了一个 KL 项”。这对理解模型行为来说远远不够。尤其是在 OPD 这类在线策略蒸馏框架里KL 的方向选择直接决定了学生模型是“覆盖教师分布”还是“聚焦教师分布”训练曲线和生成多样性都会因此产生明显差异。这篇文章我打算把反向 KL 这个细节彻底手撕一遍先对比正向 KL 和反向 KL 的数学性质再给出 OPD 的框架建模然后用 PyTorch 写出可运行的反向 KL 损失最后用一维高斯拟合的小实验观察 mode-seeking 行为。如果你正在做 Transformer 系列的模型蒸馏、LLM 对齐或生成模型微调这篇可以直接收藏。需要先说明一点OPD 在技术语境里最常见的意思是 Online Policy Distillation也就是在线策略蒸馏。它不是一个固定名字的开源仓库而是一类训练框架。核心设置是一个冻结的教师 Transformer一个正在训练的学生 Transformer每一步用学生当前策略采样输出再用教师对数似然作为训练信号。反向 KL 天然适合这种在线设置因为它的期望采样分布就是学生自己的策略。如果你在其他场景看到的是 Online Preference Distillation 或 Optimal Policy Distillation文中反向 KL 的推导、代码和调参思路同样适用。1. 核心内容速览内容维度说明技术主题Transformer / 生成模型中的反向 KL 散度核心框架OPD在线策略蒸馏Online Policy Distillation前置知识Softmax、交叉熵、自回归 Transformer 基本结构代码环境Python 3.8、PyTorch具体版本以本机环境为准核心函数reverse_kl_exact / reverse_kl_gumbel / reverse_kl_reinforce主要应用知识蒸馏、RLHF/DPO 风格策略约束、生成分布对齐是否涉及 Web API不涉及核心是训练损失函数是否涉及批量任务支持训练循环中按 batch 进行批量蒸馏显存说明取决于 batch_size、seq_len、vocab_size后文给出估算公式这篇文章不教你怎么部署模型服务也不涉及推理框架调用。它解决的是一个更底层的问题当你想让一个学生 Transformer 去模仿教师 Transformer 的分布时为什么要用“反向 KL”而不是直接用传统交叉熵或者正向 KL以及如何在 PyTorch 里把它实现成一个稳定可用的损失函数。2. 适用场景与使用边界2.1 适合谁用这个内容适合三类读者。第一类是做大模型蒸馏的算法工程师。比如有一个 7B 的教师模型想把它压缩成 2B 或者 1B 的学生模型OPD 框架下用反向 KL 可以在线获得教师对“学生当前生成结果”的反馈比一次性缓存所有教师 logits 的离线蒸馏更贴近学生当前的分布变化。第二类是在做 RLHF、DPO 或偏好对齐的工程师。RLHF 的 KL 惩罚项通常写作 KL(πθ || πref)这本身就是反向 KL。很多人在看公式时只记住了“加一个 KL 惩罚”但没有意识到这里的期望采样自当前策略 πθ而不是参考策略 π