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

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.

We 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.

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.

We will clarify this distinction in Section 2 (Background) and emphasize our contribution in Section 3.3. (Main Result):

line 158: Saunshi et al. (2025) showed that Looped Transformers can approximate standard Transformers, ..., their construction inherits this dependence, limiting parameter efficiency.

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.

W1-1. The author missed some important references.

As the work was not publicly available at the time of submission, it will be cited in the camera-ready version.

W1-2. With appropriate $k$, the looped transformer should have chance to overcome the drawback the paper discussed ...

We 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.

W3. The paper didn’t provide estimation error, which usually co-exists with approximation error ...

To 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.

W4-1. The paper considers replacing the softmax operation with hardmax, which in my view is a huge drawback.

We 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:

Before: we use the hardmax instead of the softmax
After: we use the softmax to approximate the hardmax

W4-2. The techniques of looped ReLU’s (Zhang et al., 2023) could be directly applied without much effort.

We 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:

line 127: Our question is whether the result (Zhang et al., 2023) can be extended for contextual mappings

Q2. Can your cube & ID-based analyses be improved by advanced waveless analyses? C3. The memorization of IDs might be more efficient ...

Optimality 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.

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.

We 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.

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?

Thank 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.

Reference:
[1] Csordás et al., "MoEUT: Mixture-of-Experts Universal Transformers", NeurIPS 2024.
[2] Bae et al., "Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA", ICLR 2025.

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

We 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.

line 45: Our contributions are summarized as follows: We establish the approximation rate of Looped Transformers for fixed-length continuous sequence-to-sequence functions.

W2. Regarding clarity: I found the terminology surrounding "sequence IDs" and "token IDs" confusing in Definition 3.2.

We 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:

line 175: In our proofs, we rely on the contextual token ID to describe contextual mappings.
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.
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 on the specific proof. In our case, we adopt a different construction from that of Kim et al. (2023).

Q1. In the hardmax operation, how are ties resolved? I.e., when multiple positions receive the maximum attention score.

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:

line 329: ..., where $\mathcal{O}(\delta^{d})$ arises from the case where identical tokens appear in sequences.

Q2. Line 155 (right column): where is bit complexity used? It is introduced but then doesn't appear to show up.

The 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.

Before: We evaluate bit complexity, which denotes the maximum number of bits of weights.
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.

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.

While 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:

Before: we use the hardmax function instead of the softmax function.
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).

O1. In the proof, I found it hard to track when p-norms are applied to individual vectors or across the entire input space.

We 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).

Reference:
Yun et al., "Are Transformers universal approximators of sequence-to-sequence functions?", ICLR 2020.
Kim et al., "Provable Memorization Capacity of Transformers", ICLR 2023.
Haotian Jiang and Qianxiao Li, "Approximation Rate of the Transformer Architecture for Sequence Modeling", NeurIPS 2024.
Tokio Kajitsuka and Issei Sato, "On the Optimal Memorization Capacity of Transformers", ICLR 2025.

I believe deep equilibrium models are closely related to looped transformers, but the authors don't seem to discuss them.

Thank 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:

line 77: Deep equilibrium models (Bai et al., 2019), which compute fixed points of iterative layers, are also related.

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.

Thank 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:

line 439: Beyond expressivity, improving training stability and efficiency, especially for larger numbers of loop iterations, is an important direction for future work.

Reference:
Bai et al., "Deep Equilibrium Models", NeurIPS 2019.