PaperlayerView exact revision on arXiv

Source-linked AI summary

LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation

Ahmadreza Jeddi, Marco Ciccone, Babak Taati

arXiv:2602.11451v1cs.CL

TL;DR

Looped Transformers typically fix their iteration count, leaving uncertain whether they can adapt reasoning quality to user-selected compute budgets. LoopFormer trains time- and step-conditioned trajectories with shortcut consistency, and reports strong language-modeling and zero-shot reasoning performance that scales gracefully with compute.

  • Problem

    Prior looped Transformers fix the number of iterations during training and inference, leaving flexible performance under variable compute budgets unresolved.

  • Method

    LoopFormer conditions each loop on normalized time and step size and trains variable-length trajectories with a shortcut-consistency objective aligning shorter routes to the longest route.

  • Results

    LoopFormer maintains strong language-modeling and zero-shot reasoning performance across inference budgets, with performance scaling smoothly rather than collapsing at shorter depths.

  • Takeaways & Limitations

    Loop trajectories provide a framework for budget-conditioned inference in which shorter routes remain useful and additional steps refine representations toward a shared endpoint.

Abstract

from arXiv · show

Looped Transformers have emerged as an efficient and powerful class of models for reasoning in the language domain. Recent studies show that these models achieve strong performance on algorithmic and reasoning tasks, suggesting that looped architectures possess an inductive bias toward latent reasoning. However, prior approaches fix the number of loop iterations during training and inference, leaving open the question of whether these models can flexibly adapt their computational depth under variable compute budgets. We introduce LoopFormer, a looped Transformer trained on variable-length trajectories to enable budget-conditioned reasoning. Our core contribution is a shortcut-consistency training scheme that aligns trajectories of different lengths, ensuring that shorter loops yield informative representations while longer loops continue to refine them. LoopFormer conditions each loop on the current time and step size, enabling representations to evolve consistently across trajectories of varying length rather than drifting or stagnating. Empirically, LoopFormer demonstrates robust performance on language modeling and reasoning benchmarks even under aggressive compute constraints, while scaling gracefully with additional budget. These results show that looped Transformers are inherently suited for adaptive language modeling, opening a path toward controllable and budget-aware large language models.

1 INTRODUCTION

LoopFormer addresses depth inflexibility in looped Transformers by training shortcut-consistent trajectories conditioned on time and step size. It supports budget-conditioned inference, with evolving representations and graceful performance scaling across compute budgets.

  • Motivation: Prior looped models use a single trajectory length, causing representations to collapse at shorter or longer depths and limiting flexible compute.LoopFormer instead targets a user-specified budget without retraining.
  • Method: LoopFormer conditions each loop on normalized time and step size, allowing coarse trajectories to approximate fine-grained refinement.The model evolves hidden states toward a shared t=1 endpoint.
  • Method: Shortcut-consistency training aligns shorter-route representations or logits with the longest route through a stop-gradient target.The protocol trains across families of step schedules and uses self-distillation within the loop.
  • Results: Adaptive early-exit looped models exhibit representational collapse, whereas shortcut-modulated models maintain evolving, non-degenerate states.The comparison uses anisotropy, curvature, entropy, and CKA.
  • Results: LoopFormer delivers consistent performance-per-compute gains in perplexity and zero-shot reasoning across diverse language benchmarks.Perplexity and reasoning improve smoothly with additional compute, while the model remains effective under constrained budgets.

2 RELATED WORK

Related work positions LoopFormer at the intersection of parameter-shared Transformers, dynamic computation, latent reasoning, and time-conditioned trajectory consistency. Its distinction is conditioning looped language models on both normalized time and step size while training across trajectory families.

  • Recursive / Looped Transformers: Parameter sharing offers an efficiency route to effective depth, with Universal Transformer-style recurrence and ALBERT-style weight tying reducing reliance on non-shared stacks.These approaches establish recurrent computation and parameter savings as complementary to conventional depth.
  • Recursive / Looped Transformers: Prior language-modeling work found that timestep conditioning improves looped-Transformer scaling and perplexity.LoopFormer extends this direction beyond the loop index.
  • Dynamic compute: Dynamic-compute methods reduce inference cost through early exiting, layer dropping, pruning, or sparse expert routing.These methods generally allocate computation by skipping or reallocating processing.
  • Latent reasoning: Latent-reasoning research studies reasoning in hidden states rather than through explicit chain-of-thought prompting, including looped approaches for multi-step reasoning.This literature motivates treating iterative hidden-state refinement as reasoning computation.
  • Time / shortcut modulation and conditioning: Diffusion and consistency-based methods condition transformations on time and align coarse trajectories with fine-grained solutions.LoopFormer draws on this trajectory-conditioning and shortcut-consistency framework for looped language modeling.

