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

Authors: Kevin Xu, Issei Sato
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Transformers, Looped Transformers, Expressive power, Approximation Rate

Abstract:
Looped Transformers offer advantages in parameter efficiency and Turing completeness. However, their expressive power for function approximation and approximation rate remains underexplored. In this paper, we establish approximation rates of Looped Transformers by defining the concept of the modulus of continuity for sequence-to-sequence functions. This reveals a limitation specific to
the looped architecture. That is, the analysis prompts us to incorporate scaling parameters for each loop, conditioned on timestep encoding. Experimental results demonstrate that increasing the number of loops enhances performance, with further gains achieved through the timestep encoding architecture.

Decision: reject
Reviewer scores: 6, 8, 3

Reviews and discussion:
<thread>
creation date:
1730455886952

thread_type:
official_review

<thread_object>
creator: reviewer_mxai

<content>
summary:
This paper theoretically investigates the approximation capabilities of Looped Transformers. Through theoretical analysis, it is found that scaling parameters for each loop can enhance the expressive power of Looped Transformers.

strengths:
One of the strengths of this paper is the detailed theoretical analysis, which is clearly articulated. The definitions of the notation are well presented, making the paper relatively easy to read.

weaknesses:
My main concern lies with the novelty of the work. It is well known that Looped Transformers can be viewed as a type of tied-parameters transformer, and the universal approximation theory for transformers has already been extensively studied. Moreover, the theoretical analysis presented does not yield any novel conclusions regarding the training of Looped Transformers. The finding that performance increases with the number of loops does not seem particularly innovative. Therefore, I have not yet identified any significant conclusions from this work.

My second minor concern is that the theoretical research on Looped Transformers may have lost its research value. Looped Transformers were proposed to address tasks that transformers struggle with, which RNNs can accomplish easily, while also offering parameter efficiency. However, transformers have now completely supplanted RNNs in various domains through various training techniques. In the most critical application areas for transformers, such as autoregressive models and LLMs, Looped Transformers are not commonly used, and their performance does not match that of transformer architectures.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730660506520

thread_type:
official_review

<thread_object>
creator: reviewer_c5sb

<content>
summary:
The paper presents a theoretical exploration of the expressive power of Looped Transformers, particularly regarding their function approximation capabilities. The study justifies that Looped Transformer, while offering parameter efficiency and Turing completeness, suffers from some limitations in the approximation rate, and it addresses it by introducing timestep encoding, which enhances expressive capacity by conditioning on time-dependent scaling factors.

strengths:
1.	The paper is well-written. The problem formulation, theoretical analysis, and experimental results are described clearly. 
2.	The paper's contribution is relatively significant. It derives bounds on the approximation rate of a looped transformer to demonstrate the limited approximation ability. 
3.	The paper provides strong theoretical justifications that lead to its design (of timestep encoding) on loop transformers. In particular, looped architecture is revealed to have exponential asymptotic approximation rate decay, which leads to the incorporation of timestep encoding.

weaknesses:
1.	The experimental validation is relatively limited. In particular, only one task is evaluated on the timestep encoded model compared with the baseline. The authors are encouraged to test on more tasks including basic ones (e.g. inductive heads) for better validation of the model improvement and use ablation studies for more detailed characterization.

questions:
1.	How tight is the approximation rate bound derived from the paper? It would be interesting to see empirical results that validate equation (9) or (10).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730738082654

thread_type:
official_review

<thread_object>
creator: reviewer_h7gd

<content>
summary:
This paper delves into the approximation rate of Looped Transformers for continuous sequence-to-sequence functions, offering a comprehensive mathematical analysis of the strengths and limitations of looped architectures. Specifically, it examines Looped Transformers by defining critical concepts such as sequence, contextual, and token continuity, which significantly influence the approximation rate. The authors propose a single-loop topology for contextual mapping that requires only a single-layer network for universal approximation.
This paper's novel contribution is the introduction of a new time-stepping encoding method, which addresses the approximation errors typically associated with Looped Transformers. This method incorporates a time-dependent scaling parameter, effectively discretizing variations in the target function. The study demonstrates that time-dependent Looped Transformer models can mitigate approximation errors arising from contextual and token continuity issues.
Experimental results support these claims, showing that accuracy increases linearly with the number of loops when using timestep-encoded models, in contrast to time-dependent models, which tend to saturate in training and accuracy. Additionally, the proposed approach offers resource optimization by reducing the number of layers in the neural network and decreasing the number of parameters by approximately 90% compared to previous works.

strengths:
The presented methodology introduces a novel approach that effectively addresses the known limitations of Looped Transformer model architectures. The paper provides a robust background, thoroughly reviewing previous studies in the field. It includes rigorous mathematical demonstrations that substantiate the claims, and the experimental results are consistent with these theoretical findings.

weaknesses:
The redaction can be challenging to follow in certain sections, particularly where large equations are interspersed with text (e.g., Theorem 3.7, Lemma 4.2). To enhance clarity, it would be beneficial to include comparison charts against current Looped Transformer models. These charts would help to emphasize the advantages of the proposed methodology and illustrate the behavior observed across different models.

questions:
Are you planning to scale up experiments by using different models? Is it anticipated to be feasible to apply the proposed architecture to different scenarios? Are there any future work plans?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733112613603

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Rebuttal by Authors (1/2)

comment:
Dear Area Chairs,

We have carefully addressed and respectfully refuted the questions raised by the reviewers, as summarized below. Although we have yet to receive feedback from some reviewers, we believe that all concerns have been thoughtfully considered and appropriately resolved to the best of our ability.

