On Expressive Power of Looped Transformers: Theoretical Analysis and Enhancement via Timestep Encoding
Authors: Kevin Xu
Issei Sato
Venue: ICML 2025 poster (2025)
Keywords: Transformer
Looped Transformer
Expressive power
Approximation Rate
HyperNetworks
Primary area: deep_learning->attention_mechanisms
TL;DR: We establish an approximation rate for the Looped Transformer and enhance its expressive power via timestep encoding.

Abstract: Looped Transformers provide advantages in parameter efficiency, computational capabilities, and generalization for reasoning tasks. However, their expressive power regarding function approximation remains underexplored. In this paper, we establish the approximation rate of Looped Transformers by defining the modulus of continuity for sequence-to-sequence functions. This reveals a limitation specific to the looped architecture. That is, the analysis prompts the incorporation of scaling parameters for each loop, conditioned on timestep encoding. Experiments validate the theoretical results, showing that increasing the number of loops enhances performance, with further gains achieved through the timestep encoding.

On Expressive Power of Looped Transformers:
Theoretical Analysis and Enhancement via Timestep Encoding
Kevin Xu'
Issei Sato
Abstract
To address these limitations, Looped Transformers presents promising approach. The architecture consists of fixed- size Transformer layers, in which the output is recursively fed back into the input. Looped Transformers exhibit advan- tages in parameter efficiency thanks to their weight-tying structure (Lan et al., 2020; Takase & Kiyono, 2021; Csordás et al., 2024; Bae et al., 2025), achieving performance com- parable to standard Transformers while using fewer param- eters. Additionally, they are well suited for size general- ization by adjusting the loop count based on task complex- ity (Dehghami et al., 2019; Fan al., 2024b). Their recursive structure endows them with the expressive power to emulate iterative algorithms and universal computational capabili- ties, akin to programmable computers (Giannou et al., 2023). Furthermore, their inductive bias enhances performance on reasoning tasks (Saunshi et al., 2025).
Looped Transformers provide advantages in pa- rameter efficiency, computational capabilities, and generalization for reasoning tasks. However, their expressive power regarding function approx- imation remains underexplored. In this paper, we establish the approximation rate of Looped Transformers by defining the modulus of conti- nuity for sequence-ro-sequence functions, This reveals a limitation specific to the looped archi- tecture. That is, the analysis prompts the incorpo- ration of scaling parameters for each loop, con- ditioned on timestep encodings. Experiments validate the theoretical results, showing that in- creasing the number of loops enhances perfor- mance, with further gains achieved through the timestep encoding. Code available at https: / //github... www.///wwwwww/eeeeeeoeeoeeeeeeee
In contrast, the expressive power of Looped Transform- ers and the properties unique to the looped architecture in function approximation remain unexplored. The expressive power of standard Transformers, on the other hand, has been examined extensively in prior studies. These studies show that Transformers can be universal approximators for continuous cormmutttion-qquivrrann functions on compact domains (Yun al., 2020; Kajitsuka & Sato, 2024). Further- more, their approximation rate has been analyzed by identi- fying specific properties of the target functions (Takakura & Suzuki, 2023; Jiang & Li. 2024; Wang & E, 2024), providing insights into the underlying characteristics of Transformer architectures. However, these findings cannot be directly extended due to the weight-tying constraints. Although the approximation rate of looped ReLU networks has been es- tablished only recently (Zhang et al., 2023), that of Looped Transformers remains unknown.
1. Introduction
Transformers (Vaswani et al., 2017) have become the stan- dard architecture for a wide range of machine learning tasks, including natural language processing and computer vi- sion. However, they exhibit certain limitations, particularly when applied to complex tasks. The expressive power of Transformers is theoretically constrained (Merrill & Sabhar- wal, 2023; Feng et al., 2023), and they empirically strug- gle with reasoning and planning problems (Kambhampati et al., 2024). Although chain-of-though reasoning (Wei et al., 2022) can mitigate these challenges some cases, typically relies on manually crafted prompts or costly in- termediate supervision. Moreover, Transformers encounter difficulties with length generalization (Deletang al., 2023) and require substantial computational resources as the num- ber of model parameters increases (Pope et al., 2022).
Our contributions are summarized as follows:
We establish the approximation rate of Looped Trans- formers for fixed-length continuous sgquence-to-sqqeenee functions, with respect the number of loops and three newly defined types of moduli of continuity.
'The University of Tokyo. Correspondence to: Kevin Xu wkevvnxu@gg.comm www.w/...ccctoooooooo woc.u--tooo..c.pp..... Issci Sato maato@g.ccc.u to.yo.ac.jp>.
We identify an inherent limitation of Looped Transformers and address it by proposing Timestep-Modee Modulated Looped Transformers (TMLT), which incorporate scaling parame- ters that are conditioned on timestep encodings.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