3 SHORTCUT-MODULATED LOOPED MODELS FOR LATENT REASONING

LoopFormer enables elastic-depth looped language modeling by conditioning shared Transformer iterations on trajectory time and step size. Shortcut trajectories and consistency training let users choose inference budgets while shorter routes approximate the full trajectory and longer routes refine outputs.

  • Architecture: LoopFormer repeatedly applies a shared Transformer stack while conditioning each loop on cumulative normalized time and step size.These signals modulate RMSNorm scales and gate MHSA and FFN residual connections, making each loop aware of its trajectory position and granularity.
  • Training: Shortcut conditioning and consistency training align trajectories of different lengths toward the full L-step trajectory.The consistency loss aligns per-token logits from shorter trajectories with those of the longest trajectory, alongside next-token losses for long and shortcut routes.
  • Training: Training samples shortcut lengths and schedules in addition to the maximum trajectory, exposing the model to both short and long routes.Shortcut lengths are sampled uniformly from 1 through L − 1, with schedules drawn so their step sizes sum to one.
  • Inference with elastic depth: The model supports user-chosen budgets M ≤ L by supplying a step schedule whose increments sum to one.Inference can therefore use any budget up to the maximum trained depth without retraining.
  • Inference with elastic depth: Shorter-budget routes approximate the full trajectory, while increasing the budget progressively refines next-token distributions without retraining.The design targets useful outputs at low compute and improved quality as additional steps are used.

4 EXPERIMENTS

Experiments evaluate LoopFormer across language modeling, zero-shot reasoning, inference budgets, and representation dynamics. Results show smooth scaling with depth, stronger performance than looped baselines, and less representational stagnation.

  • Latent reasoning and perplexity: LoopFormer narrows the perplexity gap to non-looped models and is competitive on reasoning, outperforming other looped variants, especially at higher budgets.Table 1 compares the (3⊗8) model at 24×, 12×, and 6× inference budgets.
  • Experimental setup: Evaluations report perplexity on three corpora and zero-shot accuracy across ten reasoning benchmarks, comparing non-looped, fixed-depth looped, and depth-elastic baselines.The experiments use FLOP-matched looped variants and evaluate budgets M ≤ L.
  • Scaling with layers and loops: Larger blocks and more iterations improve both perplexity and reasoning, while shorter trajectories remain informative and added loops provide smooth gains without collapse.These trends hold across inference budgets M ≤ L.
  • Representation dynamics: LoopFormer’s representations continue evolving across depth, unlike early-exit baselines whose flat metrics and high CKA indicate stagnation.Curvature, anisotropy, entropy, and cross-step CKA jointly characterize these dynamics.

5 DISCUSSION AND FUTURE DIRECTIONS

LoopFormer represents computation as hidden-space trajectories whose intermediate states remain useful and whose additional steps refine a shared endpoint. The approach supports budget-conditioned inference while avoiding the stagnation observed in naive early exits.

  • Discussion and future directions: LoopFormer conditions iterations on normalized time and step size and trains across trajectories with a shortcut-consistency loss.This creates trajectories with shared endpoint behavior across different lengths.
  • Discussion and future directions: Shorter routes provide useful intermediate states, while additional steps refine them toward a shared t=1 endpoint without retraining.The paper frames this as latent reasoning with elastic inference.
  • Discussion and future directions: The paper identifies global sequence-level budgeting, added multi-trajectory training overhead, and correlational rather than causal representation analyses as limitations.Future directions include instance-conditioned schedules and deeper representation-space diagnostics.

A TRAINING AND IMPLEMENTATION DETAILS

Implementation uses NanoGPT-style training on H100 GPUs with fixed data and optimization settings. LoopFormer adds time- and step-size conditioning through Fourier-feature embeddings and an AdaLN-style modulator.

  • Training setup: Models train on 4×H100 GPUs for 50,000 optimizer steps with batch size 48, context length 1024, and approximately 25B tokens.The open-source NanoGPT stack provides the reference implementation.
  • Optimization: Optimization uses AdamW, weight decay 2 × 10^-1, peak learning rate 6 × 10^-4, minimum learning rate 6 × 10^-5, cosine decay, and 4,000 warmup steps.Warmup and cosine decay mitigate occasional depth-related training instabilities.
  • Model hyperparameters: All configurations use hidden size d = 2048, 32 heads, feed-forward dimension 5120, RMSNorm, learned positional embeddings, and tied input-output embeddings.RoPE was tested but learned positions were retained for simplicity and efficiency.
  • LoopFormer conditioning: Separate embedding modules map normalized time t and step size ∆t into conditioning vectors using Fourier features followed by two-layer MLPs.The conditioning vectors are summed to form the per-iteration signal c.
  • LoopFormer conditioning: An AdaLN-style modulator converts the conditioning signal into scaling and residual-gating parameters applied inside each loop block.Zero initialization preserves the unmodulated backbone at initialization and supports stable learning.

