LoopFormer — Elastic-Depth Looped Transformers via Shortcut Modulation

Paper. Ahmadreza Jeddi, Marco Ciccone, Babak Taati. LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation. ICLR 2026. University of Toronto, Vector Institute, University Health Network. [arXiv]

0. The Picture in One Paragraph

Looped transformers — a shared block applied M times instead of M distinct layers — have become a popular way to add algorithmic and latent-reasoning inductive bias without adding parameters. But every prior design (Universal Transformers, ACT, recurrent-depth models like Huginn) fixes the loop count at training time, or samples it randomly without ensuring different depths are actually consistent with each other. LoopFormer asks a different question: can a single trained model support elastic depth at inference — pick any budget M ≤ L after training, with quality degrading gracefully rather than collapsing? The mechanism is shortcut modulation: each loop iteration is conditioned on its cumulative normalized time and step size via sine-cosine embeddings feeding a FiLM-style layer that scales norms and gates residual updates. Training adds a shortcut-consistency objective — coarse (shortcut) schedules are distilled to match the representation a fine-grained (full) schedule would have produced — so that fewer, coarser steps at inference still approximate what many fine steps would give. The result is a model that scales smoothly with added inference budget rather than requiring a separately trained model per depth.

1. The Problem — Depth as a Training-Time Hyperparameter

Weight-tied looped transformers repeat one block M times. Prior designs pick M once and bake it in:

Universal Transformers / ACT: per-token adaptive halting, but the halting policy is learned jointly with a specific training-time schedule.
Huginn (recurrent-depth): samples the recurrence depth per training step from a distribution, discarding any explicit halting mechanism, and treats depth as a test-time knob — but does not explicitly train different depths to be consistent with each other.

None of these give a single model a principled, budget-conditioned way to trade inference compute for quality at deployment time, with a guarantee that a shorter schedule is a genuine approximation of a longer one rather than an arbitrary truncation.

2. The LoopFormer Mechanism

2.1 Time and step-size conditioning

Each loop iteration i is given two scalars: the cumulative normalized time t_{i-1} (how far along the 0→1 schedule the model already is) and the step size Δ_i = t_i − t_{i-1} it is about to take. Both are turned into sine-cosine frequency embeddings, passed through small MLPs, and summed into a single conditioning vector e_i — the same style of timestep conditioning used in diffusion models, repurposed for transformer loop depth.

2.2 Shortcut modulation

The conditioning vector e_i drives a small FiLM-style MLP that produces scale parameters (γ1, γ2) for two RMSNorm layers and gates (α1, α2) applied before the attention and FFN residual connections. This lets the same shared block behave differently depending on where in the schedule — and how coarse a step — the current iteration represents.

2.3 Shortcut-consistency training

Within a batch, LoopFormer runs both a full-length route (many fine-grained steps) and a shorter “shortcut” route (a coarser step schedule) through the same weights. The shortcut route’s final representation is trained — via a stop-gradient target from the full route — to match what the full route would have produced. This is a form of self-distillation within the loop: coarse schedules learn to approximate the outcome of fine-grained ones, rather than just being truncated fine-grained runs.

2.4 Elastic-depth inference

At inference, a user picks any budget M ≤ L and any step schedule (e.g., coarser early, finer late — an ablation finding), with no retraining required. The architecture itself is a NanoGPT fork with RMSNorm in place of LayerNorm, kept close to the original for reproducibility.

3. Results

LoopFormer is evaluated across a budget sweep of loop counts K ∈ {1, 3, 6, 9, 12}, with language-modeling perplexity and zero-shot accuracy across roughly ten benchmarks — COPA, HellaSwag, LAMBADA, OpenBookQA, PIQA, RACE, Social IQA, ARC (Easy + Challenge), SciQ, and WinoGrande — compared at matched compute against vanilla non-looped transformers, fixed-depth looped baselines, and adaptive early-exit baselines.

The headline qualitative finding: LoopFormer scales smoothly and monotonically with added inference budget, whereas fixed-depth baselines that are asked to run at an off-schedule depth degrade sharply rather than gracefully. (The paper’s precise numeric perplexity/accuracy tables were not independently verifiable while drafting this review — arXiv access was intermittently blocked — so treat the exact magnitudes as reported by the authors rather than independently re-derived here.)

4. Ablations

Schedule shape matters. The best-performing schedules take coarser steps early in the loop and finer steps late — mirroring a common pattern in diffusion-model noise schedules.
Perplexity-optimal ≠ accuracy-optimal. The schedule that minimizes language-modeling perplexity is not always the one that maximizes downstream reasoning accuracy — a reminder that perplexity is an imperfect proxy for reasoning quality.
Training cost. Running both a full and a shortcut route per batch increases training cost roughly 1.5× relative to training a single fixed-depth schedule.

5. Relationship to Prior Work

Method	Depth policy	Consistency across depths
Universal Transformer / ACT	Per-token learned halting	Implicit via halting loss
Huginn (recurrent depth)	Randomly sampled per training step	Not explicitly enforced
Neural GPUs, DEQ	Fixed-point / equilibrium depth	N/A (single operating point)
LoopFormer	User-chosen budget at inference	Explicit shortcut-consistency distillation

The key differentiator from Huginn — the most directly comparable prior system, and reviewed separately on this site — is that Huginn samples depth randomly during training without an explicit mechanism ensuring a short rollout approximates a long one; LoopFormer makes that approximation an explicit training objective.

6. Why It Matters

Elastic inference-time depth, if it holds up at larger scale, turns “how much do I want to spend on this token/prompt” into a genuine runtime knob rather than a model-selection decision made before training. Combined with token-level selective iteration methods (like Think-at-Hard, reviewed alongside this post), the looped-transformer line is converging on a picture where compute allocation is both elastic at the sequence level (LoopFormer) and selective at the token level (TaH) — two largely orthogonal axes of the same broader “stop paying uniform compute for reasoning” agenda.

7. Limitations Worth Knowing

Global, sequence-level budget. The chosen depth M applies to the whole sequence, not per-token or per-instance — LoopFormer does not currently combine its elastic budget with the kind of per-token selectivity TaH introduces.
~1.5× training overhead from running both full and shortcut routes per batch.
Representation-consistency analysis is correlational, not a causal guarantee that shortcut routes always faithfully approximate full routes on out-of-distribution inputs.
Scale and exact benchmark numbers were not independently re-verified for this review beyond what search-indexed sources reported; readers should consult the primary PDF for exact figures.

8. The Takeaway for a First Reader

If you remember three things:

Prior looped transformers fix loop depth at training time; LoopFormer instead trains one model to support elastic, user-chosen depth at inference.
The mechanism is shortcut modulation (time/step-size conditioning via FiLM-style gates) plus shortcut-consistency training — coarse schedules are explicitly distilled to approximate fine-grained ones.
LoopFormer scales smoothly with added budget where fixed-depth and early-exit baselines degrade sharply off their trained schedule — depth becomes a genuine runtime trade-off knob.

References

Jeddi, A., Ciccone, M., & Taati, B. (2026). LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation. ICLR 2026. arXiv:2602.11451.
Related on this site: Think-at-Hard — token-level selective iteration, an orthogonal axis of compute allocation to LoopFormer’s sequence-level elastic depth; SOL — token-level efficiency policies for frozen LLMs, another point in the same design space of learned, granular compute allocation.