Lifted State Policy Optimization: Internalizing Self-Correction through Transition Geometry
Organizations: Hangzhou Normal University
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.
Figures & tables
| Model | Method | MATH-500 | AIME25 | GPQA | LCB | ZebraLogic | MMLU-Pro | Average | |||||||
| @1 | @16 | @1 | @16 | @1 | @16 | @1 | @16 | @1 | @16 | @1 | @16 | @1 | @16 | ||
| Qwen3.5-9B | Base | 84.8 | 93.6 | 39.6 | 57.2 | 45.8 | 61.4 | 41.9 | 48.6 | 60.2 | 73.4 | 56.3 | 68.8 | 54.8 | 67.2 |
| DPO | 85.6 | 94.1 | 40.9 | 58.3 | 46.6 | 62.3 | 43.1 | 50.1 | 61.2 | 74.3 | 57.2 | 69.7 | 55.8 | 68.1 | |
| Step-DPO | 86.4 | 94.7 | 42.7 | 60.5 | 47.5 | 63.5 | 44.4 | 52.0 | 62.4 | 75.3 | 58.1 | 70.8 | 56.9 | 69.5 | |
| GRPO | 87.6 | 95.3 | 46.0 | 65.0 | 49.1 | 66.4 | 46.6 | 55.0 | 64.4 | 77.0 | 59.8 | 72.7 | 58.9 | 71.9 | |
| GDPO | 88.0 | 95.5 | 46.6 | 65.9 | 50.0 | 67.8 | 47.4 | 55.9 | 65.1 | 77.5 | 60.6 | 73.7 | 59.6 | 72.7 | |
| Method | Net gain | Direct | Terminal |
| Base | 0.0 | 0.468 | N/A |
| DPO | 1.8 | 0.450 | N/A |
| Step-DPO | 2.5 | 0.438 | N/A |
| GRPO | 3.3 | 0.420 | 0.203 |
| GDPO | 4.0 | 0.411 | 0.201 |
| SCoRe | 3.6 | 0.411 | 0.197 |
| Method | Agg. | Latency (s) | Tokens |
| Select then SFT | 70.6 ±0.17 | 10.58 | 1218 |
| SI only | 70.8 ±0.16 | 10.74 | 1246 |
| GRPO | 70.0 ±0.18 | 12.04 | 1370 |
| SCoRe | 70.7 ±0.18 | 13.12 | 1501 |
| Reflection | 71.9 ±0.19 | 15.18 | 1705 |
| LSPO | 71.8 ±0.15 | 10.93 | 1265 |
| Variant | Agg. | Net gain | Latency (s) | Tokens | GPU hrs |
| LSPO full | 71.8 | 4.7 | 10.93 | 1265 | 2983 |
| - w/o lifting | 70.9 | 3.2 | 11.70 | 1316 | 2804 |
| - w/o progress reward | 70.0 | 1.6 | 12.35 | 1366 | 2715 |
| - w/o internalization | 70.3 | 2.4 | 12.02 | 1341 | 2566 |
| GRPO | 70.0 | 3.3 | 12.04 | 1370 | 2386 |
| DPO | 66.9 | 1.8 | 10.42 | 1196 | 1218 |
Appendix figures & tables32 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | MATH-500 | GPQA | LCB | Aggregate | Latency (s) |
| Base | 87.6 | 48.7 | 47.1 | 61.1 | 6.81 |
| DPO | 88.3 | 49.4 | 48.0 | 61.9 | 7.04 |
| GRPO | 89.2 | 50.1 | 49.4 | 62.9 | 7.18 |
| LSPO | 90.7 | 51.3 | 52.0 | 64.7 | 7.31 |
| Component | Architecture | Input / output | Key setting |
| Lift projection | 2-layer MLP, width 1024, GELU | Hidden state | Input LayerNorm |
| Energy head | 3-layer MLP, width 1024, SiLU | Pooled text, , task embedding scalar | Pairwise margin 0.20 |
| Transition block | 2 residual blocks, width 1024, SiLU | Text delta, , action embedding | |
| Action embedding | Lookup table, width 128 | Revise, Switch, Stop vector | Frozen transition conditioning |
| Value head | 2-layer MLP, width 512, SiLU | Transition state scalar value | GAE 0.95, clip 0.2 |
| Task family | Public corpus source | Samples |
| Math | open-r1/OpenR1-Math-220k 1 1 1 https://huggingface.co/datasets/open-r1/OpenR1-Math-220k | 38,000 |
| Code | BAAI/TACO 2 2 2 https://huggingface.co/datasets/BAAI/TACO | 24,000 |
| Science | allenai/sciq 3 3 3 https://huggingface.co/datasets/allenai/sciq | 18,000 |
| Logic | longface/logicLM 4 4 4 https://huggingface.co/datasets/longface/logicLM | 12,000 |
| Measure / phase | Revise (%) | Switch (%) | Stop (%) | Reach (%) |
| Overall | 49.8 | 25.9 | 24.3 | 11.8 |
| Energy decrease rate | 70.8 | 63.4 | – | – |
| Early 0–30% | 56.1 | 30.8 | 13.1 | 21.4 |
| Late 70–100% | 42.6 | 20.2 | 37.2 | 5.9 |
| Method | Agg. SD | LSPO diff. (95% CI) | Inference contract |
| Select then SFT | 70.6 0.17 | +1.2 [0.7, 1.6] | direct |
| SI only | 70.8 0.16 | +1.0 [0.6, 1.4] | direct |
| Without lifting | 70.9 0.18 | +0.9 [0.5, 1.3] | direct |
| Without progress reward | 70.0 0.19 | +1.8 [1.3, 2.3] | direct |
| Return-matched terminal reward | 69.7 0.20 | +2.1 [1.5, 2.6] | direct |
| Test-time reflection | 71.9 0.19 | – | critique + revision |
| Family | Variant | Aggregate drop | Net gain |
| Lift | dimension = 64 | 0.6 | 3.9 |
| Lift | dimension = 2 | 0.8 | 3.7 |
| Lift | No lift | 0.9 | 3.2 |
| Lift | No switch | 1.4 | 3.1 |
| Lift | Revise only | 1.7 | 2.7 |
| Lift | Switch only | 2.0 | 2.5 |
| Energy signal | Agg. | Drop | Net gain | AUC | Pairwise |
| Clean LSPO energy | 71.8 | 0.0 | 4.7 | 0.792 | 0.734 |
| Gaussian noise | 70.3 | 1.5 | 2.9 | 0.731 | 0.681 |
| Label flip or shuffle | 68.6 | 3.2 | 0.7 | 0.508 | 0.506 |
| Terminal only | 69.7 | 2.1 | 2.2 | 0.681 | 0.641 |
| Variant | Candidate pool | State space | Target selector | Controller RL | Reward signal | Transfer target |
| Full LSPO | Dynamic rollout | Lifted | Frozen | Dense policy | (SI) | |
| Return-matched terminal reward | Dynamic rollout | Lifted | Frozen | Delayed policy | Trajectory | (SI) |
| SI only on fixed pool | Fixed Full pool | Lifted | Frozen | None | None | (SI) |
| Select then SFT | Fixed Full pool | Surface text | Task verifier | None | None | Verifier-best (SFT) |
| Without lifting | Dynamic rollout | Surface text | Surface | Dense policy | (SI) | |
| Without progress reward | Dynamic rollout | Lifted | Frozen | Dense policy | (SI) |
| Input condition | Action acc. SD | NLL SD | Energy AUC SD |
| Text only | 62.8 0.7 | 0.842 0.011 | 0.681 0.010 |
| Text plus | 68.9 0.6 | 0.768 0.010 | 0.732 0.009 |
| Text plus shuffled | 60.9 0.8 | 0.874 0.013 | 0.641 0.012 |
| Text plus zero | 63.4 0.7 | 0.831 0.012 | 0.687 0.011 |
| Variant | Agg. SD | Net gain | LSPO variant (95% CI) |
| Full LSPO | 71.8 0.15 | 4.7 | – |
| Reference baseline ( Table 4 ) | |||
| Current no-lift baseline | 70.9 0.18 | 3.2 | +0.9 [0.5, 1.3] |
| Parameter-matched coordinate controls | |||
| Same-parameter no-lift MLP | 70.8 0.18 | 3.0 | +1.0 [0.6, 1.4] |
| Random , matched dim. | 70.3 0.19 | 2.8 | +1.5 [1.1, 1.9] |
| Variant | Agg. | Gain | Latency (s) | Tokens | GPU hours |
| LSPO full | 72.6 | 4.4 | 8.30 | 1097 | 2410 |
| Without lifting | 71.6 | 2.9 | 8.88 | 1141 | 2265 |
| Without progress reward | 70.8 | 1.3 | 9.38 | 1185 | 2193 |
| Without internalization | 71.1 | 2.1 | 9.13 | 1163 | 2073 |
| Outcome RL | 71.5 | 2.7 | 9.33 | 1200 | 1940 |
| DPO | 68.8 | 1.5 | 8.08 | 1048 | 990 |
| Method | Initial | Net gain | |||
| Full LSPO | 40.1 | 95.8 | 29.5 | 4.2 | +16.0 |
| Same-parameter no-lift MLP | 39.7 | 92.1 | 21.6 | 7.9 | +9.9 |
| Without progress reward | 39.8 | 89.8 | 18.7 | 10.2 | +7.2 |
| Terminal-only reward | 39.1 | 90.3 | 16.9 | 9.7 | +6.5 |
| GRPO | 39.4 | 90.9 | 18.6 | 9.1 | +7.7 |
| Score | AUC | Spearman | Kendall | Pairwise | Bins | ECE |
| Task families used for fitting include math, code, logic, and science | ||||||
| Random energy | 0.502 | 0.01 | 0.01 | 0.502 | 2 | 0.241 |
| Outcome only score | 0.681 | 0.31 | 0.22 | 0.641 | 4 | 0.132 |
| LSPO energy | 0.792 | 0.48 | 0.35 | 0.734 | 5 | 0.071 |
| Shuffled LSPO energy | 0.508 | 0.01 | 0.01 | 0.506 | 2 | 0.235 |
| Task families held out from fitting include GSM8K, ARC-C, BBH, HumanEval, and MBPP | ||||||
| Method | GSM8K | ARC-C | BBH | HumanEval | MBPP | Avg. SD |
| Base | 88.4 | 84.7 | 66.2 | 55.5 | 62.4 | 71.4 0.22 |
| DPO | 89.0 | 85.3 | 67.1 | 56.4 | 63.0 | 72.2 0.20 |
| GRPO | 90.6 | 86.7 | 69.0 | 58.6 | 64.8 | 73.9 0.22 |
| SCoRe | 91.1 | 87.1 | 69.3 | 59.1 | 65.0 | 74.3 0.24 |
| LSPO | 91.5 | 87.4 | 70.4 | 60.1 | 65.2 | 74.9 0.20 |
| Agg. SD | Gain SD | Path length | Unique ratio | ||
| 2 | 3 | 70.9 0.18 | 3.2 0.30 | 1.6 | 0.82 |
| 4 | 5 | 71.5 0.16 | 4.2 0.26 | 2.7 | 0.76 |
| 6 | 7 | 71.8 0.15 | 4.7 0.23 | 3.4 | 0.71 |
| 8 | 9 | 71.8 0.16 | 4.8 0.24 | 3.7 | 0.66 |
| Method | WritingBench L&A | Base | MT-Bench Roleplay | Base | avg. |
| Base | 7.42 0.08 | 0.00 | 7.61 0.11 | 0.00 | 7.52 |
| DPO | 7.50 0.07 | +0.08 | 7.68 0.10 | +0.07 | 7.59 |
| GRPO | 7.36 0.09 | -0.06 | 7.53 0.12 | -0.08 | 7.45 |
| SCoRe | 7.32 0.09 | -0.10 | 7.49 0.12 | -0.12 | 7.41 |
| DRO-adapted | 7.53 0.08 | +0.11 | 7.70 0.10 | +0.09 | 7.62 |
| LSPO | 7.48 0.08 | +0.06 | 7.66 0.10 | +0.05 | 7.57 |
| Subset | Spearman | Kendall | Pairwise AUC | |
| WritingBench Literature and Arts | 0.34 0.04 | 0.24 0.03 | 0.675 0.018 | +0.42 |
| MT-Bench Roleplay | 0.29 0.06 | 0.20 0.04 | 0.642 0.026 | +0.35 |
| Subset | Strongest alignment | Weakest alignment | Failure slice |
| WritingBench Literature and Arts | Instruction fit, coherence | Creativity, originality | Prompts requiring high creativity |
| MT-Bench Roleplay | Character consistency, instruction fit | Expressive richness | Playful improvisation |
| Method | Peak mem (GB) | Elapsed time (h) | GPU hrs | Pilot GPU hours | Latency (s) | Output tokens |
| DPO | 57.7 | 152.2 | 1217.6 | 146.1 | 10.42 | 1196 |
| Step-DPO | 57.7 | 203.1 | 1624.8 | 195.0 | 10.73 | 1242 |
| GRPO | 57.7 | 298.3 | 2386.4 | 286.4 | 12.04 | 1370 |
| GDPO | 57.7 | 335.6 | 2684.8 | 322.2 | 11.54 | 1324 |
| SCoRe | 57.7 | 358.2 | 2865.7 | 343.9 | 13.12 | 1501 |
| LSPO | 61.9 | 372.9 | 2983.2 | 358.0 | 10.93 | 1265 |
| Mean | Energy drop | MATH-500 EM (pp) | WritingBench L&A | MT-Bench Roleplay | |
| 0 | 0.0% | 0.0 | |||
| 1 | 19.1% | +0.2 | |||
| 2 | 32.7% | +0.6 | |||
| 4 | 43.7% | +0.4 | |||
| 8 | 51.0% | +0.2 | |||
| 16 | 57.0% | -0.4 |
RL for Language Model Reasoning
Reinforcement learning for language model reasoning uses reward-driven training to improve multi-step problem solving, tool use, and reliable answers.
See all 1,087 papers