PPO 算法在 RLHF 中的应用：让模型学会理解人类偏好
PPO 算法在 RLHF 中的应用：让模型学会理解人类偏好
前言
RLHF（Reinforcement Learning from Human Feedback）是让大模型对齐人类偏好的关键技术。PPO（Proximal Policy Optimization）是 RLHF 中最常用的强化学习算法。
我在项目中实现过完整的 RLHF 训练流程，对 PPO 算法有深入理解。今天分享 PPO 的原理和在 RLHF 中的应用。
PPO 算法原理
核心思想
PPO 的核心是在策略优化过程中保持策略的更新不要太大：
import torch import torch.nn as nn import torch.optim as optim class PPO: """PPO 算法实现""" def __init__(self, policy_net, value_net, lr=3e-4, gamma=0.99, eps=0.2): self.policy_net = policy_net self.value_net = value_net self.gamma = gamma self.eps = eps # PPO 裁剪参数 self.policy_optimizer = optim.Adam(policy_net.parameters(), lr=lr) self.value_optimizer = optim.Adam(value_net.parameters(), lr=lr) def compute_advantages(self, rewards, dones, values): """计算优势函数""" advantages = [] running_advantage = 0 for t in reversed(range(len(rewards))): if dones[t]: running_advantage = 0 running_advantage = rewards[t] + self.gamma * running_advantage - values[t] advantages.insert(0, running_advantage) # 标准化优势 advantages = torch.tensor(advantages) advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) return advantages def update(self, states, actions, old_log_probs, advantages, returns): """PPO 更新""" # 获取新的 log_probs 和 values log_probs, values = self.policy_net.get_log_probs(states, actions) values = values.squeeze() # 计算概率比率 ratio = torch.exp(log_probs - old_log_probs) # PPO 裁剪 surr1 = ratio * advantages surr2 = torch.clamp(ratio, 1 - self.eps, 1 + self.eps) * advantages # 策略损失 policy_loss = -torch.min(surr1, surr2).mean() # 值函数损失 value_loss = nn.MSELoss()(values, returns) # 总损失 total_loss = policy_loss + 0.5 * value_loss # 优化 self.policy_optimizer.zero_grad() self.value_optimizer.zero_grad() total_loss.backward() self.policy_optimizer.step() self.value_optimizer.step() return total_loss.item()RLHF 流程
步骤 1：收集人类反馈
class HumanFeedbackCollector: """人类反馈收集器""" def __init__(self): self.preferences = [] def collect(self, prompt: str, responses: list) -> int: """收集人类偏好""" # 在实际应用中，这里应该显示给人类评判者 # 简化实现：随机选择一个偏好 import random return random.randint(0, len(responses) - 1) def save_preferences(self, filepath: str): """保存偏好数据""" import json with open(filepath, "w") as f: json.dump(self.preferences, f)步骤 2：训练奖励模型
class RewardModel(nn.Module): """奖励模型""" def __init__(self, llm_backbone): super().__init__() self.backbone = llm_backbone self.reward_head = nn.Linear(llm_backbone.config.hidden_size, 1) def forward(self, input_ids, attention_mask): """前向传播""" outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask) last_hidden = outputs.last_hidden_state[:, -1, :] reward = self.reward_head(last_hidden) return reward步骤 3：PPO 微调
class RLHFTrainer: """RLHF 训练器""" def __init__(self, policy_model, reward_model, ppo_config): self.policy_model = policy_model self.reward_model = reward_model self.ppo = PPO(policy_model, reward_model, **ppo_config) def train(self, prompts, num_epochs=10): """训练主循环""" for epoch in range(num_epochs): total_loss = 0 for prompt in prompts: # 1. 生成响应 response = self.policy_model.generate(prompt) # 2. 获取奖励 reward = self.reward_model.score(prompt, response) # 3. PPO 更新 loss = self.ppo.update(reward) total_loss += loss avg_loss = total_loss / len(prompts) print(f"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}")实战示例
# 初始化模型 policy_model = load_policy_model() reward_model = load_reward_model() # 创建训练器 trainer = RLHFTrainer( policy_model, reward_model, ppo_config={ "lr": 3e-5, "gamma": 0.99, "eps": 0.2 } ) # 训练数据 prompts = [ "解释什么是机器学习", "写一首诗", "分析当前经济形势" ] # 开始训练 trainer.train(prompts, num_epochs=5)关键技巧
KL 惩罚
def compute_kl_penalty(new_log_probs, old_log_probs): """计算 KL 散度惩罚""" kl = new_log_probs - old_log_probs return kl.mean()奖励塑造
def shape_reward(raw_reward, kl_penalty, kl_coef=0.1): """奖励塑造""" return raw_reward - kl_coef * kl_penalty总结
PPO 在 RLHF 中的应用：
- 收集反馈：获取人类对模型输出的偏好
- 训练奖励模型：学习预测人类偏好
- PPO 微调：优化策略以最大化奖励
关键要点：
- PPO 通过裁剪保证策略更新的稳定性
- KL 惩罚防止策略偏离太远
- 奖励模型质量直接影响最终效果
- 需要大量高质量的人类反馈数据