Stabilizing Policy Optimization via Logits Convexity

会议: NeurIPS 2026
arXiv: 2603.00963
OpenReview: https://openreview.net/forum?id=yppWn8zxqm
代码: https://github.com/huggingface/trl
领域: 强化学习
关键词: 强化学习, 策略优化, 对数几率凸性, PPO, 大语言模型

一句话总结

针对大语言模型强化学习中策略梯度（如 PPO）容易出现梯度爆炸与训练崩溃的痛点，本文从对数几率（logits）局部凸性视角揭示了 SFT 稳定与 RL 不稳定的几何本质，并提出直接对齐闭式最优目标分布或 logits 的凸优化框架 LCO（Logits Convex Optimization），从数学上消除了负优势主导的梯度尖峰并实现高效收敛。

研究背景与动机

强化学习（RL）已成为大语言模型（LLM）实现人类偏好对齐以及激发复杂数学与逻辑推理能力的核心驱动力。然而，相较于稳定收敛的监督微调（SFT），基于策略梯度的强化学习训练过程具有众所周知的脆弱性与不稳定性。现有主流做法通常依赖广义优势估计（GAE）方差缩减、近端策略优化（PPO）或广义参考策略优化（GRPO）的重要性采样比截断，以及基于 KL 散度的策略偏移惩罚。尽管这些工程启发式手段在一定程度上缓解了发散风险，但在大规模模型微调或复杂推理场景中，剧烈的梯度震荡乃至不可逆的性能崩溃现象仍屡见不鲜，其深层数学诱因始终缺乏统一解释。

通过实证分析对比 SFT 与 PPO 的梯度演化动力学，可以发现一个反直觉的对比现象：SFT 训练过程中梯度范数随着预测置信度的提升而平滑单调衰减；而在 PPO 截断机制生效前，对于具有负优势（negative advantage）的负样本动作，其损失函数虽然在下降，参数空间与 logit 空间的梯度范数却会急剧暴增，尤其是处于中等概率（如 0.5 附近）的低概率探索动作会诱发极大的参数更新冲击。深入优化曲面可以发现，这一根本差异根源于损失函数在 logits 空间的几何性质：SFT 损失关于输出 logits 始终具备局部正半定的 Hessian 矩阵，即保证了对数几率凸性（logits convexity）；而 PPO 的剪切代理目标不仅丧失了这一凸性，其非凸区域在负优势下会形成错误的梯度指向，诱发破坏性震荡。

为了从根本上消除这一非凸曲面带来的训练不稳定，本文不再沿用启发式梯度截断，而是回归带有 KL 惩罚的闭式正则化强化学习最优解。核心 idea：将策略优化重构为向闭式最优目标对齐的凸优化框架（LCO），利用强凸的均方误差、平滑鲁棒的 log-cosh 或正向 KL 散度直接拟合最优 logits 或分布，从几何上恢复优化稳定性和梯度单调衰减界。

方法详解

整体框架

Logits Convex Optimization（LCO）的核心是将强化学习从容易诱发梯度畸变的代理目标优化，转换为对基于行为策略与优势函数计算出的最优目标进行对齐的监督式凸拟合过程。整个流水线首先通过行为策略采集状态与动作轨迹，并结合奖励模型或离散标量计算 token 级的优势估计值；接着利用正则化 RL 的一阶驻点解析解构造最优 logits 或最优动作分布目标；最终通过保持 logits 凸性的回归损失（LCO-MSE、LCO-LCH）或正向 KL 散度（LCO-KLD）更新策略网络参数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：状态序列与行为策略输出"] --> B["多粒度优势估计：稀疏采样 / DPO对数比 / 密集模型"]
    B --> C["最优目标推导：解析构建理论最优分布与Logits"]
    C --> D["对数几率凸优化：LCO-MSE / LCO-LCH / LCO-KLD"]
    D --> E["稳定策略参数更新：梯度范数随剩余误差单调耗散"]

关键设计

1. 多粒度优势估计：多源反馈信号与词表级稠密优势构建

