Prioritize the Process, Not Just the Outcome: Rewarding Latent Thought Trajectories Improves Reasoning in Looped Language Models
Abstract
Looped Language Models (LoopLMs) perform multi-step latent reasoning prior to token generation and outperform conventional LLMs on reasoning benchmarks at smaller parameter budgets. However, attempts to further improve LoopLM reasoning with reinforcement learning have failed—standard objectives such as Group Relative Policy Optimization (GRPO) only assign credit to the final latent state, creating a fundamental mismatch with the model's internal computation. To resolve this, we introduce RLTT (Reward Latent Thought Trajectories), a reinforcement learning framework which distributes reward across the full latent reasoning trajectory. RLTT provides dense, trajectory-level credit assignment without relying on external verifiers and can directly replace GRPO with negligible overhead. Across extensive experiments with Ouro-1.4B/2.6B-Thinking under identical training and inference conditions, RLTT yields statistically significant improvements over GRPO on challenging mathematical reasoning benchmarks, improving mean accuracy over MATH-500, AIME24/26, and BeyondAIME by +5.8\% on the 1.4B scale, and +10.9\% on the 2.6B scale. Despite being trained exclusively on mathematics, RLTT also transfers effectively to non-mathematical reasoning benchmarks, demonstrating the effectiveness of trajectory-level credit assignment for reinforcement learning in LoopLMs. Code is available at https://github.com/jonwill8/RLTT.git.
Lay Summary
AI models are often asked to solve difficult problems, like math questions or logic puzzles. Some models solve these problems by writing out long step-by-step explanations, but this can be slow and inefficient. Other models try to do more of the thinking internally before giving an answer. Our work improves how these “internal thinking” models are trained. Instead of only checking whether the final answer is right or wrong, our method gives the model better feedback on the thinking process it uses along the way. This helps the model learn to reach correct answers more reliably and with fewer unnecessary words. In our experiments, this approach improved performance on challenging math problems and also helped on other reasoning tasks, even though the model was only trained on math. Overall, the work shows that AI models can become better problem-solvers when we train them to improve how they think, not just what answer they give.