Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

One Ring to Rule Them All: Unifying Group-Based RL via Dynamic Power-Mean Geometry

Weisong Zhao, Tong Wang, Zichang Tan, Te Yang, Siran Peng, Haoyuan Zhang, Tianshuo Zhang, Haichao Shi, Meng Meng, Yang Yang, Xiangyu Zhu, Zhen Lei, Xiao-Yu Zhang, Xu Zhou

TL;DR

This work tackles the fixed-geometry limitation of group-based RL methods for long-form mathematical reasoning by introducing Power-Mean Policy Optimization (PMPO), which generalizes token-level aggregation with an exponent $p$ in a power-mean over token ratios. The authors show $p=1$ recovers GRPO and $p\to 0$ recovers GMPO, while adaptive per-trajectory $p$ balances aggressive credit assignment and conservative stabilization. A clip-aware ESS matching mechanism adaptively selects $p$ by matching the induced token-weight ESS to a reliability-derived target, effectively adjusting the softmax temperature in gradient updates. The approach yields state-of-the-art results on multiple math benchmarks across model scales, demonstrating improved stability and efficiency without value-function critics. This has practical impact for scalable, robust online RLHF/RLAIF on large language models performing complex reasoning tasks, with a principled way to adapt gradient concentration to signal quality.

Abstract

Group-based reinforcement learning has evolved from the arithmetic mean of GRPO to the geometric mean of GMPO. While GMPO improves stability by constraining a conservative objective, it shares a fundamental limitation with GRPO: reliance on a fixed aggregation geometry that ignores the evolving and heterogeneous nature of each trajectory. In this work, we unify these approaches under Power-Mean Policy Optimization (PMPO), a generalized framework that parameterizes the aggregation geometry via the power-mean geometry exponent p. Within this framework, GRPO and GMPO are recovered as special cases. Theoretically, we demonstrate that adjusting p modulates the concentration of gradient updates, effectively reweighting tokens based on their advantage contribution. To determine p adaptively, we introduce a Clip-aware Effective Sample Size (ESS) mechanism. Specifically, we propose a deterministic rule that maps a trajectory clipping fraction to a target ESS. Then, we solve for the specific p to align the trajectory induced ESS with this target one. This allows PMPO to dynamically transition between the aggressive arithmetic mean for reliable trajectories and the conservative geometric mean for unstable ones. Experiments on multiple mathematical reasoning benchmarks demonstrate that PMPO outperforms strong baselines.

One Ring to Rule Them All: Unifying Group-Based RL via Dynamic Power-Mean Geometry

TL;DR

This work tackles the fixed-geometry limitation of group-based RL methods for long-form mathematical reasoning by introducing Power-Mean Policy Optimization (PMPO), which generalizes token-level aggregation with an exponent in a power-mean over token ratios. The authors show recovers GRPO and recovers GMPO, while adaptive per-trajectory balances aggressive credit assignment and conservative stabilization. A clip-aware ESS matching mechanism adaptively selects by matching the induced token-weight ESS to a reliability-derived target, effectively adjusting the softmax temperature in gradient updates. The approach yields state-of-the-art results on multiple math benchmarks across model scales, demonstrating improved stability and efficiency without value-function critics. This has practical impact for scalable, robust online RLHF/RLAIF on large language models performing complex reasoning tasks, with a principled way to adapt gradient concentration to signal quality.

Abstract

Group-based reinforcement learning has evolved from the arithmetic mean of GRPO to the geometric mean of GMPO. While GMPO improves stability by constraining a conservative objective, it shares a fundamental limitation with GRPO: reliance on a fixed aggregation geometry that ignores the evolving and heterogeneous nature of each trajectory. In this work, we unify these approaches under Power-Mean Policy Optimization (PMPO), a generalized framework that parameterizes the aggregation geometry via the power-mean geometry exponent p. Within this framework, GRPO and GMPO are recovered as special cases. Theoretically, we demonstrate that adjusting p modulates the concentration of gradient updates, effectively reweighting tokens based on their advantage contribution. To determine p adaptively, we introduce a Clip-aware Effective Sample Size (ESS) mechanism. Specifically, we propose a deterministic rule that maps a trajectory clipping fraction to a target ESS. Then, we solve for the specific p to align the trajectory induced ESS with this target one. This allows PMPO to dynamically transition between the aggressive arithmetic mean for reliable trajectories and the conservative geometric mean for unstable ones. Experiments on multiple mathematical reasoning benchmarks demonstrate that PMPO outperforms strong baselines.
Paper Structure (45 sections, 2 theorems, 37 equations, 4 figures, 7 tables, 1 algorithm)

This paper contains 45 sections, 2 theorems, 37 equations, 4 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Problem Setting and Notation
  5. Fixed-Geometry Aggregations
  6. GRPO (arithmetic-mean geometry).
  7. GMPO (geometric-mean geometry).
  8. Power-Mean Policy Optimization
  9. Power Mean over Token-Level Ratios
  10. Gradient Analysis of PMPO
  11. Log-Domain Clipping and PMPO Variants
  12. Determining $p$ via ESS Matching
  13. Clip-aware target ESS via Signal Reliability.
  14. Experiment
  15. Implementation Details
  16. ...and 30 more sections

Key Result

Theorem 1.1

Let $\{x_t\}_{t=1}^n$ be a set of positive real numbers ($x_t > 0$). For any two non-zero real numbers $p_1$ and $p_2$ such that $p_1 < p_2$, the following inequality holds: where $M_p(\{x_t\}) = \left( \frac{1}{n} \sum_{t=1}^n x_t^p \right)^{\frac{1}{p}}$. Equality holds if and only if all $x_t$ are equal.

Figures (4)

  • Figure 1: (a) Gradient view comparison of GRPO, GMPO and PMPO, and (b) Adaptive $p$ via clip-aware ESS matching.
  • Figure 2: Effect of different $p$ and $\epsilon_{ess}$ on average performance on five mathematical reasoning datasets.
  • Figure 3: Training Dynamics. (Top Left and Bottom Left) Comparison of the policy gradient norm evolution across GRPO, GMPO, and PMPO over training steps based on Qwen-Math-7B and DeepSeek-R1-Distill. (Top Middle) The trajectory of the adaptive exponent $p$ (reporting both batch mean and maximum) throughout the training process. (Top Right and Bottom Middle) The curve of (average token-level entropy & reward) throughout the training. (Bottom Right) The curve of the log-importance ratio $\Delta \ell_t = \log \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}$.
  • Figure 4: Visual Illustration of the Monotonic Relationship between $p$ and ESS. We simulate token-level log-probability differences $\tilde{\Delta}\ell$ with different variances. The curves demonstrate that increasing $p$ consistently reduces the Effective Sample Size. The varying slopes highlight that the sensitivity of ESS to $p$ depends on the signal variance, necessitating the proposed ESS-matching strategy.

Theorems & Definitions (4)

  • Theorem 1.1
  • proof
  • Proposition 2.1
  • proof