PaperlayerView exact revision on arXiv

Source-linked AI summary

LoopUS: Recasting Pretrained LLMs into Looped Latent Refinement Models

Taekhyun Park, Yongjae Lee, Dohee Kim, Hyerim Bae

arXiv:2605.11011v1cs.LGcs.AI

TL;DR

LoopUS addresses the cost and instability of converting pretrained LLMs into looped models for test-time latent refinement. Its post-training framework stabilizes recursive computation and improves reasoning-oriented performance without longer generated traces or recurrent training from scratch.

  • Problem

    Existing looped LLM approaches lack principled layer selection and can require expensive recurrent training or risk hidden-state drift when adapting pretrained models.

  • Method

    LoopUS decomposes a pretrained LLM into encoder, reasoning, and decoder blocks, then stabilizes the reused reasoning block with selective gating, random deep supervision, and adaptive early exiting.

  • Results

    3.0% higher zero-shot accuracy over pretrained backbones, alongside 17.4% and 21.3% lower WikiText and LAMBADA perplexities, respectively.

  • Takeaways & Limitations

    LoopUS provides a stable and efficient post-training route for improving pretrained LLMs through controlled latent refinement across loop depths and training budgets.

  • Takeaways & Limitations

    The study does not test dedicated mathematical reasoning or longer-context, sustained multi-step reasoning because training used 1024-token contexts without math-focused adaptation data.

Abstract

from arXiv · show

Looped computation shows promise in improving the reasoning-oriented performance of LLMs by scaling test-time compute. However, existing approaches typically require either training recurrent models from scratch or applying disruptive retrofits, which involve substantial computational costs and may compromise pretrained capabilities. To address these limitations, we introduce \textbf{Looped Depth Up-Scaling} (LoopUS), a post-training framework that converts a standard pretrained LLM into a looped architecture. As a key technical contribution, LoopUS recasts the pretrained LLM into an encoder, a looped reasoning block, and a decoder. It operationalizes this latent-refinement architecture through four core components: (1) block decomposition, guided by staged representation dynamics; (2) an input-dependent selective gate to mitigate hidden-state drift; (3) random deep supervision for memory-efficient learning over long recursive horizons; and (4) a confidence head for adaptive early exiting. Collectively, these mechanisms transform a standard non-looped model into a looped form while stabilizing it against both computational bottlenecks and representation collapse. Through stable latent looping, LoopUS improves reasoning-oriented performance without extending the generated traces or requiring recurrent training from scratch. For more details, see https://thrillcrazyer.github.io/LoopUS

1 Introduction

LoopUS is a post-training framework that converts pretrained LLMs into looped latent-reasoning models by using representation-guided decomposition and mechanisms for stable, efficient recursion. It improves reasoning-oriented performance while avoiding recurrent training from scratch and extended generated traces.

  • Motivation: Test-time latent computation refines hidden states before token generation, scaling reasoning capacity without increasing model parameters or sequence length.Latent-space computation provides a complementary scaling axis through deeper internal processing.
  • Limitations: Existing pretrained-to-loop methods lack principled layer selection, suffer hidden-state drift under reuse, and require memory-intensive, unstable backpropagation through long loops.Their layers were trained for single-pass use, so repeated iteration can degrade representational fidelity.
  • Framework: LoopUS analyzes staged representation dynamics and partitions pretrained LLMs into encoder, reasoning, and decoder blocks, reusing only the middle reasoning block.Early layers rapidly transform embeddings, middle layers evolve gradually on a stable plateau, and final layers transition sharply toward decoding.
  • Mechanisms: LoopUS stabilizes and streamlines latent recursion with a selective decay gate, random deep supervision, and a confidence head for adaptive early exiting.The gate mitigates hidden-state drift, while random deep supervision avoids full backpropagation through time over long recursive horizons.
  • Empirical results: 3.0% zero-shot accuracy improvement over pretrained backbones accompanies 17.4% and 21.3% perplexity reductions on WikiText and LAMBADA, respectively.LoopUS also yields a 14.6% relative gain on TinyLlama with 17–20× fewer training tokens than existing looped baselines.

2 Background

