PPO+RLHF Training Loop with Value Model — Walkthrough
本页逐步展示在 PPO + RLHF 场景下引入 value model 时,训练循环每一步到底在做什么:4 个模型的角色、混合结构的 shaped reward、一次 rollout 的完整流程、GAE 信号沿序列衰减的数值感受,以及相对 GRPO 的对比。阅读前建议先读 ppo 和 value-model。
整体设置
训练时显存里同时存在 4 个模型:
| 角色 | 是否可训练 | 初始化来源 |
|---|---|---|
| Policy | 是 | 从 SFT checkpoint 拷贝 |
| Reference | 冻结 | SFT checkpoint 的另一份副本,用于 KL 正则 |
| Reward model | 冻结 | 之前用偏好对训练好的 Bradley-Terry RM |
| Value model | 是 | 通常从 拷贝初始化,两者架构相同 |
Policy 和 value 都是 Transformer + 一个输出头:policy 用 LM head(词表 logits),value 用 scalar head(单个标量)。两种常见架构选择:
- 分离 body: 是 的完整拷贝,有自己的 Transformer body 和 scalar head。与 policy 的梯度完全解耦,更稳定,是 TRL、DeepSpeed-Chat、LLaMA-Factory 等主流库的默认选择。代价是显存开销翻倍。
- 共享 body:policy 和 value 共享同一个 Transformer body,只挂两个独立的 head。省显存但两个任务的梯度会在 body 参数上叠加,必须小心调节 value loss 权重与两个 head 的学习率。
混合结构的 shaped reward
RLHF 的 reward 不是”只有末端有信号”,而是一个稠密 KL 惩罚 + 末端 RM 打分的混合结构。给定 prompt 和 response ,第 个 token 的 shaped reward 是:
关键观察:
- 每个 token 都有一个 per-token KL 项,作为”偏离参考策略的代价”。这一项是稠密的、方差低的。
- 只有末端 token 额外带上 ,也就是 reward model 对整条 response 的打分。这一项是稀疏的、方差高的。
- 两项都由 控制强度;典型取值 。
Value model 学的是什么
给定状态 (prompt + 已生成的前缀 ), 的目标是从 出发到 episode 结束的期望累积 shaped reward:
展开看它在学两件叠加的事:
- “如果从当前前缀继续生成,RM 最终会打多少分”。这是一个极难的估计问题,因为需要 value model 隐式地”模拟” policy 的后续生成才能给出一个合理的预测。它也是 value model 方差的主要来源。
- “从当前位置到结尾的累计 KL 偏离”。这一项相对容易学,因为每个位置都能观测到 per-token KL,属于稠密信号。
两者被”压缩”在同一个标量输出里——value model 没有办法在结构上把它们分开。这正是 LLM 场景下 value model 难训的根本:同一个输出同时被迫承担一个稀疏高方差的终端信号和一个稠密低方差的 per-token 信号。
一次 PPO+RLHF 训练迭代的完整流程
每一轮训练由 rollout 阶段 和 update 阶段 两部分组成。
Step 1:rollout
冻结当前 policy (把 拷一份叫 ),从 prompt 集合中取一批 ,对每条 prompt 用 生成完整 response 。生成时记录下面这些量:
rollout_data[i] = {
prompt: x^(i),
tokens: y^(i) = [y_1, ..., y_T],
logp_old: [log π_θ_old(y_t | x, y_<t) for t in 1..T],
logp_ref: [log π_ref(y_t | x, y_<t) for t in 1..T],
V_old: [V_ψ_old(s_t) for t in 0..T-1],
rm_score: r_φ(x, y) # 只有一个标量, 在末端
}
这一步里所有的前向都来自冻结的快照( 和 )。
Step 2:组装 shaped reward
对每条轨迹,按上面公式组装 :
for t in 1..T-1:
r_t = -β * (logp_old[t] - logp_ref[t])
r_T = rm_score + (-β * (logp_old[T] - logp_ref[T]))
注意 KL 项是用 rollout 时的 相对于 计算的,而不是用当前正在更新的 。这一点很重要:rollout 数据在所有 update epoch 里共享同一个 shaped reward 向量。
Step 3:计算 GAE advantage 与 value target
反向扫描一条轨迹计算 TD 残差和 GAE:
V_old[T] = 0 # 终止状态的 value 定为 0
for t in T, T-1, ..., 1:
δ_t = r_t + γ * V_old[t] - V_old[t-1]
Â_t = δ_t + γ * λ * Â_{t+1} # Â_{T+1} = 0
V_target[t-1] = Â_t + V_old[t-1]
RLHF 场景取 、。V_\text{target}_t 也被称作 -return,当 时就是这条轨迹从 开始的实际累积 reward。
组装好 后通常再做一次 minibatch 级的 z-score 标准化(),这是 PPO 训练稳定性的一个关键 trick。
Step 4:多轮 update
对同一批 rollout 数据做 个 epoch 的 minibatch SGD,典型 。每个 minibatch 上计算:
Policy loss(clipped surrogate,见 ppo):
其中 。注意分母的 用的是 rollout 时记录的 logp_old,不随 update 步骤变化。
Value loss(回归到 ):
L^{VF} = \hat{\mathbb{E}}_t\Big[(V_\psi(s_t) - V_\text{target}_t)^2\Big]
主流实现会再加一层 value clip:
L^{VF}_\text{clip} = \hat{\mathbb{E}}_t\Big[\max\big((V_\psi(s_t) - V_\text{target}_t)^2, (\text{clip}(V_\psi(s_t), V_{\psi_\text{old}}(s_t) - \epsilon_v, V_{\psi_\text{old}}(s_t) + \epsilon_v) - V_\text{target}_t)^2\big)\Big]
Entropy bonus:,鼓励探索。
合并:,然后反向传播更新 和 。典型 ,。
Step 5:下一轮
把 、 同步到 、,回到 Step 1 开启新一轮 rollout。
数值感受:GAE 信号沿序列衰减
设 、、、。假设 rollout 时每个位置的 per-token KL 约为 0.1(训练中期的典型值),RM 给出 。
Shaped reward 向量
Value 严重低估的情形
训练早期 value model 从 RM 初始化,但 RM 只在训练数据分布上学过”给完整 response 打分”,对不完整前缀的输出基本随机。假设它在每个位置都给出 (偏低)。
反向计算 TD 残差:
- ()
GAE 递推 :
| 20 | |
| 19 | |
| 18 | |
| 15 | |
| 10 | |
| 5 | |
| 1 |
结论:末端 token 获得 的正向梯度,但第 1 个 token 只获得 ——沿序列指数衰减。第 1 个 token 拿到的学习信号大约是末端的 倍。
这是 long-horizon credit assignment 的直接后果:GAE 再聪明也改变不了”RM 的有效信息只在序列末端”的事实。对于 RLHF 普通对齐任务 response 长度 这种情况还可接受;对于推理模型动辄 thousand-token 的 CoT, 衰减到第 1 个 token 时信号已经近似于零,这也是为什么推理模型训练更倾向于 GRPO——它把整条 response 看作一个原子动作,避免了 per-token 衰减。
Value 准确的情形(理想状态)
若 value model 经过若干轮训练后已经能准确预测终端 reward 减累计 KL,此时 。计算 TD 残差:
- (中间所有位置)
此时所有 ,policy 对这条轨迹不会有大的梯度——这是”value model 做得好”的理想状态:所有”意料之中”的样本都不贡献 policy 梯度,只有”意料之外”的样本(真正的好或坏)才会推动学习。
这套流程在实践中的几个难点
-
Value model 追不上 policy。 每一批 rollout 之间都在变,真实的 也在跟着变,但 value model 只看到了上一批 rollout 的训练样本。Policy 变化越快,value 越失准,advantage 越嘈杂。主流做法是限制每批 rollout 的 epoch 数到 2–4 次、限制 KL 增长速度、限制 rollout 规模。
-
RM 的 miscalibration 会被 value model 继承。若 RM 对某类 response 系统性偏高,value model 学到的也会偏高,advantage 就会系统性向这类 response 倾斜——这正是 奖励黑客(reward hacking) 的一条典型触发路径。见 reward-model 的 reward hacking 节。
-
Value loss 权重 的选择极敏感。 过大时 value 梯度主导整体 loss,policy 基本不更新; 过小时 value 学得慢,advantage 方差大。InstructGPT 与 TRL 默认取 。
-
共享 body 时的梯度干扰。若 和 共享 Transformer body,policy 梯度和 value 梯度会在 body 参数上叠加。由于两个任务梯度尺度差异很大,需要小心 scale value loss 或给两个 head 用不同学习率,否则训练极易崩溃。
-
Value head 初始化的不匹配。从 RM 初始化拷过来的 scalar head 是”EOS 位置打分”的分布,但现在要在每个中间位置都输出。训练早期中间位置的预测往往很嘈杂,需要几百到几千步才会稳定。一些实现会在训练前先对 value head 做一次”warm-up”——冻结 policy,只跑 rollout + 回归 若干步。
-
KL 项是否进入 value target。不同实现取法不同:有的把 KL 惩罚算进 shaped reward 从而也进入 (即 value 学”net reward”),有的只把 RM 打分算进 而 KL 惩罚单独处理。这两种选择在数学上并不等价,会导致相同代码看到不同的训练曲线。TRL 默认是前者(KL 进 reward)。
和 GRPO 的对比
GRPO 的出发点就是**“既然 value model 在 RLHF 里这么难训,那就扔掉”。对同一个 prompt 采样 条 response,用组内 reward 的经验均值和标准差做 baseline,相当于把上面 Step 3 里 value model 的功能替换成组内 Monte Carlo 估计**。对比:
| 维度 | PPO + Value Model | GRPO |
|---|---|---|
| Baseline 来源 | 每个 token 位置的 | 每个 prompt 的组内均值 |
| Baseline 偏差 | 依赖 critic 收敛度,早期有 bias | 大 时无偏,方差 |
| 每 token 区分度 | 有,GAE 给每个 token 不同 advantage | 无,整条 response 共享同一 advantage(outcome supervision) |
| 显存里模型数 | 4 | 2 |
| 每 prompt 采样数 | 1 也行 | 必须 |
| 长序列衰减 | GAE 指数衰减,早期 token 信号弱 | 无衰减,所有 token 共享同一信号 |
注意 GRPO 丢掉了 per-token advantage 的区分能力——所有 token 共享同一个 advantage 值。如果任务需要精细的 token 级信用分配(例如长程推理链里每一步贡献不同),outcome supervision 版 GRPO 反而可能不如 PPO + value model。这也是为什么 process supervision 版 GRPO 要额外引入步骤级 reward:它实质上是在”模拟” value model 对中间位置的区分作用。
反过来,对于 reasoning 任务(长 CoT + 稀疏正确性 reward),PPO + value model 的主要障碍恰恰是”长序列 × 稀疏末端 reward × GAE 指数衰减”——前面已看到 、 时第 1 个 token 只得到 38% 的信号;若 则衰减为 ,有效信号已经完全消失。在这种场景下 GRPO 的”整条 response 共享同一 advantage”反而更合理。
Comments