DeepSpeed-Chat 中PPO奖励设计的解析以及为什么只在最后一个时间步加上奖励模型的分数？

博客：DeepSpeed-Chat 中奖励设计的解析

背景介绍

在深度强化学习中的奖励设计是优化策略的核心，尤其是在基于偏好建模（Preference Model）的强化学习中。DeepSpeed-Chat 是一个开源的 RLHF（Reinforcement Learning with Human Feedback）框架，使用了 PPO 算法。其奖励函数的设计直接影响模型的学习目标与最终生成效果。

本文结合 DeepSpeed-Chat 提供的源码与理论公式，对其奖励设计进行深入解析，探讨实际代码与理论公式的对应关系。

---

奖励设计的理论基础

DeepSpeed-Chat 中的奖励函数 ( $R_t$ ) 定义如下：

$$R_t=\{\begin{array}{ll}-\text{kl\_ctl}\cdot \left(\log \frac{P(A_t|S_t)}{P_{\text{ref}}(A_t|S_t)}\right),& t\ne T\\ -\text{kl\_ctl}\cdot \left(\log \frac{P(A_t|S_t)}{P_{\text{ref}}(A_t|S_t)}\right)+R_t,& t=T\end{array}$$

其中：

• ( $\text{kl\_ctl}$ )：KL 控制系数，通常用于平衡奖励的幅度。DeepSpeed-Chat 中默认为 0.1。
• ( $P(A_t|S_t)$ )：当前策略生成的动作分布。
• ( $P_{\text{ref}}(A_t|S_t)$ )：参考策略（一般是微调模型或初始模型）的动作分布。
• ( $T$ )：生成序列的终点，只有在最后一个时间步，奖励才加上额外的 ( $R_T$ )（通常是奖励模型的分数）。

解释公式含义

1. KL 惩罚项：

   • 中间时间步 (( $t\ne T$ )) 的奖励完全由 KL 散度控制，表示当前策略与参考策略的差异。
   • 公式可以改写为：
     $$R_t=-\text{kl\_ctl}\cdot (\text{log\_probs}-\text{ref\_log\_probs}),$$
     其中：
     • ( $\text{log\_probs}=\log P(A_t|S_t)$ )
     • ( $\text{ref\_log\_probs}=\log P_{\text{ref}}(A_t|S_t)$ )

2. 终点奖励：

   • 在终点 (( $t=T$ ))，在 KL 惩罚的基础上增加奖励模型的得分 ( $R_T$ )。
   • 奖励模型得分通过裁剪函数限制在一定范围：
     $$R_T=\text{clip}(r,-\text{clip\_reward\_value},\text{clip\_reward\_value}),$$
     这里 ( $r$ ) 是从奖励模型输出的分数。

---

源码解析

在 DeepSpeed-Chat 中，这一理论设计被具体实现为以下代码：

源码链接：https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/dschat/rlhf/ppo_trainer.py#L44

def compute_rewards(self, prompts, log_probs, ref_log_probs, reward_score,
                        action_mask):

        kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)
        rewards = kl_divergence_estimate
        start = prompts.shape[1] - 1
        ends = start + action_mask[:, start:].sum(1) + 1
        reward_clip = torch.clamp(reward_score, -self.clip_reward_value,
                                  self.clip_reward_value)
        batch_size = log_probs.shape[0]
        for j in range(batch_size):
            rewards[j, start:ends[j]][-1] += reward_clip[j]

        return rewards

关于函数具体执行过程可以参考笔者另一篇博客: 详细解释DeepSpeed-Chat中ppo训练的代码：compute_rewards函数解析

1. KL 惩罚计算

kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)
rewards = kl_divergence_estimate

这部分代码对应理论中的 KL 惩罚项：

$$R_t=-\text{kl\_ctl}\cdot (\text{log\_probs}-\text{ref\_log\_probs})$$

• log_probs：当前策略生成的 log 概率。
• ref_log_probs：参考策略生成的 log 概率。
• self.kl_ctl：控制 KL 项的权重。

2. 扩展奖励序列

为了支持变长生成，每个样本的奖励需要扩展到生成的实际长度：

start = prompts.shape[1] - 1
ends = start + action_mask[:, start:].sum(1) + 1

• start：表示生成部分的起点（通常为 prompt 长度）。
• ends：每条序列的生成终点，计算方式为 start 加上生成的 token 数。

3. 奖励模型分数的裁剪

reward_clip = torch.clamp(reward_score, -self.clip_reward_value, self.clip_reward_value)

• reward_score：奖励模型的输出分数。
• self.clip_reward_value：裁剪范围，限制奖励幅度，避免过大的梯度波动。

4. 终点奖励更新

for j in range(batch_size):
    rewards[j, start:ends[j]][-1] += reward_clip[j]

这段代码在每条序列的最后一个时间步（rewards[j, start:ends[j]][-1]）添加奖励模型分数。

---

设计的意义

1. KL 惩罚控制训练稳定性：

   • KL 项的引入防止策略偏离参考策略太远，避免模型生成过多不合理的内容。
   • ( $\text{kl\_ctl}$ ) 的值调节了惩罚的力度，常设为 0.1。

