1. 策略梯度与强化学习基础回顾
在深入探讨ReMax和RLOO之前，我们需要先理解策略梯度方法的基本原理。策略梯度是强化学习中的一类重要算法，它直接优化策略函数πθ(a|s)，即给定状态s下采取动作a的概率分布。
1.1 策略梯度定理
策略梯度定理给出了目标函数J(θ)关于参数θ的梯度表达式：
∇θJ(θ) = Eπθ[∑ψt∇θlogπθ(at|st)]
其中ψt是关键权重项，决定了梯度更新的方向和幅度。不同的策略梯度算法主要区别就在于如何设计这个ψt。
1.2 常见策略梯度变体
在实践中，我们常见以下几种ψt的设计：
-
蒙特卡洛回报：ψt = ∑γ^(t'-t)r_t'
- 代表算法：经典REINFORCE
- 特点：无偏但高方差
-
优势函数：ψt = A(s_t,a_t) = Q(s_t,a_t) - V(s_t)
- 代表算法：A2C, A3C
- 特点：需要训练价值网络
-
时序差分误差：ψt = r_t + γV(s_{t+1}) - V(s_t)
- 代表算法：TD Actor-Critic
- 特点：偏差-方差折中
2. LLM强化学习的特殊挑战
将策略梯度应用于大语言模型(LLM)微调时，我们面临几个独特挑战：
2.1 动作空间特性
- 动作空间是词汇表(通常5万+token)
- 每个时间步的动作选择影响后续状态
- 序列长度可达数百甚至上千token
2.2 奖励结构特点
- 通常只在序列结束时获得单一奖励
- 中间步骤没有即时奖励信号
- 奖励函数可能非常稀疏且噪声大
2.3 计算成本考量
- 模型参数量巨大(数B到数百B)
- 每次前向传播计算成本高
- 需要高效利用有限样本
3. ReMax算法详解
3.1 核心思想
ReMax的核心创新在于baseline的选择：
b(x) = r(x, ȳ_greedy)
其中ȳ_greedy是当前策略通过greedy解码(每步取概率最高token)生成的响应。
3.2 算法流程
-
对每个prompt x：
- 生成greedy响应ȳ
- 计算baseline b = r(x, ȳ)
- 采样K个响应{y_k} (使用top-p/top-k采样)
-
对每个采样响应y_k：
- 计算优势A_k = r(x,y_k) - b
- 计算策略梯度：
∇J = ∑A_k ∇logπ(y_k|x)
-
更新策略参数θ
3.3 实现细节
在实际实现中，有几个关键点需要注意：
-
greedy解码稳定性：
- 使用beam search可能比纯greedy更稳定
- 可以加入长度归一化
-
奖励归一化：
- 不同prompt的奖励尺度可能不同
- 可考虑对优势进行批归一化
-
采样策略：
- 通常使用top-p(p=0.9)采样
- 温度参数τ=0.7左右效果较好
4. RLOO算法详解
4.1 核心思想
RLOO(Leave-One-Out)采用同prompt下其他采样响应的平均奖励作为baseline：
b(x,y_j) = 1/(M-1) ∑_{k≠j} r(x,y_k)
4.2 算法流程
-
对每个prompt x：
- 采样M个响应
- 计算每个y_j的baseline b_j = mean(r(x,y_k) for k≠j)
-
对每个响应y_j：
- 计算优势A_j = r(x,y_j) - b_j
- 计算策略梯度：
∇J = ∑A_j ∇logπ(y_j|x)
-
更新策略参数θ
4.3 实现细节
-
采样数量选择：
- 通常M=4到8
- 更多样本降低方差但增加计算成本
-
奖励归一化：
- 可考虑对每个prompt的奖励进行z-score标准化
-
并行计算：
- 可以并行生成多个采样响应
- 需要高效实现奖励计算
5. 两种方法的对比分析
5.1 计算复杂度
| 方法 | 前向传播次数 | 奖励计算次数 |
|---|---|---|
| ReMax | K+1 | K+1 |
| RLOO | M | M |
5.2 适用场景
ReMax更适合：
- 当greedy响应质量稳定时
- 需要更少奖励计算时
- 资源受限环境
RLOO更适合：
- 当采样多样性更重要时
- 可以承担更多计算时
- greedy响应质量不稳定时
5.3 实际效果比较
根据实践经验：
- ReMax通常收敛更快
- RLOO最终性能可能更好
- 两者都可以不训练critic网络
6. 实现技巧与注意事项
6.1 奖励设计建议
-
尺度一致性：
- 确保不同prompt的奖励在相似范围
- 可考虑使用sigmoid等函数归一化
-
相对性：
- 奖励函数应能区分响应质量的细微差别
- 建议使用0-1或-1到1的范围
6.2 训练技巧
-
学习率选择：
- 通常需要较小学习率(1e-6到1e-5)
- 可以使用学习率warmup
-
批大小：
- 较大批大小(32-128)有助于稳定训练
- 可以使用梯度累积
-
正则化：
- KL散度正则化防止策略偏离初始模型太远
- 可以使用entropy正则化鼓励探索
6.3 常见问题排查
-
训练不稳定：
- 检查奖励函数是否有bug
- 降低学习率
- 增加批大小
-
性能下降：
- 检查KL散度是否太大
- 可能需要更强的正则化
-
收敛缓慢：
- 检查优势计算是否正确
- 可能需要调整采样策略
7. 扩展与变体
7.1 混合baseline
可以结合两种方法的优点：
b(x,y_j) = α·r(x,ȳ) + (1-α)·mean(r(x,y_k) for k≠j)
其中α∈[0,1]是混合系数。
7.2 多目标优化
可以同时优化多个奖励：
r(x,y) = ∑w_i·r_i(x,y)
其中w_i是各奖励的权重。
7.3 离线强化学习
可以结合离线数据：
- 使用重要性采样
- 加入行为克隆项
- 约束策略更新幅度
8. 实际应用案例
8.1 对话系统微调
使用ReMax优化对话响应质量：
-
定义奖励函数：
- 流畅性(语言模型概率)
- 相关性(与上下文匹配度)
- 趣味性(基于特定指标)
-
训练流程：
- 从人类对话数据初始化
- 使用ReMax优化3000步
- 每步batch_size=64
8.2 文本摘要优化
使用RLOO提升摘要质量：
-
奖励函数组成：
- ROUGE分数
- 事实一致性
- 信息密度
-
实现细节：
- M=6个采样/提示
- 学习率5e-6
- 训练2000步
9. 性能评估方法
9.1 在线评估
-
人工评估：
- 流畅性、相关性等维度评分
- 需要足够样本量
-
A/B测试：
- 与基线模型对比
- 统计显著性检验
9.2 离线评估
-
自动指标：
- 任务特定指标(如BLEU, ROUGE)
- 多样性指标
-
模型对比：
- 与监督微调比较
- 与其他RL方法比较
10. 未来发展方向
-
更高效的baseline估计：
- 使用小模型预测baseline
- 历史奖励的滑动平均
-
混合训练策略：
- 结合监督学习和RL
- 课程学习策略
-
多任务优化：
- 同时优化多个相关任务
- 共享表示学习
在实际应用中，我发现ReMax通常在训练初期表现更好，因为greedy响应提供了一个稳定的参照点。但随着训练进行，RLOO往往能取得更好的最终性能，因为它能利用更多样本信息。一个实用的技巧是在训练初期使用ReMax，后期切换到RLOO。