★ FEATURED ARTICLE 最近在调大语言模型强化学习的时候我又一次被KL散度惩罚给坑了。Loss曲线看起来挺正常奖励也在涨但跑出来的文本却越缩越短、越长越怪最后发现罪魁祸首就是KL惩罚系数和参考模型的选择出了问题。这篇东西不是教科书式科普而是我自己在跑RLHF、GRPO这类流程时对KL散度惩罚的理解、实现细节、调参血泪和排查经验。如果你已经在跑PPO、GRPO或者正准备把SFT模型拉去做对齐训练这篇内容应该能帮你少折腾几个晚上。我默认你懂一点强化学习基本概念也大概知道PPO里的优势函数、策略梯度是怎么回事。如果你只是听说过“RLHF”想先看个全貌那也不用慌我会把KL散度惩罚为什么存在、它到底在算什么东西、怎么把代码写对、怎么调参一层层拆开讲。1. 为什么LLM强化学习离不了KL散度惩罚1.1 没有KL惩罚模型会疯狂“薅奖励”大语言模型强化学习和游戏AI、机器人控制有一个显著区别我们没有一个能给出完美真实奖励的环境。所谓“奖励模型”本质上是一个从人类标注上训练出来的打分器它也有偏见、有盲点、可以被钻空子。如果你只让策略模型最大化奖励模型的打分它会找到一堆人类看不懂但奖励很高的“捷径”。最常见的捷径包括无限堆叠重复内容、无意义地加长输出、刻意使用某些高频表达、把答案强行塞进貌似结构的模板里。更隐蔽的是如果奖励模型对“长度”有隐式偏好模型会把答案越写越长反而丢掉信息密度。这些现象在业界被统称为reward hacking奖励黑客。你如果不给策略模型一个“刹车”它迟早会跑偏到奖励模型的盲区里。KL散度惩罚就是那个刹车。它不让策略模型随心所欲地远离某个参考策略而是要求模型在“提升奖励”和“保持与参考策略相近”之间做一个权衡。参考策略通常是SFT模型也就是经过监督微调、行为还比较“正常”的模型。加了KL惩罚后模型想走偏可以但必须付出代价代价大小由你给定的系数控制。一个标准的PPO目标函数通常会写成奖励目标 原始奖励 − β × KL(新策略概率 || 参考策略概率)这个β就是KL散度惩罚的强度系数。β大了模型就不敢动β小了模型就可能放飞自我。后面我会详细讨论怎么调。1.2 参考策略不是“老模型”而是安全锚点很多人第一次接触RLHF时会把参考策略和策略模型理解成“同一个模型的两个版本”。这样做初看没错但实操里有很多细节。参考策略通常是训练初始时保存的SFT模型副本在强化学习全过程中冻结权重不参与梯度更新。它扮演的是“原本那个还算乖的模型”的角色。策略模型开始可能和参考策略一模一样随着训练推进策略模型逐渐偏离参考策略KL散度就衡量这个偏离程度。你可以把参考策略想象成风筝线盘策略模型是风筝。奖励信号是风让风筝朝着高分方向飞。KL惩罚就是牵引线防止风筝被吹到看不见的地方。线放得太长风筝容易失控线拽得太紧风筝只能在原地打转。好的RL训练就是在风力和线长之间找到让风筝飞得高且稳的位置。在具体代码实现里策略模型和参考模型通常是同一个结构的两个权重副本。每算一个batch的数据你需要分别用策略模型和参考模型算出每个token的logprob对数概率然后用它俩的差来近似KL散度。参考模型不反传梯度但必须用和策略模型完全一致的tokenizer和padding方式否则算出来的KL完全是错的。2. 核心细节解析KL散度到底该用哪个方向、怎么算才不出错2.1 前向KL和反向KL差在哪选错会出大事KL散度严格来说不是距离因为它不对称。KLP||Q和KLQ||P可能是完全不同的值含义也不一样。在大语言模型强化学习里我们通常写的是KL(策略 || 参考策略)也就是把策略模型当作P参考模型当作Q。这个选择很关键。KL(策略||参考)对“策略在参考模型认为概率很低的地方出现”很敏感。换句话说如果策略模型在某个token上给出了高概率但参考模型认为那里应该是低概率这个方向的KL会迅速变大。这种特性会让策略模型不敢在参考模型完全不确定的位置上冒险整体行为偏向保守。另一方向KL(参考||策略)的做法在部分变分推断里常见但用在RLHF里容易让模型出现“mode collapse”式的激进也就是集中到某个高分模式上丢失多样性。所以主流的开源实现里都是用前向KL即KL(pi_theta || pi_ref)。实现时不需要真的按概率分布公式去累加。对于每个token我们可以用对数概率做蒙特卡洛估计KL_single ≈ log p_ref(token) - log p_policy(token)然后对一个序列的所有有效token求平均。这里需要注意我们不能拿“整个序列生成的概率”去做一个总的KL因为RLHF的奖励信号通常落在最后但生成的概率是逐步输出的。逐token级别的KL能让模型更细粒度地保持参考模型的语言习惯、风格和知识边界。2.2 是逐token算还是整句算粒度决定手感LLM强化学习中的KL散度惩罚绝大多数情况下是per-token级计算的。也就是对序列中每个非padding的位置都算一次“参考模型logprob减策略模型logprob”然后对这些值求平均。这样做的原因有两个第一策略梯度更新本来就是作用在token级别上的。我们最终的目标是提高高优势token的概率、降低低优势token的概率KL惩罚只有对齐到同一粒度才好在每个token的loss上直接做减法。第二整句级的KL会掩盖局部风险。比如有一半token模型已经完全跑偏另一半token还非常接近参考模型平均下来KL可能看起来不高但实际生成的文本已经开始局部崩坏。per-token的KL平均虽然也会被“摊薄”但至少每个位置的差异都能进入loss配合优势函数能更精准地抑制坏位置。在实现上你通常会在得到奖励后把奖励信号转成token级别的回报然后让每个token的loss里都含有一项policy_loss_token - advantage_token * ratio_token而KL惩罚是直接作为负项加到策略目标上或者单独加到总loss里。不同框架的做法略有区别但本质一样KL惩罚必须在policy更新方向上起作用也就是说当模型在某token上偏离参考模型时即使优势为正也可能因为KL惩罚太大而得不偿失。3. 实操过程把KL惩罚正确写进训练循环并调出稳定曲线3.1 从SFT到RL的标准管线模型加载和logprob计算在动手写代码前先把训练管线理清楚。你要有四个东西策略模型policy model要被强化学习更新的模型参考模型ref model冻结的SFT基线奖励模型reward model给输出打分的模型原始prompt数据集RL训练的起点如果用的是TRL库通常只需要传入policy model和ref model奖励模型可以单独load也可以在每条样本上推理。但如果你打算自己实现或者复现论文最好手动把ref model和policy model都加载进来明确二者tokenizer完全一样。计算logprob的标准流程是对prompt和模型生成的response拼成一个序列只对response部分计算loss和logprob不更新prompt部分分别用policy model和ref model过一遍相同序列得到response段每个token的对数概率用ref logprob减去policy logprob得到per-token KL这里有一个特别容易踩的坑两个模型要严格使用相同的padding位置。如果你给ref model用的max length和padding mask和policy model不一致KL会在padding位置产生无意义的计算。我们通常会给所有涉及KL计算的token位置设置mask只保留response部分的非padding token。下面是一段简化伪代码# 伪代码只展示核心概念 def compute_kl_penalty(seq, attention_mask, policy_logits, ref_logits): # 取 response 部分的 mask去掉 prompt 和 padding response_mask get_response_mask(seq, attention_mask) # log_softmax 得到每个 token 的 logprob policy_logp log_softmax(policy_logits, dim-1) # shape: [batch, seq_len, vocab] ref_logp log_softmax(ref_logits, dim-1) # 取实际token位置的logprob policy_token_logp gather_token_logp(policy_logp, seq) ref_token_logp gather_token_logp(ref_logp, seq) # 逐token KL: logp_ref - logp_policy per_token_kl ref_token_logp - policy_token_logp # 只对response部分有效token求平均 masked_kl per_token_kl * response_mask.float() n_tokens response_mask.float().sum(dim-1) kl masked_kl.sum(dim-1) / n_tokens.clamp(min1) return kl这里用logp_ref - logp_policy算的是单点蒙特卡洛估计的KL。它不要求你显著地跑完整分布计算因为在一个实际采样到的token上这个差值的期望本身就接近真实KL方向。在PPO里我们通常希望这个值不要太大控制在个位数到十几左右。3.2 PPO和GRPO里KL项的接入方式PPO的实现里KL散度惩罚最直接的方式是把它加到总loss中。展开一点PPO会对所有响应token计算一个importance ratio并乘以优势函数同时做一个clip约束防止更新过猛。KL惩罚可以在ratio上做修正也可以作为策略目标的额外项。TRL的PPO实现中选择的经验做法是每个token的loss -优势 * ratio然后总loss里额外加一个KL惩罚项等价于在logprob更新时手动压低偏离。GRPO则有点不同。它不需要critic网络优势通过同一组采样输出之间的奖励相对值估计。TRL和很多开源复现里KL惩罚仍然保留而且被证明对稳定性至关重要。GRPO里常见的做法是给KL加一个系数β控制“当前这组样本里策略模型相对于参考模型的偏移”不要过大。有人会觉得GRPO里已经用了相对奖励是不是不用KL了我试过直接把KL系数设成0训练确实能跑但生成质量会在几千步后明显下降表现是重复率上升、句子碎片化。所以哪怕在GRPO里KL惩罚也不是装饰品。实际训练中我会把KL均值记录下来作为重要的早停指标。经验上一个长度在128到512 token的responseper-token KL均值在0到15之间比较常见。超过20就要警惕超过50基本已经疯了。你可以设定一个target_kl然后动态调整β如果最近一个窗口的KL均值大于target_kl就加大β如果KL均值远小于target_kl就适当减小β这也是很多人说的adaptive KL。它比固定β更省心因为reward scale在不同训练阶段变化很大固定β很难一直合适。3.3 超参数选择β和target_kl到底怎么设没有任何一组超参能通吃所有任务。但你可以用一个比较稳的起点β初始值0.02到0.1之间target_kl每token平均KL设置在5到10之间如果reward scale平均在1到5β用0.05左右通常不至于让KL爆炸我自己的习惯是先把奖励模型输出的reward做一下标准化让均值尽量落在0附近。不要直接拿原始reward打分去和KL惩罚相加因为不同奖励模型的scale可能差别很大。你在本地小模型上做验证时经常会出现某个奖励模型打分普遍在10以上另一个打分普遍在0.3附近。如果β不变前者会让KL惩罚完全不起作用后者会让KL惩罚重到训练停滞。做标准化思路很简单在RL初始化时先用一小批样本估算reward分布的mean和std然后对reward做类似z-score的变换再和β×KL相加。这样β的调解语义就稳定很多。下面是一个简化的loss计算公式你可以对照着理解# 每个token的最终loss per_token_kl ref_logp - policy_logp policy_gradient_loss -advantages * ratio loss policy_gradient_loss beta * per_token_kl需要注意隐藏的顺序有些实现会把KL惩罚直接放进ratio里比如把优势减掉KL再乘ratio效果类似但梯度分布略有不同。如果发现某一种写法下reward涨但KL疯涨可以试试换另一种接法往往能改善。4. 常见问题与排查技巧实录4.1 KL散度爆炸生成文本迅速退化这是最常见的翻车现场。训练进行几千步后reward可能还在涨但你打开几个样本一看全是重复词、废话或者突然变成另一种语言风格。这时候看KL曲线大概率是一条向上陡增的线。排查顺序我推荐这样确认tokenizer一致。手动让policy和ref分别tokenize同一句话比对input_id是否完全一致。确认padding mask正确。有没有把padding位置的KL也当成有效值累计。确认logprob取到的是response段而不是prompt段。有些实现误把prompt部分的KL也算了进去。降低学习率。RL training里学习率过高是KL爆炸最常见的原因之一。调大β或者调低target_kl。如果用的是adaptive KL把target_kl从10降到6往往能在几天前发现危险。另一种“伪爆炸”是优势标准化没做。如果advantage的scale突然变大ratio变化剧烈KL会被间接顶上去。这种情况你会发现KL和advantage均值同步飙升处理方式不是调β而是给advantage做更稳的标准化或者加大PPO clip范围限制。4.2 奖励不动、KL几乎为0模型“躺平”了和爆炸相对的另一种情况是KL一直很小但奖励也不涨整个训练像死水一样。很多人以为KL小就是好事代表模型没跑偏但实际上可能是“跑不动”。原因通常有三类第一β太大。KL惩罚太重模型稍微偏离参考就会受到巨大loss惩罚梯度更新朝“不偏离”的方向走直接把奖励信号淹没了。你可以把β降到原来的十分之一试试。第二参考模型太强策略模型和参考模型的logprob几乎一致。如果初始策略就是参考模型而参考模型本身已经很强模型只在很小的邻域探索奖励变化自然很小。这种情况可以适当增大KL target允许模型更大胆尝试。第三advantage计算有问题。比如你的baseline估计太强导致大多数advantage都接近0更新信号微弱。GRPO里如果你选的group内样本质量都差不多也会出现优势偏小的问题。可以检查一下生成的多样性如果beam搜索或temperature太低group内样本几乎一样优势自然没有区分度。4.3 数值问题NaN、负KL、怪异大数KL散度理论上是非负的但实际代码里你会看到负值尤其是当采用蒙特卡洛估计logp_ref - logp_policy时。因为某个token在策略模型下概率比参考模型高它就会出现负贡献。这本身不是bug但如果你把所有负值都强行clip到0就破坏了无偏估计。正确的做法是保留符号靠整体均值来反映平均偏离。NaN问题则需要重点排查log_softmax溢出当logits特别大时用朴素softmax可能溢出。用log_softmax而不是先softmax再log。奖励模型给了一个极大值导致advantage极大接着ratio极大loss也极大。给reward做截断或者缩放会有帮助。mixed precision训练下KL的均值在fp16里可能不够稳定。建议KL累加时用fp32或者单独统计。4.4 问题排查速查表症状最可能原因优先操作KL曲线持续上涨β过小 / 学习率过高 / tokenizer不一致调大β、降低LR、检查tokenizer与maskKL接近0且reward停滞β过大 / group样本多样性不足调小β、增加采样temperatureloss出现NaNlogits溢出 / reward极端值用log_softmax、reward截断、fp32统计KL生成重复严重但KL不高逐token平均掩盖局部问题观察分位KL或按token位置绘图降低采样长度reward大幅度波动reward scale未标准化对reward做z-score或固定scale训练正常但reward不涨参考模型过强允许更大KL提高β宽松度或降低β这张表是我在跑7B模型时反复用到的“第一直觉清单”。当然它们不一定百分百覆盖你的问题但能帮你快速圈定范围。5. 本地小模型上的实验策略和自研代码建议5.1 为什么我建议先用小模型验证KL逻辑很多同学一上来就想在70B上做RLHF但KL散度惩罚的调试成本非常高一次实验几百卡时起步。我自己更习惯先在一个1B或3B的本地模型上用单卡或者双卡把流程完全跑通。小模型不是“缩小版的大模型”它的reward信号和KL行为和大模型有一定差异但KL惩罚的核心问题——方向、粒度、系数、mask——完全一致。小模型上跑三小时就能看到KL曲线走向比大模型上跑三天才发现爆炸要省钱得多。本地部署的另一个优势是你可以频繁修改代码打日志、dump中间输出、画KL分布热力图。在集群上搞这些很麻烦。我通常会在小模型上先确认三件事KL均值随时间变化是否平稳、KL分布是否存在极端离群点、以及reward和KL之间是否有明显相关。确认完再去放大模型基本不会出大的稳定问题。5.2 用现成库还是手写一个可控制的实现如果你刚开始做LLM RL直接使用成熟的库肯定最快比如TRL的PPOTrainer、GRPOTrainer或者DeepSpeed的RLHF框架。这些库已经把KL惩罚、advantage估计、token mask这些基础逻辑封装好了。但缺点是出问题时你很难知道它内部到底把KL加在哪个位置、有没有做adaptive系数调整。我个人的做法是先读一遍开源库中KL相关源码确认它用的是per-token还是sequence-level然后跑一个打印日志的小实验。等你理解透了再在必要的地方魔改。我自己在落地项目时会单独抽出一个kl_penalty.py模块里面只负责输入policy logprob、ref logprob、response mask返回per-token的KL值和最终loss。这样无论上层用PPO还是GRPOKL策略都能统一更换。自研的小模块用一个简单接口就够class KLController: def __init__(self, init_coef, target_kl, methodadaptive): self.coef init_coef self.target_kl target_kl self.method method def update(self, recent_kl): if self.method fixed: pass elif self.method adaptive: if recent_kl self.target_kl * 1.2: self.coef * 1.2 elif recent_kl self.target_kl * 0.8: self.coef * 0.8 return self.coef用这个控制器每若干step更新一次β比手工改配置更平滑。训练结束后把KL变化过程和β变化过程画在同一个图上你会发现它们之间存在完好的“负反馈”对应关系。这也是判断RL训练是否健康的一个快速标准。把KL惩罚看成和reward并肩的量而不是一个附属项很多玄学问题就能解释通了。我踩过最大的坑就是把β当成固定超参以为训练一次调好就万事大吉。真正稳的做法是让β跟着KL实时走并严格盯住每一轮的平均KL、分位数KL以及KL和reward的互相关系。第一次跑通之后你再看那些开源代码的KL相关逻辑会突然觉得每个字段都很有道理。
阅读完成 · 觉得有帮助？
