Prioritize the Process, Not Just the Outcome: Rewarding Latent Thought Trajectories Improves Reasoning in Looped Language Models

AI summary

Overview

Research area: Reinforcement learning post-training for large language models, specifically latent reasoning architectures (looped transformers).

Technical level: Intermediate. The core idea is intuitive, but the paper relies on policy-gradient notation and a formal theoretical result.

Scope: This paper introduces RLTT (Reward Latent Thought Trajectories), a reinforcement learning objective that distributes reward across every internal loop of a Looped Language Model instead of only the final one, and evaluates it against GRPO on Ouro-1.4B/2.6B-Thinking across math and non-math reasoning benchmarks.

What This Paper Is About

Looped Language Models (LoopLMs) such as Ouro perform several rounds of internal "latent" computation before emitting each output token, and they outperform conventional LLMs on reasoning tasks at smaller parameter budgets. Prior attempts to improve them further with reinforcement learning have failed, because standard objectives like GRPO assign credit only to the model's last latent state, ignoring the multi-step internal computation that produced it. The goal of this work is to fix that mismatch by rewarding the whole latent thought trajectory rather than just its endpoint.

Key Contributions

1. A new RL objective, RLTT. The authors introduce Reward Latent Thought Trajectories, which replaces the single final-loop log-probability in a REINFORCE-style policy gradient with a weighted sum over all loops, so the reward signal directly shapes intermediate latent states. RLTT is designed to drop in as a direct GRPO replacement with negligible added overhead and without external verifiers.

2. Empirical gains over GRPO under matched conditions. On Ouro-1.4B/2.6B-Thinking, trained with identical MATH data, rollout budgets, optimization settings, reward functions, and advantage normalization, RLTT beats GRPO on math reasoning benchmarks.

3. Zero-shot transfer to non-mathematical tasks. Despite training only on mathematics, RLTT improves logical inference, factual recall, multi-domain QA, and code generation benchmarks.

4. A credit-assignment explanation and formal support. The paper argues RLTT works by improving reasoning efficiency, backed by a theorem (theorem A.5 in Appendix A.10) that trajectory-level credit assignment yields weakly shorter optimal decoding lengths, plus gradient signal-to-noise ratio measurements.

Main Findings

• Headline math gains over GRPO. Averaged over MATH-500, AIME24/26, and BeyondAIME, RLTT improves mean accuracy by +5.8% at the 1.4B scale and +10.9% at the 2.6B scale.

• Per-benchmark math improvements. At the 2.6B scale, RLTT gains +7.0% on MATH-500, +16.6% on AIME24, +10.0% on AIME26, and +10.0% on BeyondAIME. At the 1.4B scale, gains are +3.8% on MATH-500, +10.0% on AIME24, +3.7% on AIME26, and +6.0% on BeyondAIME.

• Non-math transfer. At the 2.6B scale, RLTT improves +0.7% on ARC-C, +3.5% on MMLU-ST, +18.7% on GPQA, and +3.3% on MBPP. At the 1.4B scale, +1.2% on ARC-C, +2.7% on MMLU-ST, +12.1% on GPQA, and +5.1% on MBPP. The authors note the GPQA gain nearly doubles GRPO performance.

• Absolute averages. Ouro-1.4B + RLTT reaches 46.0% mean math benchmark accuracy and 64.8% mean non-math accuracy; Ouro-2.6B + RLTT reaches 51.2% and 71.8% respectively. For comparison, the paper reports Qwen3-1.7B + GRPO at 27.2% (math) and 51.1% (non-math), and Qwen3-4B + GRPO at 31.7% and 58.7%.

• Statistical significance. RLTT's improvements over GRPO are significant (p < 0.05) across 7/9 benchmarks at 1.4B and 8/9 benchmarks at 2.6B. The paired t-test analysis in Appendix A.2 reports significance on all math benchmarks at 2.6B, 3/5 math benchmarks at 1.4B, and all evaluated non-math benchmarks at both scales.

• Shorter responses emerge without a brevity reward. RLTT-trained policies converge to substantially shorter responses even though the reward depends only on final-answer correctness. The authors present this as a downstream effect of aligning intermediate latent states with the final outcome.

• Training time reduction. Under the MATH training run, GRPO averaged 23.3 min/step (± 8.31) with a total of 54.42 hours, while RLTT averaged 21.1 min/step (± 9.87) with 49.05 hours, a relative cost of 0.90x and a 10.0% reduction in training time.

• Entropy drop is described as controlled confidence. RLTT shows steeper and more sustained output-token entropy reduction than GRPO; the authors argue this is confidence rather than entropy collapse and point to a Pass@k analysis in Appendix A.4.

• Robustness across decode budgets (MATH-500). Ouro2.6B-Thinking scores 38.4 / 67.8 / 75.6 / 79.6 at 1024 / 2048 / 3072 / 4096 tokens; + SFT scores 34.6 / 58.2 / 71.8 / 77.6; + GRPO scores 42.4 / 71.6 / 79.0 / 80.8; + RLTT scores 78.4 / 86.0 / 87.4 / 89.8.

• Loop-level advantage. Appendix A.6 reports RLTT outperforming GRPO at every evaluated loop count, with especially large margins in 1–2 loop regimes.

• Weighting strategy is not the source of gains. Appendix A.3 reports that RLTT's performance is largely insensitive to the precise loop-weighting strategy used to distribute credit.

• Richer gradients on hard tasks. RLTT yields statistically significant Gradient Signal-to-Noise Ratio improvements on AIME24 and BeyondAIME, described as the hardest benchmarks where credit assignment is most challenging.

Methodology in Plain English

The setup is an autoregressive model that, at every output token, runs a shared transformer block for a maximum of T_max loops, producing a "next-token distribution" at each loop. Normally only the distribution from the final loop is used for sampling; the earlier ones are treated as unobserved internal computation.

RLTT changes the learning signal in one key place. Where a standard REINFORCE-style objective multiplies the advantage by the log-probability of the token under the final loop's distribution, RLTT uses a weighted sum of the log-probabilities from every loop, with non-negative weights that add to one. This means a rollout's reward influences the gradients of intermediate latent states, not just the terminal one. A KL-divergence term against a frozen reference policy (a copy of the model from before RLVR post-training) is added to preserve general language ability.

The weights can be chosen in different ways: "Exit PDF" uses Ouro's learned exit head probability that computation halts at loop t; "Progressive" gives later loops more weight via t^α / Σ s^α; "Uniform" gives every loop equal weight 1/T_max. The main experiments use exit-probability weighting. The core algorithm is a loop over prompts: sample g rollouts, score them, normalize into advantages, record per-loop log-probabilities, then take an optimizer step — the same skeleton as GRPO.

For evaluation, everything is compute-matched: same MATH training samples, same rollout budgets, reward functions, and advantage normalization. Math evaluation is zero-shot with deterministic decoding and exact-match answer parsing. GRPO is given a 3072-token budget on all math benchmarks; RLTT gets 2048 tokens on MATH-500, 512 on GSM8K, and 3072 on the rest. Non-math benchmarks use a 2048-token budget.

A practical cost the authors flag: retaining per-loop log-probabilities makes memory scale linearly with the number of loops. Under fixed GPU memory they had to set ppo_max_token_len_per_gpu to 8192, half of GRPO's value, and compensate with

Authors’ 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.