The background motivates LoopUS through three ideas: hidden states evolve toward abstract predictive representations, looped transformers scale computation by repeated latent refinement, and gating regulates state retention and updates.

  • LLM Hidden State Representations: LLM hidden states progressively shift from local lexical processing toward abstract, prediction-oriented representations that can be refined across depth.Interpretability and representation-evolution studies support this staged view of hidden-state dynamics.
  • Looped LLMs: Looped transformers scale test-time computation by repeatedly applying the same block to refine latent representations without increasing parameter count.This paradigm complements sequence-length scaling and includes recurrent, retrofitted, latent-refinement, and adaptive-recursion formulations.
  • Deep Learning Gating Mechanisms: Gating mechanisms regulate state updates, with softmax-style gates normalizing alternatives and decay-style gates directly controlling state retention.Recent sequence models increasingly use input-dependent selective decay, including Mamba-style mechanisms.

3 Looped Depth Up-Scaling (LoopUS)

LoopUS converts a pretrained LLM into an encoder, looped reasoning block, and decoder, using selective gating to stabilize latent refinement. It trains and runs this loop with random deep supervision, confidence-based early stopping, and losses that promote predictive, non-regressive updates.

  • Architecture: LoopUS partitions a pretrained LLM into encoder E, reasoning block M, and decoder D, with boundaries selected from abrupt changes in cosine similarity across depth.The encoder produces the initial representation, the reasoning block is applied for B iterations, and the decoder maps the final state to vocabulary logits.
  • Selective Gating for Stable Loop Dynamics: The selective gate interpolates each proposed reasoning update with the previous hidden state, damping latent displacement and mitigating drift from repeatedly applying a single-pass pretrained block.A learned channel-wise decay produces α(b) ∈ (0, 1), limiting each iteration to a relaxed, contraction-like refinement step without guaranteeing a strict global contraction.
  • Adaptive Computation via Early Stopping Mechanism: A confidence head enables adaptive computation by halting refinement when q(b) reaches the predefined threshold qth, allocating additional loop steps only to insufficiently confident states.This dynamically converts pretrained transformer depth into adaptive test-time computation.
  • Random Deep Supervision for Loop Training: Random deep supervision unrolls B steps but computes gradients only for a uniformly sampled subset S of K steps, reducing memory coupling across the recursive horizon.At sampled depths, LoopUS jointly optimizes next-token prediction, monotonicity, and confidence losses; these make steps predictive, discourage regressions, and train early stopping.

4 Empirical Validation

LoopUS is empirically validated across pretrained backbones, showing consistent reasoning gains, efficient adaptation, stable test-time latent refinement, and improved behavior from its architectural components. These results indicate that LoopUS scales latent computation while preserving stable refinement dynamics.

  • Benchmark Results: Across five pretrained backbones, LoopUS reduces WikiText and LAMBADA perplexities and raises average downstream accuracy by +1.6 to +2.2 points.The evaluation spans Qwen3-1.7B, Qwen3-4B, Qwen3-8B, TinyLlama, and Phi-4.
  • Benchmark Results: LoopUS gains are task-dependent, with more consistent improvements on ARC-C, PIQA, WG, and OBQA than on MMLU and HS across model scales.The pattern suggests stronger benefits when additional latent computation can refine decisions rather than primarily retrieve broad knowledge.
  • Adaptation Efficiency: LoopUS achieves the largest average gain, ∆= + 6.3, versus ∆= −2.9 for McLeish et al. and ∆= + 3.5 for Bae et al. [11], with fewer additional training tokens.Because prior results come from corresponding papers, this is an adaptation-efficiency reference rather than a fully controlled head-to-head benchmark.
  • Test-Time Scaling: Most test-time scaling benefits appear within a few latent iterations, after which performance shows diminishing returns while remaining stable at unseen depths of 40, 80, and 100.Confidence-based stopping adaptively halts the same checkpoint after 3.39 iterations.
  • Refinement Dynamics: On “32 * 64 =”, the correct next token “2” increases from 2.17 × 10−5% at iteration 0 to 81.9% after one refinement and about 89.8% by iteration 4.The latent trajectory makes its largest move early, then contracts as token probabilities sharpen and competing candidates lose mass.
  • Ablation Study: Removing the selective gate or encoder–decoder decomposition causes convergence to a higher LLM by weakening drift control or failing to preserve the pretrained latent workspace.The decomposition separates representation extraction, latent refinement, and output decoding; without it, the loop relearns a less stable recurrent trajectory.

5 Conclusion

