Lifted state policy improves reasoning accuracy and efficiency
Lifted State Policy Optimization: Internalizing Self-Correction through Transition Geometry
Artificial IntelligenceMachine Learning
Summary
Many AI models solve problems by guessing answers step-by-step but don’t focus on the small improvements between steps. The authors created a new method called Lifted State Policy Optimization that teaches the AI to value local improvements in reasoning by assigning coordinates to answers and learning which changes lead to better results. This approach helps the AI improve its answers more efficiently, using fewer words and less time while maintaining accuracy across different subjects like math and science. Their tests show that this method works well on several AI models and tasks.
What this means in practice
- •For ai model developers: Improve large language models’ reasoning efficiency by integrating transition-based reasoning improvements to reduce token use and latency.
- •For natural language processing engineers: Develop text generation systems that internally adapt and refine answers through geometric state representations to boost accuracy on complex benchmarks.
Authors
Ren Zhuang, Ben Wang, Shuifa Sun
Abstract
Preference optimization ranks final answers and trajectory RL maximizes aggregate returns, yet neither credits the local change that turns a flawed reasoning state into a better one. We introduce Lifted State Policy Optimization (LSPO), which establishes this finite, cost-adjusted transition as the primary unit of learning. Each textual answer is augmented with a learned continuous coordinate, forming a lifted state that an energy function evaluates and orders. A shaped reward drives energy descent net of edit and step costs, and an internalization loss absorbs selected low-energy answers into the direct generator. On Qwen3.5-27B, the resulting single-pass policy reaches 71.8% aggregate accuracy across mathematics, science, code, and logic benchmarks, matching test-time reflection at 71.9% with 28% lower latency and 25.8% fewer tokens. The pattern holds across dense and mixture-of-experts backbones, and matched controls confirm that transition credit and the lifted coordinate each contribute, as does answer internalization. Our code is available at https://github.com/Venomrko/lspo.