2. 终点奖励增强奖励信号：

   • 奖励模型的分数只在终点加入，避免生成中间部分奖励过多、导致不必要的策略偏移。
   • 裁剪函数进一步稳定了训练过程。

3. 支持变长序列：

   • start 和 ends 的计算允许批次内的序列长度动态变化，提高模型的生成适应性。

---

总结

DeepSpeed-Chat 中的奖励设计核心是结合 KL 惩罚和奖励模型分数，并通过裁剪与动态长度扩展确保稳定性与灵活性。这种设计在实际中平衡了模型生成的多样性和与参考策略的接近程度，是 RLHF 实现中的一种经典方案。

为什么只在最后一个时间步加上奖励模型的分数 ( $R_T$ )?

在 RLHF（Reinforcement Learning with Human Feedback）的实现中，这种设计主要是基于以下原因：

---

1. 奖励模型分数的语义性

奖励模型分数通常表示整个生成序列的质量，而不是某个具体时间步的表现。例如：

• 奖励模型的训练目标通常是对整个生成的句子或段落进行打分，而不是评估生成过程中每一步的质量。
• 在这种情况下，将奖励模型分数直接赋值给整个序列显然不合理。因此，它被设计为只加到终点时间步 ( $t=T$ ) 的奖励中，以反映最终生成结果的整体质量。

这种设计确保了奖励模型的分数主要作用于整个生成序列，而不会误导中间时间步的策略。

---

2. 稳定性和训练效率

在 RLHF 中，奖励信号是通过策略网络和参考模型（Ref Model）的 KL 惩罚控制生成的。如果每个时间步都引入额外奖励分数，可能会导致以下问题：

• 梯度不稳定：如果奖励模型的分数过大或波动较大，分散到所有时间步可能导致梯度爆炸或策略训练方向不稳定。
• 训练效率低：将奖励分数扩展到每个时间步可能让策略网络对局部优化更敏感，难以收敛到对全局质量更优的策略。

因此，将奖励分数局限于最后一个时间步，可以减少梯度噪声，提高训练的稳定性和效率。

---

还能进行怎样的优化？

虽然当前设计有其合理性，但仍存在改进的空间：

1. 平滑奖励分布

• 方法：将奖励模型的分数 ( $R_T$ ) 按权重分散到多个时间步，而不是只加到终点。例如：
  $$R_t=w_t\cdot \text{clip}(r,-c,c)$$
  其中，( $w_t$ ) 是一个递增的权重函数，例如 ( $w_t=\frac{t}{T}$ )。
• 优点：
  • 更平滑的奖励信号，策略网络能够感知中间时间步对最终生成的贡献。
  • 提高对中间时间步生成质量的关注。

2. 引入细粒度奖励信号

• 方法：在训练奖励模型时，设计更细粒度的评分机制，对序列中的每个子片段或时间步进行单独打分。这需要将奖励模型扩展为生成过程的局部评价器。
• 优点：
  • 提升奖励的时间步相关性。
  • 有助于避免生成过程中局部错误对全局奖励的影响。

3. 动态权重分配

• 方法：使用策略网络的注意力机制或其他特征，将奖励模型的分数动态分配到不同时间步。例如：
  $$R_t={\alpha }_t\cdot R_T$$
  其中 ( ${\alpha }_t$ ) 由策略网络生成，表示时间步 ( $t$ ) 的重要性权重。
• 优点：
  • 更灵活地调整不同时间步的奖励。
  • 利用上下文信息提高奖励分布的合理性。

---

其他框架如何实现？

1. OpenAI 的 PPO-RLHF

• 设计类似：在 GPT-3.5 等模型的训练中，OpenAI 的 PPO-RLHF 实现也将奖励模型分数仅作用于终点时间步，而中间时间步主要通过 KL 惩罚调整。
• 原因：保持简单性，避免对中间时间步的奖励分配导致策略过拟合局部时间步。

2. Anthropic 的 Constitutional AI

• 差异点：Anthropic 的 Constitutional AI 框架尝试通过规则约束引导生成，某些规则可以动态作用于生成过程中的不同时间步，而不是仅在终点处理整体奖励。
• 特点：利用规则的奖励信号更注重生成过程中的逻辑性和一致性。

3. 自监督强化学习（Self-Supervised RL）

• 设计不同：一些强化学习方法尝试在每个时间步直接引入奖励信号，例如将对生成结果的打分信号回传到每个时间步的策略中。
• 挑战：这种方法可能导致策略对局部奖励信号过于敏感，难以优化全局生成质量。

---

总结

只在最后一个时间步加上奖励模型的分数是一种合理的折中设计，兼顾了奖励信号的语义性和训练的稳定性。然而，改进方向包括：

1. 平滑奖励信号分布；
2. 细粒度奖励建模；
3. 动态调整不同时间步的奖励权重。

不同框架根据应用需求和模型特点，采用了类似的折中策略，但也在尝试探索更精细的奖励设计，以提高生成质量和优化效率。

参考

[1] 图解大模型RLHF系列之：人人都能看懂的PPO原理与源码解读

后记

2024年12月12日19点47分于上海，在GPT4o大模型辅助下完成。