Beyond the Mean: Multi-Moment Policy Optimization for LLM Reasoning

2026-08-03

AI Summary

This work addresses a critical limitation in existing large language model (LLM) inference strategy optimization, which typically focuses solely on the mean of failure probability while neglecting its distributional characteristics, thereby compromising robustness. To overcome this, the paper introduces Multi-Moment Policy Optimization (MMPO), the first framework to explicitly model failure probability as a random variable and jointly optimize its higher-order moments to comprehensively shape its distribution. The authors further develop a general moment transformation theory that unifies diverse optimization objectives under a common formalism. Extensive experiments across five mathematical reasoning benchmarks and LLMs of varying scales demonstrate that MMPO substantially outperforms strong baselines, significantly reducing the expected truncation time required to obtain the first successful response.

Abstract

Reinforcement learning has become a central paradigm for improving the reasoning capabilities of large language models. Existing methods generally aim to reduce the failure probabilities induced across problems. In this paper, we introduce a moment-based perspective on policy optimization for LLM reasoning by treating the failure probability of a randomly sampled problem as a random variable and characterizing optimization objectives through its moments. Under this perspective, many existing methods optimize only a single moment of the failure-probability distribution, leaving its broader distributional structure largely uncharacterized. We propose \textbf{M}ulti-\textbf{M}oment \textbf{P}olicy \textbf{O}ptimization (MMPO), a novel policy optimization framework that jointly minimizes multiple moments of the failure-probability distribution. MMPO admits a direct operational interpretation as minimizing the expected truncated time required to obtain the first successful response. Beyond MMPO, we further develop a general moment-transformation framework that systematically induces different moment profiles and provides a unified view of a broader family of policy optimization objectives. Experiments across five mathematical reasoning benchmarks and models of different scales demonstrate that MMPO consistently outperforms strong baselines. We hope this moment-based perspective offers new insights into the design of policy optimization objectives for LLM reasoning.

Problem

• policy optimization
• LLM reasoning
• failure probability
• moment-based perspective
• distributional structure

Innovation

• multi-moment optimization
• policy optimization
• failure probability distribution
• reinforcement learning
• LLM reasoning