PaperlayerView exact revision on arXiv

Source-linked AI summary

LoopRPT: Reinforcement Pre-Training for Looped Language Models

Guo Tang, Shixin Jiang, Heng Chang, Nuo Chen, Yuhan Li, Huiming Fan, Jia Li, Ming Liu, Bing Qin

arXiv:2603.19714v1cs.CL

TL;DR

LoopLMs need reinforcement signals aligned with reasoning that unfolds across latent recurrent steps rather than only final output tokens. LoopRPT supplies those signals through hard-token selection, EMA-teacher step-wise rewards, and noisy latent rollouts, improving accuracy-computation trade-offs and enabling earlier effective exits across Ouro scales.

  • Problem

    Output-token RL rewards are structurally mismatched with LoopLMs’ implicit multi-step reasoning, where sparse feedback complicates credit assignment across latent iterations.

  • Method

    LoopRPT reframes next-token prediction as reasoning and trains hard tokens with EMA-teacher step-wise rewards, noisy latent rollouts, and joint exit-policy and representation optimization.

  • Results

    LoopRPT consistently improves per-step quality and accuracy-computation trade-offs across Ouro scales, including higher hard-token accuracy, fewer average steps, and downstream gains.

  • Takeaways & Limitations

    The results support hard-token-focused reinforcement pre-training as an approach for learning efficient latent reasoning in looped language models.

Abstract

from arXiv · show

Looped language models (LoopLMs) perform iterative latent computation to refine internal representations, offering a promising alternative to explicit chain-of-thought (CoT) reasoning. However, existing reinforcement learning (RL) paradigms primarily target output tokens, creating a structural mismatch with looped architectures whose reasoning unfolds implicitly. In this work, we propose LoopRPT, a reinforcement pre-training framework tailored for LoopLMs. By reframing next-token prediction as a next-token reasoning task, LoopRPT assigns reinforcement signals directly to latent steps using an EMA teacher reference and noisy latent rollouts. This formulation enables RL to directly shape intermediate representations, compressing effective reasoning into fewer iterations. We instantiate LoopRPT on the Ouro architecture across multiple model scales. Results demonstrate that LoopRPT consistently improves per-step representation quality, achieving Pareto dominance in accuracy-computation trade-offs. Notably, significant gains on hard tokens indicate that LoopRPT enhances early-stage reasoning rather than merely encouraging premature exits. Our findings highlight reinforcement pre-training as a principled paradigm for learning efficient latent reasoning in LoopLMs.

1. Introduction

LoopLMs offer latent iterative reasoning as an alternative to token-based CoT, but standard RL provides sparse output-level feedback that mismatches implicit computation. LoopRPT addresses this gap with latent-step reinforcement signals and improves efficiency and downstream performance.

  • LoopLMs refine representations recurrently in latent space, compressing multi-step reasoning without proportionally increasing parameter count or emitting intermediate tokens.
  • Standard RLVR rewards final output tokens, creating credit-assignment and stability challenges for LoopLMs with implicit reasoning and adaptive exits.
  • Looped models achieve higher accuracy than standard LLMs across next-token reasoning difficulty levels, motivating RL designed for latent recurrence.
  • LoopRPT selects hard tokens, uses EMA-teacher step-wise rewards, and applies noisy latent rollouts to optimize latent representations and exit behavior.
  • Experiments on Ouro-1.4B and Ouro-2.6B show improved per-step quality, better accuracy-computation trade-offs, fewer average steps, and gains on GSM8K, MBPP, and HumanEval.
  • LoopRPT is presented as the first reinforcement pre-training framework for LoopLMs, achieving Pareto-dominant accuracy-computation trade-offs across scales.

2. Preliminaries

LoopLMs perform recurrent latent computation for next-token prediction, while RPT reframes prediction as rewarded reasoning over sampled continuations. This supplies reinforcement learning with scalable corpus-level signals beyond likelihood maximization.

  • LoopLMs update a latent state through a parameter-shared backbone for multiple iterations, then derive the next-token distribution from the final state.
  • Standard next-token prediction maximizes log-likelihood but treats predictions independently and does not explicitly model latent reasoning for difficult tokens.
  • RPT samples an output continuation and receives a verifiable reward based on the ground-truth continuation.
  • By using intrinsic corpus rewards, RPT scales reinforcement learning to large datasets and incentivizes next-token reasoning rather than simple likelihood maximization.

3. Method

