Segment-Aligned Policy Optimization for Multi-Modal Reasoning

2026-05-02

AI Summary

Existing reinforcement learning approaches for large language models optimize policies either at the token or full-sequence level, which poorly aligns with the natural stepwise structure of multimodal reasoning, leading to credit assignment difficulties and training instability. This work proposes Segment-Aligned Policy Optimization (SAPO), which elevates the policy update unit to semantically coherent reasoning segments for the first time. SAPO formulates a segment-level Markov decision process and introduces tailored mechanisms for value estimation, advantage computation, and importance sampling. Evaluated across multiple reasoning benchmarks, SAPO substantially outperforms existing token-level and sequence-level methods, achieving higher accuracy while simultaneously improving training stability and consistency in value estimation.

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. Codes and models will be released to ensure full reproducibility.

Problem

• reinforcement learning
• multi-modal reasoning
• credit assignment
• policy optimization
• reasoning structure

Innovation

• Segment-Aligned Policy Optimization
• reasoning segments
• multi-modal reasoning
• reinforcement learning
• credit assignment