Research map / Architectures and efficiency / Attention and transformers

On Expressive Power of Looped Transformers: Theoretical Analysis and Enhancement via Timestep Encoding

ICML 2025 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Derives an approximation-rate bound for Looped Transformers on continuous permutation-equivariant sequence-to-sequence functions using three moduli of continuity, identifies a looped-architecture limitation, and proposes Timestep-Modulated Looped Transformer (TMLT) with per-loop scaling conditioned on timestep encodings.

Problem

The expressive power and function-approximation rate of Looped Transformers (weight-tied Transformers applied recursively) for continuous sequence-to-sequence functions is unknown, and standard Transformer and looped ReLU results do not directly apply.

Approach

Introduces modulus of sequence, contextual, and token continuity to bound approximation error of looped Transformers in terms of loop count; shows weight-tied feed-forward layers add error from contextual/token variation, and fixes this with per-loop scaling vectors generated from sinusoidal timestep encodings through an MLP and time-dependent RMSNorm and gating.

Key results

• On reasoning tasks, looped Transformer accuracy rises with loop count (e.g., Sudoku 65.6 at r=16 to 87.9 at r=32), and timestep encoding raises Sudoku to 90.2 at r=32 and ED(60) from 47.7 to 88.3 at r=100.
• In-context learning MSE drops from 8.6e-3 (standard TF L=12) and 1.4e-2 (looped r=12) to 1.7e-3 with timestep encoding.
• WikiText-103 test perplexity improves from 20.5 (TF L=12) and 20.6 (looped r=24) to 19.6 with timestep encoding.

Measured effects reported in the paper

Surprising findings

• Looped Transformers with fixed O(d) parameters can match or approach the approximation of standard Transformers when loops are increased, with the error bound shrinking as loop count grows.
• Without timestep encoding, looped Transformer error depends on contextual and token continuity, which can amplify approximation error beyond the standard sequence continuity.

Limitations

• Analysis restricted to single-layer looped Transformers with fixed sequence length and two heads with head size 1.
• Bit complexity of the construction is O(delta^{-(N+1)d}), so the number of loops grows exponentially in dimension and sequence length.
• Only expressivity is studied; estimation performance and training stability are left as future work.

Experimental setup

Datasets: Edit distance (synthetic) LCS (synthetic) Sudoku dataset Countdown decision tree in-context learning WikiText-103

Benchmarks: Sudoku Countdown LCS ED (edit distance) in-context decision tree learning WikiText-103

Metrics: test accuracy MSE perplexity

Baselines: Standard Transformer (TF) Looped Transformer without timestep encoding

Builds on: Looped Transformer Transformers as universal approximators (Yun et al.) Looped ReLU networks approximation rate adaptive instance normalization RMSNorm

Closest papers by meaning

• Stability and Generalization in Looped Transformers (NeurIPS 2026 · similarity 0.89)
• Exact Expressive Power of Transformers with Padding (NeurIPS 2025 · similarity 0.87)
• Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping (NeurIPS 2026 · similarity 0.87)
• Looped Transformers for Length Generalization (ICLR 2025 · similarity 0.86)
• A Formal Framework for Understanding Length Generalization in Transformers (ICLR 2025 · similarity 0.86)
• Chain of Thought Empowers Transformers to Solve Inherently Serial Problems (ICLR 2024 · similarity 0.86)