LoopRPT trains looped models with dense, hard-token-focused rewards over latent iterations, using an EMA teacher, difficulty-aware computation penalties, noisy rollouts, and step-weighted representation learning. Its objective jointly improves early prediction and adaptive exiting while preventing premature collapse.

  • Core framework: The method selects hard tokens by entropy, compares latent-step predictions with an EMA teacher, and jointly optimizes exits and representations through noisy rollouts.
  • Exit-step distribution: Each latent step produces next-token logits and an exit-gate probability, while cumulative exit probabilities determine a threshold-based exit step.
  • Step-wise rewards: The step-wise reward combines accuracy gain against the teacher reference with a difficulty-aware time penalty, providing dense supervision aligned with early exiting.
  • Representation learning: LoopRPT normalizes per-step rewards into advantages and weights representation learning toward latent steps that produce higher rewards.
  • Noisy latent rollouts: Gaussian perturbations create on-policy latent rollouts, whose rewards are normalized across groups and used for policy-gradient learning on hard-token positions.
  • Optimization: The total objective combines policy-gradient, representation, entropy-regularization, and KL-style trust-region terms, followed by an EMA teacher update.

4. Experimental Setup

Experiments train Ouro models with OMNI-MATH and evaluate next-token reasoning across entropy-defined difficulty splits, adaptive exits, and a broad set of downstream benchmarks.

  • Training Data: OMNI-MATH provides 4,428 competition-level mathematical problems, with 200 examples held out for validation.
  • Models: The study evaluates Ouro-1.4B and Ouro-2.6B, initializing teacher and student models from the same base model.
  • Next-token Reasoning: Tokens are divided into easy, medium, and hard splits using entropy thresholds of 0.5, 1.0, and 1.5.
  • Inference Settings: Ouro models are tested with maximum-loop execution and adaptive early exit, while Qwen3-1.7B is tested in vanilla and explicit-CoT modes.
  • Downstream Evaluation: End-task evaluation covers MMLU, MMLU-Pro, BBH, ARC-C, HellaSwag, Winogrande, GSM8K, MBPP variants, and HumanEval variants.

5. Experimental Results

LoopRPT improves next-token reasoning accuracy and efficiency across model scales, transfers gains to downstream tasks, and benefits from its complementary training components.

  • Main Results: +3.58 Hard Peak accuracy on Ouro-2.6B accompanies a reduction in average steps from 3.51 to 2.28.
  • Downstream Benchmarks: LoopRPT delivers consistent downstream gains, including MBPP+ +2.91 and GSM8K accuracy rising from 81.76 to 85.36 at 2.6B.
  • Exit Dynamics: LoopRPT increases early-step exits while preserving dominant final-step exits, with larger models shifting notably toward step-3 exits.
  • Stabilization and Trajectory Diversity: Removing Gaussian noise lowers accuracy and increases average exit steps, while removing policy-gradient or KL stabilizers causes performance drops and higher adaptive-inference latency.
  • Representation Shaping and Efficiency: Removing representation and entropy shaping produces the largest accuracy-efficiency degradation, while removing hard-token selection hurts Hard-split accuracy.

6. Analysis and Discussion

Analysis attributes LoopRPT’s trade-off gains to better intermediate representations, especially for hard tokens, while showing robustness across hyperparameters, depths, and model scales.

  • Hyperparameter Sensitivity: LoopRPT follows a predictable accuracy-compute trade-off, with moderate KL and larger GRPO groups improving adaptive performance before diminishing returns.
  • Accuracy–Computation Trade-offs: Across three difficulty buckets and model scales, LoopRPT achieves higher accuracy with fewer average computation steps than Ouro.
  • Per-step Representation Quality: LoopRPT improves per-step prediction accuracy at every latent step, with especially pronounced step-1 and step-2 gains for hard tokens.
  • Interpretation: The results indicate that LoopRPT improves intermediate representations rather than merely encouraging earlier exits, compressing effective computation without sacrificing predictive performance.
  • Forced-depth Evaluation: Forced-depth accuracy need not increase monotonically with depth, but LoopRPT outperforms the baseline across all buckets and depths.

7. Related Work

Related work spans reinforcement learning for language-model training and recurrent or adaptive-depth architectures that perform computation through halting, shared depth, or memory.

  • Reinforcement Learning for Language Models: Modern language-model reinforcement learning includes preference feedback, policy-gradient updates, AI feedback, rule-based supervision, and direct preference optimization.
  • Looped Language Models and Adaptive Depth: Adaptive-depth sequence models include dynamic halting, Universal Transformers, sparse shared-depth variants, and learned pondering mechanisms.
  • Looped Language Models and Adaptive Depth: Complementary recurrent approaches operate over segments or explicit memory for long-context modeling.

8. Conclusion

