Math-05.信息论-20.RLHF与KL约束

发表于 | 分类于
RLHF 流程、奖励模型、PPO 中 KL 惩罚、DPO 目标及参考策略约束的信息论解读。

本页从信息论视角梳理 RLHF（Reinforcement Learning from Human Feedback，人类反馈强化学习）中的 KL 约束：为何要对齐参考策略、公式如何出现。

段末注释：RLHF 用人类偏好训练奖励模型，再以强化学习微调语言模型；KL 惩罚 $D_{\mathrm{KL}}(\pi_\theta | \pi_{\mathrm{ref}})$ 防止策略偏离预训练模型过远导致能力遗忘。

系列入口：00.系列规划 | 前置：03 交叉熵与 KL、10 语言模型

1. RLHF 三阶段（D1）

图 1　SFT → RM → PPO/DPO

阶段 目标 损失/目标（概览）
1. SFT 监督微调 token CE（10 LM）
2. RM 奖励模型 偏好对 $(y_w, y_l)$：Bradley–Terry / ranking loss
3. RL 对齐 最大化奖励 + 约束 PPO + KL；或直接 DPO

$\pi_{\mathrm{ref}}$ 通常为 SFT 后策略，冻结作参考。

2. 带 KL 约束的优化目标（D3）

图 2　奖励 − β·KL

理想形式（每 prompt $x$，生成 $y$）：

$$
\max_\theta \ \mathbb{E}{x \sim D,, y \sim \pi\theta(\cdot|x)}\left[ r(x,y) - \beta , D_{\mathrm{KL}}(\pi_\theta(\cdot|x) | \pi_{\mathrm{ref}}(\cdot|x)) \right]
$$

  • $r(x,y)$：奖励模型分数
  • $\beta > 0$：KL 系数——大则更接近 $\pi_{\mathrm{ref}}$，小则更追奖励
  • $D_{\mathrm{KL}}$：在同一 $x$ 下，两个输出分布（整个序列或逐 token）的 KL

信息论含义：在「人类奖励」与「不浪费预训练信息」之间权衡；KL 是「相对参考多编了多少 bit」。

3. PPO 中的 KL（D7）

近端策略优化（Proximal Policy Optimization，PPO）在 RLHF 中常用：

$$
\mathcal{L}_{\mathrm{PPO}} = \mathbb{E}\left[\min\left(r_t(\theta)\hat{A}_t,\ \mathrm{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_t\right)\right]
$$

其中 $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\mathrm{old}}(a_t|s_t)}$。

额外 KL 惩罚或约束：

  • 在 reward 中加 $-\beta \log\frac{\pi_\theta}{\pi_{\mathrm{ref}}}$（token 级）
  • 或监控 $D_{\mathrm{KL}}(\pi_\theta | \pi_{\mathrm{ref}})$ 超阈值 early stop

防止 reward hacking 与 模式崩溃（重复、胡言乱语）。

4. DPO：闭式偏好目标（D3–D7）

图 3　DPO 无需显式 RM 采样

Direct Preference Optimization（DPO）将上述 KL 约束 RL 问题重参数化为偏好损失。对偏好 $(y_w \succ y_l)$：

$$
\mathcal{L}{\mathrm{DPO}} = -\mathbb{E}\left[\log \sigma\left(\beta \log \frac{\pi\theta(y_w|x)}{\pi_{\mathrm{ref}}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{\mathrm{ref}}(y_l|x)}\right)\right]
$$

  • 隐式含 KL 约束（$\beta$ 与 RL 中角色类似）
  • 无需单独 PPO 循环，实现简单
  • $\pi_\theta$ 与 $\pi_{\mathrm{ref}}$ 的 log 概率差 = 对数似然比，与信息论紧密相关

5. 与 CE / PPL 的关系（D6）

指标 阶段
CE / PPL SFT、评估通用能力
偏好 win rate RM / DPO 评估
$\mathbb{E}[D_{\mathrm{KL}}(\pi_\theta | \pi_{\mathrm{ref}})]$ 监控遗忘；过大→能力掉，过小→对齐不足

酶学/蛋白 LM：若有 SFT + 偏好数据（活性排序、专家评分），同样可用 LoRA + DPO（Math-03/20）。

6. 局限与工程（D8）

图 4　局限

问题 说明
$\beta$ 难调 需 grid search / 自适应 KL target
奖励模型偏差 RM 错则 RL 放大
KL 估计方差 token 级 Monte Carlo 噪声
长度偏差 长回复 RM 分高
仅 KL 不够 还需 repetition penalty、SFT mix

7. 概念代码（D12）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import torch
import torch.nn.functional as F

def token_kl(logits_theta, logits_ref):
    """逐 token KL(pi_theta || pi_ref)，对 vocab 维。"""
    log_p = F.log_softmax(logits_theta, dim=-1)
    log_q = F.log_softmax(logits_ref, dim=-1)
    p = log_p.exp()
    return (p * (log_p - log_q)).sum(dim=-1)  # (batch, seq)

def dpo_loss(logp_theta_w, logp_theta_l, logp_ref_w, logp_ref_l, beta=0.1):
    """标量 log prob 已对序列求和。"""
    logits = beta * ((logp_theta_w - logp_ref_w) - (logp_theta_l - logp_ref_l))
    return -F.logsigmoid(logits).mean()

# 示意
logp_tw = torch.tensor(-2.0); logp_tl = torch.tensor(-5.0)
logp_rw = torch.tensor(-2.5); logp_rl = torch.tensor(-4.0)
print("DPO loss:", dpo_loss(logp_tw, logp_tl, logp_rw, logp_rl).item())

8. 小结

RLHF = 奖励最大化 + KL 锚定 $\pi_{\mathrm{ref}}$；DPO 把 KL-RL 化为偏好 CE。Math-05 信息论 全系列至此贯通 CE → KL → MI → RLHF。

系列导航：00 规划 | Math-04 优化

-------------本文结束感谢您的阅读-------------