Looped Transformers

A transformer architecture that applies one shared block repeatedly over many loop iterations instead of stacking many distinct layers.

Looped transformers reuse one shared transformer block many times in a row. Each loop iteration reads the current hidden states, runs the same attention and feed-forward weights again, and writes back updated states before the model makes a final prediction.

At a glance

Released

May 2024

Authors

Liu Yang, Kangwook Lee, Robert D. Nowak, et al.

Optimizes

  • Parameter Efficiency
  • Iterative Inference
  • In Context Learning

What It Is

Looped transformers are a transformer architecture that creates depth by repeating one shared block instead of stacking many separate layers. The model embeds input tokens into hidden states, then applies the same transformer block—self-attention followed by a feed-forward path—over and over for a fixed loop count. Each pass refines the hidden states in place. After the last loop iteration, a prediction head reads the final states and produces the output.

Why It Exists

A standard transformer grows its parameter count every time a new layer is added, because each layer keeps its own weights. Looped transformers ask whether depth can come from repeated computation on shared weights instead. That pattern can reduce parameter count while still giving the model many refinement steps. Researchers study looped transformers partly because repeated passes over the same block may help a model carry out iterative in-context learning behavior with fewer total parameters.

How It Works

Input context tokens are embedded into hidden states that carry information across positions. A loop counter sets how many times the shared block runs. On each iteration, the block applies self-attention so tokens can exchange information, then applies a feed-forward network to refine each position. Residual connections carry the previous state forward so each loop updates rather than replaces the representation. After each pass, the updated hidden states feed back into the same shared block until the loop counter reaches its limit. When the loop count is exhausted, the final hidden states feed a prediction head. The same block weights are used on every iteration; only the hidden states change from pass to pass. In the ICLR 2024 study, training sets a maximum loop count and a truncated loss window over recent loop outputs, while inference can use a different loop count; their linear-regression experiments report that a model trained for one loop budget can still reach a stable fixed-point solution when unrolled for more loops at test time.
Looped transformer compute flow
Hidden state computation path
Loop iteration control
Residual connection
Loop count input
Shared-weight note

Math Or Compute Schema

A standard depth-L transformer uses L distinct blocks with separate parameters. A looped transformer applies one shared block L times to the evolving hidden state, then reads the final state with a prediction head. The formulas below contrast layer-specific depth against shared-block iteration and show how the final prediction is produced.
Standard depth-L transformer stack
h(l)=Blockl(h(l−1)),l=1,…,Lh^{(l)} = \mathrm{Block}_l\big(h^{(l-1)}\big),\quad l = 1,\ldots,L
h(l)h^{(l)}
Hidden states after layer l.
ll
Layer index from 1 to L.
LL
Total number of distinct layers.
Blockl\mathrm{Block}_l
Transformer block with layer-specific parameters.
Looped transformer with shared block and prediction head
h(ℓ)=Block(h(ℓ−1)),ℓ=1,…,L;y^=Head(h(L))h^{(\ell)} = \mathrm{Block}\big(h^{(\ell-1)}\big),\quad \ell = 1,\ldots,L;\quad \hat{y} = \mathrm{Head}\big(h^{(L)}\big)
h(ℓ)h^{(\ell)}
Hidden states after loop iteration ℓ.
ℓ\ell
Loop iteration index from 1 to L.
LL
Loop count for repeated block application.
Block\mathrm{Block}
Shared transformer block reused on every iteration.
y^\hat{y}
Model prediction after the last loop iteration.
Head\mathrm{Head}
Output projection that maps final hidden states to predictions.
h(L)h^{(L)}
Hidden states after the final loop iteration.

Compared To Nearby Modules

A standard transformer stack assigns fresh weights to every layer, so depth and parameter count grow together and each layer runs once per forward pass. A looped transformer keeps one shared block and sets depth through a loop count L, so stored parameters stay tied to that block while compute depth comes from repeated passes. The same attention and feed-forward weights are reused on every iteration, which trades duplicated parameters for extra forward compute. Where a standard stack hands states to a new layer map at each step, a looped design can keep refining the same states with one shared map, which is the intuition behind fixed-point or convergence-style behavior when more loops are unrolled at test time. The shared block still contains ordinary attention and feed-forward submodules; the architectural difference is parameter sharing, iteration count, compute reuse, and iterative refinement rather than a replacement for those building blocks.
Comparison dimensionStandard transformer stackLooped transformerShared transformer block
How depth is createdL distinct layers stacked in sequence, each with its own weightsOne shared block applied L times in a loop over the same statesSingle attention and feed-forward block used as the repeated unit
How parameters scale with depthParameter count grows with layer count because every layer stores separate matricesParameter count stays tied to one block while loop count adds compute depthWeights are shared across loop iterations rather than duplicated per layer
What repeated passes doEach layer transforms states once on the forward passEach loop iteration refines the same hidden states with reused weightsSelf-attention and feed-forward run inside every loop pass
How loop or layer depth is chosenDepth equals layer count; every layer runs once per forward passLoop count L sets how many times the shared block runs; the studied setup can train and test with different countsNot applicable; describes the reusable unit inside each loop pass
Iterative refinement or fixed-point roleStates pass through distinct layer maps without revisiting earlier weightsRepeated application of one block can refine states toward a stable fixed point when unrolled further at test time in the reported linear-regression experimentsSelf-attention and feed-forward execute inside every loop iteration

Example Architectures

In the ICLR 2024 study by Liu Yang et al., looped transformers trained from scratch matched or outperformed standard depth-matched transformers on several in-context data-fitting tasks, including linear regression, sparse linear functions, decision trees, and two-layer neural networks. On the reported linear-regression setup, a one-block looped model reached error close to a twelve-layer standard transformer and to a least-squares solver while using roughly one twelfth of the transformer's parameters. Those experiments use controlled prompts and synthetic function classes rather than large production language-model stacks, so they show where looped depth helps in research benchmarks rather than claiming a universal deployment pattern.

Limitations And Tradeoffs

Each loop iteration adds another full pass through attention and feed-forward compute, so inference can cost more wall-clock time than a shallow parameter-matched model even when total weights are smaller. The ICLR 2024 training recipe depends on loop count, a truncated loss over recent loop outputs, and schedule choices that can make quality sensitive to how many iterations are used at train time versus test time. Their out-of-distribution loop-count tests on linear regression suggest some robustness when unrolling beyond the training budget, but that evidence stays within the studied tasks and does not extend every looped design to broad out-of-distribution behavior. The paper's benchmarks are controlled in-context learning setups, not general production transformers, so looped transformers remain a research architecture for iterative refinement and parameter sharing rather than a drop-in substitute for every standard depth-stacked design.

Why It Still Matters

Looped transformers matter because they separate depth from parameter count in a concrete, testable way and give readers a named pattern for iterative in-context learning. The ICLR 2024 results provide a direct comparison point against standard depth-stacked transformers on data-fitting tasks that already appear in in-context learning literature, so the looped idea can be evaluated on mechanism and tradeoffs instead of only through a paper abstract.

Tags

References

  1. Yang, Liu, et al. "Looped Transformers are Better at Learning Learning Algorithms." ICLR, 2024, https://arxiv.org/abs/2311.12424.

Diffusion Transformer Block

Previous Page

Mamba Selective State-Space Module

Next Page

On this page

What It IsWhy It ExistsHow It WorksMath Or Compute SchemaCompared To Nearby ModulesExample ArchitecturesLimitations And TradeoffsWhy It Still MattersRelatedTagsReferences