LoopUS is a post-training framework that converts pretrained LLMs into looped latent-refinement models while preserving standard interfaces. It improves reasoning performance across model scales and consistently reduces perplexity.

  • Framework: LoopUS recasts pretrained LLMs into looped latent-refinement models through encoder–reasoning–decoder decomposition.The framework is designed as a post-training conversion rather than a from-scratch recurrent model.
  • Core components: The framework combines a selective gate, random deep supervision with stepwise detachment, and a lightweight confidence head for adaptive stopping.These components support latent looping, memory-efficient supervision, and adaptive computation.
  • Outcomes: Across diverse model scales, LoopUS improves pretrained backbones while preserving standard interfaces, enhancing reasoning performance and consistently reducing perplexity.The reported improvements apply across multiple pretrained model sizes without changing the standard model interface.

A Experimental Details · A.1 Backbones and Training Data · A.2 Optimization Details

The experiments evaluate LoopUS across five pretrained backbones using a 3B-token FineWeb-Edu training setup. Optimization combines distributed, memory-efficient training with specified loop-specific supervision and evaluation configurations.

  • A.1 Backbones and Training Data: Experiments use Qwen3-1.7B, Qwen3-4B, Qwen3-8B, TinyLlama, and Phi-4 backbones.Main experiments stream 3B tokens from FineWeb-Edu CC-MAIN-2025-26 with sequence length 1024.
  • A.1 Backbones and Training Data: The layer assignments vary by backbone: Qwen3-1.7B uses encoder layers 0–1 and decoder layer 27, while Qwen3-4B uses 0–1 and 35.Qwen3-8B and Phi-4 use encoder layers 0–5 with decoder layers 35 and 39, respectively; TinyLlama uses encoder layer 0 and decoder layer 21.
  • A.1 Backbones and Training Data: LoopUS selects shallow encoder and late decoder layers while reusing the intermediate transformer layers as the looped latent workspace.The architecture separates a shallow encoder from a late decoder and repurposes the entire middle block.
  • A.2 Optimization Details: Training uses Accelerate, distributed data parallelism, and gradient checkpointing throughout.The reference recipe uses AdamW, 5 × 10−5 learning rate, one epoch, bf16 mixed precision, FlashAttention-2, cosine scheduling, and 300 warmup steps.
  • A.2 Optimization Details: The reference script uses 8 dataloader workers, pinned-memory loading, logging every 50 steps, checkpointing every 5000 steps, and retains at most 3 checkpoints.These settings are part of the stated optimization and run-management configuration.
  • A.2 Optimization Details: Loop-specific training uses 20 reasoning steps, deep supervision on 5 loop positions per example, a 0.55 stopping threshold, and the all stopping mode.These parameters define the loop configuration in the training setup.
  • A.2 Optimization Details: The released script optionally runs checkpoint-time lm-evaluation-harness evaluation on WikiText with a limit of 200 samples.This evaluation option is enabled in the released script.

A.3 Evaluation Details

The evaluation uses five pretrained backbones across standardized training and inference configurations. It reports zero-shot perplexity and accuracy across language-modeling and reasoning benchmarks.

  • Backbones and Training: Five pretrained backbones are evaluated: Qwen3-1.7B, Qwen3-4B, Qwen3-8B, TinyLlama, and Phi-4.Training runs use NVIDIA L40S, RTX PRO 6000, or H200 GPUs depending on the backbone.
  • Backbones and Training: Unless otherwise stated, training uses FineWeb-Edu for 3B tokens with context length 1024, AdamW, cosine learning-rate decay, and bf16 mixed precision.The passage also specifies a default LoopUS setting beginning with B=, but the supplied text is truncated.
  • Evaluation Protocol: Zero-shot evaluation uses lm-evaluation-harness, reporting perplexity on WikiText and Lambada and accuracy on seven additional benchmarks.The accuracy benchmarks are MMLU, HellaSwag, ARC-Easy, ARC-Challenge, PIQA, WinoGrande, and OpenBookQA.
  • Evaluation Protocol: Standard inference uses a maximum recursion budget of 8, a stopping threshold of 0.6, and KV caching for autoregressive decoding.The main experiments are organized to test three claims, but the supplied passage truncates the first claim after language-mode.

A.4 KV-Cache Implementation for Autoregressive Inference