LCO 框架构建最优目标的前提是获得各 token 位置的优势信号 \(A(s_t, a_t)\)。针对不同任务与算力约束，设计了三类不同粒度的优势估计策略。最基础的是针对环境只给出当前采样 token 反馈的稀疏采样估计，非采样词元优势直接置零。针对偏好对齐场景，框架进一步引入基于 DPO 对齐模型与参考模型对数概率差的比值优势：

\[A(s_t, a_t) := \log \frac{\phi_{\text{DPO}}(a_t|s_t)}{\phi_{\text{ref}}(a_t|s_t)}\]

由于语言模型在最大因果熵强化学习框架下等价于 Q 值与 V 值的能量分布，该对数比通过 Bellman 方程直接解构为瞬时奖励与状态价值转移差，从而为词表上的全部候选 token 提供了稠密、平滑的连续标量指导，避免了单点采样的剧烈梯度方差。

2. 最优目标推导：解析构建理论最优分布与无偏 Logits

标准正则化强化学习目标是在最大化期望优势的同时，施加系数为 \(\beta\) 的 KL 散度以约束策略偏离参考策略 \(\pi_{\text{old}}\)。通过拉格朗日乘子法求解该有约束极值问题，可知其在时间步 \(t\) 存在唯一的全局最优闭式解：