***

> Q1: (Reviewer MxAi) My main concern lies with the novelty of the work. It is well known that Looped Transformers can be viewed as a type of tied-parameters transformer, and the universal approximation theory for transformers has already been extensively studied. Moreover, the theoretical analysis presented does not yield any novel conclusions regarding the training of Looped Transformers. The finding that performance increases with the number of loops does not seem particularly innovative. Therefore, I have not yet identified any significant conclusions from this work.

A1: **Our paper is the first to establish the approximation rate of Looped Transformers, introducing three types of modulus of continuity newly defined for sequence-to-sequence functions**. Due to the structural constraints imposed by weight-tying, which limit their flexibility, **universal approximation theories for Transformers cannot be directly applied**.

***

> Q2: (Reviewer MxAi) Moreover, the theoretical analysis presented does not yield any novel conclusions regarding the training of Looped Transformers. The finding that performance increases with the number of loops does not seem particularly innovative. Therefore, I have not yet identified any significant conclusions from this work.

A2: **Beyond establishing the approximation rate, our work identifies a specific limitation inherent to the looped architecture**. To address this, we introduce the incorporation of scaling parameters for each loop, implemented with conditioning on timestep encoding (see Sec.4). We demonstrate the **enhancement via timestep encoding** both theoretically and experimentally.

***

> Q3: (Reviewer MxAi) My second minor concern is that the theoretical research on Looped Transformers may have lost its research value…. In the most critical application areas for transformers, such as autoregressive models and LLMs, Looped Transformers are not commonly used, and their performance does not match that of transformer architectures.

A3: Looped Transformers hold their research value because of the following three points.

(1). **Extensive Practical Applications and Research Interest**: Looped Transformers have been practically studied for efficient and competitive capabilities [1, 2] and remain an active area of research. Previous works have studied their Turing completeness [3], generalization ability [4, 5], inductive bias for reasoning tasks [6], ability of learning iterative algorithms for in-context learning [7, 8, 9], and potential to simulate graph algorithms [10]. (All the reference is shown in rebuttal to Reviewer MxAi)

We also emphasized (2) **the distinctiveness from RNNs** and (3). **the potential advantages for complex tasks**.

***

> Q4 (Reviewer H7Gd): The redaction can be challenging to follow in certain sections, ...

A4: To enhance clarity, we add the chart of our timestep encoding methodology as Figure 2.

***

> Q5: (Reviewer H7Gd) Are you planning to scale up experiments by using different models? (Reviewer C5sB) The experimental validation is relatively limited.

A5: **We expanded 5 tasks**: dynamic programming problems (subset sum, knapsack, LCS), in-context learning (decision tree), and language modeling (Wikitext-103).
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:**
This paper studies expressive ability and approximation rate of the looped transformers. 

**Strengths:**
This paper provides a theoretical analysis on expressive abilities of the looped transformers, claiming that it has revealed limitations specific to the looped architecture. On the basis of the theoretical findings, this paper furthermore proposes introduction of timestep encoding with the aim of improving the approximation rate.

**Weaknesses:**
The validity of the arguments in the theoretical part of this paper is questionable. See below.

**Reasons:**
I realized in the very final phase of reviewing that there appear several serious flaws in Theorem 3.7 and its proof, which constitutes the main contribution of this paper. Unfortunately, it seems that these flaws have not been noticed by any reviewers, most likely because they are not asked to read supplementary materials. Given these flaws, I would judge the technical quality of this paper not to reach the level of the conference, and would refrain from recommending acceptance.

The following part details those flaws.

**On $r$:** It is claimed that the statement of Theorem 3.7 holds "for any $r\in\mathbb{N}$" (line 203), but it cannot. (a) If $r$ is less than or equal to $N$, then $\delta=((r-N)/2)^{-(N+1)d-1}$ is not positive, causing problems. (b) If the for-any claim be valid, then one might let $r=1$, implying that the statement holds for non-looped transformers, making the discussion on looped transformers meaningless.

**On Equation (10):** It is claimed that the conditions $\delta^{-1}+2\delta^{-(N+1)d}=r-N$ (line 366) and $\delta^{-1}\cdot2\delta^{-(N+1)d}\ge r-N$ (line 367) on $\{\delta,N,d,r\}$ are equivalent, but this claim is not correct. It makes the development here, as well as the whole statement of Theorem 3.7, invalid.

**On $\epsilon$:** As shown in Propositions A.1 and A.2, compositions of ReLU functions approximate the rectangular and step functions, only in the limit $\epsilon\to0$. In the proof of Lemma A.4, however, this was not treated appropriately. The formula in line 1055, as well as the claim of Lemma A.4, is valid only in the limit $\epsilon\to0$. 

One also observes that some weights will diverge in the limit $\epsilon\to0$: see the formula in lines 1087-1092. The unboundedness of the weights may be problematic, especially in contrast to the well-known Barron-type universal approximation theorems, where the universal approximation property holds with parameters whose values are bounded (see, e.g., Caragea, Petersen, Voigtlaender, Ann. Appl. Probab. 33(4): 3039-3079, 2023). In any case, the authors should discuss $\epsilon$-dependence of their approximation rate, as well as dependence of bounds of the weights on $\epsilon$.

In addition to the points mentioned above, there are numerous points which would require revision, although each of them might not be fatal. I would like to encourage the authors to proofread their manuscript much more carefully. These are listed in **Additional comments** section below due to the space constraint.