Beyond the Mean: Multi-Moment Policy Optimization for LLM Reasoning
2026-08-03 • Artificial Intelligence
Artificial Intelligence
AI summaryⓘ
The authors propose a new way to improve how large language models (LLMs) solve problems by focusing on the chances of failure as a random variable with multiple characteristics, called moments. Instead of just reducing average failure rates, they introduce a method called Multi-Moment Policy Optimization (MMPO) which minimizes several aspects of failure probability, aiming to get successful answers more efficiently. They also develop a general framework to understand different optimization goals based on these moments. Tests on math reasoning tasks show that their approach works better than previous methods.
reinforcement learninglarge language modelspolicy optimizationfailure probabilitymoments (statistics)mathematical reasoningmoment transformationoptimization objectivesMMPOrandom variable
Authors
Yijun Zhang, Yule Xie, Jiaxin Ding, Xin Ding, Fan Xu, Haoxiang Zhang, Luoyi Fu
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.