LoopUS implements autoregressive inference with separate encoder, per-refinement-depth reasoning, and decoder KV caches, enabling incremental decoding instead of recomputing the full prefix. With B = 8, this caching produces substantial speedups across Qwen3 model sizes, especially for longer generations.

  • Cache organization: LoopUS stores separate DynamicCaches for the encoder, each reasoning refinement depth, and the decoder, alongside a seen_tokens counter.Each refinement depth maintains its own KV cache because hidden representations evolve across recursive refinement steps.
  • Two-phase inference: During prefill, LoopUS processes the prompt once to populate all caches; during decoding, it updates cached state for each new token without recomputing the full prefix.Absolute positions supplied through cache_position keep rotary position IDs and causal masks aligned with the full prefix.
  • Cache organization: Each refinement depth requires its own cache because sharing one cache would mix keys and values from different latent trajectories and break the actual recurrence.The block parameters are shared, but the hidden state entering each depth differs after every refinement step.
  • Decoding speed: At 1024 generated tokens with B = 8, KV caching yields 1.64×, 2.31×, and 2.49× speedups for Qwen3-1.7B, Qwen3-4B, and Qwen3-8B, respectively.The benchmark measures token-generation throughput over five repeated runs on an NVIDIA L40S GPU.

A.5 Halting and Recursion Diagnostics · B Dynamical and Geometric Interpretation of LoopUS · C Additional Qualitative Example

LoopUS is evaluated through halting and recursion-depth diagnostics, interpreted rigorously as a gated latent fixed-point iteration, and illustrated qualitatively as smooth, progressively contracting refinement trajectories. The analysis distinguishes these supported fixed-point and surrogate-energy views from stronger diffusion or explicit-energy-based interpretations.

  • A.5 Halting and Recursion Diagnostics: Figure 18 compares threshold, hidden-state-convergence, and CDF-based halting using a fixed Qwen3-1.7B LoopUS checkpoint.The evaluation uses the task set {MMLU, HellaSwag, ARC-Easy, PIQA, WinoGrande, Lambada, WikiText}, maximum recursion budget B = 10, batch size 16, maximum length 1024, seed 2026, and runtime-stat logging.
  • A.5 Halting and Recursion Diagnostics: A separate Qwen3-4B study varies inference depth to identify performance gains before saturation or over-refinement trade-offs.The remaining evaluation settings are held fixed, with results reported in Figure 5 in Section 4.4.
  • B Dynamical and Geometric Interpretation of LoopUS: LoopUS reuses a near-fixed-point latent workspace and converts it into an iterative refinement process represented as a diagonally preconditioned relaxed fixed-point iteration.With recursion-depth index b and P^(b) = Diag(α^(b)), the rigorous interpretation is latent-state search for h⋆ satisfying h⋆ = M(h⋆).
  • B Dynamical and Geometric Interpretation of LoopUS: Under a scalar-potential assumption, the LoopUS update becomes diagonally preconditioned gradient descent, with sufficiently small gated steps decreasing the potential.If Φx is L-smooth and ∥P^(b)∥2 ≤ 1/L, fixed points of M coincide with stationary points of Φx.
  • B Dynamical and Geometric Interpretation of LoopUS: LoopUS explicitly trains decoder-induced surrogate energy to be approximately non-increasing, while the confidence head halts refinement when predicted additional benefit becomes small.The monotonicity loss penalizes positive increments in Ex and mildly rewards negative increments; halting trends and smooth contraction-like trajectories support this interpretation.
  • B Dynamical and Geometric Interpretation of LoopUS: Monotonicity loss, stepwise detachment, and selective gating favor small, progressively improving latent edits while avoiding globally entangled backpropagation through all loop positions.Each supervised depth is optimized primarily as a local correction to the current latent state.
  • B Dynamical and Geometric Interpretation of LoopUS: The diffusion and explicit-EBM analogies are limited: LoopUS injects no noise, learns no reverse diffusion process or schedule, uses no score-matching objective, and does not optimize a standalone output-space scalar energy.The precise characterization supplied is an implicit latent optimizer with an amortized refinement process.
  • C Additional Qualitative Example: Qualitative visualizations show large early representational corrections followed by smaller stabilizing adjustments and smooth, low-dimensional PCA trajectories with progressive contraction across LoopUS backbones.Figure 12 presents a representative generation trajectory, while Figures 13–15 visualize trajectories for multiple backbones.

