Segment-Aligned Policy Optimization for Multi-Modal Reasoning

Lei Gao, Zhuoming Li, Mengxi Jia, Jiakang Yuan, Hongbo Sun, Hao Sun, Xuelong Li

Abstract: Existing reinforcement learning approaches for Large Language Models typically perform policy optimization at the granularity of individual tokens or entire response sequences. However, such formulations often misalign with the natural step-wise structure of reasoning processes, leading to suboptimal credit assignment and unstable training in multi-modal reasoning tasks. To bridge this gap, we propose Segment-Aligned Policy Optimization (SAPO), a novel reinforcement learning paradigm that treats coherent reasoning steps, rather than tokens or full sequences as fundamental units of policy update. SAPO introduces a step-wise Markov decision process abstraction over reasoning segments, accompanied by segment-level value estimation, advantage computation, and importance sampling mechanisms that are semantically aligned with reasoning boundaries. Experiments on representative reasoning benchmarks demonstrate that SAPO consistently outperforms token-level and sequence-level policy optimization methods, achieving significant accuracy improvements while exhibiting better training stability and value estimation consistency. Our work underscores the importance of aligning reinforcement learning updates with the intrinsic structure of reasoning, paving the way for more efficient and semantically grounded policy optimization in complex reasoning tasks. Code is available at https://github.com/Graysonicc/SAPO.

Lay Summary: Existing reinforcement learning methods for LLMs optimize at the token or sequence level, leading to misaligned credit assignment for step-wise reasoning. We propose Segment-Aligned Policy Optimization (SAPO), which performs policy updates over coherent reasoning segments through segment-level value estimation and optimization. SAPO improves reasoning performance and training stability across representative benchmarks.

Link To Code: https://github.com/Graysonicc/SAPO

Primary Area: Reinforcement Learning->Deep RL

Keywords: Reinforcement Learning, MLLM Reasoning

Submission Number: 5733