强化学习优化语言模型数学推理能力实践

1. 项目概述

在CS336课程第五次作业中，我们将探索如何通过强化学习技术提升语言模型在数学推理任务上的表现。这项作业聚焦于三个核心环节：构建零样本基线模型、实施监督微调(SFT)以及应用专家迭代(EI)和组相对策略优化(GRPO)算法。我们将使用Qwen 2.5 Math 1.5B作为基础模型，在MATH竞赛数学数据集上进行实验验证。

提示：本实验需要至少2块GPU（一块用于策略模型训练，一块用于vLLM推理评估），建议使用H100等高性能显卡以获得最佳体验。

2. 核心组件解析

2.1 语言模型作为策略的实现

在强化学习框架下，我们将语言模型视为策略函数πθ，其核心操作包括：

python复制class LanguageModelAsPolicy:
    def __init__(self, model_path):
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    def sample_action(self, state):
        """从当前状态采样下一个token"""
        inputs = self.tokenizer(state, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = self.model(**inputs)
        probs = torch.softmax(outputs.logits[0, -1], dim=-1)
        return torch.multinomial(probs, num_samples=1).item()
    
    def log_prob(self, state, action):
        """计算给定状态下动作的对数概率"""
        inputs = self.tokenizer(state, return_tensors="pt").to(device)
        action_id = self.tokenizer.encode(action, add_special_tokens=False)[0]
        with torch.no_grad():
            outputs = self.model(**inputs)
        logits = outputs.logits[0, -1]
        return F.log_softmax(logits, dim=-1)[action_id].item()

2.2 轨迹生成与奖励计算

轨迹生成是强化学习中的关键环节，我们实现了专门的RolloutWorker类：

python复制class RolloutWorker:
    def __init__(self, policy, reward_fn, max_steps=1024):
        self.policy = policy
        self.reward_fn = reward_fn
        self.max_steps = max_steps
    
    def generate_rollout(self, prompt):
        state = prompt
        trajectory = []
        for _ in range(self.max_steps):
            action = self.policy.sample_action(state)
            new_state = state + self.policy.tokenizer.decode([action])
            reward = 0  # 中间步骤奖励为0
            
            trajectory.append({
                'state': state,
                'action': action,
                'reward': reward,
                'new_state': new_state
            })
            
            state = new_state
            if '</answer>' in new_state:  # 终止条件
                break
        
        # 计算最终奖励
        if trajectory:
            final_reward = self.reward_fn(prompt, state)
            trajectory[-1]['reward'] = final_reward
        
        return trajectory

3. 关键算法实现

3.1 监督微调(SFT)实现细节

监督微调采用标准的交叉熵损失，但有几点特殊处理：

1. 响应掩码(Response Mask)：确保只对响应部分计算损失
2. 梯度累积：解决显存限制问题
3. 动态批处理：提升训练效率

核心训练代码如下：

python复制def sft_train_step(batch, model, optimizer, grad_accum_steps=4):
    model.train()
    inputs = batch['input_ids'].to(device)
    labels = batch['labels'].to(device)
    masks = batch['response_mask'].to(device)
    
    outputs = model(inputs, labels=labels)
    logits = outputs.logits
    
    # 计算掩码损失
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    shift_mask = masks[..., 1:].contiguous()
    
    loss_fct = nn.CrossEntropyLoss(reduction='none')
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), 
                   shift_labels.view(-1))
    loss = (loss * shift_mask.view(-1)).sum() / shift_mask.sum()
    
    # 梯度累积
    loss = loss / grad_accum_steps
    loss.backward()
    
    if (step + 1) % grad_accum_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()
    
    return loss.item() * grad_accum_steps

3.2 专家迭代(EI)算法优化

我们在基础EI算法上做了三点改进：

1. 多样性采样：采用温度采样增加轨迹多样性
2. 精英保留：保留历史高质量样本
3. 动态过滤阈值：逐步提高过滤标准

算法核心流程：

python复制def expert_iteration(policy, train_data, reward_fn, n_iter=5):
    elite_pool = []
    
    for iter in range(n_iter):
        # 1. 采样生成轨迹
        rollouts = []
        for prompt in train_data:
            trajectories = [generate_rollout(policy, prompt) for _ in range(G)]
            rewards = [reward_fn(traj[-1]['new_state']) for traj in trajectories]
            rollouts.extend(zip(trajectories, rewards))
        
        # 2. 过滤保留高质量样本
        filtered = [traj for traj, reward in rollouts if reward > 0]
        elite_pool.extend(filtered)
        
        # 3. 动态调整过滤阈值
        if len(elite_pool) > 2000:
            elite_pool = sorted(elite_pool, key=lambda x: x[1], reverse=True)[:2000]
        
        # 4. 监督微调
        sft_train(policy, elite_pool)
        
        # 5. 评估验证集性能
        eval_results = evaluate(policy, valid_data)

4. 实验配置与优化

4.1 超参数设置

经过多次实验验证，我们确定了以下最优超参数组合：

4.2 性能优化技巧

1. 内存优化：

   • 使用梯度检查点技术
   • 启用FlashAttention-2
   • 采用BF16混合精度

2. 计算加速：

   • 使用vLLM进行高效推理
   • 实现异步数据加载
   • 优化KV缓存管理

3. 训练稳定性：

   • 梯度裁剪(1.0)
   • 学习率预热(500步)
   • 动态批处理策略

5. 常见问题与解决方案

5.1 格式错误分析

在零样本基线测试中，我们观察到三类典型错误：

1. 格式正确但答案错误：

   • 原因：数学推理能力不足
   • 解决方案：增加数学专项预训练

2. 格式错误但答案正确：

   • 原因：模型不遵循模板
   • 解决方案：强化模板遵从微调

3. 格式与答案均错误：

   • 原因：基础能力缺陷
   • 解决方案：更换更强基础模型

5.2 训练不稳定问题

现象：损失值剧烈波动
解决方法：

• 减小学习率
• 增大批大小
• 加强梯度裁剪
• 添加更多正则化

现象：验证性能下降
解决方法：

• 早停机制
• 模型平均
• 增加验证频率
• 调整温度参数

6. 实验成果与结论

通过系统实验，我们得出以下结论：

1. 监督微调效果：

   • 完整数据集上验证准确率达18.7%
   • 过滤后数据(保留60%)准确率提升至22.3%

2. 专家迭代优势：

   • 5次迭代后准确率提升至25.1%
   • 模型熵值从3.2降至2.4，显示置信度提高

3. GRPO表现：

   • 最终测试准确率达到28.9%
   • 相比基线提升近3倍

从实际训练中获得的几点关键经验：

1. 初始阶段应优先确保格式正确性
2. 数学推理能力提升需要分阶段进行
3. 奖励设计对最终性能影响显著
4. 小模型也能通过精心调优获得不错效果