D Hidden-State Trajectory Analysis · E Halting Strategies

LoopUS exhibits a recurring three-phase hidden-state trajectory across multiple pretrained backbones, supporting the generality of its staged representation dynamics. Its halting strategies trade off training alignment, latent-trajectory geometry, and cumulative confidence evidence under synchronized batched inference.

  • D Hidden-State Trajectory Analysis: Across six additional backbones, hidden-state PCA trajectories and layer-to-layer distance profiles reproduce the qualitative structure observed for Qwen3-1.7B.Figures 16 and 17 show that the analyzed trajectory pattern is not unique to Qwen3-1.7B.
  • D Hidden-State Trajectory Analysis: The recurring pattern has three phases: rapid early-layer motion, a smooth middle-layer regime, and a sharper transition near the final layers.PCA plots show a middle-layer arc or plateau followed by a pronounced turn, while distance profiles show reduced consecutive-layer change in the middle block.
  • E Halting Strategies: LoopUS compares threshold-based, convergence-based, and CDF-based inference-time halting rules while keeping the checkpoint fixed.Each strategy selects an exit depth under a maximum recursion budget B.
  • E Halting Strategies: The three rules encode distinct inductive biases: training alignment for threshold halting, latent geometry for convergence halting, and more aggressive evidence aggregation for CDF halting.Threshold is most tightly matched to the training objective, convergence diagnoses contractive refinement behavior, and CDF aggregates multiple confidence values.
  • E Halting Strategies: Threshold halting stops when the confidence head exceeds a fixed threshold, directly matching its supervision to post-update token accuracy.The rule treats confidence as evidence that the current latent state is sufficiently predictive and applies conservative batch-wide stopping.
  • E Halting Strategies: Convergence halting exits when consecutive last-token hidden states differ by less than a preset tolerance ϵ, without using the confidence head.It assumes marginal latent-state change indicates insufficient expected decoder improvement and synchronizes stopping across the batch.
  • E Halting Strategies: CDF halting reinterprets stepwise confidence as a hazard rate and accumulates exit evidence across refinement steps.Because LoopUS trains confidence with a binary post-update accuracy target rather than Ouro’s full exit-distribution objective, this is a comparative heuristic.

F Limitations and Future Work

The study identifies open questions about LoopUS across modalities, architectures, training scales, reasoning domains, post-training methods, and model-development stages. Future work includes evaluating multimodal and hybrid models, broader reasoning coverage, integration with instruction and preference tuning, diffusion LLMs, and pretraining-native designs.

  • Extension beyond text-only language modeling: LoopUS has only been evaluated on text-only, decoder-only language models, leaving the stability of its encoder–reasoning–decoder decomposition for multimodal models untested.The method is formulated as latent-space refinement and is not tied to a particular input modality in principle.
  • Heterogeneous and hybrid model architectures: Modern heterogeneous architectures may not provide a uniform reusable middle-layer block, leaving the optimal recursion policy for LoopUS uncertain.Relevant architectures include gated delta networks, sparse attention, mixture-of-experts routing, state-space modules, and other hybrid sequence mechanisms.
  • Scaling to larger and more diverse training regimes: LoopUS remains untested under substantially larger corpora, longer contexts, curriculum schedules, and more diverse data mixtures beyond the moderate post-training budget.Scaling the training recipe is especially relevant because the confidence head and monotonicity objective may be affected by broader training regimes.
  • Dedicated math and long-context reasoning coverage: The study lacks dedicated mathematical reasoning tasks, math-focused training corpora, and contexts beyond the released 1024-token setup.Compute constraints prevented additional math-oriented datasets during adaptation, so the paper does not yet test dedicated math and long-context reasoning coverage.
  • Integration with instruction tuning and preference optimization: Future work should integrate LoopUS with instruction tuning, long-context adaptation, and preference optimization methods such as reinforcement learning from feedback or GRPO-style objectives.The current evaluation treats LoopUS as post-training adaptation for base models and does not assess these practical development stages.
  • Future architectural directions: Future directions include applying LoopUS to diffusion LLMs and incorporating looped latent refinement natively during pretraining rather than retrofitting pretrained backbones.Diffusion LLMs use iterative continuous-latent denoising, while pretraining-native LoopUS could expose recurrence to broader intermediate-state distributions.
Loading 2605.11011v1…