\[\pi^*(a_t|s_t) = \frac{\pi_{\text{old}}(a_t|s_t)\exp\left(\frac{A(s_t, a_t)}{\beta}\right)}{\sum_{a'_t \in \mathcal{V}} \pi_{\text{old}}(a'_t|s_t)\exp\left(\frac{A(s_t, a'_t)}{\beta}\right)}\]

通过消去配分函数中的 softmax 归一化项，可以推导出诱导该最优分布的最优 logit 向量为 \(z^*(s_t, a_t) = z_{\text{old}}(s_t, a_t) + \frac{A(s_t, a_t)}{\beta}\)。进一步从收敛速度界分析证明，在实践中对优势向量施加标准均值归一化（\(\frac{1}{|\mathcal{V}|}\sum A = 0\)）后，该零常数平移的最优 logits 正好能使初始空间的距离误差达到全局极小，从而保证参数沿最优梯度轨迹最快收敛。

3. 对数几率凸优化：三种具备凸几何特性的对齐损失函数

为消除 PPO 目标在非凸区产生的梯度爆炸，LCO 提出了三套严格保持 logits 局部凸性的损失函数。第一种是针对最优 logits 的均方误差回归损失 \(\mathcal{L}_{\text{LCO-MSE}}\)，其 Hessian 矩阵为常数矩阵 \(\frac{2}{|\mathcal{V}|}\bm{I}\)，具备全局强凸性并享有线性收敛速度；第二种是引入双曲对数余弦函数的 \(\mathcal{L}_{\text{LCO-LCH}}\)，其在误差较小时表现为二阶 MSE，在误差较大时平滑过渡为一阶线性惩罚，其 Hessian 为对角正定矩阵，有效防御了异常优势噪声导致的局部梯度激增：

\[\mathcal{L}_{\text{LCO-LCH}} = \mathbb{E}_t \left[ \frac{1}{|\mathcal{V}|} \sum_{a_t \in \mathcal{V}} \log \cosh \left( z_\theta(s_t, a_t) - z^*(s_t, a_t) \right) \right]\]

第三种是基于概率分布匹配的正向 KL 散度损失 \(\mathcal{L}_{\text{LCO-KLD}}\)，其最小化最优策略 \(\pi^*\) 与当前参数化策略 \(\pi_\theta\) 间的相对熵。理论推导证明这三种损失的参数空间梯度范数均受关于当前损失值的单调递增连续函数严格上界约束（如 \(\mathcal{L}_{\text{LCO-KLD}}\) 的梯度范数满足 \(\|\nabla_\theta \mathcal{L}\| \le \sigma_{\max}\sqrt{2\mathcal{L}}\)），保证了随着模型逼近最优解，更新步长自然耗散衰减，彻底消除了负样本引起的梯度尖峰。

损失函数 / 训练策略

在训练实施上，默认采用 \(\beta = 1.0\)，模型学习率设定为 \(5\times 10^{-6}\)。模型在预热阶段首先在高质量响应集上进行 1 个 epoch 的轻量 SFT 预热（1,000 条样本），以保证基础指令遵循能力与输出 logits 处于平滑区间。采样生成采用温度 0.6 与 top-p 0.95。针对 GRPO 等组采样对比场景，单输入生成 4 条响应计算奖励统计量，而 LCO 则既支持密集奖励模型提供的词表分布指导，也支持单一响应配合采样词元稀疏优势优化。

实验关键数据

主实验

主实验在数学推理任务（MATH500、AMC23、MinervaMath）、文本阅读理解（QA-Feedback）与通用指令遵循（AlpacaEval 2.0）三大核心基准上展开，基座模型覆盖 Qwen-2.5-3B、Qwen-3-4B、Llama-3.2-3B 及 Mistral-3-3B。

模型基座	优化方法	MATH500 (Pass@1)	AMC23 (Pass@1)	MinervaMath (Pass@1)
Qwen-2.5-3B	SFT	41.60 ± 0.72	27.50 ± 1.44	10.29 ± 0.71
Qwen-2.5-3B	REINFORCE	54.60 ± 0.53	40.15 ± 1.25	12.35 ± 0.93
Qwen-2.5-3B	PPO	55.40 ± 1.02	39.20 ± 2.01	13.97 ± 0.81
Qwen-2.5-3B	GRPO	57.79 ± 0.34	41.75 ± 1.25	14.89 ± 0.87
Qwen-2.5-3B	DAPO	57.05 ± 0.31	40.20 ± 1.46	13.23 ± 0.78
Qwen-2.5-3B	GSPO	56.80 ± 1.38	38.75 ± 2.01	14.87 ± 0.71
Qwen-2.5-3B	LCO-MSE (本文)	59.20 ± 1.75	40.20 ± 1.58	15.12 ± 0.43
Qwen-2.5-3B	LCO-LCH (本文)	61.40 ± 2.01	39.50 ± 2.37	15.38 ± 0.16
Qwen-2.5-3B	LCO-KLD (本文)	60.80 ± 0.51	42.65 ± 1.63	16.71 ± 0.97
Qwen-3-4B	PPO	67.80 ± 0.83	47.50 ± 1.25	20.95 ± 0.28
Qwen-3-4B	GRPO	67.60 ± 1.23	51.95 ± 1.92	21.75 ± 0.19
Qwen-3-4B	LCO-KLD (本文)	73.20 ± 1.34	55.50 ± 1.88	23.95 ± 0.31

在 AlpacaEval 2.0 上，Qwen-3-4B 下 LCO-KLD 取得了 29.05% 的胜率与 32.93% 的长度控制胜率（LC WR），相比 PPO（26.79% / 27.20%）与 GRPO（25.37% / 26.38%）取得一致领先。

消融实验

消融实验重点验证了在受限的稀疏优势反馈（仅采样动作获得优势，非采样动作为 0）以及规则奖励模型（正确 +1，错误 -1）下的稳定性与效果：

实验设定 (Qwen-3-1.7B)	优化算法	MATH500 (Pass@1)	AMC23 (Pass@1)	MinervaMath (Pass@1)
DPO-RM (稀疏反馈)	PPO	55.20 ± 1.32	32.20 ± 2.31	15.31 ± 1.97
DPO-RM (稀疏反馈)	LCO-MSE	61.30 ± 1.23	38.10 ± 1.10	16.13 ± 1.38
DPO-RM (稀疏反馈)	LCO-LCH	64.70 ± 1.93	37.50 ± 1.03	15.43 ± 0.98
DPO-RM (稀疏反馈)	LCO-KLD	63.25 ± 2.01	39.50 ± 1.83	16.98 ± 1.34
规则 RM (稀疏反馈)	GRPO	56.20 ± 1.31	34.50 ± 2.27	14.37 ± 1.02
规则 RM (稀疏反馈)	LCO-KLD	61.40 ± 1.31	37.10 ± 2.13	16.54 ± 1.34

与在策略知识蒸馏方法对比（Qwen-3-4B 教师模型监督）中，LCO-KLD 在 MATH500 上达到 73.20%，显著超过 MiniLLM（66.30%）与 GKD（67.70%），展现了引入优势驱动对齐相比纯粹分布模仿的性能突破能力。

关键发现

• 负样本梯度爆炸是 PPO 崩溃的主因：实证曲线显示，PPO 在训练步数达到 6K 步后出现显著的熵激增、采样动作概率跳水与梯度范数振荡；而在负优势下，PPO 截断机制无法阻止梯度在中间概率区间的非理性放大。
• 强凸损失带来更高的样本效率：在训练收敛效率对比中，LCO-KLD 仅需 PPO 约一半的训练样本量即可达到相同准确率；而强凸性质更显著的 LCO-LCH 其样本利用效率接近 LCO-KLD 的 3 倍，验证了理论推导中线性收敛半径的优势。
• 分布外泛化能力提升：在 MMLU 与 MMLU-Redux 评估中，LCO-KLD 分别取得 72.11% 与 76.71% 的高分，显著高于 PPO（67.87% / 73.93%），表明优化过程避免了因过拟合局部非凸梯度尖峰而破坏模型的泛化表征。

亮点与洞察

• 切中数学本质的稳定性视点：跳出了传统研究在重要性采样截断比或方差缩减上的工程修补，从神经正切核理论（NTK）与输出 logits Hessian 矩阵的正定性切入，首次从对数几率凸性（Logits Convexity）层面严谨解释了 SFT 与 RL 的稳定性鸿沟。
• 巧妙衔接闭式最优解与监督回归：将策略梯度转变为追踪动态解析极值 \(\pi^*\) 或 \(z^*\) 的过程，不仅规避了非凸优化曲面，更证明了当优势均值归一化为零时，直接加性调整 logits 恰好达到最紧收敛界，兼顾理论优雅性与工业落地极简性。

局限与展望

• 依赖参考策略与优势的估计质量：虽然 LCO-LCH 具备抗噪声能力，但若奖励模型或基准策略严重漂移，闭式推导的最优目标本身会带有系统偏差，此时优化器会迅速拟合至局部次优解。
• 词表级稠密运算的显存与计算开销：LCO-KLD 需要在整个词表维度上计算正向 KL 散度与概率分布，当模型词表极大（如超过 128k）且上下文较长时，全词表矩阵运算相较于仅计算单 token 交叉熵的 PPO 会带来额外的显存吞吐压力。未来可探索局部 top-k 截断或稀疏近似下的保凸对齐。

相关工作与启发

• vs PPO / GRPO：PPO 与 GRPO 通过硬性截断比率 \(r_t(\theta) \in [1-\epsilon, 1+\epsilon]\) 强行截断梯度，但其内部活跃区间在负优势下依然是非凸曲面，无法阻断中间概率动作引起的梯度暴增；LCO 直接构建强凸/凸损失曲面，梯度天然随误差衰减。
• vs GKD / MiniLLM：在策略知识蒸馏将学生分布对齐固定教师模型分布，缺乏探索环境自适应的能力边界；LCO 虽采用监督形式的目标对齐，但其目标是由动态更新的优势函数即时修正得到的强化学习全局最优解，能够突破固定教师的性能天花板。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 首次从 logits 局部凸性角度解构 RL 训练不稳定根源，理论视角极具原创性与启发性。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 大模型家族、3 种主流任务、稀疏与规则 RM 扩展以及详细的梯度与熵动力学监控。
• 写作质量: ⭐⭐⭐⭐⭐ 理论推导闭环严密，从 NTK 假设、Hessian 判定到收敛速度证明一气呵成。
• 价值: ⭐⭐⭐⭐⭐ 为大模型强化学习（尤其是大模型长思维链推理训练中的防崩塌）提供了全新且极易部署的基础算法范式。

相关论文

• [NeurIPS 2026] Flow Matching Policy Optimization with Mirror Descent and Entropy Constraints
• [NeurIPS 2026] Policy Optimization in Hybrid Discrete-Continuous Action Spaces via Mixed Gradients
• [ICLR 2026] BAPO: Stabilizing Off-Policy Reinforcement Learning for LLMs via Balanced Policy Optimization with Adaptive Clipping
• [NeurIPS 2026] Performance-Driven Policy Optimization for Speculative Decoding with Adaptive Windowing
• [NeurIPS 2026] Amortized Molecular Optimization via Group Relative Policy Optimization