B.1 COMPUTATIONAL OVERHEAD OF LOOPFORMER

LoopFormer’s dual-trajectory training adds computational overhead, but inference retains baseline-like FLOPs while supporting variable loop budgets. The expected training cost is approximately 1.5× fixed-loop training, with an observed wall-clock slowdown of about 1.3×.

  • LoopFormer computes a full L-loop trajectory and a sampled shorter M-loop trajectory for language-modeling and shortcut-consistency objectives.The shorter trajectory samples M uniformly from {1, ..., L −1}.
  • C(L) = Cio + LC1 decomposes fixed-loop cost into loop-independent overhead and the cost of L shared-stack loops.Cio covers embedding and unembedding, while C1 is the cost of one loop through the shared K-block stack.
  • 1.5× is the approximate training-FLOP overhead of LoopFormer relative to fixed-loop or vanilla training in the considered regimes.The additional cost comes from computing the sampled shortcut trajectory.
  • 1.5× training FLOPs and 1.3× wall-clock time were observed against fixed-loop or vanilla baselines under the same token budget and iteration count.The wall-clock measurement used 4×H100 GPUs with identical batch size, optimizer, and data.
  • Inference has computational FLOPs similar to vanilla or fixed-loop baselines while remaining robust under truncation and refining representations with more loops.The same shared-parameter model supports elastic-depth inference.

B.2 COMPARISON WITH VANILLA TRANSFORMERS UNDER A SIMILAR PARAMETER BUDGET

Under a similar parameter budget, LoopFormer matches a vanilla Transformer at low inference compute and improves consistently as more loops are allocated. At higher budgets, it substantially outperforms the vanilla baseline while retaining depth elasticity.

  • LoopFormer matches the 3-layer vanilla Transformer at low inference budgets and improves consistently across (3 ⊗2), (3 ⊗4), and (3 ⊗8).The comparison reports perplexity and downstream language-task accuracy.
  • At higher loop budgets, LoopFormer substantially outperforms the vanilla Transformer under a similar parameter budget.The comparison uses a shared stack with K = 3 layers.
  • LoopFormer retains depth elasticity while its performance scales with additional inference compute.The table compares three selectable inference budgets against one 3-layer vanilla baseline.

B.3 COMPARISON UNDER FLOPS-MATCHED TRAINING

Under FLOPs-matched training, LoopFormer uses fewer training iterations and approximately 8B fewer Pile tokens than the baselines. Despite this reduced token budget, it matches TMLT and remains competitive on reasoning across budgets.

  • LoopFormer is trained for 34k iterations versus 50k for the baselines to match total training FLOPs.All other hyperparameters remain unchanged.
  • ∼8B fewer Pile tokens are seen by the compute-matched LoopFormer than by the 50k-iteration baselines.
  • Despite the reduced token budget, LoopFormer matches TMLT and remains competitive on reasoning across budgets.The model retains depth-elastic behavior in the FLOPs-matched comparison.

B.4 PYTORCH PSEUDOCODE FOR LOOPFORMER

LoopFormer’s shared stack applies loop-conditioned Transformer blocks repeatedly, using time and step-size information to modulate attention and MLP residual branches. The same K-block stack is reused for trajectories of different lengths.

  • Each LoopFormerBlock uses AdaLN-style modulation conditioned on a vector derived from time t and step size ∆t.The modulation produces gates and scales for the attention and MLP residual branches.
  • A SharedBlock stacks K distinct Transformer blocks and applies them sequentially within each loop iteration.
  • The shared K-block stack is reused across loops and repeated for M or L iterations during training or inference.This reuse supports variable compute budgets.
  • The implementation accepts token states x, a loop-conditioning vector c, and a configuration defining the shared stack depth.
  • AdaLN-Zero initialization starts the modulation from identity updates.
  • The block updates token states through gated, scaled attention and feed-forward residual branches.The implementation adds each modulated branch back to x.
Loading 2602.11451v1…