LoopRPT combines hard-token selection, step-wise latent supervision, noisy rollouts, and EMA references to improve looped reasoning while reducing inference computation. Its theoretical analyses connect entropy selection to informative gradients and noisy rollouts to smoothed objectives.

  • Conclusion: LoopRPT improves next-token reasoning under maximum-loop and adaptive early-exit execution while reducing average inference steps and improving downstream math and code benchmarks.The conclusion also reports better-calibrated early-exit behavior and identifies larger scales, broader mixtures, and distribution-shift robustness as future directions.
  • Method: The framework assigns dense, step-wise rewards to latent recurrence and uses an EMA teacher as a stable reference for exit and reward baselines.The teacher supplies a reference step and baseline while latent steps receive reinforcement signals.
  • Method: LoopRPT applies reinforcement-style losses to top-ρ high-entropy tokens, focusing training on positions selected by teacher uncertainty.Teacher entropy is computed from the final latent step, and subsequent losses use only selected positions.
  • Entropy selection: Under the teacher-as-label-uncertainty assumption, higher teacher entropy yields a larger lower bound on expected gradient energy.The result provides a theoretical basis for treating high-entropy tokens as more informative training examples.
  • Noisy rollouts: Noisy latent rollouts induce a Gaussian-smoothed on-policy objective whose score-function gradient is estimated without backpropagating through the reward.Monte Carlo estimation with independent noisy rollouts is unbiased for the smoothed objective’s gradient.

A.3. EMA Teacher as a Provably Slowly Moving Reference

The EMA teacher is analyzed as a slowly moving, low-pass-filtered reference whose exit decisions remain stable under sufficiently small parameter drift. This supports joint optimization of exit policies and intermediate representations without continual target chasing.

  • EMA teacher: The EMA teacher lags the student through a geometrically weighted sum of recent updates, acting as a low-pass filtered version of the student.The lag decomposition is expressed in terms of recent student parameter increments.
  • Reference-step stability: A teacher reference step remains unchanged when teacher parameter drift is below the exit-CDF margin scaled by the Lipschitz constant.The stability condition is ∥¯θ_n −¯θ_n−1∥≤m/LΠ.
  • EMA teacher: When student and prior teacher parameters are close, EMA suppresses teacher drift by a factor of 1 −ϕ.The bound follows directly from the EMA update.
  • Connection to LoopRPT: Consequently, EMA provides a slowly moving reference and stabilizes the jointly optimized exit policy and intermediate representations.The paper connects reference-step invariance and suppressed teacher drift to reduced target chasing.
  • Training iteration: LoopRPT’s training iteration separates dense reward construction from gradient updates that optimize noisy-rollout exits and step-weighted next-token learning.The algorithm’s two phases operate on selected high-entropy tokens.

F.1. Training Hyperparameters and Infrastructure

The experiments train LoopRPT on OMNI-MATH with Ouro-1.4B and Ouro-2.6B models using a distributed, mixed-precision training pipeline and Gaussian latent-noise sampling.

  • Infrastructure: Training uses 3 epochs, sequence length 4096, AdamW, cosine learning-rate decay, warmup, distributed data parallelism, mixed precision, and gradient checkpointing.These infrastructure and optimization choices are listed as the training setup.
  • Evaluation: The evaluation compares LoopRPT-trained Ouro models with the original model and similarly scaled Qwen, Gemma, and LLaMA model series.The evaluation spans knowledge, reasoning, mathematics, science, coding, and multilingual capabilities.

G. Training Dynamics Visualization

Training-dynamics figures track validation accuracy, average exit steps, and student behavior relative to the EMA teacher. The reported curves show improving validation performance alongside progressively fewer reasoning steps.

  • Observed dynamics: The training curves indicate that LoopRPT improves validation performance while progressively reducing required reasoning steps.The passage interprets this pattern as consistent with better early-exit calibration.

H. Case Studies

LoopRPT shows stronger constraint adherence and reasoning across general, mathematical, and coding cases than the base Ouro model. Training dynamics also indicate alignment with the EMA reference while using fewer reasoning steps.

  • Case-study scope: The qualitative cases span general reasoning, mathematical problem-solving, and code generation, highlighting error correction and multi-step dependency tracking.The comparison uses MMLU-Pro, GSM8K, and coding examples.
  • General reasoning & factual consistency: LoopRPT maintains factual consistency in interdisciplinary reasoning, unlike base Ouro traces that contain hallucinations or logical fragility.In biology, the base model contradicts itself by categorizing birds as ectotherms before later mentioning their endothermic nature.
  • Mathematical logic & constraint satisfaction: LoopRPT tracks separate variables and global set constraints in GSM8K cases where the base model makes mapping and set-neglect errors.The examples include scheduling, finance, and fractional-calculation problems.
  • Algorithmic invariants in coding: LoopRPT adheres to coding-task constraints by handling boundary conditions and algorithmic invariants that often cause off-by-one errors in the base model.The coding analysis covers MBPP and HumanEval cases.
  • Training dynamics: The student progressively aligns with the EMA teacher reference while converging to fewer reasoning steps, indicating improved early-exit calibration.The training-dynamics plots show rolling-mean student and teacher steps with within-window variability bands.
Loading 2603.19714v1…