Stabilizing Policy Optimization via Logits Convexity
Authors: Hongzhan Chen, Tao Yang, Yuhua Zhu, Shiping Gao, Xiaojun Quan, Ting Yao
Organizations: School of Computer Science and Engineering, Sun Yat-sen University, China · Shanghai Innovation Institute · Wechat Search, Tencent Inc, China · Shenzhen Loop Area Institute
While reinforcement learning (RL) has been central to the recent success of large language models (LLMs), RL optimization is notoriously unstable, especially when compared to supervised fine-tuning (SFT). In this work, we investigate the stability gap between SFT and RL from a gradient-based perspective, and show that the convexity of the SFT loss with respect to model logits plays a key role in enabling stable training. Our theoretical analysis demonstrates that this property induces favorable gradient directionality during optimization. In contrast, Proximal Policy Optimization (PPO), a widely adopted policy gradient algorithm utilizing a clipped surrogate objective, lacks this stabilizing property. Motivated by this observation, we propose Logits Convex Optimization (LCO), a simple yet effective policy optimization framework that aligns the learned policy with an optimal target derived from the original RL objective, thereby emulating the stabilizing effects of logits-level convexity. Extensive experiments across multiple model families show that our LCO framework consistently improves training stability and outperforms conventional RL methods on a broad range of benchmarks.
Figures & tables
Figure 1: (a) Gradient norm during training for SFT and PPO. (b) Pass@1 results of SFT and PPO on the MATH500 benchmark..
Figure 2: (a) Gradient norm during training for LCO-KLD. (b) Pass@1 results of LCO-KLD on MATH500.
Figure 3: Training dynamics of SFT. (a) Gradient norm decreases as training progresses while target token probability on training sample increases. The magnitudes of both target token logit gradient (b) and the non-target token logit gradient (c) diminish as training progresses.
Figure 4: Training dynamics of PPO for actions with negative advantages. (a) Gradient norm increases before being clipped, while sampled action probabilities decrease. The magnitudes of both sampled action logit gradient (b) and non-sampled action logit gradient (c) increase before being clipped. The results show a correlation between the behavior of parameter-space gradients and logit-space gradients.
Figure 5: Training dynamics of LCO-KLD for actions with negative advantages. Unlike training dynamics of PPO, LCO‑KLD exhibits stable gradient updates that scale with the remaining error and diminish progressively as the model converges.
Figure 6: Training dynamics of LCO-KLD. (a) Gradient norm for actions with positive advantages. (b) Gradient norm for actions with negative advantages. Unlike training dynamics of PPO, these results demonstrate a stable gradient update of LCO-KLD during training.
Figure 7: Training efficiency of LCO-LCH, LCO-KLD, and PPO, evaluated by policy accuracy on the MATH500 evaluation set.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Training dynamics of PPO for actions with positive advantages. (a) Gradient norm decreases as training progresses while sampled action probabilities increase. The magnitudes of both sampled (b) and non-sampled action logit gradient (c) decrease.
Figure 9: Training dynamics of LCO-KLD for actions with positive advantages. We show the dynamics of gradient norm (a) , logit gradient for sampled (b) and non-sampled (c) actions. LCO‑KLD exhibits stable gradient updates as the model converges.
Post-training paradigms for Large Language Models (LLMs), primarily Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL), face a fundamental dilemma: SFT provides stability (low variance) but suffers from high fitting bias, while RL enables exploration (low bias) but grapples with high gradient variance. Existing unified optimization strategies often employ naive loss weighting, overlooking the statistical conflict between these distinct gradient signals. In this paper, we provide a rigorous theoretical analysis of this bias-variance trade-off and propose DYPO (Dynamic Policy Optimization), a unified framework designed to structurally mitigate this conflict. DYPO integrates three core components: (1) a Group Alignment Loss (GAL) that leverages intrinsic group dynamics to significantly reduce RL gradient variance; (2) a Multi-Teacher Distillation mechanism that corrects SFT fitting bias via diverse reasoning paths; and (3) a Dynamic Exploitation-Exploration Gating mechanism that adaptively arbitrates between stable SFT and exploratory RL based on reward feedback. Theoretical analysis confirms that DYPO linearly reduces fitting bias and minimizes overall variance. Extensive experiments demonstrate that DYPO significantly outperforms traditional sequential pipelines, achieving an average improvement of 4.8% on complex reasoning benchmarks and 13.3% on out-of-distribution tasks. Our code is publicly available at https://github.com/Tocci-Zhu/DYPO.
Taojie Zhu, Dongyang Xu, Ding Zou +4
Shenzhen International Graduate School, Tsinghua University · Intelligent System Department, Zhongxing Telecom Equipment (ZTE) · Institute of Advanced Interdisciplinary Studies, Chongqing University of Posts and Telecommunications
Reinforcement learning (RL) has become a cornerstone for improving the reasoning capabilities of large language models (LLMs), but the need for on-policy data substantially limits training efficiency. Reusing off-policy data through importance sampling (IS) can improve efficiency but introduce considerable instability. Hence, algorithms such as PPO and GRPO widely adopt IS-ratio clipping to stabilize training. However, training stability and gradient estimate are mainly determined by the product of IS ratio and advantage. To further stabilize training, we propose ACPO, which clips the product of the IS ratio and the advantage, leading to more stable gradient estimates. We also establish a connection between ACPO and gradient clipping in policy mirror descent (PMD), which is a standard technique to stabilize optimization process, and prove the convergence of clipped-PMD under the standard RL setting. Experiments on widely used mathematical reasoning benchmarks show that ACPO consistently outperforms PPO and GRPO in both accuracy and training efficiency, delivering 4-6 percentage points gains on standard math benchmarks, with Qwen3-8B+PPO. Hence, ACPO is a practical and effective alternative to conventional IS-ratio clipping for RL post-training of LLMs.
Reinforcement learning (RL) has become an effective post-training paradigm for long-horizon large language model (LLM) agents. However, we find that the resulting policies can be sensitive to various policy perturbations, such as hidden-state noise, pruning, and quantization. In this work, we study how to improve perturbation robustness during policy optimization. We first introduce the notion of a perturbation robust policy and analyze conditions under which perturbed policy updates preserve stable monotonic improvement. Based on this analysis, we introduce Stable Perturbation-Robust Policy Optimization (SPrPO), which applies adaptive and sensitivity-aware perturbations during RL training. We evaluate SPrPO on ALFWorld and WebShop and conduct systematic experiments across multiple perturbation types and scales, showing improved perturbation robustness while maintaining stable policy optimization.
Pengxin Wang, Yuanzhe LI, Yuxin Ren +2
Department of Electrical and Computer Engineering University of Arizona