超越稳定性-探索困境:面向大语言模型策略优化的环境正则化 Beyond the Stability-Exploration Dilemma: Environmental Regularization for LLM Policy Optimization
把KL正则从响应侧移到查询侧,稳住训练环境漂移又不牺牲探索
前置知识
Policy-KL(动作侧策略KL正则)
RLHF/RLVR训练中对当前策略 $\pi_\theta$ 与冻结参考策略 $\pi_{\theta_0}$ 的响应分布施加KL散度惩罚(如k3估计器),约束每步更新不偏离SFT参考太远,防止策略崩溃与奖励投机。它作用于动作(响应)侧,代价是消耗探索预算。
ERPO的出发点正是用查询侧QKL替代这一动作侧正则,理解Policy-KL的作用与代价才能体会作者所说的稳定性-探索两难。
GRPO(组相对策略优化)
对同一查询采样K个响应,用组内奖励归一化得到优势 $A^{GRPO}_\theta(q,o^{(k)})=\frac{g(q,o^{(k)})-\mathrm{mean}(g)}{\mathrm{std}(g)}$,无需价值网络,是DeepSeek-R1等模型的主力RLVR算法。
本文全部实验基于GRPO流水线(EasyR1框架),ERPO以即插即用方式改造其损失。
RLVR(可验证奖励强化学习)
用可自动验证的结果(数学答案正确性、代码测试通过等)代替人工偏好标注作为奖励信号训练LLM,数学推理是其最典型场景,通常要求模型用包裹推理、答案放入\boxed{}。
本文训练与评测全在RLVR数学推理设定下:MATH Level 3-5训练,六个数学基准评测。
策略诱导查询分布与环境漂移
论文定义 $\rho_\theta(q)=P_\theta(q)$,即把查询q当作自终止token序列的自回归似然 $\ell_\theta(q)=\sum_t \log P_\theta(x_t|x_{<t})$。随参数更新,$\rho_\theta$ 相对RL前参考 $\rho_{\theta_0}$ 漂移,EnvShift(θ)=KL(ρ_θ‖ρ_θ0),构成输入侧非平稳环境并放大梯度方差。
这是全文核心概念,ERPO的两个组件(Query-KL与查询重加权)都围绕约束和利用这一漂移设计。
重要性重加权
把从经验分布 ρ_train 采样的期望改写为按 $w^*(q)=\rho_{\theta_0}(q)/\rho_{train}(q)$ 加权的期望,从而近似直接从参考分布优化目标;均匀采样N条数据时 $w^*(q)=N\rho_{\theta_0}(q)$。实践中用有界代理(如归一化后裁剪到[0,2])控制方差。
ERPO用数据集静态的参考导出权重 w_B(q) 重加权每查询梯度,是降低估计方差、驯服高低温波动的第二支柱。
研究动机
主流RLHF/RLVR策略优化(TRPO/PPO、DPO、GRPO等)把训练稳定性完全寄托在动作侧Policy-KL上:它约束响应分布、同时消耗探索预算;一旦去掉又没有任何漂移控制。作者用实验揭示被忽视的输入侧问题:即使训练语料固定,随θ更新,模型赋予每条训练查询的自回归似然也在演化,策略诱导查询分布 ρ_θ 相对RL前参考 ρ_θ0 无约束漂移。Figure 1显示固定Policy-KL预算下,批估计Query-KL在整个GRPO训练中持续上升而Policy-KL几乎持平,查询级KL比策略级高一个数量级。后果是长程训练中优化噪声与分布偏移累积:GRPO在240步(约15个epoch)内低温采样尚稳定,但约400步后高温采样性能率先崩塌并向全温度蔓延;同时出现严重reward hacking——训练平均准确率77.3%而评估仅69.8%/70.1%,Train-Eval平均gap达6.47%,Step-240时评估准确率从约75%骤降至58.4%。
本文的目标是本文目标是打破稳定性-探索两难:在不给响应分布施加任何直接KL约束的前提下,依然显式控制训练环境的漂移,同时获得稳定性与响应侧探索自由。具体做法是把正则化对象从动作侧的 $\pi_\theta(o|q)$ 移到输入侧的查询分布:引入Query-KL(QKL)惩罚 $R_{query}(\theta)=KL(\rho_\theta\|\rho_{\theta_0})$ 约束查询分布漂移;同时用数据集静态、参考模型导出的每查询权重 w_B(q) 偏置每次更新,使其偏向参考分布下典型的查询,降低估计器方差、驯服高温解码下的长尾敏感。工程上要求估计器无关、即插即用于GRPO/PPO/REINFORCE风格流水线,且不增加任何额外前向传播。
与已有工作不同的是,已有LLM工作或形式化提示分布(EVA、Align-Pro)、或对训练数据重加权(StablePrompt、WPO),而主流RLHF PO只关注对SFT参考的动作侧Policy-KL,没有方法直接约束查询分布。本文的独特切入是把经典RL中环境非平稳(初始状态分布或转移核漂移)与模仿学习的covariate shift教训(DAgger/AggreVaTe的动机)迁移到LLM-RL:选择模型自身的查询序列似然 $P_\theta(q)$ 作为可正则化的输入侧统计量,而非外生数据分布 ρ_train。由此自然导出两个配合的机制——约束漂移的QKL与向参考先验对齐的查询重加权,并从理论上证明QKL梯度与响应score function结构解耦。
核心方法
直觉:RL训练语料虽然固定,但有效训练环境并不固定——模型对每道题的似然随参数更新而变,等价于RL中初始状态分布漂移,会放大梯度方差并让训练与评估行为脱节。ERPO把参考环境定义为RL前模型诱导的查询分布 ρ_θ0,然后做两件事: 在RL目标上加查询级正则 $J_{ERPO}(\theta)=J_{RL}(\theta)-\alpha R_{query}(\theta)$,其中 $R_{query}(\theta)=KL(\rho_\theta\|\rho_{\theta_0})$; 在小批量经验损失中用数据集静态权重重加权每查询贡献:$\hat{\mathcal{L}}_{ERPO}(\theta)=-\frac{1}{m}\sum_{q\in B}w_B(q)\bar{g}_\theta(q)+\alpha\hat{R}_{query}(\theta)$。内层PG估计器可以是GRPO组相对优势、PPO裁剪代理或REINFORCE,ERPO只作用于外层查询循环,实验默认 $\alpha=10^{-2}$。
核心创新是结构解耦(Proposition 1):$\nabla_\theta R_{query}(\theta)=\mathbb{E}_{q\sim\rho_\theta}[(\ell_\theta(q)-\ell_{\theta_0}(q))\cdot\nabla_\theta\ell_\theta(q)]$,梯度严格只流经查询对数似然的梯度 $\nabla_\theta\ell_\theta(q)$,而策略梯度估计器使用的响应score function $\nabla_\theta\log\pi_\theta(o|q)$ 完全不出现在QKL项中——因此QKL对响应分布没有直接梯度压力,与Policy-KL既管行为又吃探索预算有本质区别。第二点是零开销设计:$\ell_{\theta_0}(q)$ 训练前用参考模型对全数据集算一次并缓存为常数表,$\ell_\theta(q)$ 直接从PG估计器本来就做的每查询前向读取,不增加任何额外前向传播。查询权重则是对理想重要性权重 $w^*(q)=N\rho_{\theta_0}(q)$ 的有界单调代理。
方法步骤详情
流程四步: 预计算:用冻结参考模型 $\pi_{\theta_0}$ 对约8.5K条MATH Level 3-5查询计算 $\ell_{\theta_0}(q)=\sum_t\log P_{\theta_0}(x_t|x_{<t})$ 并缓存;记负对数概率 $s_i=-\log p_{\theta_0}(q_i)$,取 $\tilde{w}(q_i)=\bar{s}/s_i$,裁剪到[0,2]得 w(q_i)。 每步前向:从现有PG前向读出当前 $\ell_\theta(q)$,按k3式批代理计算 $\hat{R}_{query}(\theta)$。 查询重加权:外层权重从 1/m 换成 w_B(q)/m;内层每题采样K=8或16个响应,优势用GRPO组内归一化。 联合优化:最小化 $\hat{\mathcal{L}}_{ERPO}=-\frac{1}{m}\sum_q w_B(q)\sum_o u_\theta A^*_\theta\,\nabla_\theta\log\pi_\theta+\alpha\hat{R}_{query}$,即用QKL替换原Policy-KL项。
技术新颖性
技术新颖性有三层。视角新颖:首次把LLM-RL的不稳定归因到输入侧环境非平稳,并给出可操作的度量 EnvShift(θ)=KL(ρ_θ‖ρ_θ0),区别于动作侧信任域/KL惩罚(TRPO/PPO/GRPO),也区别于只优化提示本身的EVA/Align-Pro和只重加权数据的StablePrompt/WPO。结构新颖:Proposition 1给出闭式梯度证明QKL与响应分布解耦,使不用动作侧KL也能约束训练成为可能——消融显示仅把Policy-KL换成Query-KL平均提升15.9%。工程新颖:缓存参考似然+复用PG前向实现零额外前向的即插即用,且与GRPO/PPO/REINFORCE/DAPO/RLOO全部兼容(Table 6中DAPO+10.24%、RLOO+2.28%)。重加权设计明确声明是对理想密度比的有界单调逼近而非无偏估计,务实且可解释。
实验结果
在MATH Level 3-5(约8.5K题)上训练Qwen2.5-7B/32B(EasyR1,rollout批512/更新批128/240步),在AIME24/25、AMC、MATH500、Minerva、OlympiadBench六基准评测(温度0.1-1.5聚合)。Table 1:Avg@32平均从GRPO的0.274升至0.336(+6.2%),Pass@1从0.275到0.332(+5.69%),MATH500 Avg@32为0.677 vs 0.528(+14.9个百分点)。消融(Table 2):仅换成Query-KL时T≤1.0均分80.90 vs 68.80,平均+15.9%;32B下T=1.5为80.80 vs 25.20。长程训练(Figure 5/6,至960步):GRPO约400步后高温先崩后蔓延全温度,ERPO退化小、高温反升。Table 4/5:Train-Eval gap 6.47%→3.14%(↓51%);Table 3显示Query-KL把查询KL从0.9679压到0.0041,是主要增益来源;重加权主责稳定训练(熵最低0.2782)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 六数学基准平均(AIME24/25、AMC、MATH500、Minerva、OlympiadBench) | Mean Avg@32(温度0.1-1.5聚合) | 0.336 | GRPO 0.274(Base 0.143) | +6.2%,单项最高+14.9% |
| MATH500 | Avg@32 | 0.677 | GRPO 0.528 | +14.9个百分点(最大单项增益) |
| AIME24 / AIME25 | Avg@32 | 0.218 / 0.110 | GRPO 0.174 / 0.072 | +4.6 / +3.8 个百分点 |
| MATH500(Qwen2.5-32B) | Pass@1,T≤1.0均值 / T=1.5 | 84.60 / 80.80 | GRPO 81.62 / 25.20 | +2.98;高温1.5下+55.6个百分点 |
| MATH500(Qwen2.5-Math-7B,每题16采样) | Pass@1,T=1.5 | 56.20 | GRPO 10.60 | +45.6个百分点 |
| 训练-推理一致性(reward hacking量化) | Train-Eval平均gap | 3.14% | GRPO 6.47% | 下降51% |
| DAPO+ERPO(算法迁移) | Pass@1,T≤1.0 | 78.4 | DAPO 68.16 | +10.24 |
| RLOO+ERPO(算法迁移) | Pass@1,T≤1.0 | 79.56 | RLOO 77.28 | +2.28 |
局限与改进
作者承认:实验限于数学推理基准与Qwen系模型,向指令遵循、对话、代码生成、多语言设定的迁移未验证;未对正则系数α做彻底搜索;训练中估计查询似然的质量与计算成本随数据选择机制和模型规模而变;ERPO在极长训练下仍可能出现熵突增式崩溃,只是推迟和减轻。我的补充观察: ρ_θ(q)定义为查询的自回归序列似然并非严格意义上的查询分布——查询是外部数据而非模型采样,该KL更像正则化统计量; 工作假设A1(保持ρ_θ与参考对齐可保留泛化能力)只有经验支持,无理论保证; w(q)是 $\ell_{\theta_0}(q)$ 的单调代理而非密度比 $N\rho_{\theta_0}(q)$,丢掉了绝对比例信息; 7B模型参数空间存在近似正交子空间的论证(Figure 4附近)也仅是经验推测。
独立分析的弱点
其一,超参敏感且默认α未必最优:Table 2中α=5e-3时低温0.1仅53.80(GRPO为66.80),α=5e-2才全面占优,说明需按模型规模调α,改进方向是把α做成对EnvShift的在线自适应惩罚(拉格朗日松弛式控制)。其二,权重设计启发式:w(q)=clip(s̄/s_i,0,2)丢弃密度比的绝对比例,且对极难/超长查询(s_i很大)权重趋近0,可能系统性欠训练难题,可改用 $\exp(\beta\ell_{\theta_0})$ 型软权重或分位数归一。其三,缓存参考似然表在数百万级数据集上内存与一次性算力开销可观,可改为滚动估计或抽样缓存。其四,评测协议把温度0.1-1.5聚合,与常见单温度报告不可直接比较;脚注亦承认GRPO在多次实验中学不出格式,基线强度存疑,应补充DAPO原版(无KL)与Dr.GRPO变体等更强基线。
未来方向
作者提出的方向:推广到指令遵循、对话、代码生成与多语言设定;对α及超参做更系统的分析。基于成果可延伸的: 用反馈控制视角让α自动维持目标查询KL水平,消除手动调参; 利用 w_B(q) 天然给出的参考典型度排序做课程学习与数据筛选,与主动采样调度结合; 把QKL引入DPO族偏好优化的输入侧正则,检验假设A1在非可验证奖励场景是否成立; 理论上刻画查询似然漂移与reward hacking的联系——目前仅有Appendix C/D的经验证据(prompt与response的NLL在95%样本上相关系数接近1); 探索多参考环境(多个 $\rho_{\theta_0}$ 混合)以支持多领域联合训练;研究查询KL与熵崩溃因果链的长程动力学。
复现评估
复现条件较好:代码开源于 https://github.com/AlibabaResearch/ERPO,基于公开的EasyR1框架;数据为公开MATH数据集Level 3-5(约8.5K题);模型为公开的Qwen2.5-Math-7B与Qwen2.5-32B。训练配置完整:最大序列3K token,每题8个采样@温度1.0,rollout批512、更新批128、共240步,token级loss,KL系数/α=0.01;评测用vLLM并报告TP1/TP2两种张量并行结果(作者指出TP会影响性能)。7B版在单节点多卡(约8×A100级)应可复现;32B需更多算力且论文未给出确切GPU配置与训练时长,是主要不确定项。方法本身改动小(预计算 $\ell_{\theta_0}$、换外层权重、加 $\alpha\hat{R}_{query}$ 三步),总体难度中等,主要成本在于对齐其多温度聚合评测协议的开销。
论文图表
GRPO训练全程中,批估计的Query-KL(深色曲线)持续上升,而Policy-KL(浅色曲线)几乎保持低位,说明只约束动作分布无法稳定查询过程,查询级KL比策略级高一个数量级。
这是全文动机的核心实证证据:直观展示策略诱导查询分布 ρ_θ 相对参考 ρ_θ0 的漂移完全未被Policy-KL控制,直接引出环境正则化的必要性。
附录D:对8K+训练问题在温度1.0下分别计算prompt与response的NLL直方图与相关性,95%样本(prompt NLL<300)内两者相关系数接近1,低概率响应集中在正样本中并推高训练熵。
为输入侧正则化的合理性提供统计依据:查询似然与响应似然强相关,约束查询分布的漂移能间接抑制梯度方差与熵失控。