Technical Summary: Stabilizing Recurrent Dynamics for Test-Time Scalable Latent Reasoning in Looped Language Models

1. Problem Statement

Looped Language Models (LoopLMs) represent a promising paradigm for enhancing Large Language Models (LLMs) through depth recurrence with shared parameters. This architecture aims to emulate human-like latent reasoning processes, offering improved computational efficiency by avoiding the expansion of context windows or parameter counts during inference. However, the authors identify a critical failure mode in existing LoopLMs: unreliable test-time scaling.

While the theoretical goal is that increased recurrent iterations lead to progressively refined latent representations, empirical observations (e.g., on the Ouro-1.4B model) reveal that performance often peaks at a specific iteration depth and subsequently collapses as recurrence continues. The paper attributes this to an inherent, previously overlooked trade-off between effectiveness (the ability to propagate information and perform deep reasoning) and stability (the boundedness of the latent trajectory).

• Internal Normalization (e.g., Pre-Norm): Maintains information flow (effectiveness) but allows hidden states to grow exponentially, leading to trajectory divergence and performance collapse.
• External Normalization (e.g., Post-Norm): Ensures bounded states (stability) but often restricts the system to shallow reasoning or poor fixed points, failing to support deep latent reasoning.

Standard remedies such as auxiliary Prelude/Coda layers, L2 regularization, or random loop sampling alone are found to be insufficient to resolve this deadlock.

2. Methodology: STARS Framework

To address the instability of latent dynamics, the authors propose STARS (STAbility-driven Recurrent Scaling), a unified training framework. The core insight is that effective latent reasoning requires the hidden states to converge toward asymptotically stable fixed points while maintaining effectiveness.

2.1 Theoretical Foundation

The approach is grounded in dynamical systems theory. The authors model the recurrent transformation as a discrete-time dynamical system $h^{(t+1)} = \Phi_\theta(h^{(t)})$. According to the Lyapunov Linearization Theorem, the local stability of a fixed point is determined by the spectral radius ($\rho$) of the system's Jacobian matrix $J$.

• If $\rho(J) < 1$, the fixed point is asymptotically stable (perturbations decay).
• If $\rho(J) \geq 1$, the system may oscillate or diverge.

2.2 Jacobian Spectral Radius Regularization (JSRR)

STARS introduces a regularization term that constrains the spectral radius of the Jacobian matrix during training to ensure $\rho(J) < 1$.

• Efficiency Challenge: Direct eigenvalue computation for large-scale models is computationally prohibitive.
• Solution: The authors employ a lightweight estimation scheme using single-step power iteration combined with Jacobian-vector products (JVP). This avoids constructing the full Jacobian matrix while providing an efficient estimate of the spectral radius ($\rho(J) \approx \|Jv\|_2$).
• Implementation: The regularization is applied as a proxy at specific iterations $t$ during the loop, rather than waiting for the exact fixed point, which is computationally infeasible to locate during training.

2.3 Integration with Random Loop Sampling

To ensure global stability across the inference trajectory rather than just at a single point, STARS integrates JSRR with random loop sampling. During training, the number of recurrent iterations $t$ is sampled from a distribution (e.g., Log-Normal). The final training objective is a weighted combination of the standard supervised fine-tuning (SFT) loss and the JSRR loss:
$$\mathcal{L}_{STARS} = \mathbb{E}_{t \sim P} \left[ (1 - \lambda) \cdot \mathcal{L}_{SFT}^{(t)} + \lambda \cdot \mathcal{L}_{JSRR}^{(t)} \right]$$
This encourages the model to maintain a small spectral radius across the entire support set of latent states encountered during iterative inference.

3. Key Contributions

1. Diagnostic Analysis: The paper provides a systematic diagnostic study of LoopLM latent trajectories, identifying the irreconcilable trade-off between stability and effectiveness caused by normalization placement (Internal vs. External).
2. STARS Framework: The proposal of a training framework that enforces asymptotic stability via Jacobian Spectral Radius Regularization, mathematically compelling the model to converge to stable fixed points.
3. Efficient Regularization: The development of an efficient, memory-aware spectral radius estimator using single-step power iteration and JVP, making the approach scalable for large language models.
4. Empirical Validation: Demonstration that STARS achieves reliable test-time scaling on both synthetic arithmetic tasks and complex mathematical reasoning benchmarks, significantly mitigating performance degradation at deeper recurrence levels compared to baselines.

4. Experimental Results

The authors evaluated STARS on two primary tasks:

4.1 Multi-digit Addition (Synthetic Task)

• Setup: Randomly initialized Transformers trained on 4-digit addition.
• Result: Models trained with STARS achieved 100% accuracy consistently across all recurrent steps. The latent states converged to stable fixed points with fast convergence rates, demonstrating fully reliable test-time scaling. In contrast, baselines exhibited performance collapse as iterations increased.

4.2 Mathematical Reasoning (Real-world Task)

• Setup: Fine-tuning the pre-trained Ouro-1.4B model on the NuminaMath-1.5 dataset.
• Baselines: Standard Ouro-1.4B and Ouro-1.4B with standard SFT.
• Results:
  • Peak Performance: STARS achieved a peak average accuracy of 74.18% on GSM8K (at 4 recurrent steps), outperforming the SFT baseline (70.46%).
  • Scaling Stability: When scaling to 8 recurrent steps, the SFT baseline's performance collapsed to 52.97% (a 20.47% drop from its peak). In contrast, STARS degraded by only 8.26%, maintaining an average accuracy of 65.55%.
  • Generalization: STARS showed superior robustness across other benchmarks (MATH500, ASDiv, SVAMP, AMC23), consistently outperforming the baselines at deeper recurrence depths.

5. Significance and Claims

The paper claims that STARS resolves the fundamental instability in current LoopLM architectures, enabling reliable test-time scalable latent reasoning. By conceptualizing reasoning as an iterative process of uncertainty reduction that must converge to stable fixed points, the authors demonstrate that it is possible to simultaneously achieve high effectiveness and rigorous stability.

The significance lies in providing a mathematically grounded training framework that allows models to leverage increased computational resources (via deeper recurrence) without suffering from the performance collapse observed in prior LoopLMs. The authors position this as a methodological advancement for improving the stability and scaling behavior of recurrent language model architectures, specifically for reasoning tasks, without introducing new applications or datasets.