SimKO方法优化LLM推理多样性：解决概率过度集中问题
1. 项目背景与问题定义
在大型语言模型（LLM）的推理能力优化领域，强化学习与可验证奖励（Reinforcement Learning with Verifiable Rewards, RLVR）已成为主流方法。其核心思想是通过验证机制（如数学问题求解器）自动判断模型生成答案的正确性，并据此提供奖励信号来优化模型参数。然而，当前RLVR方法普遍存在一个关键缺陷：随着训练进行，模型会逐渐收敛到少数"安全"的推理路径上，导致输出多样性下降。
这种现象具体表现为：模型在Pass@1（单次采样正确率）指标上持续提升，但在Pass@K（K次采样中至少一次正确的概率）指标上却出现下降。例如，在MATH500数据集上，使用GRPO方法训练的Qwen2.5-Math-7B模型，Pass@1从基准模型的25.8%提升到41.7%，但Pass@256却从76.4%降至76.1%。这种探索能力的退化严重限制了模型应对复杂、多解问题的能力。
2. 核心发现：概率过度集中现象
2.1 概率分布动态分析
通过对训练过程中token级概率分布的跟踪，我们发现RLVR方法会导致明显的**概率过度集中（Over-Concentration）**现象。具体表现为：
Top-1主导：在解码过程中，排名第一的候选token逐渐吸收绝大部分概率质量（如图1右所示，GRPO训练的模型top-1概率高达0.93）
其他候选被抑制：排名2-5的候选概率被压缩到极低水平（GRPO下top-2概率仅0.06，top-5仅0.00）
动态演变规律：如图3所示，随着训练步数增加：
- Λ(1)（top-1平均对数概率）快速上升并接近1
- Λ(2)/Λ(3)则骤降至10^-8~10^-10量级
2.2 与Pass@K的关联性
关键发现是：概率集中程度与Pass@K性能呈负相关。通过对比不同训练方法发现：
- PSR（仅正样本强化）加剧集中效应，Pass@256下降3.6%
- NSR（仅负样本强化）缓解集中效应，Pass@256提升4.2%
- 集中程度每增加0.1，Pass@256平均下降2.3%
这种现象的根源在于RLVR的梯度更新机制：
# 标准GRPO的梯度计算 grad = advantage * (π_current/π_ref) * ∇logπ当正确样本的top-1候选获得正向奖励时，其概率会被进一步放大，形成"富者愈富"的马太效应。
3. SimKO方法设计
3.1 整体架构
SimKO（Simple Pass@K Optimization）通过非对称梯度重分配机制解决概率集中问题，其核心组件包括：
- 关键token识别：基于熵阈值筛选"分叉token"（H(π)>τ）
- 正样本处理：对正确响应实施top-K标签平滑
- 负样本处理：对错误响应加强top-1惩罚
3.2 关键技术实现
3.2.1 分叉token检测
使用熵作为决策点检测指标：
def is_forking_token(prob_dist, τ=0.8): entropy = -sum(p * log(p) for p in prob_dist) return entropy > np.percentile(all_entropies, τ*100)实验表明，80%分位的τ值能有效捕捉关键决策点（如数学推理中的运算符选择）。
3.2.2 正样本的top-K平滑
对正确响应，将梯度分配到top-K候选：
def topk_label_smoothing(logits, K=4, α=0.01): topk_probs = torch.topk(softmax(logits), K) smoothed_target = (1-α)*one_hot + α/K*topk_sum return cross_entropy(smoothed_target, logits)这形成概率"高原"而非尖峰（如图4b）。
3.2.3 负样本的top-1惩罚
对错误响应的top-1候选施加λ倍（λ=1.1）更强惩罚：
if is_incorrect and is_top1: ratio *= λ # 增强梯度幅度4. 实验验证
4.1 数学推理任务
在MATH500、AIME等6个数学基准测试中，Qwen2.5-Math-7B模型表现：
| 方法 | Pass@1 | Pass@256 | ΔPass@256 |
|---|---|---|---|
| Base | 25.8 | 76.4 | - |
| GRPO | 41.7 | 76.1 | -0.3 |
| SimKO | 43.4 | 80.5 | +4.4 |
| KL-Cov | 42.5 | 79.0 | +2.9 |
特别在AIME25难题上，Pass@256提升12.5%（51.8%→64.6%）。
4.2 逻辑推理任务
在Synlogic和BBH数据集上，SimKO展现出更强泛化性：
- Synlogic：Pass@128达55.0%，比GRPO高5.6%
- BBH：保持92.0%的Pass@128，仅下降2.2%
4.3 消融实验
关键参数影响分析（Qwen2.5-Math-7B）：
| 参数 | 最优值 | Pass@1 | Pass@256 |
|---|---|---|---|
| α | 0.01 | 43.4 | 80.5 |
| τ | 80% | 42.8 | 79.3 |
| K | 4 | 43.4 | 80.1 |
5. 实施建议与注意事项
5.1 工程实践要点
- 梯度计算优化：
# 使用stop-gradient技巧保持重要性采样无偏 topk_ratio = (ratio.detach()/topk_ratio.detach()) * topk_ratio- 内存管理：
- 仅需缓存top-K logits（K=4时额外显存<5%）
- 相比全词汇表操作，速度损失<3%
- 超参数设置：
# 推荐配置 math_task: α: 0.01 λ: 1.1 τ_quantile: 0.8 logic_task: α: 0.005 λ: 1.05 warmup_steps: 505.2 常见问题排查
- Pass@1下降：
- 检查α是否过大（建议≤0.05）
- 验证τ是否过低（导致非关键token被修改）
- 收敛不稳定：
- 逻辑任务需50-100步warmup
- 适当减小λ（1.05-1.15区间）
- 效果不显著：
- 确认模型是否已过拟合（查看Λ(1)是否>0.9）
- 尝试增大K（3→5）
6. 扩展应用方向
- 多模态推理：在图文联合推理任务中，SimKO可保持解题路径多样性
- 代码生成：提升Beam Search时候选程序的有效性
- 对话系统：避免陷入重复性回应模式
关键洞见：概率集中本质是RLVR的固有偏差，SimKO通过梯度层面的干预，实现了比熵正则化更精细的探索控制。在实际部署中发现，该方法对7B以上大模型效果尤为显著，可能与小模型的表达能力限制有关。