Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stabilizing Policy Optimization via Logits Convexity

Hongzhan Chen, Tao Yang, Yuhua Zhu, Shiping Gao, Xiaojun Quan, Ting Yao

TL;DR

This work investigates the stability gap between SFT and RL from a gradient-based perspective, and shows that the convexity of the SFT loss with respect to model logits plays a key role in enabling stable training.

Abstract

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.

Stabilizing Policy Optimization via Logits Convexity

TL;DR

This work investigates the stability gap between SFT and RL from a gradient-based perspective, and shows that the convexity of the SFT loss with respect to model logits plays a key role in enabling stable training.

Abstract

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.
Paper Structure (59 sections, 7 theorems, 71 equations, 9 figures, 6 tables)

This paper contains 59 sections, 7 theorems, 71 equations, 9 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Problem Setup
  4. Supervised Fine-Tuning
  5. Proximal Policy Optimization
  6. Gradient Dynamics
  7. Gradient Dynamics of SFT
  8. Gradient Dynamics of PPO
  9. On the Convexity of Logits
  10. LCO Objectives
  11. Regression-based LCO Objectives
  12. Distribution-based LCO Objective
  13. Advantage Estimation
  14. Sampled Action Sparse Estimation
  15. Log Probability-Based Dense Estimation
  16. ...and 44 more sections

Key Result

Proposition 4.1

For a given behavioral policy $\pi_\text{old}$ and advantage function $A(s_t,a_t)$, the optimal policy $\pi^*$ that maximizes eq:ppo-objective-with-kl-penalty at time step $t$ is: A specific solution for the optimal logits $z^*(s_t,a_t)$ that induces $\pi^*$ is obtained by direct advantage-based adjustment: where $z_\text{old}(s_t,a_t)$ is the logit of behavioral policy $\pi_\text{old}$ correspo

Figures (9)

  • 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.
  • ...and 4 more figures

Theorems & Definitions (11)

  • Proposition 4.1
  • proof
  • Definition 4.3: Logits Convexity
  • Proposition 4.4: Gradient Directionality
  • Lemma 4.5
  • Lemma 4.6
  • Proposition 4.7: Gradient Norm Upper Bounds
  • proof
  • Proposition 4.8: Convergence Bound of $\mathcal{L}_\text{LCO-MSE}$
  • Proposition 4.9: Convergence Bound of $\mathcal{L}_\text{LCO-LCH}$
  • ...and 1 more