On Expressive Power of Looped Transformers
2. Background
2.2. Looped Transformer
This section defines the Transformer and Looped Trans- former architectures, reviews related work, and examines prior theoretical studies on the function approximation capa- bilities of Transformers and weight-tied networks, thereby clarifying the research question addressed in this paper.
Looped Transformer with single layer is represented as:
L2 TF Lt,
where L2 and L1 represent token-wise affine linear layers, and TF denotes the composition of TF applied times.
While we focus on single-layer as (Dehghani et al., 2019; Lan et al., 2020; Yang et al., 2024: Fan et al., 2024b), they can also be implemented with multiple layers as (Csordás et al., 2024; Bae et al., 2025; Saunshi et al., 2025).
Vectors are represented by lowercase boldface letters e.g., v, and matrices are denoted by uppercase bold- face letters e.g., X. The i-th element of a vector o denoted by Vi, and the (i, j-th element of a matrix X is denoted by Xij. The n-th column of a matrix X is denoted by X... Given input sequence X [21 T2, ................................ xNN] RT*NN N where x; E Rm, and a function f: R + Ie" the token- wise application of f is denoted by the bold symbol f i.e.
Notations:
Overview of Previous Work The recursive structure was introduced into Transformers (Dehghani al., 2019), where the number of loops can be adaptively adjusted, allowing for size generalization (Fan et al., 2024a). Looped Transformers are closely related to weight-tying Transformers (Lan et al., 2020; Takase & Kiyono, 2021), achieving performance com- parable to standard Transformers using fewer parameters. Deep equilibrium models (Bai al., 2019), which compute fixed points of iterative layers, are also related. In addition, the recursive structure enables the model to emulate iterative algorithms, including basic computational primitives (Gi- annou et al., 2023) and learning algorithms (Giannou et al., 2024; Yang et al., 2024). Furthermore, recent studies have demonstrated that Looped Transformers exhibit an inductive bias towards reasoning tasks (Saunshi et al., 2025). To im- prove performance, more sophisticated architectures, such mixture-of-exxper (Csordás al., 2024) and relaxed weight-tying (Bae et al., 2025), have been introduced.
f(x) [f f(2)). f(xx) R=XN,
For E [1, oo), the p-norm, denoted by I1-Ilp- represents the entry-wise p-norm. This norm applies to both vectors and matrices e.g., || The L -norm of a function is defined
for E [1, xc) as:
M = (IIIX)))))X 1/P
where Q represents the domain of the function f.
2.1. Transformer Architecture
Given an input sequence X E IRT composed of N token embedding of dimension size m, the self-attention layers with h heads and head size and the feed-forward layer with width size q. are defined as follows:
2.3. Theoretical Analysis on Expressive Power
We review related work and summarize the comparisons between our problem setting and previous studies in Table 1.
Attn(X)= = Σ w Xo, (W x) XTWxXX)
i=1
FF(X..)) = W2on W X + by) + b2,
Universality of Transformers The universal approxima- tion theorem for fully connected neural networks (Cybenko, 1989; Hornik al., 1989) shows that networks of sufficient size can approximate certain classes of functions with arbi- trarily low error. For Transformers, the target function class extends to seguence-to-sqqeeee compute contextual mapping of the input, which requires capturing the entire sequence and computing the token em- bedding within context (Yun et al., 2020), formulated as:
where W Rmxx Wi, Wi: W/ E Ra Wi E Raxm W₂ E Tm by E R9, b2 E Rm are parameters, RR denotes ReLU function, and a, denotes softmax oper- ator applied to the columns of the matrix.
functions. Transformers
Transformer block TF. R + R is defined by
X' = X + At),
( X' FF(X')
where FF represents wike FF. In other words,
Definition 2.1 (Yun et al., 2020). Consider finite set L C RdxN A map CM L - ÆIxN defines a contextual mapping if the map satisfies the following:
TF= = (id fid+FD FF) o (id + Attn),
where id denotes the identity mapping, where omit the domain of definition for simplicity. For the analysis of ex- pressive power in Section 3, we exclude layer normalization and our constructive proof relies on the softmax function to approximate the hardmax function as in previous stud- ies (Yun et al., 2020; Kim et al., 2023)
1. For any L E L,, the N entries in CM(L) are all distinct.
2. For any L, L' E IL, with L + L', all entries of CM(L) and CM(L') are distinct.
Prior studies have shown that Transformers can compute contextual mappings, enabling memorization (Kim et al.,
2

---

On Expressive Power of Looped Transformers
Table 1. Comparisons of our problem setting with related work the theoretical analysis of function approximation.
Paper
Model Type
Function Class
Approximation Rate
Looped (Weight- -Tying)
Yarotsky (2018)
FFN
Continuous functions
X
Yun et al. (2020)
Transformer Continuous seq-to-seq functions
x
X
Takakura & Suzuki (2023)
Transformer
y-smoooth infinite-enggth
x
Kajitsuka Sato (2024)
Transformer Continuous seq-to seq-to-seq functions
X
Jiang & Li (2024)
Transformer
Temporal coupled functions
X
Wang & E (2024)
Transformer
Long but sparse memories
X
Zhang et al. (2023)
FFN
Continuous functions
Ours
Transformer
Continuous seq-to-seq functions
2023) and universal approximation (Yun et al., 2020; Kajit- suka & Sato, 2024).
require contextual mappings. For a srquence-to-sqqueee function / RdxN -> Rdxx,, the modulus of continuity can be generalized as:
For Looped Transformers, as the fixed parameters of a single Transformer layer are used, the results of previous studies cannot be directly applied. This leads to the question: Can Looped Transformers compute contextual mappings? and Are they universal approximators?
wy (6) = sup (I || x s =6}
We investigate whether this modulus of continuity alone can determine the approximation rate.
Approximation Rate of Transformers Beyond the uni- versality, the approximation rate provides deeper insights into the characteristics of models (Barron, 1993; Yarotsky, 2018). This rate is derived as an upper bound of eтror in terms of the properties of the target functions and the com- plexity of the networks. For Transformers, recent studies have investigated these rates and the nature of the target functions (Takakura & Suzuki, 2023; Jiang & Li, 2024; Wang & E, 2024). Specifically, they have shown conditions under which Transformers can overcome the curse of dimen- sionality (Takakura & Suzuki, 2023) and revealed structures in target functions that Transformers can effectively approx- imate (Jiang & Li, 2024; Wang & E, 2024).
For Looped Transformers, it has been shown that they can represent standard Transformers, although their parameter count grows with both the desired approximation accuracy and the sequence length (Saunshi et al., 2025). Moreover, no existing work has established their approximation rate.
3. Approximation Rate Analysis
In this section, we establish the approximation rate of Looped Transformers. We define three types for the modu- lus of continuity in Section 3.2 that determine the approxi- mation rate. The main results are presented in Section 3.3, followed by a proof sketch in Section 3.4.
Our study focuses on understanding the architectural proper- ties of Looped Transformers, particularly in comparison to standard Transformers. To this end, we explore the approxi- mation rate and investigate the properties of target functions that determine their approximation errors.
3.1. Preliminaries
The target function class of our analysis is continuous func- tions that Transformers can represent. Specifically, these are cotmmntttioo-qquivariii
functions, defined as follows:
Definition 3.1 (Yun et al., 2020). A function f RRdxN
Expressive Power of Weight-Tied Neural Networks Re- cently, it has been shown that single fixed-size networks can serve as universal approximators in parameter-efficient manner; that is, the parameter count depends solely the input dimension, not the approximation error (Zhang et al., 2023). Furthermore, the approximation rate of weight-tied ReLUU networks has been established with respect to the number of loops and the modulus of continuity of con- tinuous functions (Zhang al., 2023). The modulus of continuity for g R2 IR. and > 0 is defined as:
R is said be permutation equivariant if f(XP) f(X) holds for any permutation matrix P. Let FPE
denote the set of continuous functions, defined on Ω that are permutation equivariant.
We evaluate both the number of parameters and the bit com- plexity, the maximum number of bits required to represent the network's 's weights (Vardi et al., 2022; Kim et al., 2023).
In our proofs, we introduce IDs for tokens, sequences, and tokens within sequences as theoretical constructs to formal- ize contextual mappings.
(6) = sup (1((z) g(x')l ): || x < 5}.
Our question is whether the results can be extended to sequence-to-sequee functions and Transformers, which
Definition 3.2. A token ID is a unique integer assigned to each token. A sequence ID uniquely identifies each sentence.
3

---

On Expressive Power of Looped Transformers
A contextual loken ID uniquely identifies specific token within specific sentence. We denote the set of contextual token IDs as KC o.I,..., K 1, with corresponding embeddings yı E E for each & E K.
3.3. Main Result
The result establishes the approximation rate of Looped Transformers in terms of the number of loops and the three types of moduli of continuity of the target function.
This notion is defined in Kim et al. (2023), to which we refer for further details, for constructive proofs of contextual mappings. The actual construction of contextual token IDs may vary depending the specific proof. In our case, we adopt a different construction from that of Kim et al. (2023).
Theorem 3.6. Given function J E FFE((O, 1] 1jdxN), N, there exists a Looped Transformer, composed of TF. R(1T449)NN + (1(dd+9)xNN with two heads, head size 1, and width size of 49d + 25, and two affine linear maps L1 : Rdd + +R.17 and L2 : TR. 17d99 + Rd s.L.
L2 o TFor o Li -1-
3.2. Definition of modulus of Continuity
S (Nd)² (w)) (BVC)) + ( wy(6V/N))
As briefly mentioned in the preliminary discussion, we de- fine the modulus of continuity in Eq. 2.3 as:
+ +O(N2/P + ((MM8)- 1/p)
Definition 3.3 (Modulus of Sequence Continuity). Given
sequence woaennneett-seeeeeee to-sequence continuous function f R RdxN the modulus of sequence continuity is defined by:
where 6 = N -11//))1)((()
where M is the maximun absolute value of the model pa- rameters, and the bit complexity is O O(8-(N+1dd)
(6) = sup { [I(()) [Ip:)
|| - X'12 S 83.
We omit the subscript p for simplicity. This quantifies how the output sequence shifts relative to differences in input, hence referred to as sequence continuity.
Theorem 3.6 shows that increasing the number of loops reduces the approximation error. Under infinite- -eenite-rreiiion weights, this leads universal approximation theorem.
We found that this alone is insufficient to determine the approximation rate of Looped Transformers, in contrast to the case of ReLU networks (Zhang et al., 2023). Informally, this issue arises because Transformers compute contextual mappings. We notably identified two additional types of modulus of continuity, defined as follows.
Corollary 3.7 (Universality). The hypothesis space of Looped Transformers, defined by
H ( o TF o -L :2:00,1 10,1 + [0, 1jdxN m, m Cd, h h s = r MEN, W }
where C is positive constant, W denotes the flattened set of all weights in the network, and twe represents the total number of these weights, is dense in FPE(10, 1)dxN) w.r.tthe LP norm.
Definition 3.4 (Modulus of Contextual Continuity). Given a sequence-to-sequenee continuous function : R + Rd N. the modulus of contextual continuity is defined by: (llf(X):. ):m f(XX)(()) dimlle
sup X.X
This approximation analysis highlights the characteristics of Looped Transformers, including both their capabilities and limitations, as summarized below:
X X 8, X., = }-
Definition 3.5 (Modulus of Token Continuity). Given a sequence-to-sqquence continuous function f RaxN RdXNN, the modulus of token continuity is defined by: (X')s.llz
While the number of parameters remains fixed (d), in- dependent of the desired approximation accuracy and the sequence length, the error can be reduced by increasing the number of loops.
lok (5) == sup (II((X):. X.X,X'
1 -on-X' 1 S6 X = X', (vg + })).
Looped Transformers, even with weight-tied self-attention using a hard-max function, can compute contextual map- pings and become universal approximators.
The modulus of contextual continuity quantifies the variation in the contextual token embeddings induced by perturba- tions of context. For example, consider the sentences: (1) "I write papers" and (2) "You write books". It measures the difference in the contextual token embedding of the same word 'write' within different contexts.
The approximation rate depends three types of conti- nuity, with contextual and token dependencies unique to Looped Transformers; these dependencies are not present in standard Transformers or looped ReLU networks.
On the other hand, the modulus of token continuity quantifies the variation in the output embedding caused by perturba- tions to the token itself within the same context such as (1) write papers" and (2) "I draft papers".
Our contribution lies in establishing the approximation rate with respect to the number of loops, based on novel moduli of continuity that are unique to Looped Transformers.
4

---

On Expressive Power of Looped Transformers
from Step
Sequence from Step 2
Kim al., 2023) cannot be directly applied to Looped Trans- formers due the following distinctions:
11 paper
Beethoven writes music.
Yun et al. (2020) employed both sparse and uniform at- tention mechanisms, whereas Looped Transformers limited to single fixed attention layer.
draft papers.
write
music
2 draft
write papers.
Kim al. (2023) used N layers store N parameters required for representing the target function, whereas Looped Transformers have fixed parameter size.
2: token ID
B1
B1,1
Notably, found that Looped Transformers with N-loops can compute contextual mapping. Let z E 10. 1 represent sequence consist of N ordered and distinct token IDs, satisfying ø > 22 > ... NN. The network then maps 2 to a sequence ID through an inner product with u (8
write
Beethoven writes music.
K: Contextual tokea ID
Figure 1. The networks construct contextual token IDs by combin- ing token IDs with sequence IDs.
(8-d(N 10-d(N-1 -X--1 8-d 1). which satisfies
Furthermore, the additional dependency can amplify the ap- proximation error, revealing limitation inherent Looped Transformers. A detailed discussion of this issue, along with improvement methods, is provided in Section 4.
> 2
This guarantees that the network assigns distinct sequence IDs for different z. Combined with token IDs, the network computes contextual mapping. The key idea is that the network requires only 8-4 to represent u, allowing to be implemented with Looped Transformers.
3.4. Proof Sketch
This section presents proof sketch, emphasizing distine- tions from prior studies and challenges unique to the looped architecture. The formal proof is provided in Appendix A.
Step Function Value Mapping. The network maps the contextual token IDs into the target embeddings in token-wise manner, using K loops to sequentially map h = 0, 1....., K - 1 to ÿ Rd. which approximates yk: in each iteration. In our constructive proof', we design both the set contextual token IDs and their ordering.
The basic strategy involves approximating the continu- ous target function I with a piecewise constant function f. which approximated by the network, denoted by f. For 8-1 N, 8-1 Z 2, the input space [0, 1] dddx is divided into discretized cubes with width 8, denoted by (QB) BE (o..,,.....1111111
Weight-tied herd-forarrd networks cannot map accurately, and the error only be bounded by the maximum differ- ence between adjacent contextual token embeddings, i.e.
Fach cube is assigned rep- resentative point XB E QB. and the piecewise constant function f is then defined as:
[ yb)al S max (Bk - -1-l KEE
(4)
f(x) = (()) where B satisfies X E QB.
holds for k E K. and == 1, ......d. d.
The approximation with networks consists of three steps. First, the network assigns token ID to each token. Second, it assigns a sequence ID. The combination of the token ID and sequence ID) constitutes the contextual token IDs as in Fig. 1. Finally, these are mapped to embeddings that represent the output of the target function at each token.
Generally, the following inequality holds, for æ E R
max b < ||> I m max |xy].
That is, by controlling the p-norm, 1/ yk- ppp. the error in Eq. 4 can bounded. We require K. to be designed such that the differences between neighboring contextual token embeddings are bounded w.r.t.he p-norm.
Step Token-wise Quantization. The network uses the feed-forward network to assign each input token, denoted by X, X,, token ID, denoted by z, in a token-wiss manner.
To illustrate our idea, consider the following sentences:
X E 10
8- 10.1....,5
8-4-13 1}.
(2)
(1) I write papers. I write papers. (different token ID with same sequence ID)
Step 2. Contextual Mapping. The network, given N to- ken IDs computes their sequence ID. We notice that the result of previous studies on Transformers (Yun al., 2020;
(2) I write papers. You write books. (same token ID with different sequence ID)
5

---

On Expressive Power of Looped Transformers
yx (model)
linear interpolation
cocont (8vNd)
lok (8)d)
a cont (Bvnd) + to (8vd)
write
He writes papers,
Mozart writes music.
She talks
drinks
K
papers.
papers.
at you.
coffee.
Figure 2. Approximation error and modulus of continuity. The linear interpolation technique reduces the error by a factor of 1/8
We found that none of the moduli of continuity, defined in Section 3.2, alone can bound the difference between "write and 'papers' in (I), contrast, the error of 'write' in (2) can be bounded by the contextual continuity, w Thus, we designed contextual token IDs such that, basically, identical similar tokens with different sequence IDs are positioned adjacent to each other, as shown in Fig. 2. To reduce errors in corner cases, linear interpolation is applied; further details are provided in Appendix A. This allows us to obtain the following error bound.
where ((N2/p84/P) arises from the case where identical tokens appear in sequences, and O (((M0)-1dN) 1/p) (P) re- sults from the restriction on the norm of weights.
Considering the error within cubes in Eq. I, we obtain
J [11(()- s(x)) lI) S wy((vNd)).
(10)
Since, generally, the norm of sequences be bounded by the maximum norm of the token-wisg wiiee vectors
max yyu- yk 1llo y S w/ (ōv + wi cont (ōv Nd). (ōvNd). K'EK
(6)
ll I (Nv) max (()),,))
(11)
3,a
Substituting yı- - into Eq. 5, with Eq. 4 and Eq. 6, the following result holds:
the error between f and can be bounded by
/ ((x(x) < (Nd)# max lüw yve (12)
lix yk)aa ( + (ōVNNd)
(7)
Substituting æ = UN ys into Eq. 5, and using Eq. 7 and Eq. 12, we obtain:
for i = l....,d d and k E K.
Concatenated into Single Transformer Layer In the final construction, we show that the composition of the three sub-networks from Steps 1, 2, and can implemented within a single Transformer block. While our proof strategy follows Zhang et al. (2023), their approach necessitates an additional layer. In contrast, show that single Trans- former block suffices, detailed in Appendix A.
/ |1 f(X)llpdX
(13)
S (Nd) (wp + (om
We then express in terms of the number of loops r to determine the approximation rate, We use 8-1 loops for Step 1, N loops for Step 2, and 28 2 1 loops for Step with 1 loop required to connect each step i.e.
Deriving Approximation Rate Lastly, we analyze the approximation error of our construction and establish the approximation rate in terms of the number of loops.
8-1 + 20 - (N+1)d N.
(14)
Now, å can be bounded in terms the number of loops as:
With the triangle inequality, we obtain the following:
ò + 28 2NV+)
N 5-1 28 (N-(NN1)d 2 N (8-1 2)
(15)
[1f(x) || fllo= S / - f(X)1
(8)
(16)
/l1((x) ((X))) / [1F(X) ((00)f((x1)) f(x
N
11((N+1)d41)
(17)
+0(N3/18d/n) +o aN((
By combining Eq. 10 and Eq 13 with Eq. 9, and substituting Eq. 17, we obtain Theorem 3.6.
(9)
6

---

On Expressive Power of Looped Transformers
4. From Theory to Practice: Introducing Timestep Encoding
4.2. Timestep-Moddulated Looped Transformer
We employ timestep encodings to condition scaling parame- ters on the loop index (timestep). This method is inspired by adaptive instance normalization (Peebles & Xie, 2023).
The theoretical result in Section 3 highlights limitation of the looped architecture. We show that variant of architec- ture can overcome this limitation.
To condition on timesteps, frequency embeddings are pro- cessed through a two-layer MLP with hidden size matching the Transformer block and SiLU activation. Let TE(t) E Fdd denote timestep embeddings, defined as:
4.1. Motivation
Limitation Specific to Looped Transformers The ap- proximation rate in Theorem 3.6 includes two additional moduli of continuity, which can lead to increased errors, reflecting a limitation inherent to Looped Transformers.
TE(t) = W3 SW4 PE(() + ba) + b3.
where W3, W4 E Raxd b3. by E Rdd, and PF(() E Rd de- notes the timestep encoding function that maps the timestep into d-dimensional embedding, s.t.
We can identify the cause of additional dependency in the error in Eq. 4, caused by weight-tied feed-forward networks. This can be formalized as follows:
PE(t);: = sin(t/10000
Lemma 4.1. Given yk E Rdd for k = 0,1,...,
K -1 yy - yk- for = fori=1,-. 1,...,d
with
PE(f)2141 cos(t/10000
We use the root mean square layer normalization (RM- SNorm) (Zhang & Sennrich, 2019), which is widely used in several recent LLMs (et al., 2023; Team, 2024), defined as:
there exists a feed-forward layer FF., 122dd + Igldd with a width size of 18d, and two affine linear maps L1 R + 1222d and L2 1R.2d - TRd s.I.
( (id FF)O FF) o (k())) (000-yy)))) y:), 1 Si for = I,...,d and ls 0, 1,....k K
(18)
T RMS(#)) where RMS(w)
1
where a E Rd a gain parameter for rescaling. We define dime-dependent RMSNorm, denoted by RMSN,
This shows that large variations in the target function may lead to approximation errors, raising the question of whether inequality in Eq. 18 can be replaced with equality.
RMSN( 1) = a(I)) C RMS(x)
Improving Approximation Rate of Looped Transformers To eliminate this dependency, we introduce time-dependent parameters. Specifically,, we modify the feed-forward layers by adding scaling vector for each loop step as follows:
where (1) Rd is time-dependent parameter generated by network. With scaling parameters, the time-dependent Transformer block is defined as follows:
FF(X) (1) FF( for the 1-th loops,
X' X + (t O Attn (RMSN, (X, t)), TF t) X' (t @ FF(MMSN, (X', 2X' ()),
where @ is element-wist product,t E N is the loop index, and (1) E I& is the scaling parameter for each loop. This method is analogous to Hypernetworks (Ha et al., 2016).
where 7: (t), (t) E RR are time-dependent parameters applied token-wise, as well as RMSNorm.
With the definition
The time-dependent vector parameters are generated as:
(id + (id+noFN) FF)" = (id+ 2
a (1), 0z (1), (1) W, -ILUU(E((() + bs:
we show that this model can memorize labels exactly.
Theorem 4.2. Given yv E Rd for k 0, 1,.... K there exists a feed-forward layer FF R.4d + 4dd with a width size 6d, nnt) TR. fort = 1, ...,, K 1, and two affine linear maps L1= IR
where W2 E Rdxd and by E TRd.
5. Experiments
+ 1d and L2 1R gd s.t.
This section presents experimental results supporting our theoretical findings. We used Looped Transformers with varying numbers of loops, both with and without timestep encoding, and compared to standard Transformers. We assess approximation capabilities based on test evaluation, as we observe a strong correlation between train and test performance. The details are provided in Appendix C.
(L2 o sid + O FF) (K-))
Li (k) yk); O,
for I, and &- 0, 1,
K 1.
The proof is provided in Appendix B. For implementation, adding parameters per loop increases the total parameter count proportionally. Thus, we introduce timestep encoding.
7

---

On Expressive Power of Looped Transformers
Table 2. Test accuracy for reasoning tasks.
Performance improves as the numher of loops increases..
TF
Looped TF
w/ Timestep Encoding
Task
L=6
r=4
r=8
r=16
r=32
r=4
r=8
r=16
r=32
Sudoku
0.0
0.0
0.0
65.6
87.9
0.0
0.0
62.0
90.2
Countdown
53.8
28.3
52.7
81.0
88.1
33.2
54.4
80.2
90.5
L=12
r=5
r=10
r=50
r=100
r=5
r=10
r=50
r=100
LCS (60)
70.0
66.0
81.8
98.6
96.9
68.5
80.5
99.3
97.1
LCS (100)
39.8
39.6
45.1
93.5
98.2
36.7
45.6
98.1
98.6
ED (40)
54.2
41.4
57.9
85.4
90.4
44.8
63.5
94.5
96.1
ED (60)
41.4
23.8
32.6
47.3
47.7
26.6
38.9
57.3
88.3
5.1. Problem Setting
5.2. Results
evaluate the model on two types tasks. The first consists of reasoning problems known to be challenging for standard Transformers. These used to examine whether increasing the number of loops and incorporating timestep encodings can enhance performance. The second includes core Transformer benchmarks, such as in-context learning and language modeling.
The results in Table demonstrate that increasing num- ber of loops improves performance reasoning tasks, with higher loop counts significantly outperforming standard Transformers. Furthermore, incorporating timestep encod- ings leads to additional gains; in particular, for the edit distance task with input size = 60, the model with loop counts r = 100 achieves significantly better performance when timestep encodings are incorporated.
5.1.1. REASONING TASKS
Dynamic Programming method for solving complex problems by breaking them down into simpler sub-problems. We use edit distance (ED) and longest common subsequence (LCS) tasks with varying input lengths. Each task has 106 train samples and 1 test samples.
Table 3. MSE (L) on the in-context learning task.
TF L=12
Looped r=12
w/ Timestep r=12
Test
8.6e-03
1.4e-02
1.7e-03
Sudoku is a constrained satisfaction problem that involves filling x9 grid with digits from to 9, such that each digit appears exactly once in every row, column, and predefined 3 x 3 sub-grid. The grid is flattened into sequence rep- resentation. Unlike (Yang et al., 2023), we use the dataset from (Radcliffe, 2020), sampling 3M instances for training and 100K for testing.
Table 4. Perplexity (1) on the WikiText-103 dataset.
TF L=12
Looped r=24
w/ Timestep r=24
Train Test
15.9
17.1
15.9 19.6
20.5
20.6
As evidenced by the results in Table 3 and Table 4, the use of timestep encodings Ieads to performance gains both in-context learning and language modeling.
Countdown is a game which given set of input numbers must be combined using basic arithmetic operations reach target number (Yao al., 2023; Gandhi et al., 2024). We consider cases with 4 input numbers and target numbers ranging from 10 to 100, where 10% of the target numbers are reserved for evaluation. We generate 5M samples for training and 1K samples for testing.
6. Conclusion
We establish the approximation rate of Looped Transform- ers with respect to the number of loops and the moduli of continuity of the target function. Our analysis reveals a limitation of Looped Transformers, which addressed by timestep encodings. To the best of our knowledge, this study is the first to investigate the function approximation capa- bilities of Looped Transformers. Extending the analysis to multiple layers, varying input lengths, and characterizing op- timal memorization capacity presents promising avenues for future research. Beyond expressivity, investigating estima- tion performance and enhancing training stability constitute important challenges moving forward.
5.1.2. IN-CONTEXT AND LANGUAGE MODELING
The in-context learning problem is learn the function class from given sequence, which investigated with Looped Transformers ( (Yang et al., 2024) without timestep encodings. We use decision tree functions. For the language modeling task, we use the WikiText- 103 (Merity et al., 2017) dataset, containing over 100 million tokens from Wikipedia articles. Details are Appendix C.2 and Appendix C.3.
8

---

On Expressive Power of Looped Transformers
Acknowledgements
al., H. T. Llama: Open and efficient foundation lan- guage models, 2023. URL httpts://axxiiv..ooto abs abs/2302.13971
This work was supported by JSPS KAKENHI Grant Num- ber JP24H00709.
Fan, Y., Du Y., Ramchandran, K., and Lee, K. Looped transformers for length generalization, 2024a. URL
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must specifically highlighted here.
Fan, Y., Du Y., Ramchandran, K., and Lee, K. Looped transformers for length generalization, 2024b. URL
https://arxie http://axxxvv..r/abeee// 2409. 15647.
Feng, G., Zhang, B., Gu, Y., Ye, H., He, D., and Wang. L. Towards revealing the mystery behind chain of thought: A theoretical perspective. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https: httt:::://ppeeeevieee www./nne//eteee forum? id-qHrADgAdYu.
References
Bae, S., Fisch, A., Harutyunyan, H., Ji, Z.,, Kim, S., and Schuster, T. Relaxed recursive transformers: Effective pa- rameter sharing with layer-wise loRA. In The Thirteenth International Conference Learning Representations, 2025. URL httpttt:/ttteeeeeeeeee http///ppppeeeeeee//e/epee net/ forum? id=WwpYSOkkCE
Gandhi, K., Lee, D. H. J., Grand, G., Liu, M., Cheng, W., Sharma, A., and Goodman, N. Stream of search (sos): Learning to search in language. In First Con- ference on Language Modeling, 2024. URL https: / //openreview net/ htt////rrmmmi//////// wwwrrummdddcoopppcmo jMOVL. JmOV..
Bai, S., Kolter, J. Z., and Koltun, V. Deep equilibrium mod- cls. Advances in neural information processing systems, 32, 2019.
Barron, A. R. Universal approximation bounds for super- positions a sigmoidal function. IEEE Transactions on Information theory, 39(3):930 945, 1993.
Garg, S., Tsipras, D., Liang, P., and Valiant, G. What can transformers learn in-context? a case study of sim- ple function classes. In Oh, A. H., Agarwal, A., Bel- grave, D., and Cho, K. (eds.), Advances in Neural In- formation Processing Systems, 2022. URL https: / //openreview.oo net ftrumoirrdi///// FINZJ2e0et.
Bartlett, P. Maiorov, V., and Meir, R. Almost linear ve dimension bounds for piecewise polynomial networks. Advances in neural information processing systems, 11, 1998.
Giannou, A., Rajput, S., Sohn, J.-Y., Lee, K., Lee, J. D., and Papailiopoulos, D. Looped transformers as programmable computers. In Proceedings of the 40th International Conference on Machine Learning, /ttptuttsssppppoeeeeeeeiioi http:///coooceeeseeeeeee/ mlr mlr.
Csordás, R., Irie, K., Schmidhuber, J., Potts, C., and Man- ning, C. D. MoEUT: Mixture-of-expere universal trans- formers. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024. URL https:
2023. URL htth:tts:/ v202/giannou2333 html.
Giannou, A., Yang, L., Wang, T., Papailiopoulos, D., and Lee, D. How well can transformers emu- late in-context newton's method? arXiv preprint arXiv:2403.03183 2024. URL https://arxis org/abs 03183.
Cybenko, G. Approximation by superpositions of a sig- moidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, Universal transformers. In In- ternational Conference on Learning Representations, 2019. URL tattts::/opeereeet id Hdd-Hy2R.R
Ha. D., Dai, A., and Le, Q. V. Hypernetworks. arXiv preprint arXiv: 1609966009, 2016.
net. forumm
Horni,, K., Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. Neu- ral networks, 2(5):359-366, 1989.
Deletang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, I.. K., Catt, E.,, Cundy, C., Hutter, M., Legg. S., Veness, J., and Ortega, P. A. Neural networks and the chomsky hierarchy. In The Eleventh International Confer- ence on Learning Representations, 2023. URL ht tps: / //openreview net/ forum? http://oommmiddoo/tAtataeo d=w.wd
Jiang, H. and Li, Q. Approximation rate of the transformer architecture for sequence modeling. arXiv preprint arXiv:2305.18477 2024. URL https: //rxivv. org/abs/2300

---

On Expressive Power of Looped Transformers
Kajitsuka, T. and Sato, I. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? In The Twelfth International Conference on Learning Representations, 2024. URL https:// openrevi ew, ntt//fruu forum? 1=n.Jnky5944.
Saunshi, N., Dikkala, N., Li, Z.., Kumar, S., and Reddi, S. J. Reasoning with latent thoughts: On the power of looped transformers. In The Thirteenth International Conference on Learning Representations, 2025. URL https: openreview, net / wtw.ummmdddiinne@e
Kambhampati, S., Valmeckam, K., Guan, L., Verma, M., Stechly, K., Bhambri, S., Saldyt, L. P., and Murthy, A. B. Position: LLMs can't 't plan, but can help planning in LLM-modulo frameworks. In Forty-first International Conference on Machine Learning, 2024. URL https: wtwpilewweee/fooee wwt T- E.
Takakura, and Suzuki, T. Approximation and estimation ability of transformers for sequence-to-seqeene func- tions with infinite dimensional input. In Proceedings of the 40th International Conference on Machine Learning, 2023, URL https: httpt:::/pprccceciinngnto mlr press/ v202 2.02/takakura250 html.
//openreview.co
Kim, J. Kim, M., and Mozafari, B. Provable memo- rization capacity of transformers. In The Eleventh In- ternational Conference on Learning Representations, 2023. URL 1 1d=8JCg5xJCTPR.
Takase, S. and Kiyono, Lessons on parameter shar- ing across layers in transformers. arXiv preprint arXiv: (tti::0/04000222. 2021.
Team, G. Gemma 2: Improving open language models at a practical size, 2024. URL httpt:://aaxi...ottt abs 2408..48..0 00118.
Lan, Z., Chen, M., Goodman, S., Gimpel, K., Sharma, P., and Soricut, R. Albert: A lite bert for self- supervised learning of language representations. In International Conference on Learning Representations, 2020. URL httptttt/////c hwwph:ttt////eeeee http://///p/eeveeeeeeeeeeeeeeeee http:/iwwwweeeeeeeeeeeeeeeeeeeee id-HleA7AEtvV..
Vardi, G., Yehudai, G., and Shamir, O. On the optimal memorization power of reI. neural networks. In In- ternational Conference on Learning Representations, 2022. URL https:/ http://///ppeeeeeeeeeeoomemmmeep wttpeeeiiwwwetefffrrreeeeee forum? id-MkIPtnjeYTV.
Loshchilov, I. and Hutter, F. Fixing weight decay regular- ization in adam, 2018. URL https:://pperreeeeeeeee //ppnrevvew. net/ forum?
Vaswani, A., Shazeer, N. M., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Neural Information Process- ing Systems, 2017.
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference on Learning Representationss 2017. URL https:// 72udxe.
Wang. M. and E, W. Understanding the expressive power and mechanisms of transformer for sequence modeling. In The Thirty-eighth Annual Conference on Neural In- formation Processing Systems, 2024. URL https: 7Rd5 jngV.
Merrill, W. and Sabharwal, A. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531- 545, 2023.
// wwwpneeviw..eeeee.eeeeeeeeeeeeee net htt///orummmmi
Wei, J., Wang, X., Schuurmans, D., Bosma, M. brian ichter, Xia. F., Chi, E. H. Le, Q. V., and Zhou, D. Chain of thought prompting elicits reasoning in large language models. In Oh, A. H.,, Agarwal, A., Belgrave, D., and Cho, K. (cds.), Advances in Neural Information Processing Systems, 2022. URL https: opppreveeevvveeeeeeeeeeeee net./ forum? waat-dddddioeesss_ ViQ.MeSB_BB..
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF Interna- tional Conference Computer Vision, pp. 4195-4205, 2023.
Pope, R., Douglas, S.. Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference, 2022. URL hhttps://arxii/ wtrpiii.orrra.oooo abts/211111 05102.
Yang, L., Lee, K., Nowak, R. D., and Papailiopoulos, D. Looped transformers are better at learning learn- ing algorithms. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: //openreview.oo net www.mmiddd-nuooooooo. HHHRRxODDX..
Radeliffe, D. G. 3 million sudoku puzzles with ratings, 2020.
Yang. 7., Ishay, A., and Lee, J. Learning solve constraint satisfaction problems with recurrent transformer. In The Eleventh International Conference on Learning Represen- tations, 2023. URL https: openreview net/ forum d=udNhDCr220N..
Radford, A., WU, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
10

---

On Expressive Power of Looped Transformers
Yao, S., Yu, D., Zhao, J., Shafran, I., Griffiths, T. L.,, Cao, Y., and Narasimhan, K. R. Tree of thoughts: Deliberate problem solving with large language models. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https: / http://ppeeevvveeeee www.rrevvieeeeee.e n Forum?id-5Xclecxe
Yarotsky, D. Optimal approximation of continuous func- tions by very deep relu networks. In Conference on learn- ing theory, pp. 639-649. PMLR, 2018.
Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S., and Kumar, S. Are transformers universal approximators of sequence-to-sequence functions? In International Confer- ence on Learning Representations, 2020. URL https: //ppereevew. net / forum? i By
Zhang, B. and Semnrich, R. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Zhang, S., Lu, J., and Zhao, H. On enhancing expressive power via compositions of single fixed-size relu network. In International Conference Machine Learning, pp. 41452-41487. PMLR, 2023.

---

On Expressive Power of Looped Transformers
A. Proofs for Theorem 3.6
The main theorem incorporates a restriction on the norm of weights, leading to errors when approximating discontinuous functions, such as step functions with ReLU or hardmax functions with softmax. We first establish the approximation rate assuming that weights can take arbitrary precision real values, as outlined below. Then, account for the bit complexity of bounded weights complete the proof of the main Theorem 3.6.
Theorem A.1. Given a function f (P. 1] r > N. there exists a Looped Transformer, composed of TF R(LTd49)xX (111=)))) with Iwo heads, head size width size of 49d + 25, and two affine linear maps :::R Rd
R 17d49 and L2 :R II17449
IRd s.I.
|| o TF
OFF o L1 fll (Nd) (W + (6V wyllvv O(N2/18/PP)
(19)
where = ((r-N)/2) 1/(N+1)6++)
A.1. Proof of Theorem A.1
Proof. Since any continuous function can be approximated by piecewise constant function with arbitrarily small errors, we
approximate f E FFE((0, )j with piece-wise constant function f: [o. 1 1 + Rd We choose 5-1 E N, 1-. 2 2,
determining how finely the input is divided: we divide the input space [0, 1] i]dxN into 8-discretized cubes, denoted by (QB)
for B {0, ..... 1 JdxN defined by
QB = x = 1j X [B, (Bi, + 1)6), i. d, n2 1.2....,NN }
(20)
Each cube Qs is associated with a representative point XB, defined as the vertex of Qs with the minimal & norm. Then, we define the piecewise constant function f for X [0, 1jdxN as
f(x)= f(XB).
(21)
Since we can bound the error within each cube, we have:
max Xe
(11((x) s(X)lp) ((X)Il) (oy().
(22)
Our construction consists of three steps approximate f. as outlined below.
1. The network, with 8-1 - loops, maps the input space [0, 1] token-wise to the coordinates 3 E {0, 1...., .. 8-1 1 of discretized cubes, and then bijectively maps these coordinates to integers, representing token IDs in the set 10......, 8-d 1}. using 8 base system; for example, d = 2 and 8-1 3, then coordinates (31, 32) = (2, 1) are mapped to the integer 3 1 30 7.
2. The network, with N loops, computes contextual mapping from the set of N distinct token IDs into the of contextual token ID. Contextual token IDs refer to token IDs assigned each token within the context of sequence ID.
3. The network, with 26 (NN+ - loops, approximately maps contextual token IDs into the output embeddings of each token in token-wise manner. To achieve small approximation error, the network has to be designed so that neighboring IDs correspond to similar output token embeddings. Furthermore, dummy indices are used to reduce the error.
The details for each step are provided below.
Step 1. Token-wise Quantization. The input space for each token æ E [0, 1]d are divided into 5-discretized cubes denoted
by (Qa) for G E (0, 1,..., 8-1 - 1]d, defined as
QB { a0 E [0, 112 as E [3,5, (3, 1)6), for all i = 1, i.,, 1,,
(23)
12

---

On Expressive Power of Looped Transformers
By Lemma A.5, there exists feed-forward layer FF(1), Réd RRdd of width size g 7d, and two affine linear maps
c Rd + Rdd and C²1) Rod - Rd such that
(x) B s.s. Q33 2 (id + FF(U) 0 111 (x)
(24)
In addition, we need to bijectively map the d-dimensional vector f to an integer token ID, denoted by 2. We use a 5-1-base system: we define the vector (66-1) IRd as
u16 1) = (6-(d-1) (d(22) ,...., 10, ,))
(25)
and define as
2 = ui6 1/B3 E {0, 1,...,6.6 .....,,- -1.
(26)
To implement this, we define an affine linear map c/²²' 5d + R via
(x) = uib 1) 1)(22(()) (z).
T
(27)
Thus, we have
(c o (FF()))((-111)) (x)),, = =u(s 11B = s.t. x E QB.
(28)
We establish an upper bound on the maximum distance in input space between adjacent token IDs to derive the approximation error for the following steps. Define the input cubes corresponding to each token ID z as follows:
Q2 = x E j0, [0,11= x, E [8,6, (G, + 1)6) for = 1,2..... d s.t. 2 = =u 36
(29)
Then we have
Vvd, Vd.
if E (1. 1}, 3 0.
8-1
max s.mE EE EQ.
2 2'112 S
(30)
Informally, in this based representation, the least significant digit corresponds to the index of the d-th dimension, B., As the token ID increments sequentially, the index in the d-th dimension increases as 0, 0, 8 1, while the higher-order digits remain unchanged. Consequently, consecutive token IDs correspond to tokens that are "similar" in the d-dimensional space, with a maximum distance of dvd. However, when a carry occurs, the higher-order digits may change significantly, leading to cases where tokens that are not adjacent in input space become adjacent in their indices. In such cases, the distance is only bounded by vd.
Step Contextual Mapping. The networks, with N-loops, map the list of N token IDs, denoted by z E (0, 1, 1,....5 8-d_ 11a, into sequence IDs bijectively. Combined with token IDs, the network computes contextual mapping.
We consider only the case where all N input tokens are distinct, disregarding other cases, as they be treated as negligible when 8 is small. The number of subsets in which at least one of the N tokens is duplicated is given by
N(N - 1), 8 3-N--1)d
(8-4)N- (o-d. 18)..-)... N 1)
<
(31)
2
when 8 is sufficiently small. The volume of these subsets is bounded by
C5- CN-(N-)
N-Nd Nd = (NV8d).
(32)
Thus, the error with respect to the IP norm is bounded by ((N2/18//0)
Let LL denote the set of N distinct token IDs, i.e.
TLg == {z E (0, 1, ,.... 8-d 13 2, + zy for all i #110 j.
(33)
13

---

On Expressive Power of Looped Transformers
Due to permutation equivariance, we can assume without loss of generality that elements of E Ls are ordered, i.e. 21 >>>3 22> zN- Define ((6-)) == (8-(N-1)d ..... 8-d 1)T, which satisfy
1 d)² - 18 مد ا > 1, for any 2, =' E Tlyy with 2 # z'.
T
(34)
This mapping, "16 416-1z2. represents z in --_ -- base system, Then, we define sequence ID for 2 E L as:
N Σ 2n0
s(z) = u/6 2)²
(N-n)d
(35)
n=1
By Lemma A.6, there exists a Transformer block TF R5XN + R5xN with single head, head size 8 = 1, and width
size q = 3, and two affine linear maps 2/ IR + RS and C/2)) 12 + RR such that
2 (TF
2 = 5(2) 1T-
(36)
Furthermore, we have to add dummy indices to alleviate the approximation error caused by the looped architecture in Step
3. Recall that B E {0. 1...... 8-1 11dxNN represents the coordinates of the inputs. Let Z E {0. 10.1.....6 8-1 1xxx denote the ordered coordinates of B where the tokens are ordered by their token IDs, i.e., (u/6-1) ZB), (4/6-1) ZB) x. in other words, u (6--13 ZB. Recall that we consider only the case where all N input tokens distinct. By redefining the sequence ID of Eq. 35 for B instead of z. sequence IDs in 8 g-d- base can be rewritten in the 8- base system as follows:
s(B) = u ((--a)( a5 13(2(161)) ZB²
(37)
N
ZBB) in 8- (((N-n)d+((-s))
(38)
1=1 x=1
Then, define extended sequence IDs as:
Svalid (B) 2ss (ZB)a.
(39)
and the dummy sequence IDs as:
SJummy
(B)
Svatia (B) b.
(40)
for be E 18-1,8-11 +
1,..,,
28-1 - 1}. Then, define each set as follows:
Syalid == (Svatial (B) : B E {0, 10,1,...,, 5-1 - 13dxN },
Summmy
{ suummy
(B) 13 {0, 1,...,8-1 11 MxxN be 16- 28- 1}}.
+1,
Recalling that (ZB),,N E .0...... 8-1 13. we observe that
Syated O Saummy 0,
(41)
Syana U Saummyy {0, 1,...,2.. 28 Nd 1}.
(42)
We define the input cubes for each valid sequence ID s E Svalid as follows:
Qs = x 1],1 X E QB, s.t. s 8BS(B)}
(43)
Analogous to Eq. 30, we have
if Bf E (1,2...,, -1.
EES
max X'EQ...E.....
X - X' <
VNd
if B.,N
(44)
To implement this, we slightly modified (F(2), By Corollary A.7, there exists a Transformer block (2) R8XN - R8XN with two heads, head size s = 1, and width size q 5, and two affine linear maps e222) IR2 - R and c(2) R + R.t.
c (TF(2) oN (020) ([
TT
1) = (2s(z) Za.v) I/.
(45)
Zd.
14

---

On Expressive Power of Looped Transformers
Step 3. Token-wise Function Value Mapping. In Steps and 2, the network receives token ID and extended sequence ID input for each token, collectively forming contextual token ID. With 20 28-(N+1)d loops, the network approximately maps contextual token IDs to the output token embeddings of the target function.
To construct contextual token IDs, we define bijective afline linear mapping 2) N22 N follows:
c (z, = 220 Nd + 8,
(46)
where 2 represents a token ID, defined Eq. 26, and 8 represents an sequence ID. Recall that & {0, (0.1..... 8- 3- sequence IDs are less than 28 Na so informally, it's as if we are adding another digit, 2, the most significant digit in based system. Define the set of contextual token IDs as:
1) and
<3,) :2 20..,,....,, 88...,.--111 13, Svalid j
{ (2,5):
Kvalid
(47)
and dummy contextual token ID
{ :2E E 10, a-- 13. Sdummy
K-uummy
(48)
From Eq. 41 and Eq. 42, the following holds:
Kvvaid Kdummy = 0,
(49)
IC == Kvalid U Kdummy {0, 1, 1,..,,,,8" 28 28-1N+114 -
(50)
We now define the target output embedding for each ID. Let ys Id denote the contextual token embedding corresponding each contextual token ID, defined as follows:
f(XB).n lin interp (Yuear (N.)Ca)))) Qnear + (ro))) k near (ks, Kvatia), 8-1)
co (Zn, (BB) k
for ks Kvalid for E Kduumyy
yk
(51)
where the nearest functions defined
near+ (a
arg min |a b/, near mear"a,, (a = arg min |a b/, beS
(52)
b6S
and function lin interp is defined by
li interp(a,,,,, b,t, = a (b a).
(53)
n
The illustration of y is shown Fig.
Thanks to our design of K,, the error between neighboring contextual token embeddings can be bounded follows. There types of error: the variation induced by contextual perturbation and the variation induced by token perturbation, both. The examples of each pattern are shown in Fig. 2, such that
(1) I draft papers. : I write papers.
(perturbation of token)
(2) He writes papers. Mozart writes music. (perturbation of context)
(3) He drinks coffee.
He drinks coffee.
(perturbation of both token and context)
Recall that there are two types of adjacency that generally have small errors, with a few instances causing large errors when 'carryover' occurs, as stated in Eq. 30 and Eq. 44. The key point our design IC is that when a large variation occurs, linear interpolation inevitably takes place to smooth out the steep changes between adjacent indices.
15

---

On Expressive Power of Looped Transformers
Thus for token ID in Eq. 30, if small variation of token ID, with same context, in input space, 8vd. occurs, the emror c llyk yk- in the output contextual token embedding can be bounded by the modulus of token continuity as
ek (óvd).
(54)
In contrast, if large variation in token input space, vd, occurs in the token input space, the error can be bounded using linear interpolation with 8-1 intermediate points
ok S 8Wlt (Vd).
(55)
The same holds for sequence IDs in Eq. 44. That is, since the variation in context is bounded by sequence variation, the difference in adjacent contextual token IDs caused by perturbations in context is bounded as
ek S (dv (öv
or ek S bu
cont (S (VNd).
(56)
for k = O......,A K 1.
Since that
w tok (n S n wf ui tok (t)
(57)
for any E andd E [0, oo), with follows that (note that it holds with opposite inequality due to d < 1)
8w (VNd) S (6 and 8wggont (v Nd) S (6
(58)
Considering the maximum difference when both token and context perturbations occur, we have, with the triangle inequality,
max yy - -1l- (wi)t (Vd) wi (VNd)) S wii (5Vd) + wf"h "o (5VNd).
(59)
Generally, the following inequality holds, for any vector Rd
max [ S ||0||| p S d max /m/l-
(60)
Substituting x = yk- yk-1 into the above inequality results in
max (3 yk-11)l
S 13/k yk-1llp-
(61)
By Lemma 4.1, there exists a feed-forward layer FF R12d c( R + gl2d and cb3) gl2d + Rd such that
R12d of width size 18d and two affine linear maps
(c o (id
(K-1) (l) yk: FF(3) FF(3))
max (2 KeR
ye-1))bb
(62)
for i = 1,2,..... d and h: = 0, 0, K - 1, where K = 28-+41+
Let yb E IRd be defined as ÿJ = (3) o (id + FF( (K-1) ((3) (k). Then we have
1 yb)a s max (y -1)l
(63)
Kek
< max ek
| y
because of Eq. 61
(64)
< jotk (8vd) d) + w/
cont (ov because of Eq. 59,
(65)
for 1,2,..., d and h: 0, 0,.....,
16

---

On Expressive Power of Looped Transformers
Concatenated into Single Transformer Layer
Define the input space for each step
(0)
E RIXX
X E R
x E R
X (3) R2dxx
(66)
where ( act as counter. Define Attn = 5 via
1117499xN p(1T419)x with two heads, head size 1, and width size
x
01XN O N A 012dxN
X
Attn
(67)
(
where 2t denote the self-attention
layer of TF(2) and mMX denote m x N zero matrix. Let
2u
R,
21 E R
3' R
denote the token-wise input space. Define FF:
with the impulse function in Proposition A.4, as: 11744 1774+
30
(m1)
æ1
FF
(2 (a:2) + impulso (8-1 (f' (f'( mo) FF (x²) + impulse 1NN) (x2), In)
(68)
222
U3
where (2 denotes the feed-forward layer of TF(2), two linear maps L' Rdd R and L" IR. - I2ddd are defined
respectively as follows:
('(x) -( ([
C (x)
1)
(69)
(x)x) = 2 (c)) ( (x2)))
(70)
and impulse refers to the dimension-wiz wise application of impulse function. Note that serves the role of counter.
Each step should always be zero or set to an appropriate initial value the beginning the corresponding loop iteration. the bias term causes it to deviate from this value before iteration starts, the offset can be subtracted in advance to compensate.
As shown in Proposition A.4, the impulse function requires 2 ReLU functions per dimension and 2 ReLU functions for the corresponding loop iteration. Since the total dimension of and is 12d + 8, this results in an additional width size of 24d + 16 + 4 = 24d + 20. Since the width size is 7d for FF(1), (), 5 for T2( and 18d for FF(3), resulting in a total width size of 49d + 25.
Define two affine linear maps L1 RRdd_ - R17d+9 and L2 RI74+9 - R via
(x) (0, -( (x). (x)) 012448 T)", L2 x2, xs) C/ (xx)).
(71)
Thus, the network, consisting of three steps, is defined as:
f(x) L2 TF (x)
(72)
where 5-1 +N N 25- (N++)) and TF. (1 (17d49)) consists of Attn and FF.
Deriving Approximation Error
Generally, the following inequality holds.
12
n
Zz for x; 2 0 and p 2 1.
(73)
Substituting x; = [1 f(X)lpp into the above inequality results in
|| - fller = (111)(0)
1/p s -11x())
(74)
17

---

On Expressive Power of Looped Transformers
Also, generally, the following inequality holds, for IRm
max - - ||>|| S m max |zy].
(75)
With the triangle inequality, we can bound the approximation error as
I1 - fl - / [1((x) ((x)
because of Eq. 74
(76)
/ [1((x)- -(x))lpxxx + / [1f(X) (X)))xdx
(77)
S (Nd)² max K'EK.
Lkk- yel+ / [1f(X) f(X)lxdx + ((N2/88//)) because Eq. 75
(78)
(Nd)² max
Ük- yel wy(bv Nd) + O(N2/P84/P) because of Eq. 22
(79)
KEE
(Nd)
(w) + w wf + Nd) ( because of Eq. 65.
(80)
Then, 8 is expressed in terms of the number of loops
8- N + 28 (N+1)e
+> 8-1 1 20 (N+1)d -
(81)
+ 5- 25 (N+1)d 2r-N
(82)
-1/ -11((N+1)441)
N
# <
(83)
2
Thus, we have completed the proof of Theorem A.1.
A.2. Proof of Theorem 3.6
Extending the construction in Theorem A. we then account for the boundedness of the weights and bit complexity.
Proof. Due to the use of bounded weights to approximate discontinuous functions, there inevitably exist regions where quantization errors arise in Step 1 of our construction. We define these regions, for < 6,
8-1
2(10, 1jaaN 5),,) x E (0, 1J, Ji,n such that Xi, E U (kō e, kō)
(84)
k=1
That is, Ω consists of all inputs for which at least one coordinate X lies within an c-neighborhood of quantization discontinuity point kō for some k E (1,....8-11.
Outside the region St the quantization function is piecewise constant and can be precisely approximated using bounded weights. By the construction Lemma A.5, the maximum magnitude of weights required is proportional 1/c. We now estimate the Lebesgue measure of Ω. For each coordinate (i, n)., the
8-1
U (k: e, k:) A=1
a union of 8-1 intervals, each of length e, so the total measure of this set is èe, Since the input X has d x N coordinates, and since we consider the event that at least one coordinate lies in bad region, we apply the union bound for measures to obtain meas(!?) < dNV x 8-'e. Substituting the bound on meas(S)), and using the fact that the maximum magnitude of the weights M satisfies M = 1/c, we have mcas([2) S dN x 3-1M-1 Consequently, the I--moorm of the quantization error is bounded by o ((N)))-'dN)11 1/p)
When replacing hardmax with softmax, it is required that the error in step 2 remains sufficiently small so that it does not affect step 3. Specifically, step function is used, in Lemma 4.1 for step 3, map the index, defined for > 0
if x 2 0, if f C,
step, (z)
(85)
18

---

On Expressive Power of Looped Transformers
The construction of Lemma 4.1 I this function for indices & K. obtained from step step, (k- for 0, 0.1... K 1. Thus, the error in step 2 does not affect step if the perturbed indices, denoted by ki satisfies
step, (k + -11 i) step, (k- i) for all i 0, 1 1....A. K 1
(86)
To estimate the error introduced by replacing hardmax with softmax in step 2, we revisit the construction of Lemma A.6. Specifically, we extract and reformulate the key components necessary for this estimation. In particular, we consider a simplified form of the attention computation in Eq. 135, denoted by 9n IR x RN -> RR which is defined as
gH a) = Varg max; a, When hardmax in Eq. 135 is replaced with softmax, the function can be expressed as
(87)
exp(Aa;) 8:: exp(Xa;)
gs(u, a, 1) Σ
(88)
where A > Note that limx->oe (s a, A) = gu (v, a). According to the construction of Lemma A.6, such as Eq. 147, the domains, denoted by Lo and Asg are restricted on
L = {z E {0, 1, 10,,,,..,,,-- 5-d-3 11 2 mj for all i j},
(89)
As = {a E R a, E 8 - 1} or a, <0, O, if a, a, 2 0, then a; * a, for all i + j. Bi s.t. aj > 0).
(90)
We impose the following two additional assumptions on Ls and E As
arg max ai ISSS
Un
max U/-
(91)
ISVSN
Under these assumptions, for any finite we have
< gs( g2(0,a) a)
(92)
and
exp(daa) Σ/=1 exp(da;)
gs( a, 1)
Vns
(93)
Thus we have
gn(v,a) a) - gs(v, a, Un gs(v, a. A)
(94)
Σn= exp()a,) exp(>an)
Un
(95)
E exp(>a;)
Ejen exp(>a,)
Σ itn exp(Aa,) + exp(Aaa) Un
(96)
Σ jonn exp(>aj)
exp(Nam) Un
(97)
Un Σ exp (X(a) - -an) s+n
(98)
S N exp(->)
because a, an -1
(99)
< 8 ' exp(-A).
(100)
Thus, if we aim to bound the enror within e > 0, À must satisfies
8-dN
8 N exp(-A) < A log
(101)
19

---

On Expressive Power of Looped Transformers
From Eq. 46, the error of contextual ID 1s can be bounded in terms of the emror of IH (v, a) (s( a, 2) as:
k < (u(s-a))" (elx) S eN--
(102)
where 41 = (6-(N-1))) 11dd 8- 1)T and I E IR denotes the all-ones vector.
Since Eq. 86 holds 0 < k- - k < 1, it follows that
log (6--NN) (NN-(N-11)) = -No- 1,
(103)
and Eq. 86 holds. Thus, the bit complexity of À can be bounded by
0 log log(6-NN
(104)
while ensuring that no error occurs in step 3 when using the softmax function instead of the hardmax function.
The bit complexity at each step of the computation can be analyzed as follows. In Step 1 and Step 2, the bit complexity is bounded by Olog )) reflecting the cost of maintaining precision within error ó. contrast, Step incurs significantly higher cost, with bit complexity bounded by (228 +1)d),, due to the need to evaluate higher-order terms accurately,
As result, the overall bit complexity of the Looped Transformer is dominated by Step 3 and can be bounded by
Ommax(log log(6 0 Na N2), 2, 20- = O(8 100-00008110100 (N((+))d
(ō + 0).
(105)
With Theorem A.I, the proof of Theorem 3.6 is completed.
A.3. Approximating Discontinuous Piecewise Functions
We define three utility functions using ReLU activations. Since the target function is discontinuous, there negligible 'trifling regions' introduced by the bounded weights of the networks.
Proposition A.2 (Rectangular function). Given 1 12 E R. there exist four ReLU functions that approximate the following function, denote by rect, IR + for 0 < < t2 - t1. such that:
[1
if E tt - -c], i& (-00, 41 - -e] [|2, (c),
rect (x))
(106)
which is represented by
rect(t..2..)) (x) OR (=fftte)
oR (=ي + +oR (=ett
RR( - - 1.
(107)
Note that maximum magnitude of the weights is 1/c and 'trifling regions' are (t - (tt) 11,----11 Proposition A.3 (Step function). There exist four ReL U functions that can approximate the following function, denote by step: R + R, for > O, such that:
step, step,(x)= (x)=
(1 0
if 2
(108)
ifx < 0 -
which is represented via
step.(x) = RR + 1) on (2)
(109)
Proposition A.4 (Impulse function). Given 0 E N, there exist four ReLU functions that can approximate the following function, denote by impulse, IR x N R x E [-M, and E N, such that:
if = 0, otherwise.
impulsee(ar,
(110)
0
which is represented via
impulsc,(r, :0=n( + 2t 8 1/2)) 2Mon(t - 0 1/2))
(111)
- - 2M 8- 1/2)) + 2Mon(t - 0 - 1/2).
20

---

On Expressive Power of Looped Transformers
(kō + k),
k
hx
(kō - E, 0]
0
k6
kō -
k6
kō
kō +
Figure 3. An illustration of hz())).
A.4. Step 1. Token-wise Quantization
We aim to construct quantization function g [0, 1] - (o........ 6--jd, for e < 6, for each dimension as
g(z) (g(zz). g(()))
where g(æ) = k. if [k6, (I. + 1)8 for k 0, 0,...,ot1-- - 1. (112)
This function g TR + IR can be expressed as
g(x) rect)(,((44)))))0)))
(113)
for any n E and R. The illustration of h2 (x: k r0c)( (2) is shown Fig The key idea that (x) can be represented with a single function the form h h(kz k2, k). Lemma A.5 implement h(kz, 12, k) with a feew-fo layer and perform the summation through skip connection.
Lemma A.5. Given any 8-1 E N and x E Rd, there exists a feed-for forward layer FF Rod + Rd of width size q = 7d with the maximun magnitude of the weights 1/e, and two affine linear maps L1 Rd + Rddd and L2 : Rdd + Rd s.t.
(C2 (id FF) 06 0 (z)))) if E [k6, (k: + 1)8 e). for k = 0, 0,..,, 8-1
(114)
for any i = 11,,....,
Proof. On the basis of Proposition A.2, define function (x(2)) k- rect+(x) via
(x) = ORR ( kd + +))
RR( ( (tr - ko)) - -o ( kg 1))
(115)
-oR( ( (-x + kō + 1 --6)) k,
which satisfies
k
if kō, (k + 1)8 -c]
h4 (x) =
(116)
if E (-o0, kō - o] U [(A: + 1)0, 20),
For any x E [kō, (k + 1)5 <) for k = 0, 1....
5 .
1, we have
1-1
Σ ==0
hx(x) = (x(x)=) =k.
(117)
Define function h: 3 - R to represent hy via
h(kx,, 1kkx,k2,4 k2,K) k) on( tiz - 3'8 + k)
166 F kr 1-d
+ +
on
(118)
onl
kz + k + k 12-)) - -on(k) (x))))
21

---

On Expressive Power of Looped Transformers
Define &
En(ee)
& kx, 12, k x,
(119)
==
Then, construct feed-forward layer FF R R with skip connection such that
(id + FF)+FF)(AA (Ex) = (id + FF) (krr h2 k, x; Σ n((()))))
(120)
(kk + 1)x, (k+ + 1)?, ++1, 1, x,,
hn(())
(121)
==0
Ek+1-
(122)
via
kx
kx
k2
k2
(id FF)
k
k
x
x
k 1-1
(2)
R 1
(x(
2
100
0
0
0
11111
0
20000
kr
TUTO
0
0
0
0
1.2
10000
0
0
0
OR
1000
k
(123)
0
0
0
0
0
x
0
0
-1
0
h2(1))
0
0
kw
x
k-2
2k
k
(124)
Σ == (1(
k-1
le (3)
fx +
1.2 + 21 + 1
ke +
(125)
Ei ha(x))
+ ha (xr)
1)x
+ 1)2
++1
(126)
E h4(x))
Then, define two affine linear maps L1 RR. - R and L2 R R by
(z) = 20,0,0,2,00)
22. = 23.
(127)
Thus, have
L2 0 (id + FF) o6 1-13 o fz(x) L2 (id + FF) of8 '-11(Eo)
(128)
L2(E6-1)
(129)
1-1
hz(
(130)
1=0
For d-dimensional inputs, we need d-times more parameters.
22

---

On Expressive Power of Looped Transformers
A.5. Step 2. Contextual Mapping
The network takes token IDs imputs, denoted by z E {0, 8.....,844111 distinct. The network maps token IDs into sequence ID using the inner product with the vector R defined as = (8- (N-1)d 8- (N 22d ..... 8 , 1) i.e.
13% We only consider cases where all token IDs
s(2) == T
(131)
Due to permutation equivariance, we can assume without loss of generality that elements of z E Ls are ordered, i.e., 21
22 zv. Then the map s satisfies
u 1=== u-=uul 2' if #
(132)
In other words, represent in 8 8-d-base system. The network computes u in the form of Σ 8- 8-(N-)) The network computes ( = Σi= 8-(k-x)dd 2; in each loop, and after N-loops, outputs s(N) = auuz z. To implement this, the self-attention layer selects 2 in the k-th loop iteration. We design the key, query, and value weights to select the maximum token ID. The feed-forward layer post-processes the token ID such way that it selected, then is replaced with negative value prevent selection in subsequent iterations, i.e., the post-processed token IDs for the k-th loop are
0 if isk
(k)
s.l.
otherwise.
We focus on self-attention layers that employ the hardmax function.
Lemma A.6. Consider the set of distinct indices corresponding to the d-dimensional 6-discretized regions of N tokens, i.e.
Lg == (= E (0,1,..., 13 2, + 2j for all + j}.
(133)
There exists function s N - TR composed of Transformer block TF. RRON + R with the hardmax function, single head, head size = 1, and width size
3, and two affine linear maps L1 R - R and L2 RR. - R, such that
L TF ==(x))= (u=z) 2) =
for any E LLs, where u = (N-(N-1)d 8- 2N- 1).
Proof. Due permutation equivariance, we assume without loss of generality that elements of E L,g ordered, iee.
NN. Define u E N as = (8- (N-1) ......-..
8-d 1) which satisfy
21
T z u 2 1, x + x' for any 2, E Lg.
(134)
We construct Transformer block TF: p85x
p with single head and head size 1 such that, for any z E Ls.
T T
TEON
IN
(135)
Og
0g
(u 2) I
where 02, I E N denote the all-zero and all-one vectors, respectively. For z E Lgg, define two series for k = 0, ... N as:
k for 1, I,..,,,, e R
(1) S.
(136)
z;
otherwise,
8 (k)
8 (k- (R-1)d 2; E R.
(137)
23

---

On Expressive Power of Looped Transformers
While (4)) is not uniquely determined, any vector that satisfies the conditions is accepted 2()) The series g(k) satisfies
R+1
(k+1)
(K+ 1)d 2i
(138)
ó (k+ (k+1 (++1=1)d + 2
(139)
1=1
8-d 8- (1--)) 2x+1
(140)
8 s 2k+1,
(141)
for k 0, ... 1. Note that (NN u² z. Define single-head self-attention Attn ROXN R such that
a N
Attn
O
(142)
(Varg az)
IN
T
ON
where RV,, and RN denotes arbitrary vectors, via the weight parameters
0
Wo
0
Wy [10000) Wx [ 1
1000
We
0 (143)
Define FF R > R of width size q 3 via:
J1
0
J22
M
0
422
FF
J3
0
0
OR
0
0
233
+
(144)
J44
0
-1
0
0
274
J55
0
1
8-d
MoR(12 + E)
(145)
-oR(xs)
(8 ()))((s) + (n(
where 0 < < 1 and M > 5 1_1
For 22 E 10.1......6
8-d d 1) U (xr 50 and E {0, 1, 1,.,1 1} with 2 we have
Rg > D2. 2 <<< if 24 F2.
= L1
52- MσR(Zz I4 +
(146)
This post-processes the token ID in such a way that if it is selected, then it is replaced with negative value i.e.
0 if isk+1, 2 = 2; otherwise,
(k) Mon(2.1)
2
2 + s.I.
(147)
2 (k+1)
for 1 1...NN.
24

---

On Expressive Power of Looped Transformers
We confirm that the Transformer block TF R + RRxx composed of Attn and FF, satisfies, for k -,,..., N
I (I(h)) 1 (s(R))T
TT
(2()))
TF
(id FF) o (id Aut))
1 (s()))
(148)
T
(x
(id + FF)
I
(149)
2k+1 IV (s(R))
T
0
(2
(M 2k 1 + x1
0
(150)
Zk+1 11
-x+1 1
(s()))
(8-a (
(1((())
1 o F (x
T
(2(K+1))
1
because Eq. 147
(151)
0
(s(k)
1
1 N
T
(2((+1))T
because of Eq. 141.
(152)
O
(s(k+)))
1
Define two affine linear maps L1 R - IR5 and L2 respectively. Thus, we have
R + IR via (x) = (x, x, 1,0,0) and (z 22; 23. F4, 26) = 36
L2 TFON L1(xT)= =uu 2) (u)2) IN-
Combining the sequence ID with token ID, the network computes contextual mapping. Corollary A.7. There exists Transformer block TF2 188.N + RSXN with the hardmax functiionn two heads, head size = 1, and width size 5, and two affine linear maps Li R²a + R and L2 IR - IR s.f.
C2 (2(2)
([
Zd.: ))
TT
Li
(2u Za,N) IN for any E L.g.
Proof. Define self-attention Attn 88N 88N such that
(Varg max, a,
1
Attn
(153)
0
Za
0
Za.N I
0 N
25

---

On Expressive Power of Looped Transformers
where v, R-V, and R- denotes arbitrary vectors, via the weight parameters
0
0
1 [1
0], W [ 1 0], W [
0 10100
0
(154)
and
WO
W W = 0
0 0100],
W 10010 ] (155)
where denotes sequence of zeros.
Define FF: R - IR of width size 5 via:
0
21
0
0
0
31
0000
0
0
42
0
100
12
0
43
0
1000
0
J33
FF
27
24
0
0
8-d
0
-1
OR
0
F4
175
0
0
0
0
J6
0
0
0
-1
277
0
0
-1
.tS
(156)
0 Mon(xz
n4 +
0
-o((0)
(8-4 d ()x)( tOR(TA)
(157)
0
-oR(2)) oR(27) (R(TS)
where 0 < 1 and M 5-"-1 Then, define two affine linear maps L1 IR.. - R and L2 R -> R respectively:
(1(11, (I1, 01
12(11,02, 24, Fs. Ne. 27, rs) 2xg 26
(158)
From Lemma A.6, the corollary holds for this construction.
A.6. Step 3. Function Value Mapping with Bit Extraction
We employ a bit extraction technique (Bartlett et al., 1998). as used (Zhang et al., 2023) for weight-tied ReLU networks, to approximately memorize the label set. Given E N input indices & = 1, 1, ., with associated values y1, V2, VK E R. the network approximately memorizes the differences yi V:-1 using their base-2 representation, Since the binary
26

---

On Expressive Power of Looped Transformers
representation is limited (0. 1), the differences y; Vi-1 must be rescaled by factor max; 1y V/-1
a. Vi
(159)
where [a] == max(n: S 2, n E Z). Then, the difference b; = a, a;-1 satisfies E {,,,,, 1 and can be represented
using two binary values C,,d, {0, follows:
by
dis
(160)
and we have
0k = ag +
b;
ag
+
Σ 1=0
G
Σ 1=0
dj for == =0,,,....
(161)
1=0
dj can be realized by composition of single feed-forward layer. Thus, the network can approximate yi using ea,, denoted as Dr. with the following accuracy:
Lemma A.9 and Lemma 4.1 show that Σ -i=o- and E Ei=o
Us - vol = |c
yi
c
yi = <ll V/ - s
(162)
For a d-dimensional input-output pair, we construct the networks for each dimension i.e.
ÿ
g*)
(163)
The basic strategy of our lemma and proof follows Lemma D.I from Zhang et al. (2023), as shown below and Proposition 3.2. However, their result cannot be directly applied here, as it requires depth-2 networks.
Proposition A.8 (Lemma D.1 in Zhang et al. (2023)). Given any r E N+, there exists a function FF : RR3d - RR with width and depth 2, utilizing two affine linear maps Li IR² + IR and L2 R + RR such that for any 0,, 02,000,,0.. 0, {0, 1}. the following holds:
C2 OF Ci, bin 0.0,0.00...... 0,) ..)))))))))) Σe 02 for k 1,2,...,
(164)
where bin 0.0,02-0-02... 02. 0, denote the binary representation of 0 Ei=1 0,2
We found that the loop unrolling technique allows us to reduce the number of layers from 2 to 1 by replacing k-+ ReI.U(ReT.U(<>*)) with (a*+1,
ReLU(s** x**). Although our method makes the weights dependent on 01, 02. ... Or. this does not present an issue for our construction in function approximation. Specifically, 02, .......9. O,
fixed for each target function, and the role of the network is to learn the weights tailored to that single function.
Lemma A.9. Given 01, .2..0....,, 0, {0, 1) for some N+ there exists a feed-forward layer FF R R of width size 9 and two affine linear maps L1 R IR. and La IR6 TR. s.t.
C2 o (id (id+ (k) Σ
for k 1,2,....
(165)
where the bit complexity is bounded by O(r).
Proof. From Proposition A.3, we have a function (x), for e > 0, defined by
step,(x) on + 1) - (
(166)
as shown in Figure 4, and it satisfies
1 if 2 0, 0 if < < e.
(x)p,
(167)
27

---

On Expressive Power of Looped Transformers
(0, 1)
step;
(-E, 0)
0
-2
-1
1
Figure 4. An illustration of step,(r]
Define B1 for 0, 0,11. ,,....r as
B = bin 0.0 fr.
(168)
where bin 0.0, .-. 0, denote the binary representation of 0 1,2 0,2 and go = 0. If we set < 2",, it follows that
0: = step, (bin 0.0.... ...... 2 step, (3r - }
(169)
(170)
implying, for = 0, 0. 1,
Bt 2,31 A1
(171)
= 23: step, (3: - -12
(172)
For all / 1, ,r. since the product ry satisfies
my = max (0, + 1}
(173)
= oR( + 1),
(174)
for x. y E (0, 1}, follows that
0
01
(175)
1=k+1
Σ step, (k 1)
(176)
Σ on 0: step, -1) -1)
(177)
Σ oR stepp (B: - 1)) + step, (k- 11) -1 -1).
(178)
To compute the right-hand side, we require two nested ReLU functions. By employing loop unrolling, we precompute T(3: }) and [(k 1) the previous iterations, reducing the requirement to a single layer.
Define & for = 0, 1, 1,..., 1
E: (k: 16-11 1, Br, Bi+1, step, (3: } step. (k- - 1), sum(t))
(179)
where sum(l) Σ on (topp - }) + stop. i) - -).
Note that we have B441 the 1-th loop to precompute step, (31+1 - 1)) and step, ((I: - (1 + 1)) for the (1 + 1)-th loop.
28

---

On Expressive Power of Looped Transformers
Define FF: R IRo with width size of 9 such that
k -
B1
Bi+1
(id + (E) = (id FF)
(180)
step. (3
2
step. (k
sum(l)
k
O
-1
0
B141 step. (3t
D
-1
0
-1
0
-1
-1
step, (A
0
0
00
1.
sum(1)
0
0
0
1
1-
0
0
0
B1
0
1/e
0
-1/(2)) +1
B[41
dR
0
1/e
0
step. (B) - =
+
(2ee2)
(181)
0
0
step, (A: 1)
1/e
0
0
0
1/E+1
sum(()
1/c
0
0
0
- 1/6
0
0
1
1
0
-1
-1
k -
on(B1) (stepp (B: 11)
31
rR((142) (on( 1/2 + onl (1111122))
B1+1
(stepp, (3) - 11) + Rnl 1421-122 + 1) onl (
(182)
step, (Bt - 2
(ntespp (k- 1)) an( k- k-(41) 1) rzl (((l+1)
step. (k- sum(1)
OR step, (k- 1) + step, (3: )
(1+1)
k (1 + 1)
23, (3: 1)
3++1
23 step, (31+1
B142
E111
(183)
step, (B1+1 2) step, (k- (1 1)) sum(1 +
step, (51+1 - 2)
step. ((k - (1 + 1))
sum(1 + 1)
Define L1 R R and Le R + R via
(k) = (k, Bo, Br, 0,0,0) = Eo. [2 24, 24, = Z6.
(184)
respectively. The lemma holds for this construction.
Then, we prove Lemma 4.1 with Lemma A.9.
Lemma 4.1. Given E Rdd for k 0, 1,...., K - with
18 yk U1-1)a Se for I,....d,
there exists a feed-forward layer FF R 12d + 12d with width size of 18d, and two affine linear maps L R + 12d
and C2 : 2d + Rd s.f.
(CC2 (id FF)O(K-)) (k))) - yk),1 S sn
(18)
for 1,... d and &s 0, I,..,, K 1.
29

---

On Expressive Power of Looped Transformers
Proof. We prove this for the where d 1, considering yı E IR for 0,1, K Define
a2
Mk
for k ....,, 1,
(185)
where [a]] max n E x, ] and set
by = ag ax-1 for 1,,,...., K 1.
(186)
Since by {-1,0,1 there exist O E {0, l and dy E {0,1} such that
hy = Ch dk
for h: = 1,2...., K K-1. 1.
(187)
Thus, we have
k
ak = do +
Σ
C4
d
for any & E (1, K --11 1}
(188)
1=1
1=1
From Lemma A.9, there exist FF())FFF()
FF
R + R of width size 9 and affine linear maps C2 R - and
(c)) ((() R - R s.t.
C2 (id + FF(0))
( (m-1) (K) Σ cr. C22 (id + +F((d) (
o(m-1) - (1) Σ di.
(189)
for k = 0,0,, 1...... 1. Then, define FF R212 IR 12 of width size 18, for az y E R
x 1) = FF (c)) (z) ].
FF (c) (x)
(190)
Define L1 IR > 12 and L2 1212_ TR.
c (x)
L2 ([ x 1) = ca -2xx)= (z)
(z(z) =
(191)
2 (x)
(x)
We can confirm that
L2 o (id + FF) (k)
(192)
L2 (id + (K-) (1
c (k) c1500 (A:)
(193)
Co ([
(id + FF (F (K--)) o c(r) (k)
( (K( c)d (k) 1)
(194)
(id FF
6 <ao +
Σ
Σ di)
=ea.
(195)
i=1
1=1
Thus we have
LC2 o (id + FF) (K- o Lik) yk =leas - yk S E.
(196)
For d-dimensional inputs, we need d-times more parameters.
30

---

On Expressive Power of Looped Transformers
B. Role of Teeendep Scaling Parameters
show that time-dependent scaling parameters overcome the limitations inherent to the looped architecture and eliminate the dependence of the modulus of continuity. We use the architecture defined in Section 4
(F - (1() ④ (F(x)) for the 1-th loops,
(197)
The following lemma demonstrates that time-dependent scaling parameters can exactly map indices to output vectors.
Theorem 4.2. Given yh 12d for k 0. ..... K size of 6d. (t) E R4d for t I..... K and two affine linear maps L1 R. - 1R and C2 R - R2 s.t.
there exists a feed-forward layer FF R4d R2dd with width
I(C2 (id+ @ FF)o(K-)) (1()) yk:); = 0,
for = I,....d and lo 0, 1,..., K- 1
Proof. We consider the case when 1, where yn R for K 0,1,..... 10.1....,4 1. We update ye as follows:
+ Uk te,
(198)
where is chosen such that none of the yt values are zero.
Next, define n!)) R2 as:
(11 = (0, 1, 91
y11 T
for 1,2,.... K
(199)
V1-1
By Proposition A.4, we have, a E [-M, .] and t E N,
impulse (x, 1) = oR( + 2M (1 - 2MoR(( + 1/2) - -on(x + 2M(1 1/2)) + 2Monn 1/2)
(200)
0
x
ift= 0,
(201)
otherwise,
where M > max. yk.
Let k E {0, 1...., K- 1} be the input index that specifies which yk to extract. Define
s(1) Σ impulsep (v,, - i).
(202)
1=0
for = 0, 1....... K 1, which satisfies
s(K yk-
(203)
Define & E Rd via
& = ks 1 1, U1, s())) T
(204)
for = 0, 1,..,, K-1 1.
31

---

On Expressive Power of Looped Transformers
Then, define FF R² R. of width size y via:
(id + v()
& (É
-
0
0
0
0
0
-1
0
0
0
0
0
2M
0
M
(()
1
1
0
0
0
On
2M
0
+--1
MM
(205)
0
0
1
-1
-2M
2M
1
0
0
1/2
0
1
0
0
-1/2
0
-1
(R(M--1) -O(R-((-i
0
OR(M-I + 2M - 1) + 1/2))
(206)
g1
2MoRR( 11 (k-1) 1/2)
Mt
o 2M (k: - 1/2)) + 2MoR - 1/2)
0
0
1
-1
VI
1
④
(207)
21
Dc-1
31-1
51
impulse, ( k 1))
32-1
k
k
(208)
y1
x(1)
for ...,, K- 1. Thus we have
(id + - 1) o ....o (id + +1) O (E.) ÊK-1
(209)
Then, define two affine linear maps L1 IR + 4 and Le R + Ry
(x) = (k, ki yo. 0), f2(21, 23 x4) == 24 -
(210)
We can extend this to d-dimensional input by using times more parameters, by applying the above to each dimension
C. Details of Experiments
This appendix section provides additional details the experiments for each task.
C.1. Reasoning Tasks
C.I.1. PROBLEM SETTINGS
Longest Common Subsequence (LCS) is the longest common to given set of sequences. We use problems with input lengths of 60 and 100. Two sequences are sampled uniformly from the alphabet.
Edit Distance (ED) problem, also known Levenshtein distance, is find the minimum cost required to change sequence into the other. We adopted the problem setting and data generation approach from Feng et al. (2023), but applied larger input lengths. The costs for insertion, deletion, and replacement were set to 2, 2, and 3, respectively. They generate instances of the edit distance problem as shown Algorithm The first string is randomly selected, while the second is generated two ways: (1) a random string yielding a large edit distance, and (2) a corrupted copy of the first string, resulting in small edit distance.
32

---

On Expressive Power of Looped Transformers
Algorithm 1 ED Data Generation from Feng al. (2023)
1: Input: Length of the First String 7
2: Input: Alphabet
{a, b...2)
3: Output: Sequence 81, 82 Sample uniformly from (3,4..100) T Sample letters from V 81 Sample letters
uniformly from T' Sample p uniformly from [0, 1]
4: if 0.4 then
5: Sample I uniformly from (n 3, 2.... 2}
6: $2 < Sample 7 letters uniformly from T
7: else
8: while len(s2) not in [n 3, do
9:
f- 81 $2
10:
for i < 1 to do
11:
Sample p uniformly from {0, 1s..)e(((()) - 1}
12:
Sample uniformly from T
13:
Randomly conduct one of the followings: pop 82 [p], substitute s2[p] with 7, insert into 82 s2[p]
14:
end for end while
end if
Sudoku We use the dataset from Radeliffe (2020), which contains over million Sudoku puzzles. The puzzles are flattened, with representing blank grids. The input sequence is formatted as:
00000000000000000000000000000000
Countdown To generate the dataset, we randomly sampled the target and input numbers for each instance. Pairs that have no solution were excluded. For tokenization, we assigned unique tokens to each number and symbol. The target sequence represented as:
58
84
48
62
96
62
58
48
12
84
12
96
The model learns to predict the target sequence from the input:
58
84
48
62
96
0
0
0 0
C. 1. TRAINING AND TEST ACCURACY CORRELATION FOR EDIT DISTANCE
Given that our study focuses on function approximation capabilities, one might question whether it appropriate to rely on test evaluations, which are influenced by generalization. Here, we confirm that there is strong correlation between training and test results, validating this approach. Figure demonstrates strong positive correlation between training and test accuracy, enabling the evaluation of approximation power through test accuracy.
Train Accuracy Number of Loops
Test Accuracy Number of Loops
0.9
w/o time depnedncy
w/o time depnedncy
0.8
w/ timestep encoding
0.8
w/ timestep encoding
Accuracy
Test Accuracy
0.7
0.6
0.6
0.5
0.4
0.4
20
40
100
20
40
80
100
Number Loops
Number of Loops
Figure Training and test accuracy for the edit distance task with sequence length of 60.
33

---

On Expressive Power of Looped Transformers
C.1.3. MODEL AND TRAINING CONFIGURATION
We used Looped Transformers with attention heads and 256-dimensional embedding. The AdamW optimizer (Loshchilov & Hutter, 2018) was used with 31 = 0.9, B2 0.999, weight decay of 0.01, and linear learning rate decay scheduler starting at Ir 10 and ending at 0, with 5 warm-up steps. Training consisted of 50 epochs for reasoning, 200K steps for in-context learning, and 100K iterations for language modeling, using batch size of 64. For time-dependent models, (tt) and (c()) were initialized as zero and one vectors, respectively, following Peebles & Xie (2023). The input embeddings are added at each loop iteration. Furthermore, for Sudoku and in-context learning tasks, the output of each intermediate loop is incorporated into the loss as (Yang et al., 2023).
C.2. In-Context Learning
We followed the setting of Yang et al. (2024). The problem is to learn the function class from a given sequence composed of the pairs of input a, and output values (a,)). The input for model (ar)), .Tk [(xx) Fress), and model learns to predict (s)s) The model trained on (*xx and its performance is evaluated on (( using the squared error.
We use 4 decision trees with 20-dimensional inputs. Each function in this class is represented by full binary tree with 16 leaf nodes. Non-leaf nodes are associated with specific input coordinates, while leaf nodes are assigned target values. To evaluate f(z), the tree is traversed from the root, moving to the right if the coordinate value is positive and to the left otherwise. Inputs and leaf node values are sampled from N)O, and the coordinates for non-leaf nodes are drawn uniformly at random. Our training setup follows the approach of Yang et al. (2024). Following the curriculum training approach of Garg et al. (2022); Yang et al. (2024), we progressively increase the task dimensionality from to 20 in steps of 1 every 5000 steps, while the sequence length increases from 26 to 101 in increments of 5 over the same interval.
C.3. Language Modeling
Tokenization is performed using byte-pair encoding, following GPT-2 (Radford et al., 2019). The Transformer model is based on the GPT-2 decoder architecture (Radford et al., 2019). The baseline standard Transformer model consists of 6 layers, 8 attention heads, and an embedding size of 512. The Looped Transformer has a layer, 12 attention heads, and hidden dimension of 768, which were chosen match the parameter size of the baseline. We initialize y(t) zero vectors and (t) as one vector for dime-dependent models. The AdamW optimizer (Loshchilov & Hutter, 2018) is used with 31 = 0.9, 32 = 0.95, a weight decay of 1 x 10-4, and a learning rate schedule with 2000 warmup steps. The maximum learning rate is set to 2 x 10 and decays to x 10 using a cosine schedule. Training is conducted for 100k iterations with a batch size of 48 and a block size of 1024.
D. Disccusion
Multiple Layers A natural question is whether our analysis, which focuses on single-layer Looped Transformers, can be extended to multi-layer architectures. We restricted our analysis to a single layer in order to highlight a key strength of Looped maansoormers-namlyy their universality as function approximators even with just one layer. A more specific question is whether multi-layer Looped Transformers can overcome potential limitations inherent to the single-layer design. While it conceivable that deeper architectures but with fixed-depth feedforward layers may achieve better approximation accuracy, this remains an open question. The difficulty lies in the fact that such improvements are not captured in terms of asymptotic order but rather in constants, which are harder to analyze theoretically. For instance, if we allow a logarithmic number of layers depending on the desired approximation precision, then even our current construction may overcome the limitations of the looped architecture. However, this deviates from our main objective, which is to characterize the approximation rate solely in terms of the number of loops.
Additional Experiments To assess the model's sensitivity to input continuity, we designed a perturbed version of the WikiText- 103 dataset, where 10% of the tokens were randomly replaced. We trained models with and without timestep encoding and evaluated both their memorization performance and continuity behavior. Continuity was measured by applying small perturbations to the input and quantifying the change in output embeddings. The model with timestep encoding showed improved memorization (cross-entropy loss reduced from 4.32 to 4.18) and significant reduction in continuity coefficients (from 130.6 to 21.5). These results suggest that timestep encoding not only enhances stability under perturbations but also enables more faithful input-output mappings, thereby improving both robustness and learning efficiency.
34


Reviews:
[{"review_id": "evi01hWnRX", "reviewer": "Reviewer_s53h", "summary": "The authors conduct theoretical analysis on the expressive power and approximation rate of looped transformers, a variant of transformers with weight tying across layers and unbounded recurrence. The analysis led to some limitations of looped transformers on function approximation, especially in memorizing input-output pairs. The authors also propose a way to fix this by adding a scaling mechanism conditioned on the loop index, whose effectiveness is empirically validated.\n\n## update after rebuttal\nThe rebuttal helps complement the draft and addresses some of my concerns. I will maintain my original score, which leans positive overall.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> I believe deep equilibrium models are closely related to looped transformers, but the authors don't seem to discuss them.\n\nThank you for the suggestion. We agree that Deep Equilibrium Models are related to Looped Transformers. We will add the following sentence to the related work section:\n \n**line 77:** Deep equilibrium models  (Bai et al., 2019), which compute fixed points of iterative layers, are also related.\n\n> W1. While it's good to understand the expressive power of looped transformers, it's a bit unclear how/whether the ideal weight configurations could be learned through gradient-based optimization. Looped transformers are notorious hard to train stably, and there are various tricks that people have been using such as only computing the gradient at the last several recurrence steps, etc. It would be great if the analysis could provide insights into how to better train these models.\n\nThank you for the thoughtful comment. We did an additional analysis, specifically, we empirically observed that the trained models tend to exhibit smaller continuity coefficients with timestep encodings. Since smaller continuity coefficients (such as Lipschitz constants) are often associated with more stable training dynamics, our results suggest that the proposed approach may also contribute to improved trainability of Looped Transformers. We will include this point as an important direction for future work:\n\n**line 439:** Beyond expressivity, improving training stability and efficiency, especially for larger numbers of loop iterations, is an important direction for future work.\n\n**Reference:**  \nBai et al., \"Deep Equilibrium Models\", NeurIPS 2019.", "claims_and_evidence": "Most claims are theoretical and the correctness depends on the proof. Another claim on how adding timestep encoding should improve performance is empirically validated.", "methods_and_evaluation_criteria": "Yes the methods and evaluation make sense.", "theoretical_claims": "I could understand the high level proof strategies which make the conclusions plausible, but didn't check the proof correctness due to limited time and expertise in theoretical analysis.", "experimental_designs_or_analyses": "The experiments are on tasks repeated adopted in prior work, and the settings are sound.", "relation_to_broader_scientific_literature": "The paper mostly extends the current theoretical understanding of looped transformer on its parameter efficiency and approximation abilities, and how they compare with standard transformers.", "essential_references_not_discussed": "I believe deep equilibrium models are closely related to looped transformers, but the authors don't seem to discuss them.", "other_strengths_and_weaknesses": "While it's good to understand the expressive power of looped transformers, it's a bit unclear how/whether the ideal weight configurations could be learned through gradient-based optimization. Looped transformers are notorious hard to train stably, and there are various tricks that people have been using such as only computing the gradient at the last several recurrence steps, etc. It would be great if the analysis could provide insights into how to better train these models.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qKpiLsqPNv", "reviewer": "Reviewer_7dbK", "summary": "This paper studies Looped Transformers from a theoretical perspective. Its first main contribution is to prove a universal approximation property, showing that a wide class of sequence functions can be represented by Looped Transformers. The main technical contribution here, compared to existing universal approximation results for transformers, is to show that this can be achieved with coupled parameters, i.e., looping a single fixed transformer layers. A second contribution is to propose timestep encoding, which shows improved empirical performance on a few tasks.\n\n## update after rebuttal\n\nThe authors have promised to address the points I raised.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> W1. Unlike a lot of other literature, the present paper follows Yun et al in only studying functions on fixed-length inputs. This restriction can be fine, but I think this should be made more transparent earlier in the paper\n\nWe respectfully point out that the study of **function approximation** and memorization commonly assumes fixed-length inputs (Yun 2020, KIm 2023, Kajitsuka 2024, Jiang 2024). On the other hand, the prior work cited by the reviewer focuses on computational complexity rather than function approximation. Nevertheless, we agree that this distinction is important. Accordingly, we will revise the introduction to make this assumption more explicit.\n\n**line 45:** Our contributions are summarized as follows: We establish the approximation rate of Looped Transformers for **fixed-length** continuous sequence-to-sequence functions.\n\n> W2. Regarding clarity: I found the terminology surrounding \"sequence IDs\" and \"token IDs\" confusing in Definition 3.2. \n\nWe agree that the terminology in Definition 3.2 could be made clearer. We will explicitly clarify in the definition that the sequence IDs and token IDs are integers. Furthermore, we will explain the motivation: they are defined for a constructive proof to describe contextual mappings as Kim et al. (2023). Specifically, we will revise the definition and add the following explanation as:\n\n**line 175:** In our proofs, we rely on the *contextual token ID* to describe contextual mappings.  \n**Definition 3.2.** A *token ID* is a **unique integer** assigned to each token. A *sequence ID* uniquely identifies each sentence. A *contextual token ID* uniquely identifies a specific token within a specific sentence.  \nThis notion is defined in Kim et al. (2023), to which we refer for further details, for constructive proofs of contextual mappings. The actual construction of contextual token IDs may vary depending on the specific proof. In our case, we adopt a different construction from that of Kim et al. (2023).\n\n> Q1. In the hardmax operation, how are ties resolved? I.e., when multiple positions receive the maximum attention score.\n\n**Such cases result in errors, but their errors diminish as approximation improves.** As you rightly noted, when multiple positions have identical values, the hardmax operation cannot distinguish between them. In our proof, we explicitly evaluate the measure of such regions and include them in the error term, $\\mathcal{O}(\\delta^{d})$. As the discretization becomes finer ($\\delta\\to 0$), the proportion of these regions decreases and eventually becomes negligible. To clarify this point, we will add the following sentence:\n\n**line 329:** ..., where $\\mathcal{O}(\\delta^{d})$ arises from the case where identical tokens appear in sequences.\n\n> Q2. Line 155 (right column): where is bit complexity used? It is introduced but then doesn't appear to show up.\n\nThe bit complexity is part of the result stated in Theorem 3.6. However, we agree that its role is not clearly connected at the point of introduction. To improve clarity, we will revise the text (line 155) to explicitly indicate where the bit complexity appears in the subsequent analysis.\n\n**Before:** We evaluate bit complexity, which denotes the maximum number of bits of weights.  \n**After:** We evaluate bit complexity ... weights. **Our theoretical results show how the approximation rate depends on both the number of parameters and the bit complexity**.\n\n> C1. One issue is that the theoretical analysis uses the hardmax operation (line 106, left column), which is a nontrivial abstraction and may have expressiveness implications. I think this should be mentioned when describing the paper's overall contributions (abstract, introduction) as in other work making such assumptions.\n\nWhile theoretical analysis with the hardmax operation is a common approach in function approximation (Yun et al., 2020; Kim et al., 2023), we agree that it constitutes a strong assumption. To clarify this, we will revise as: \n\n**Before:** we use the hardmax function instead of the softmax function.  \n**After:** Our constructive proof **relies on the softmax function to approximate the hardmax function** as in previous work (Yun et al., 2020; Kim et al., 2023).\n\n> O1. In the proof, I found it hard to track when p-norms are applied to individual vectors or across the entire input space. \n\nWe appreciate the comment. We will clarify, in the camera-ready version, the distinction between vector $p$-norms and function $L^p$-norms, using $|f|_{L^p}$ for the latter, as defined in Equation (2).\n\n**Reference:**  \nYun et al., \"Are Transformers universal approximators of sequence-to-sequence functions?\", ICLR 2020.  \nKim et al., \"Provable Memorization Capacity of Transformers\", ICLR 2023.  \nHaotian Jiang and Qianxiao Li, \"Approximation Rate of the Transformer Architecture for Sequence Modeling\", NeurIPS 2024.  \nTokio Kajitsuka and Issei Sato, \"On the Optimal Memorization Capacity of Transformers\", ICLR 2025.", "claims_and_evidence": "The claims made in the Abstract, Introduction, and Conclusion are supported. I believe the stated theorems to be correct.\n\nOne issue is that the theoretical analysis uses the hardmax operation (line 106, left column), which is a nontrivial abstraction and may have expressiveness implications [1]. I think this should be mentioned when describing the paper's overall contributions (abstract, introduction) as in other work making such assumptions [e.g. 2].\n\n[1] Yang et al, Simulating Hard Attention Using Soft Attention, arXiv 2024\n\n[2] Amiri et al, Lower Bounds for Chain-of-Thought Reasoning in Hard-Attention Transformers, arXiv 2025", "methods_and_evaluation_criteria": "Reasoning tasks used for evaluation make sense.", "theoretical_claims": "I read through the proofs in the appendix, and checked the individual steps of proof of Theorem A.1 in detail. The later proofs are plausible to me, but I did not check every single step.", "experimental_designs_or_analyses": "I checked the soundness of the experiments on reasoning tasks as described in Appendix C.", "relation_to_broader_scientific_literature": "The paper is related to both the theoretical literature on Looped Transformers, and literature on the universal approximation property of transformers. It brings these together by showing the latter kind of property under the constraints imposed by the looped transformers architecture.", "essential_references_not_discussed": "None, as far as I see.", "other_strengths_and_weaknesses": "Weaknesses:\n\n- Of note, unlike a lot of other literature on the expressive power of transformers [e.g. 1,2,3], the present paper follows Yun et al in only studying functions on fixed-length inputs, without any regard to performance across inputs of potentially unbounded length. This restriction can be fine, but I think this should be made more transparent earlier in the paper (e.g., abstract, intro).\n\n- Regarding clarity: I found the terminology surrounding \"sequence IDs\" and \"token IDs\" confusing in Definition 3.2. At this point, a reader may wonder: Is a Token ID something like a word embedding? or a positional embedding? or something different? And what exactly is a sequence ID -- what kind of object even is it -- a number, a tensor, or something else? I think the reader should be given a clearer idea at this point.\n\n[1] Strobl et al, What Formal Languages Can Transformers Express? A Survey, TACL 2024\n\n[2] Hahn, Theoretical limitations of self-attention in neural sequence models, TACL 2020\n\n[3] Merrill and Sabharwal, The Expressive Power of Transformers with Chain of Thought, ICLR 2024", "other_comments_or_suggestions": "line 88, right column: \"can multiple graph\" a word such as \"perform\" is missing after \"can\"\n\nline 192, right column: \"Looped Transformers is defined by\" -- delete the word \"is\", as there is also a verb in the second-to-last line of the Corollary.\n\nSection 5.1: The abbreviations LCS and ED and the mapping between the tasks listed in the text is not made explicit. The reader needs to refer to the appendix to figure this out.\n\nline 571: \\delta^{-1} \\in mathbb{N}, \\delta \\geq 2 -- should the second \\delta be \\delta^{-1}?\n\nIn the proof, I found it hard to track when p-norms are applied to individual vectors (as in equation 49) or across the entire input space (as in equation 46, or in the line below equation 59). Can the authors make the distinction more explicit?\n\nLine 1302: \"combined\" --> \"combining\"", "questions_for_authors": "In the hardmax operation, how are ties resolved? I.e., when multiple positions receive the maximum attention score.\n\nLine 155 (right column): where is bit complexity used? It is introduced but then doesn't appear to show up.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CFLkMPVKTX", "reviewer": "Reviewer_YeAT", "summary": "The paper theoretically investigates the expressive power of Looped Transformers, neural architectures that recursively reuse a single Transformer block, which are appealing for their parameter efficiency and generalization capabilities. The authors define a new modulus of continuity tailored to sequence-to-sequence functions to establish the approximation rate of Looped Transformers. Their analysis identifies a specific limitation of the looped architecture due to weight-tying constraints. To address this, they propose introducing scaling parameters conditioned on timestep encoding, thereby enhancing the expressive power of the architecture. Empirical results across various reasoning tasks and standard Transformer benchmarks validate the theoretical findings, demonstrating improved performance with increased loop counts and further gains achieved through timestep encoding.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1 (W1). Could you elaborate on whether and how the increase in computational complexity (due to timestep encoding) affects practical deployment scenarios compared to standard Transformers?\n\nThank you for the insightful comment. While timestep encoding adds slight computational overhead, the increase in parameters is minimal, and we observed only a minor impact on throughput. Prior work [1,2] has also explored architectural modifications to enhance expressiveness without sacrificing efficiency, and we believe our design follows this practical direction.\n\n**Reference:**  \n[1] Csordás et al., \"MoEUT: Mixture-of-Experts Universal Transformers\", NeurIPS 2024.  \n[2] Bae et al., \"Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA\", ICLR 2025.", "claims_and_evidence": "- Claim: Looped Transformers can universally approximate continuous permutation-equivariant sequence-to-sequence functions.\n  - Evidence: Proven through a theorem (Theorem 3.6), supported by mathematical proof utilizing newly defined modulus of continuity.\n\n- Claim: Looped Transformers face inherent limitations due to their weight-tying nature, impacting their expressive power.\n  - Evidence: Demonstrated analytically (Lemma 4.1), which highlights limitations in approximating specific target mappings exactly due to recurrent feed-forward layer constraints.\n\n- Claim: Incorporating timestep encoding with scaling parameters mitigates the limitations and improves approximation capabilities.\n  - Evidence: Theoretical justification (Theorem 4.2) shows exact memorization capability with time-dependent scaling, and experiments validate improved performance on reasoning and standard Transformer tasks.", "methods_and_evaluation_criteria": "The methods and evaluation criteria proposed in the paper are suitable. The authors use well-established theoretical frameworks (modulus of continuity, universal approximation results) and practical machine learning benchmarks (reasoning tasks, dynamic programming problems, Sudoku, in-context learning, language modeling) to validate their theoretical insights.", "theoretical_claims": "Theoretical Claims:\n- Looped Transformers are universal approximators for permutation-equivariant sequence-to-sequence functions (Corollary 3.7).\n- Introducing timestep encoding and scaling parameters can overcome certain theoretical limitations specific to Looped Transformers (Theorem 4.2).\n\nThe theoretical claims appear correct. The proofs provided are logically structured, detailed, and sound. The innovative use of three separate moduli of continuity to evaluate approximation capabilities is well justified.", "experimental_designs_or_analyses": "Experiments included:\n\n- Reasoning tasks (Sudoku, Countdown, Edit Distance, Longest Common Subsequence).\n\n- In-context learning task.\n\n- Language modeling on WikiText-103 dataset.\n\nThe experimental design is sound, with clear benchmarks appropriate for evaluating reasoning capabilities and general transformer-related tasks. The experimental analyses demonstrate clear performance improvements due to the proposed timestep encoding.", "relation_to_broader_scientific_literature": "The paper relates to broader scientific literature through its theoretical examination of expressive power, aligning closely with previous works on sequence-to-sequence function approximation for Transformers (Yun et al., 2020). It extends existing results to weight-tied Looped Transformers, providing new insights into their capabilities and limitations. The introduction of timestep encoding aligns with existing literature on conditional scaling and adaptive methods in neural network architectures.", "essential_references_not_discussed": "The paper references most relevant prior works thoroughly.", "other_strengths_and_weaknesses": "Strengths:\n\n- Clearly stated theoretical contributions with rigorous proofs.\n\n- Novel and insightful identification of a limitation in Looped Transformers.\n\n- Introduction of timestep encoding is a simple yet effective conceptual innovation, clearly validated by experiments.\n\nWeaknesses:\n\n- Additional computational overhead due to timestep encoding may diminish some practical advantages of looped architectures regarding parameter efficiency.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "Could you elaborate on whether and how the increase in computational complexity (due to timestep encoding) affects practical deployment scenarios compared to standard Transformers?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jukyq4xvmE", "reviewer": "Reviewer_7KaS", "summary": "The paper study the limitation of hardmax 1-layer  $r$-times looped transformer with fixed limited embedding size, extending Zhang et al. (2023) that studied looped ReLU. Furthermore, the paper suggest adopted layer-dependent scaling for looped transformer to mitigate the suggest limitation.\n## update after rebuttal\n### *After Discussion*\n\n*I remain unsatisfied and disappointed following in-depth discussions with the authors. In all, I think the paper needs a big revision for submission and currently missing many important parts.*\n\n1. ***Downplaying Fixed-Width Architecture**. In the title, abstract, contribution bullet points, and Table 1, the authors—intentionally or not—downplay / hide their use of a **fixed-width** transformer architecture, which, though adopted in some prior work, isn’t widely celebrated.*\n2. ***Missing Comparison over the Proposed ‘Limitation’**. This is particularly striking given their second bullet point highlighting the limitations of looped transformers. They stress that a 1-layer looped transformer’s reliance on specific continuity coefficients drives its error, yet **they fail to compare this with non-looped transformers** of equivalent bit complexity (e.g., to what degree might non-looped designs mitigate this dependency?).*\n3. ***Unsatisfactory Response w.r.t. a Related Work**. I cited Saunshi et al. (2025), which shows a 1-layer looped transformer overcoming the limitation the authors note (without explicitly fixing width). Their rebuttal emphasized parameter efficiency—an outcome they later concede wasn’t satisfactorily achieved and is not their story—and offered a flawed excuse for overlooking this key work on looped transformer expressivity. They claimed unavailability despite ICLR submissions being anonymously accessible for citation and reading well in advance. Since Saunshi et al. (2025) undermines author’s claim ‘looped transformer’s expressive power remains underexplored’ in their abstract and main story, it should not be absent.*\n4. ***Missing Parts in Experiments**. Following my suggestion, they designed language tasks with varying continuity coefficients tied to their core theories in the rebuttal. However, these should have been conducted earlier, across more datasets and disturbance conditions, and **should have been** detailed in Section 5’s first part—yet they’re absent. Their further response on comparison fairness is unprofessional: I don’t expect a reduced base model parameter count as they conjectured, but the base + timestep total should match the without-timestep setup. Given their focus on this **limitation**, experiments with non-looped and k-layer looped counterparts are needed to fairly assess if these overcome it—perhaps all underperform with 1-layer looped + timestep. Moreover, the extent to which their encoding boosts parameter efficiency at equal performance levels warrants discussion but is currently absent.*\n5. ***Failure of Deeper Discussions**. Top-conference rebuttals value discussions on potential topics and potential improvement of current techniques, as these enrich the Discussion Section or appendix, fostering its impact and welcoming impactful future-up work. I explicitly noted my suggestions—like k-layer looped transformers, wavelet analysis, and memory capacity—extend beyond their current scope, not expecting polished analyses now. While the authors acknowledge these as critical, they adopt a dismissive stance, avoiding deeper heuristic and brainstorming-style exploration based on current knowledge. Instead, they waste characters distancing these from their contributions—which I already understand—rather than engaging constructively as somebody eager to explore non-triviality. They admit their cube and ID-based analyses could improve, potentially easing model issues, yet resist a brainstorming-style discussion. Though they expressed “optimality is non-trivial,” I struggle to detect sufficient enthusiasm for making their techniques more realistic or elegant.*", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1 (W2). The theoretical results are based on constraining the hidden embedding size of transformer $(17d+9)×N$, which is very tricky. In fact, Saunshi et al. (2025) has shown that a 1-layer transformer looped can simulate the non-looped counterpart at the cost of a larger embedding size, still, the looped transformer is more parameter-efficient.\n\nWe respectfully point out that, **in contrast to Saunshi et al. (2025), our results achieve parameter efficiency** by showing that **the parameter complexity is $\\boldsymbol{O(d)}$**. The construction by Saunshi et al. (2025) depends on the number of layers and width of the target non-looped transformer. Since the number of layers or width in a non-looped transformer typically grows with the approximation error or sequence length, the corresponding looped transformer inherits the dependence, which is not parameter-efficient.\n\n> **Theorem 5.2 (Saunshi et al., 2025):** For any transformer with $L$ layers, at most $R$ distinct, $d$ embedding size, $d_{FF}$ hidden dimension for MLP ... loops a 1-layer transformer block for $L$ times, with $d + R + 2$ embedding size, $R{h_{FF}} + O(L)$ hidden dimension.\n\nWe will clarify this distinction in Section 2 (Background) and emphasize our contribution in Section 3.3. (Main Result):\n\n**line 158:** Saunshi et al. (2025) showed that Looped Transformers can approximate standard Transformers, ..., their construction inherits this dependence, limiting parameter efficiency.\n\n**line 207:** **The parameter complexity is $\\boldsymbol{O(d)}$**, depending solely on the input dimension, not the approximation error or sequence length, highlighting the parameter efficiency.\n\n> W1-1. The author missed some important references.\n \nAs the work was not publicly available at the time of submission, it will be cited in the camera-ready version.\n\n> W1-2. With appropriate $k$, the looped transformer should have chance to overcome the drawback the paper discussed ...\n\nWe respectfully point out that achieving universal approximation **even with a single-layer** is a non-trivial theoretical contribution. Exploring how efficiently multiple layers fall under the study of non-looped Transformers. In contrast, our work investigates how the approximation rate depends on the number of loops. We will add as future work.\n\n> W3. The paper didn’t provide estimation error, which usually co-exists with approximation error ...\n\nTo the best of our knowledge, estimation error bounds are not commonly discussed in studies on approximation (Yun 2020, Kim 2023, Kajitsuka 2024, Jiang 2024). This is likely due to a difference in problem settings: estimation error is analyzed in regression. We will add this as a future direction.\n\n> W4-1. The paper considers replacing the softmax operation with hardmax, which in my view is a huge drawback. \n\nWe use the softmax function to **approximate** (not replaced) the hardmax function by scaling input (Yun et al., 2020; Kim et al., 2023). To avoid confusing, we will revise as:\n\n**Before:** we use the hardmax instead of the softmax  \n**After:** we use the softmax to **approximate** the hardmax\n\n> W4-2. The techniques of looped ReLU’s (Zhang et al., 2023) could be directly applied without much effort. \n\nWe respectfully point out that the techniques from ReLUs (Zhang et al., 2023) cannot be directly applied to Transformers due to fundamental differences in architecture and target function class, that is, transformers had to compute **contextual mappings**. It required newly defined continuity and techniques. We will clarify as:\n\n**line 127:** Our question is whether the result (Zhang et al., 2023) can be extended **for contextual mappings**\n\n> Q2. Can your cube & ID-based analyses be improved by advanced waveless analyses? C3. The memorization of IDs might be more efficient ...\n\nOptimality is non-trivial. Existing proof techniques (e.g., Kajitsuka 2025) do not directly apply to looped architectures, and new methods are needed. We will include this as an important direction for future work.\n\n> Q3. Can you design some language tasks that different tasks depend on different continuity coefficients to validate theoretical claims? No experiments showed that they suffer from continuity coefficients and cannot validate the author’s claim that timestep alleviates the specific limitation.\n\nWe designed a perturbed WikiText-103 task (10% token randomly replacement) to test sensitivity to continuity. We further analyzed the continuity coefficients of the trained models by applying perturbations to the input and measuring the change in output embeddings. We found that timestep encoding, with only about 5% increase in parameters, improved memorization and consistently reduced continuity coefficients. Without it, the model was less stable and failed to capture accurate input-output relationships.\n\n|   | CE Entropy Loss (Train) | Continuity Coefficients |\n|--|--|--|\n| w/o | 4.32 | 130.6 |\n|  w/ Timestep | **4.18** | **21.5** |", "claims_and_evidence": "The paper’s claim regarding the limitations of looped transformer is not well-illustrated by experiments, but the effectiveness of the layer-dependent scaling strategies is validated by experiments on case studies.", "methods_and_evaluation_criteria": "The proposed Timestep Encoding method suggest layer-dependent scaling for loop transformer with specific design motivated by prior work. The author validated it through three case studies.", "theoretical_claims": "The paper claimed that their hardmax 1-layer  $r$-times looped transformer with fixed limited embedding size’s expressivity would depend on certain continuity coefficients, and their proposed layer-dependent scaling would mitigate this.", "experimental_designs_or_analyses": "1. No experiments showed that looped transformer really suffer from certain continuity coefficients. In the author experiments, since the width is the same, the comparisons between L=6 to r=8, or L=12 and r=10 show that TF is not obviously superior to Looped TF.\n2. The effectiveness of layer-dependent encoding suits common intuition, which essentially cost more parameter to make each layer slightly distinct. This cannot validate the author’s claim that it alleviate the specific limitation.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "Saunshi et al. (2025), a paper published in ICLR 2025, has shown that looped transformer can simulate non-looped counterpart while requiring a larger NN width.", "other_strengths_and_weaknesses": "**Strengths**\n\n1. The paper’s logic flow is clear, with good illustrations.\n2. The paper cast a table to compare different studies.\n\n**Weakness**\n\n1. The author missed some important reference such as Saunshi et al. (2025). In these studies as well as in real-world practice, looped transformers adopted $k$-layer transformer looped $L$ times, instead of 1-layer transformer looped $L$ ($r$ is the paper’s notation) times considered in this paper. With appropriate $k$, the looped transformer should have chance to overcome the drawback the paper discussed, and the looped transformer should be parameter efficient than non-looped counterpart, similar to literature like Saunshi et al. (2025).\n2. The theoretical results based on constraining the hidden embedding size of transformer ($(17d+9)\\times N$), which is very tricky in my view. The limitations of looped transformer (depending on continuity coefficients) that the author proved might be relaxed when the hidden embedding is not constrained. In fact, Saunshi et al. (2025) has shown that a 1-layer transformer looped $L$ ($r$ is the paper’s notation) times can simulate the non-looped counterpart at the cost of a larger embedding size. But still, looped transformer is more parameter-efficient.\n3. The paper didn’t provide estimation error, which usually co-exist with approximation error result in similar well-established theoretical studies like Takakura and Suzuki (2023). \n4. The paper considers replacing transformer’s softmax operation with hardmax, which in my view is a huge drawback. In fact, by appropriately choosing the QK matrice, the transformer could equate looped ReLU network, and the techniques in prior study that studied looped ReLU’s drawback could be directly applied (Zhang et al. (2023)) with some improvements without much effort. The potential method to use softmax can refer to my first point in “Other Comments Or Suggestions” section.", "other_comments_or_suggestions": "1. Considering the actual softmax operation would make the analysis more realistic, despite introducing additional technical challenges. However, it remains feasible with sufficient technical effort, as demonstrated in Lemma C.1–2 of Takakura and Suzuki (2023).\n2. As discussed in Weakness 1, the discussions over $k$-layer $L$ looped transformer should be included to complement the real-world scenario.\n3. The memorization of IDs might be more efficient using random feature transformation or even NTK (similar approaches adopted in Nichani et al. (2025)). In this merit, how far did your looped transformer from the Optimal Storage Capacity is worth discussing.", "questions_for_authors": "1. Please respond to Weakness 2.\n2. Can your cube & ID-based analyses (which induces $\\sqrt{d}$ and the design of $\\mathcal{K}$) be improved by advanced waveless analyses? Some case studies?\n3. See “Experimental Designs Or Analyses” section. Can you design some language tasks that different tasks depend on different scale of your defined continuity coefficients to show the effectiveness of your theoretical claims?\n\n*Reference*\n\nZhang et al. (2023). On enhancing expressive power via compositions of single fixed-size ReLU network. In Proceedings of the 40th International Conference on Machine Learning, 2023.\n\nSaunshi et al. (2025). Reasoning with Latent Thoughts: On the Power of Looped Transformers. In International Conference on Learning Representations (ICLR2025), 2025.\n\nTakakura and Suzuki (2023). Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Input. In International Conference on Machine Learning, pages 26517–26582. PMLR, 2023.\n\nNichani et al. (2025). Understanding Factual Recall in Transformers via Associative Memories. In International Conference on Learning Representations (ICLR2025), 2025.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper focuses on understanding the approximation power of looped transformers. In particular, it focuses on looped transformers where a single layer is looped $r$ times. It establishes approximation rates for the looped transformers in terms of multiple notions of modulus of continuity for the target function. The paper then highlights a limitation of Looped transformers and proposes timestep encoding to scale different loops. The paper presents empirical results to validate the utility of the proposed timestep encoding.

Overall, the paper studies a timely problem and makes non-trivial contributions toward improving the understanding of looped transformers.

The reviewers raised multiple valid questions/comments during the review process. Some of these comments were centered around improving the presentation which the authors have adequately addressed. The reviewers also brought up a recent work by Saunshi et al. 2025 that also explores the approximation power of looped transformers by arguing how a looped transformer can simulate a non-looped transformer. The authors pointed out how their problem setting differs from the one covered by Saunshi et al. The authors are encouraged to include a detailed comparison with this work. Furthermore, as requested by one of the reviewers, the authors should include a discussion on those looped transformers that loop a block of k-layers multiple times and comment on the challenges in extending their analysis to this setting. The authors are also encouraged to include the new experimental results in their response with a detailed discussion. The author must also incorporate all the changes they had promised to the reviewers during the review process.
Decision: Accept (poster)