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

## Abstract

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. 1 1 1 Code for our paper is available at https://github.com/kevin671/timestep-encoding

## 1 Introduction

Transformers (Vaswani et al., ) have become the standard neural network architecture for a wide range of machine learning tasks, such as natural language processing and computer vision.
Looped Transformers have an architecture composed of fixed-size Transformer layers, where the output is fed back into the input. Dehghani et al. () first introduced the idea; they aimed to add recursive inductive bias into Transformers to address tasks that recurrent neural networks (RNNs) handle more easily.
Looped Transformers are also related to weight-tying Transformers (Lan et al., ; Yang et al., ).
They demonstrated that Looped Transformer achieve performance comparable to the standard Transformer, while using less than $10$% of the parameter counts.
While earlier studies focused on a limited number of loop iterations, recent studies have explored the potential of a large number of loops. Giannou et al. () theoretically demonstrated that the recursive structure of Looped Transformers allows them to function as Turing machines. Yang et al. () empirically demonstrated that increasing the number of loop iterations improves performance on more complex tasks.
In summary, Looped Transformers are not only more parameter-efficient but also potentially more expressive than standard Transformers.

The expressive power of standard Transformers has been extensively explored, showing that they can approximate continuous permutation-equivariant functions on compact domains (Yun et al., ; Kim et al., ; Takakura & Suzuki, ; Jiang & Li, ; Kajitsuka & Sato, ). In addition, their approximation rate has been studied: Takakura & Suzuki (); Jiang & Li () established it by identifying the properties of the target functions, revealing the types of functions that Transformers can approximate effectively.
In contrast, the expressive power of Looped Transformers in function approximation remain unexplored. Moreover, the approximation rate and the appropriate properties of the target functions for Looped Transformers have yet to be investigated.
Recently, Zhang et al. () first explored the expressive power of *looped* models. They established an approximation rate for ReLU networks for continuous functions in terms of loop counts and modulus of continuity.
Still, this remains unclear for the case of Looped Transformers.

In this paper, we derive the approximation rate of Looped Transformers for continuous sequence-to-sequence functions by defining the concept of *sequence continuity*, *contextual continuity*, and *token continuity*.
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.

## 2 Background

We define the Transformer architecture in Section 2.1 and Looped Transformers in Section 2.2. We then introduce related work of theoretical analysis on function approximation power of Transformers in Section 2.3.

### 2.1 Transformer architecture

Given an input sequence ${\bm{X}}\in\mathbb{R}^{d\times N}$, composed of $N$ token embedding of dimension size $d$, the self-attention layers with $h$ heads and head size $s$, and the feed-forward layer with width size $q$, are defined as follows:

$$ $\displaystyle{\mathrm{Attn}}({\bm{X}})$ $\displaystyle=\sum_{i=1}^{h}{\bm{W}}_{i}^{(O)}\left({\bm{W}}_{i}^{(V)}({\bm{X} })\right)\sigma_{S}\left[\left({\bm{W}}_{i}^{(K)}({\bm{X}})\right)^{\top}\left ({\bm{W}}_{i}^{(Q)}({\bm{X}})\right)\right]\in\mathbb{R}^{d\times N},$ (1) $\displaystyle{\mathrm{FF}}({\bm{X}}_{:,i})$ $\displaystyle={\bm{W}}^{2}\cdot{\sigma_{R}}({\bm{W}}^{(1)}\cdot{\bm{X}}_{:,i}+ {\bm{b}}^{(1)})+{\bm{b}}^{(2)}\in\mathbb{R}^{d},$ (2) $$

where ${\bm{W}}_{i}^{(V)},\,{\bm{W}}_{i}^{(K)},\,{\bm{W}}_{i}^{(Q)}\in\mathbb{R}^{s
\times d},{\bm{W}}_{i}^{(O)}\in\mathbb{R}^{d\times s},{\bm{W}}^{(1)}\in\mathbb
{R}^{q\times d},{\bm{W}}^{(2)}\in\mathbb{R}^{d\times q},{\bm{b}}^{(1)}\in
\mathbb{R}^{q},{\bm{b}}^{(2)}\in\mathbb{R}^{d}$ are parameters, ${\sigma_{R}}$ denotes ${\mathrm{ReLU}}$ function, and $\sigma_{S}$ denotes a softmax operator applied to the columns of the input matrix.

The Transformer block ${\mathrm{TF}}({\bm{X}}):\mathbb{R}^{d\times N}\to\mathbb{R}^{d\times N}$ is defined by:

$$ $\displaystyle{\bm{X}}^{\prime}$ $\displaystyle={\bm{X}}+{\mathrm{Attn}}({\bm{X}}),$ (3) $\displaystyle{\mathrm{TF}}({\bm{X}})$ $\displaystyle={\bm{X}}^{\prime}+\bm{{\mathrm{FF}}}({\bm{X}}^{\prime}).$ (4) $$

where $\bm{{\mathrm{FF}}}$ represent token-wise ${\mathrm{FF}}$. In other words

$$ ${\mathrm{TF}}=(\mathrm{id}+\bm{{\mathrm{FF}}})\circ(\mathrm{id}+{\mathrm{Attn} }),$ (5) $$

where $\mathrm{id}$ denote the identity mapping. For simplicity, we omit the domain of definition.

For the analysis of expressive power in Section 3, we use the hardmax operator $\sigma_{H}$ instead of $\sigma_{S}$ as in previous studies (Yun et al., ; Kim et al., ). The reason is that the softmax function can be interpreted as an approximation of the hardmax function; also softmax can be made arbitrarily close to the hardmax by scaling the temperature.
For simplicity, we also exclude Layer Normalizations from the analysis in Section 3.

### 2.2 Looped Transformer

Looped Transformers feed output back into input, defined as

$$ $\bm{\mathcal{L}}_{2}\circ{\mathrm{TF}}\circ\cdots\circ{\mathrm{TF}}\circ\bm{ \mathcal{L}}_{1},$ (6) $$

where $\bm{\mathcal{L}}_{2}$ and $\bm{\mathcal{L}}_{1}$ denote token-wise linear layers.
Let ${\mathrm{TF}}^{\circ r}$ denote the $r$-times composition of ${\mathrm{TF}}$. We refer to $\bm{\mathcal{L}}{2}\circ{\mathrm{TF}}^{\circ r}\circ\bm{\mathcal{L}}_{1}$ as a Looped Transformer with $r$-loops.

Dehghani et al. () first introduced the idea of adding recursive inductive bias into Transformers.
Looping Transformers are also related to parameter-efficient weight-tying Transformers (Lan et al., ).
Giannou et al. () theoretically showed that the recursive structure of Looped Transformers enables them to function as Turing machines.
Yang et al. () demonstrated that increasing the number of loop iterations can improves performance on some tasks.
Other recent works (Back De Luca & Fountoulakis, ; Gao et al., ; Giannou et al., ) have theoretically analyzed the emulation power of certain algorithms or training convergence, particularly focusing on in-context learning. The expressive power of Looped Transformers in function approximation or their approximation rate remain unexplored.

### 2.3 Theoretical analysis on function approximation

The function approximation of neural networks is well-studied.
The universal approximation theorem for fully connected neural networks, as demonstrated by Cybenko (); Hornik et al. (), shows that networks of sufficient size can approximate certain classes of functions with arbitrarily low error.
Transformers have been shown to be universal approximators of sequence-to-sequence functions (Yun et al., ; Takakura & Suzuki, ; Jiang & Li, ; Kajitsuka & Sato, ), and their memorization capacity has also been studied (Kim et al., ).
In particular, Jiang & Li () established an approximation rate by identifying the appropriate properties of the target sequence-to-sequence function.
Recently, Zhang et al. () revealed that even single fixed-size networks can be universal approximators. They explored the surprising potential of composition (loop) and derived the approximation rate in terms of the number of loop counts and modulus of continuity of the target function.

## 3 Approximation Rate of Looped Transformers

We establish the approximation rate of Looped Transformers by defining the modulus of continuity for continuous sequence-to-sequence functions. First, we introduce the preliminary in Section 3.1.
Then, we present and explain three types of continuity for sequence-to-sequence functions, which determine the approximation rate.
In Section 3.3, we present our main results on approximation rate, along with some implications. In Section 3.4, we provide a proof sketch with a detailed explanation, outlining three steps for approximation and deriving the approximation rate.

### 3.1 Preliminary

Transformers are permutation-equivariant, so we define the target function class as follows.

###### Definition 3.1 ( Yun et al. ( 2020 ); Kim et al. ( 2023 ) ) .

A function $f:\mathbb{R}^{d\times N}\to\mathbb{R}^{d\times N}$ is said to be permutation equivariant if for any permutation matrix ${\bm{P}}$, we have $f({\bm{X}}{\bm{P}})=f({\bm{X}}){\bm{P}}$.
Let $\mathcal{F}_{\mathrm{PE}}([0,1]^{d\times N})$ denote the set of permutation equivariant and continuous functions.

To approximate sequence-to-sequence functions, networks need to map each token within the context of sequences, formulated as contextual mapping.

###### Definition 3.2 ( Yun et al. ( 2020 ); Kim et al. ( 2023 ) ) .

Consider a finite set ${\mathbb{L}}\subset\mathbb{R}^{d\times N}$.
A map $\mathrm{CM}:{\mathbb{L}}\to\mathbb{R}^{1\times N}$ defines a contextual mapping if the map satisfies the following:

- 1.
For any ${\bm{L}}\in{\mathbb{L}}$, the $N$ entries in $\mathrm{CM}({\bm{L}})$ are all distinct.
- 2.
For any ${\bm{L}},{\bm{L}}^{\prime}\in{\mathbb{L}}$, with ${\bm{L}}\neq{\bm{L}}^{\prime}$, all entries of $\mathrm{CM}({\bm{L}})$ and $\mathrm{CM}({\bm{L}}^{\prime})$ are distinct.

Let $\left\lVert\cdot\right\rVert_{p}$ denote the entry-wise $L^{p}$-norm for a vector for any $p\in[1,\infty)$.

###### Definition 3.3 (Norm of function) .

We define the $L^{p}$-norm for a function $f$ on $[0,1]^{d\times N}$ by:

$$ $L^{p}([0,1]^{d\times N})\coloneq\Big{(}\int\left\lVert f({\bm{X}})\right\rVert _{p}^{p}d{\bm{X}}\Big{)}^{1/p}.$ $$

###### Definition 3.1 ( Yun et al. ( 2020 ); Kim et al. ( 2023 ) ) .

###### Definition 3.2 ( Yun et al. ( 2020 ); Kim et al. ( 2023 ) ) .

###### Definition 3.3 (Norm of function) .

### 3.2 Definition of continuity for sequence-to-sequence functions

The modulus of continuity of a continuous function $g:[0,1]^{d}\to\mathbb{R}$ can be defined as

$$ $\sup\big{\{}|g({\bm{X}})-g({\bm{X}}^{\prime})|:\|{\bm{X}}-{\bm{X}}^{\prime}\|_ {2}\leq t,\ {\bm{X}},{\bm{X}}^{\prime}\in[0,1]^{d}\big{\}}.$ $$

It can be extended for sequence-to-sequence functions $f:[0,1]^{d\times N}\to\mathbb{R}^{d\times N}$ as follows.

###### Definition 3.4 (Modulus of Sentence Continuity) .

Given a sequence-to-sequence continuous function $f:[0,1]^{d\times N}\to\mathbb{R}^{d\times N}$, the modulus of *sentence continuity* is defined by:

$$ $\omega_{f}(t)\coloneqq\sup\big{\{}\|f({\bm{X}})-f({\bm{X}}^{\prime})\|_{p}:\|{ \bm{X}}-{\bm{X}}^{\prime}\|_{2}\leq t,\ {\bm{X}},{\bm{X}}^{\prime}\in[0,1]^{d \times N}\big{\}}.$ $$

We explain what this continuity represents and why it is referred to as *sequence continuity* with illustration. We consider the following two sentences: (1) I read books; and (2) He writes music.
The sequence continuity measures how much the overall meaning of a sequence changes:

$\text{`I read books'}\ \overset{\Delta}{\longleftrightarrow}\ \text{`He writes
music'}$,

in proportion to the differences between the input sequences, measured by comparing each token:

$\text{`I'}\ \overset{\Delta}{\longleftrightarrow}\ \text{`he'}$,

$\text{`read'}\ \overset{\Delta}{\longleftrightarrow}\ \text{`write'}$, and

$\text{`book'}\ \overset{\Delta}{\longleftrightarrow}\ \text{`music'}$.

We found that this concept is insufficient to derive the approximation rate of Looped Transformers, primarily because one key characteristic of the Transformer architecture is the sharing of parameters across all tokens. Specifically, the feed-forward layers are applied in a token-wise manner. In other words, Transformers output token embeddings in the context of sequences for each token.

This observation leads us to define two additional forms of continuity: *contextual continuity* and *token continuity*, which we found to determine the approximation rate of Looped Transformers.

###### Definition 3.5 (Modulus of Contextual Continuity) .

Given a sequence-to-sequence continuous function $f:[0,1]^{d\times N}\to\mathbb{R}^{d\times N}$, the modulus of *contextual continuity* is defined by:

$$ $\omega^{cont}_{f}(t)\coloneqq\sup_{n,{\bm{X}},{\bm{X}}^{\prime}}\big{\{}\|f({ \bm{X}})_{:,n}-f({\bm{X}}^{\prime})_{:,n}\|_{p}:\|{\bm{X}}-{\bm{X}}^{\prime}\| _{2}\leq t,\ {\bm{X}}_{:,n}={\bm{X}}^{\prime}_{:,n},\ {\bm{X}},{\bm{X}}^{ \prime}\in[0,1]^{d\times N}\big{\}},$ $$

###### Definition 3.6 (Modulus of Token Continuity) .

Given a sequence-to-sequence continuous function $f:[0,1]^{d\times N}\to\mathbb{R}^{d\times N}$, the modulus of *token continuity* is defined by:

$$ $\displaystyle\omega^{tok}_{f}(t)\coloneqq\sup_{n,{\bm{X}},{\bm{X}}^{\prime}} \big{\{}\|f({\bm{X}})_{:,n}$ $\displaystyle-f({\bm{X}}^{\prime})_{:,n}\|_{p}:\|{\bm{X}}_{:,n}-{\bm{X}}^{ \prime}_{:,n}\|_{2}\leq t,$ $\displaystyle{\bm{X}}_{:,m}={\bm{X}}^{\prime}_{:,m}\text{ for any }m\neq n,\ { \bm{X}},{\bm{X}}^{\prime}\in[0,1]^{d\times N}\big{\}},$ $$

The modulus of *contextual continuity* measures the variation in the output of token embeddings induced by a perturbation of context. For example, we consider the following three sentences:

(1) I write papers,

(2) You write books, and

(3) Mozart writes music.

The output embedding of the second token, ‘write’, should be similar in sentences (1) and (2) due to their similar context. In contrast, a larger variation in context, as seen in sentence (3), can induces a significant variation in the output of token embedding.

On the other hand, the modulus of *token continuity* measures the variation in the output embedding caused by perturbations to the token itself within the same context. For instance, consider the following sentences:

(1) I write papers, and

(2) I draft papers.

In this example, both sentences have the same context, but the variation in the verb (‘write’ vs. ‘draft’) reflects a perturbation in the token itself. The modulus of *token continuity* quantifies how this change influences the output embeddings.
A small modulus of *token continuity* means that the output embeddings of ‘write’ and ‘draft’ are expected to be similar.

###### Definition 3.4 (Modulus of Sentence Continuity) .

###### Definition 3.5 (Modulus of Contextual Continuity) .

###### Definition 3.6 (Modulus of Token Continuity) .

### 3.3 Main Result

The following main theorem demonstrates the approximation rate of Looped Transformers in terms of the modulus of continuity and the number of loops.

###### Theorem 3.7 (Approximation rate) .

Given a function $f\in\mathcal{F}_{\mathrm{PE}}([0,1]^{d\times N})$, for any $r\in\mathbb{N}$, there exists a Looped Transformer; $\mathrm{TF}:\mathbb{R}^{(24d+1)\times N}\to\mathbb{R}^{(24d+1)\times N}$ of single head, head size $s=1$, and width size $q=99d+8$, and two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{24d+1}$ and $\mathcal{L}_{2}:\mathbb{R}^{24d+1}\to\mathbb{R}^{d}$ such that

$$ $\big{\|}\bm{\mathcal{L}}_{2}\circ\mathrm{TF}^{\circ r}\circ\bm{\mathcal{L}}_{1 }-f\big{\|}_{L^{p}([0,1]^{d\times N})}\leq\omega^{tok}_{f}(\delta\sqrt{d})+ \omega^{cont}_{f}(\delta\sqrt{Nd})+\omega_{f}(\delta\sqrt{Nd})+\mathcal{O}( \delta^{d}),$ $$

for $\delta=\big{(}(r-N)/2\big{)}^{-1/((N+1)d+1)}$, where $\bm{\mathcal{L}}_{1}$ and $\bm{\mathcal{L}}_{2}$ represent the token-wise applications of $\mathcal{L}_{1}$ and $\mathcal{L}_{2}$, respectively.

Thus, Looped Transformers are universal approximators.

###### Corollary 3.8 .

The hypothesis space of Looped Transformers $\mathcal{H}$, defined by

$$ $\mathcal{H}\coloneqq\Big{\{}\bm{\mathcal{L}_{2}}\circ{\mathrm{TF}}^{\circ r} \circ\bm{\mathcal{L}_{1}}:\,r\in\mathbb{N},\,\bm{\mathcal{L}}_{2}\textnormal{ and }\bm{\mathcal{L}}_{1}\textnormal{are token-wise affine linear maps}\Big{\}},$ $$

are dense in $\mathcal{F}_{\mathrm{PE}}([0,1]^{d\times N})$
in terms of the $L^{p}([0,1]^{d\times N})$ norm.

Here, we outline several aspects of Theorem 3.7.

- •
A function with small modulus of continuities, e.g., robust to contextual perturbations, is suited for approximation by Looped Transformers.
- •
The total parameter count is $\mathcal{O}(d)$, independent of both $\delta$ and $N$, highlighting the parameter efficiency of Looped Transformers.
- •
The optimal approximation rate of ReLU networks of size $n$ is $\mathcal{O}(\omega_{f}(\mathcal{O}(n^{-2/d})))$ for continuous functions on $[0,1]^{d}$ (Yarotsky, ); the exponential rate is unavoidable.

###### Theorem 3.7 (Approximation rate) .

###### Corollary 3.8 .

### 3.4 Proof Sketch

We provide a proof sketch that highlights differences from prior studies and discusses associated difficulties. The formal proof is provided in Appendix A.

#### Approximation with Piecewise Constant Function.

We approximate $f\in\mathcal{F_{\mathrm{PE}}}$ with piece-wise constant function $\bar{f}:[0,1]^{d\times N}\to\mathbb{R}^{d\times N}$.
Specifically, for $\delta^{-1}\in\mathbb{N}$, we divide the input space $[0,1]^{d\times N}$ into $\delta$-discretized cubes, denoted by $\{Q_{\bm{\mathcal{B}}}\}_{\bm{\mathcal{B}}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d
\times N}}$. Each cube is associated with a representative $\hat{{\bm{X}}}_{\bm{\mathcal{B}}}\in Q_{\bm{\mathcal{B}}}$. Define a piecewise constant function $\bar{f}$ for ${\bm{X}}\in[0,1]^{d\times N}$ as

$$ $\bar{f}({\bm{X}})=f(\hat{{\bm{X}}}_{\bm{\mathcal{B}}})\quad\text{where }\bm{ \mathcal{B}}\text{ satisfies }{\bm{X}}\in Q_{\bm{\mathcal{B}}}.$ $$

We can bound the error within each cube as: $\|\bar{f}({\bm{X}})-f({\bm{X}})\|_{p}\leq\omega_{f}(\sqrt{\delta^{2}+\cdots+
\delta^{2}})\leq\omega_{f}(\delta\cdot\sqrt{Nd})$ for any ${\bm{X}}\in[0,1]^{d\times N}$.
We involve three steps to construct $\bar{f}$. The first and second steps map the input ${\bm{X}}$ to the coordinates of the discretized input space, involving $\bm{\mathcal{B}}$. The third step approximately maps these coordinates to the target embeddings.

#### Step 1. Token-wise Quantization.

The network, with $(\delta^{-1}-1)$ loops, token-wise maps the input space into indices (the proof is provided in Section A.4). Then it maps them to an integer, referred to as a *token ID*:

$$ ${\bm{X}}_{:,n}\in[0,1]^{d}\to\bm{\beta}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d}\to z \in\{0,1,\dotsc,\delta^{-d}-1\}.$ (7) $$

The key idea behind our proof follows Zhang et al. (); however, we cannot directly apply it here due to the need to account for skip connections. Additionally, we have to consider a bijective mapping of $\bm{\beta}$ to a token ID in the $\delta^{-1}$-base system for the next step of contextual mapping.

#### Step 2. Contextual Mapping.

The network performs contextual mapping, which maps $N$ token IDs to a *sequence ID* in the set of $\{0,1,\dotsc,\delta^{-Nd}-1\}$.
Previous studies (Yun et al., ; Kim et al., ) use multiple layers for constructions; however, these results do not extend to Looped Transformers for two reasons.

- (1)
Yun et al. () used both sparse and uniform attention, whereas Looped Transformers are limited to a single fixed attention layer.
- (2)
Kim et al. () used $N$-layers to store $N$ parameters, whereas fixed-size $\mathcal{O}(d)$ Looped Transformers cannot store $N$ weight components.

While these considerations indicate some limitations of Looped Transformers, we found that Looped Transformers with $N$-loops can perform contextual mapping. The proof strategy follows Kim et al. ();however, we have to update for single Transformer block.
Let ${\bm{z}}\in\{0,1,\dotsc,\delta^{-d}-1\}^{N}$ represent a sequence of $N$ ordered and distinct token IDs, where ${\bm{z}}_{1}>{\bm{z}}_{2}>\cdots>{\bm{z}}_{N}$.
The networks map the set of token IDs into sequence ID via inner product with ${\bm{u}}\coloneqq(\delta^{-d(N-1)},\dotsc,\delta^{-d},1)\in\mathbb{R}^{N}$ i.e.

$$ $\mathrm{CM}({\bm{z}})\coloneq{\bm{u}}^{\top}{\bm{z}},$ $$

which satisfy

$$ $\left\lVert{\bm{u}}^{\top}{\bm{z}}-{\bm{u}}^{\top}{\bm{z}}^{\prime}\right \rVert>1,\quad\text{if }{\bm{z}}\neq{\bm{z}}^{\prime}.$ $$

Thus $\mathrm{CM}$ is a contextual mapping. The key point is that the network only needs to store $\delta$ in order to represent ${\bm{u}}$, allowing it to be implemented with Looped Transformers. Details are provided in Section A.5

#### Step 3. Token-wise Mapping.

The network token-wise maps the coordinates of discretized regions approximately to the target token embedding.
From Steps 1 and 2, each token in the input sequence is assigned a token ID with a sequence ID, where the sequence ID is consistent across all tokens. The combination of the token ID and sequence ID determines the coordinates, referred to as the *contextual token ID*. Notably, the design of the *contextual token ID* plays a crucial role in Looped Transformers.

This comes from the constrain of looped architecture.
Let ${\mathcal{K}}$ denote the set of contextual token IDs, where each index is associated with a distinct cube $\bm{\mathcal{B}}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d\times N}$ in the discretized space.
For each $k\in{\mathcal{K}}$, let ${\bm{X}}_{k}$ denote the cube associated with $\bm{\mathcal{B}}$, i.e., ${\bm{X}}_{\bm{\mathcal{B}}}$.
Looped Transformer $\tilde{f}:\mathbb{R}^{d\times N}\to\mathbb{R}^{d\times N}$ can approximate piecewise constant function $\bar{f}$ with the error of

$$ $\|\tilde{f}({\bm{X}}_{k})_{:,n}-\bar{f}({\bm{X}}_{k})_{:,n}\|\leq\max_{n,k^{ \prime}\in{\mathcal{K}}}\|\bar{f}({\bm{X}}_{k^{\prime}})_{:,n}-\bar{f}({\bm{X} }_{k^{\prime}-1})_{:,n}\|_{p}\quad\text{for any }k\in{\mathcal{K}},$ (8) $$

for $n={0,1,\dotsc,N}$ (details are provided in  Lemma 4.1).
This requires us to design ${\mathcal{K}}$ so that $\|\bar{f}({\bm{X}}_{k^{\prime}})_{:,n}-\bar{f}({\bm{X}}_{k^{\prime}-1})_{:,n}
\|_{p}$ is small, i.e., the outputs of neighboring contextual token IDs are similar.
We explain the core idea of this design with illustration.
Compare the output embeddings of the underlined tokens:

- (1)
I write papers.  ;  I write papers.   (different token ID with same sequence ID)
- (2)
I write papers.  ;  You write books.   (same token ID with different sequence ID)

The contextual continuity $\omega^{cont}_{f}$ of Definition 3.2 ensures that the outputs embeddings of ‘write’ in (2) are similar. However, none of the continuity properties provide guarantees that the outputs embeddings of ‘write’ and ‘papers’ in (1) are similar.
Thus, we design the contextual token ID so that the same token with different sequence IDs comes next to each other except for unavailable corner cases (details in Section A.4).

#### Consolidation into Single Looped Transformer.

At the end of the construction, we demonstrate that the composition of the three sub-networks from Steps 1, 2, and 3 can be realized with a single Transformer block. The proof strategy follows Zhang et al. (); however we cannot directly apply it because they requires an additional layer. We found that a single Transformer block is sufficient.

#### Deriving Approximation Rate.

Figure: Figure 1: Approximation error and modulus of continuity. The linear interpolation technique reduces the error by a factor of $1/\delta^{-1}$.
Refer to caption: extracted/5911203/figures/contextual_token_ids.png

We estimate the approximation error of our construction and establish the approximation rate.

In Step 2, we consider only the case where all $N$ input tokens are distinct, disregarding other cases. These cases can be treated as negligible when $\delta$ is small. The number of subsets where one of the $N$ tokens is duplicated is

$$ ${(\delta^{-d})}^{N}-\delta^{-d}\cdot(\delta^{-d}-1)\cdots(\delta^{-d}-N-1)<C \delta^{-(N-1)d},$ $$

where $C$ is a constant. The volume of these subsets is $C\delta^{-(N-1)d}/\delta^{-Nd}=C\delta^{d}$, thus the error with respect to the $L^{p}$ norm is $\mathcal{O}(\delta^{d})$.

In Step 3, we can bound the approximation error as follows:

$$ $\displaystyle\|\tilde{f}({\bm{X}}_{k})-\bar{f}({\bm{X}}_{k})\|_{p}\leq\max_{n} \|\tilde{f}({\bm{X}}_{k})_{:,n}-\bar{f}({\bm{X}}_{k})_{:,n}\|_{p}\leq\max_{n,k ^{\prime}\in{\mathcal{K}}}\|\bar{f}({\bm{X}}_{k^{\prime}})_{:,n}-\bar{f}({\bm{ X}}_{k^{\prime}-1})_{:,n}\|_{p}.$ $$

There are two types of error for the right-hand side term: the variation induced by contextual perturbation and the variation induced by token perturbation.
We illustrate this with examples of each pattern, as show in Fig. 1. Specifically, we consider the following three cases:

- (1)
I write papers.  ;  I write books.   (*small* perturbation of context)
- (2)
You write books.  ;  Mozart write music.   (*large* perturbation of context)
- (3)
Beethoven writes music. ; I draft papers.   (perturbation of both token and context)

The error in each case can be bounded with corresponding modulus of continuity:

- 1.
$\omega^{cont}_{f}(\delta\sqrt{d})$
- 2.
$\omega^{cont}_{f}(\sqrt{Nd})\to\ \delta\omega^{cont}_{f}(\sqrt{Nd})$
- 3.
$\omega^{tok}_{f}(\sqrt{d})+\omega^{cont}_{f}(\sqrt{Nd})\to\delta\big{(}\omega^
{tok}_{f}(\sqrt{d}))+\omega^{cont}_{f}(\sqrt{Nd})\big{)}$

where $\to$ represents the use of *linear interpolation techniques* to reduce the error with extra $\delta$-loops (see Appendix A).
Since that $\omega_{f}^{cont,tok}(n\cdot t)\leq n\cdot\omega_{f}^{cont,tok}(t)$ for any $n\in\mathbb{N}$ and $t\in[0,\infty)$ with $\delta<1$, we can then derive the upper bound for the three terms:

$$ $\displaystyle\max_{n,k^{\prime}\in{\mathcal{K}}}\|\bar{f}({\bm{X}}_{k^{\prime} })_{:,n}-\bar{f}({\bm{X}}_{k^{\prime}-1})_{:,n}\|_{p}$ $\displaystyle\leq\max\Big{\{}\omega^{cont}_{f}(\delta\sqrt{d}),\ \delta\omega^ {cont}_{f}(\sqrt{Nd}),\ \delta\big{(}\omega^{tok}_{f}(\sqrt{d})+\omega^{cont}_ {f}(\sqrt{Nd})\big{)}\Big{\}}$ $\displaystyle\leq\max\Big{\{}\omega^{cont}_{f}(\delta\sqrt{d}),\ \omega^{cont} _{f}(\delta\sqrt{Nd}),\ \omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}( \delta\sqrt{Nd})\Big{\}}$ $\displaystyle\leq\omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}(\delta \sqrt{Nd}),$ $$

With the triangle inequality, we have approximation error as

$$ $\displaystyle\|\tilde{f}-f\|_{L^{p}([0,1]^{d\times N})}$ $\displaystyle\leq\|\tilde{f}({\bm{X}})-f({\bm{X}})\|_{p}\cdot 1$ $\displaystyle\leq\|\tilde{f}({\bm{X}})-\bar{f}({\bm{X}})\|_{p}+\|\bar{f}({\bm{ X}})-f({\bm{X}})\|_{p}$ $\displaystyle\leq\max_{n,k^{\prime}\in{\mathcal{K}}}\|\bar{f}({\bm{X}}_{k^{ \prime}})_{:,n}-\bar{f}({\bm{X}}_{k^{\prime}-1})_{:,n}\|_{p}+\|\bar{f}({\bm{X} })-f({\bm{X}})\|_{p}+\mathcal{O}(\delta^{d})$ $\displaystyle\leq\omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}(\delta \sqrt{Nd})+\omega_{f}(\delta\sqrt{Nd})+\mathcal{O}(\delta^{d}).$ (9) $$

Then, $\delta$ is expressed in terms of the number of loops $r$ to determine the approximation rate. We use $\delta^{-1}-1$ loops for Step 1, $N$ loops for Step 2, and $2\delta^{-(N+1)d}-1$ loops for Step 3, with 1 loop required to connect each step. Thus we have

$$ $\displaystyle r=(\delta^{-1}-1)+1+(N)+1+\big{(}2\delta^{-(N+1)d}-1\big{)}$ $\displaystyle\Leftrightarrow\delta^{-1}+2\delta^{-(N+1)d}=r-N$ $\displaystyle\Leftrightarrow\delta^{-1}\cdot 2\delta^{-(N+1)d}\geq r-N$ $\displaystyle\Leftrightarrow 2\delta^{-(N+1)d-1}\geq r-N$ $\displaystyle\Leftrightarrow\delta\leq{\Big{(}\frac{r-N}{2}\Big{)}}^{-1/((N+1) d+1)}.$ (10) $$

From Eq. 9 and Eq. 10, we can derive Theorem 3.7.

#### Summary.

Our contribution is to establish an approximation rate for Looped Transformers by identifying the continuity of sequence-to-sequence functions.
Additionally, for technical construction, we demonstrate that a single looped Transformer block is sufficient for contextual mapping, and while Zhang et al. () requires three feed-forward layers of Looped ReLU networks for universal approximation, we achieve this with just one layer.

## 4 From Theory to Practice: Introducing Timestep Encoding

The theoretical result in Section 3 reveals a limitation specific to the looped architecture. The analysis prompts us to incorporate *time-dependent* scaling parameters for each loop, conditioned on the timestep encoding.

### 4.1 Motivation

#### Limitation Specific to Looped Architecture.

Theorem 3.7 demonstrates that the approximation rate of Looped Transformers depends on the modulus of three types of continuity. The dependency on sequence continuity arises from the approximation of continuous functions with piecewise constant functions, which is, in some sense, unavoidable.

In contrast, we found that the dependencies on contextual continuity and token continuity are specific to the looped architecture.
Referring to previous results (Yun et al., ; Kim et al., ), there is no such dependency when approximating with standard Transformers.
The dependency on the extra term indicates larger approximation errors, which limits the approximation capabilities of Looped Transformers.

On the basis of our construction, we identify the cause of dependency in Step 3: Token-wise Mapping, specifically due to Lemma 4.1, as show in below.

###### Lemma 4.1 .

Given ${\bm{y}}_{k}\in\mathbb{R}^{d}$ for $k=0,1,\dotsc,m-1$ with

$$ $|({\bm{y}}_{k}-{\bm{y}}_{k-1})_{i}|\leq\varepsilon_{i}\quad\textnormal{for $k= 1,2,\dotsc,m-1,$}$ $$

there exist feed-forward layer $\mathrm{FF}:\mathbb{R}^{14d}\to\mathbb{R}^{14d}$ of width size $20d$ and two affine linear maps ${\mathcal{L}}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{14d}$ and ${\mathcal{L}}_{2}:\mathbb{R}^{14d}\to\mathbb{R}^{d}$ such that

$$ $|\big{(}{\mathcal{L}}_{2}\circ(\mathrm{id}+\mathrm{FF})^{(m-1)}\circ{\mathcal{ L}}_{1}(k)-{\bm{y}}_{k}\big{)}_{i}|\leq\varepsilon_{i}\quad\textnormal{for $k= 0,1,\dotsc,m-1$},$ $$

for any $i=1,2,\dotsc,d$.

Lemma 4.1 implies that large variations in the target function, represented by discretized points ${\bm{y}}_{k}$, increase the approximation error for Looped Transformers. In particular, when the outputs of the target function at nearby points differ significantly, there are no guarantees of a small approximation error.

###### Lemma 4.1 .

#### How Can We Improve the Approximation Rate of Looped Transformers?

On the basis of the above considerations, we aim to alleviate or remove the dependency of the approximation rate on contextual continuity and token continuity.
The analysis prompts us to incorporate *time-dependent* parameters for each loop. We introduce a variant of the feed-forward layers, where added scaling vector parameter, dependent on the loop index (timestep), defined as follows:

$$ ${\mathrm{FF}}({\bm{X}})\to\bm{\eta}(t)\odot{\mathrm{FF}}({\bm{X}})\quad\text{ for the $t$-th loops},$ $$

where $\odot$ is an element-wise product, $t\in\mathbb{N}$ denotes the loop index (timestep), and $\bm{\eta}(t)\in\mathbb{R}^{d}$ represents a *time-dependent* scaling parameter.
This kind of *dynamic* scaling vector parameters is also used by HyperNetworks (Ha et al., ) for RNN to enhance the expressive power.

We demonstrate that the time-dependent Looped Transformer can avoid approximation errors caused by both contextual continuity and token continuity.
That is, we can replace Lemma 4.1 with the following Theorem 4.2, which shows that time-dependent models can always precisely approximate a target function.

###### Theorem 4.2 .

Given ${\bm{y}}_{k}\in\mathbb{R}^{d}$ for $k=0,\dotsc,m-1$,
there exist feed-forward layer $\mathrm{FF}:\mathbb{R}^{4d}\to\mathbb{R}^{4d}$ of width size $6d$ and $\bm{\eta}(t)\in\mathbb{R}^{4}d$ and two affine linear maps ${\mathcal{L}}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{4d}$ and ${\mathcal{L}}_{2}:\mathbb{R}^{4d}\to\mathbb{R}^{d}$ such that

$$ ${\mathcal{L}}_{2}\circ(\mathrm{id}+\bm{\eta}(m-1)\odot\mathrm{FF})\circ\cdots \circ(\mathrm{id}+\bm{\eta}(1)\odot\mathrm{FF})\circ{\mathcal{L}}_{1}(k)={\bm{ y}}_{k}.$ $$

The proof is provided in Appendix B.
While adding parameters for each loop is effective, the number of parameters increases with the number of loops. Therefore, we introduce time-dependent scaling methods, which are discussed in the following subsections.

###### Theorem 4.2 .

### 4.2 Timestep Encoding

Although various architectures for the time-dependent methods can be considered, we consider generating timestep encodings, which were first proposed by Dehghani et al. () for Looped Transformers. While their approach used timestep encoding as a simple addition to the input, recent studies on diffusion models (Peebles & Xie, ) used *time-dependent* standard Transformers by regressing layer normalization parameters from the timestep encodings. Building on these insights, we adopt this more advanced approach.

To condition on timesteps, we create frequency embeddings and then process them through a two-layer MLP, with equal dimensional to the Transformer block hidden size, and the SiLU function. Let $\operatorname{TE}(t)\in\mathbb{R}^{d}$ denote timestep embeddings defined as:

$$ $\operatorname{TE}(t)={\bm{W}}^{(TE)}_{1}\cdot\mathrm{SiLU}({\bm{W}}^{(TE)}_{2} \cdot\operatorname{PE}(t)+{\bm{b}}^{(TE)}_{1})+{\bm{b}}^{(TE)}_{2},$ $$

where ${\bm{W}}^{(TE)}_{1}\in\mathbb{R}^{d\times d},{\bm{W}}^{(TE)}_{2}\in\mathbb{R}^
{d\times d},{\bm{b}}^{(TE)}_{1}\in\mathbb{R}^{d},{\bm{b}}^{(TE)}_{2}\in\mathbb
{R}^{d}$ are parameters, and $\operatorname{PE}(t)\in\mathbb{R}^{d}$ is defined as

$$ $\displaystyle\operatorname{PE}(t)_{2i}=\sin(t/10000^{2i/d}),\quad\operatorname {PE}(t)_{2i+1}=\cos(t/10000^{2i/d}).$ $$

The root mean square layer normalization (RMSNorm) (Zhang & Sennrich, ) has been employed in several recent large language models (LLMs), such as LLaMA (et al., ) and Gemma (Team, ).
RMSNorm, denoted by ${\mathrm{RMSLN}}$, is defined as

$$ $\displaystyle\begin{split}{\mathrm{RMSLN}}({\bm{x}})=\bm{\alpha}\odot\frac{{ \bm{x}}}{\text{RMS}(\mathbf{{\bm{x}}})},\quad\text{where}~{}~{}\text{RMS}( \mathbf{{\bm{x}}})=\sqrt{\frac{1}{d}\sum_{i=1}^{d}{\bm{x}}_{i}^{2}}.\end{split}$ (11) $$

where $\bm{\alpha}\in\mathbb{R}^{d}$ is a gain parameter for rescaling.

Time-dependent RMSNorm is defined as:

$$ ${\mathrm{RMSLN}}({\bm{x}},t)=\bm{\alpha}(t)\odot\frac{{\bm{x}}}{\text{RMS}( \mathbf{{\bm{x}}})}$ $$

where $\bm{\alpha}(t)\in\mathbb{R}^{d}$ is a time-dependent parameter generated by a *hypernetowrk*.
We also add vector parameters for scaling output; Time-dependent Transformer block is defined as follows.

$$ $\displaystyle{\bm{X}}^{\prime}$ $\displaystyle={\bm{X}}+\bm{\gamma}_{1}(t)\odot{\mathrm{Attn}}(\bm{{\mathrm{ RMSLN}}}_{1}({\bm{X}},t)),$ (12) $\displaystyle{\mathrm{TF}}({\bm{X}},t)$ $\displaystyle={\bm{X}}^{\prime}+\bm{\gamma}_{2}(t)\odot\bm{{\mathrm{FF}}}(\bm{ {\mathrm{RMSLN}}}_{2}({\bm{X}}^{\prime},t)),$ (13) $$

where $\bm{\gamma}_{1}(t),\bm{\gamma}_{2}(t)\in\mathbb{R}^{d}$ are time-dependent parameters applied token-wise, and $\bm{{\mathrm{RMSLN}}_{1}}$ and $\bm{{\mathrm{RMSLN}}_{2}}$ represent token-wise applications of ${{\mathrm{RMSLN}}}_{1}$ and ${{\mathrm{RMSLN}}}_{2}$, respectively.

We use the SiLU function and weight parameters to generate time-dependent vector parameters:

$$ $\bm{\alpha}_{1}(t),\bm{\alpha}_{2}(t),\bm{\gamma}_{1}(t),\bm{\gamma}_{2}(t)={ \bm{W}}^{(H)}\cdot{\mathrm{SiLU}}(\operatorname{TE}(t))+{\bm{b}}^{(H)},\\$ $$

where ${\bm{W}}^{(H)}\in\mathbb{R}^{4d\times d},{\bm{b}}^{(H)}\in\mathbb{R}^{d}$ are parameters.

## 5 Experiments

We conducted experiments to validate our analysis by training Looped Transformers with $10$, $50$, and $100$ loops, both with and without time dependency of timestep encodings. The results show that increasing the number of loops improves expressive power in both models, and the enhancement is observed through the timestep encoding.

### 5.1 Experimental Setup

#### Problem Setting.

We use an edit distance problem. We followed the problem setting in Feng et al. (), where the costs for insertion, deletion, and replacement were set to $2$, $2$, and $3$, respectively. We set the input string length to $60$. We use $10^{6}$ training samples and perform evaluation on $10^{3}$ samples. The loss function was cross-entropy, and the evaluation metric was accuracy. Details on the data generation are provided in Appendix C.

#### Model Architecture and Training Configuration.

We used Looped Transformers based on GPT-2 decoder models (Radford et al., ), with $4$ attention heads and a $256$-dimensional, replacing the stacked multiple Transformer blocks with looped single Transformer block. We used the AdamW optimizer (Loshchilov & Hutter, ) with $\beta_{1}=0.9,\beta_{2}=0.999,\text{lr}=10^{-4}$, and $\text{weight decay}=0.01$, training for $50$ epoch with batch size $64$.
We train Looped Transformers with $10,50,100$-loops, with and without time dependency.
For time-dependent models, we initialize $\gamma(t)$ as zero vectors and $\alpha(t)$ as one vectors, following Peebles & Xie () and Bachlechner et al. ().

Figure: Figure 2: Training and test accuracy of Looped Transformers with $10$, $50$, and $100$ loops, both with and without time dependency, on the edit distance task.
Refer to caption: x1.png

### 5.2 Results.

Figure 2 presents the training and test accuracy results for models with timestep encoding and without time dependency. In both cases, training and test accuracies improve as the number of loops increases ($100$ loops $>$ $50$ loops $>$ $10$ loops).

The accuracy gains for the Looped Transformer without time dependency are minimal between $50$ and $100$ loops. In contrast, the timestep encoding models show a consistent improvement in accuracy without saturating as the number of loops increases. Despite these differences, both cases support the analysis in Theorem 3.7, which demonstrates that the approximation rate improves as the number of loops increases.

In addition, when comparing models without time dependency to those with timestep encoding, we observe a consistent performance improvement in all models with timestep encoding. This validates the analysis in Section 4, which highlights the limitations of the looped structure and suggests that incorporating scaling parameters can mitigate these limitations, thereby enhancing the expressive power.

## 6 Conclusion

In this paper, we have established approximation rates for Looped Transformers by defining the concept of the modulus of continuity for sequence-to-sequence functions. Our analysis reveals a limitation specific to the looped architecture, which led us to incorporate a time-dependent scaling parameter. This research is the first to explore the function approximation power of Looped Transformers and serves as a foundation for future studies.
While our results require few assumptions about the target functions, making them broadly applicable, future work could explore identifying specific function classes or properties to derive alternative approximation rates.
Additionally, designing an optimal time-dependent architecture for Looped Transformers to achieve greater enhancement could be a promising direction for future research.

## Appendix A Proofs of Function approximation

We present the formal proof of Theorem 3.7. In Section A.2, we provide the whole proof relying on lemmas, introduced in the following subsections. In Section A.4, Section A.5, and Section A.6, we provide the key lemmas. The overview of proof outline is shown in Fig 3.

Figure: Figure 3: Overview of proof sketch.
Refer to caption: extracted/5911203/figures/proof_sketch.png

### A.1 Notations

- •
The bold notation for functions indicates that they are applied in a token-wise manner.
- •
We define the nearest functions as follows:
$\displaystyle\operatorname{near^{+}}(a,\mathcal{S})$
$\displaystyle\coloneqq\operatorname*{arg\,min}_{b\in\mathcal{S},b>a}\left|a-b\right|$
$\displaystyle\operatorname{near^{-}}(a,\mathcal{S})$
$\displaystyle\coloneqq\operatorname*{arg\,min}_{b\in\mathcal{S},b<a}\left|a-b\right|$
The function $\operatorname{near}(a,\mathcal{S})$ identifies the element in the set $\mathcal{S}$ that is closest to $a$. The function $\operatorname{near^{+}}(a,\mathcal{S})$ finds the closest element greater than $a$, while $\operatorname{near^{-}}(a,\mathcal{S})$ identifies the closest element less than $a$.
- •
Given two values $a$ and $b$, and the number of divisions $n$, we can define the linear interpolation output at an index $t$ (where $t$ ranges from 0 0 to $n$) as follows:
$\operatorname{lin_interp}(a,b,t,n)\coloneqq a+\frac{t}{n}(b-a)$

### A.2 Proof of Theorem 3.7

See 3.7

###### Proof.

Since any continuous function can be approximated by a piecewise constant function with arbitrarily small errors, we approximate $f\in\mathcal{F_{\mathrm{PE}}}$ with piece-wise constant function $\bar{f}:[0,1]^{d\times N}\to\mathbb{R}^{d\times N}$.
We choose $\delta^{-1}\in\mathbb{N}$, which determines how finely the input is divided; then, we divide the input space $[0,1]^{d\times N}$ into $\delta$-discretized cubes, denoted by $\{\bm{Q}_{\bm{\mathcal{B}}}\}$ for ${\bm{\mathcal{B}}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d\times N}}$ defined by

$$ $\bm{Q}_{\bm{\mathcal{B}}}\coloneqq\Big{\{}{\mathbf{X}}\in[0,1]^{d\times N}:{ \mathbf{X}}_{i,n}\in\big{[}\bm{\mathcal{B}}_{i,n}\delta,\bm{\mathcal{B}}_{i,n} \delta+1\big{)},\quad i=1,2,\dotsc,Nd\Big{\}}.$ $$

Note that we do not consider the *trifling regions* in Zhang et al. () as we use the $L_{p}$ norm rather than the uniform norm.

Each cube is associated with representative $\hat{{\bm{X}}}_{\bm{\mathcal{B}}}\in Q_{\bm{\mathcal{B}}}$.
We can define piecewise constant function $\bar{f}$ for ${\bm{X}}\in[0,1]^{d\times N}$ as

$$ $\bar{f}({\bm{X}})\coloneq f(\hat{{\bm{X}}}_{\bm{\mathcal{B}}})\quad\text{where }\bm{\mathcal{B}}\text{ satisfies }{\bm{X}}\in\bm{Q}_{\bm{\mathcal{B}}}.$ $$

We can bound the error within each cube as: $\|\bar{f}({\bm{X}})-f({\bm{X}})\|_{p}\leq\omega_{f}(\delta\sqrt{Nd})$ for any ${\bm{X}}\in[0,1]^{d\times N}$.

Our construction consists of three steps to approximate $\bar{f}$. The first and second steps map the input ${\bm{X}}$ to the coordinates of the discretized input space. The third step approximately maps these coordinates to the target embeddings.
The outline of three step is:

- 1.
The network, with $(\delta^{-1}-1)$-loops, maps the input space $[0,1]^{d}$ token-wise to the coordinates $\bm{\beta}\in\{0,1,\dots,\delta^{-1}-1\}^{d}$ of divided cubes, and then bijectively maps these coordinates to an integer, representing token IDs in the set $\{0,1,\dots,\delta^{-d}\}$, using a $\delta^{-1}$-base system.
- 2.
The network, with $N$ loops, performs a contextual mapping from the set of distinct $N$ token IDs into the set of *extended sequence ID*. Specifically, the network using Step1 and Step2 maps the discretized coordinates ${\bm{\mathcal{B}}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d\times N}}$, for each dimension $d$ and each token $N$, representing it in a $\delta^{-1}$ system with $Nd$ digits ($N$ tokens are ordered). Furthermore, *dummy indices* are needed to reduce the approximation error in the next step.
- 3.
The network, with $\delta^{2(N+1)d}-1$ loops, approximately maps *contextual token IDs* to the output embeddings of each token in a token-wise manner. *Contextual token IDs* refer to token IDs assigned to each token within the context of a sequence ID. To achieve a small approximation error, the network has to be designed so that neighboring IDs correspond to similar output token embeddings.

The details for each steps are provided below.

###### Proof.

#### Step 1. Token-wise Quantization.

The input space for each token ${\bm{x}}\in[0,1]^{d}$ are divided into $\delta$-discretized cubes, denoted by $\{Q_{\bm{\beta}}\}$ for ${\bm{\beta}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d}}$, defined as

$$ $Q_{\bm{\beta}}\coloneqq\Big{\{}{\bm{x}}\in[0,1]^{d}:{\bm{x}}_{i}\in\big{[}\bm{ \beta}_{i}\delta,\ \bm{\beta}_{i}\delta+1\big{)},\quad i=1,2,\dotsc,d\Big{\}}.$ $$

By Lemma A.4,
there exist a feed-forward layer ${\mathrm{FF}}^{(1)}:\mathbb{R}^{5d}\to\mathbb{R}^{5d}$ of width size $q=7d$, and two affine linear maps $\mathcal{L}^{(1)}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{5d}$ and $\mathcal{L}^{\prime(1)}_{2}:\mathbb{R}^{5d}\to\mathbb{R}^{d}$ such that

$$ $\mathcal{L}^{\prime(1)}_{2}\circ\big{(}\mathrm{id}+{\mathrm{FF}}^{(1)}\big{)}^ {\circ(\delta^{-1}-1)}\circ\mathcal{L}^{(1)}_{1}({\bm{x}})=\bm{\beta}\quad \textit{s}.\textit{t}.\quad{\bm{x}}_{i}\in\big{[}\bm{\beta}_{i}\delta,\ \bm{ \beta}_{i}\delta+1\big{)}$ $$

for any $i=1,2,\dotsc,d$.
Let $\bm{\mathcal{B}}$ denote the output of the token-wise application to the input ${\bm{x}}\in[0,1]^{d\times N}$.

In addition, we need to bijectively map the $d$-dimensional vector $\bm{\beta}$ to an integer *token ID*, denoted by $z$. We use a $\delta^{-1}$-base system; define the vector ${\bm{u}}_{(\delta^{-1})}\in\mathbb{R}^{d}$ as

$$ ${\bm{u}}_{(\delta^{-1})}\coloneqq(\delta^{-(d-1)},\delta^{-(d-2)},\dotsc, \delta^{-1},1)^{\top},$ $$

and define $z$ as

$$ $z\coloneq{\bm{u}}_{(\delta^{-1})}^{\top}\bm{\beta}\in\{0,1,\dotsc,\delta^{-d}- 1\}.$ $$

To implement this, we define $\mathcal{L}^{(1)}_{2}:\mathbb{R}^{5d}\to\mathbb{R}^{d}$ with $\mathcal{L}^{\prime(1)}_{2}:\mathbb{R}^{5d}\to\mathbb{R^{d}}$ via

$$ $\mathcal{L}^{(1)}_{2}({\bm{x}})={\bm{u}}_{(\delta^{-1})}^{\top}\mathcal{L}^{ \prime(1)}_{2}({\bm{x}}).$ $$

Thus, we have

$$ $\left(\bm{\mathcal{L}}^{(1)}_{2}\circ({\mathrm{FF}}^{(1)}_{1})^{\circ{(\delta^ {-1}-1)}}\circ\bm{\mathcal{L}}^{(1)}_{1}({\bm{x}})\right)_{n}={\bm{u}}_{( \delta^{-1})}^{\top}\bm{\beta}=z\quad\textit{s}.\textit{t}.\quad{\bm{x}}_{:,n} \in Q_{\bm{\beta}},$ $$

for $n=1,2,\dotsc,N$.

We define the input cubes for each token assigned to $z$ as follows:

$$ $Q_{z}\coloneqq\Big{\{}{\bm{x}}\in[0,1]^{d}:{\bm{x}}_{i}\in\big{[}\bm{\beta}_{i }\delta,\ \bm{\beta}_{i}\delta+1\big{)}\quad\text{for }i=1,2,\dotsc,d\quad \textit{s}.\textit{t}.\quad z={\bm{u}}_{(\delta^{-1})}^{\top}\bm{\beta}\Big{\}}.$ $$

We can confirm that

$$ $\|{\bm{x}}_{z}-{\bm{x}}_{z-1}\|_{2}\leq\delta\sqrt{d}\quad\forall{\bm{x}}_{z} \in Q_{z},\forall{\bm{x}}_{z-1}\in Q_{z-1}\quad\textit{s}.\textit{t}.\quad\bm{ \beta}_{d}=1,2,\dotsc,\delta^{-1}-1.$ $$

While, we have

$$ $\|{\bm{x}}_{z}-{\bm{x}}_{z-1}\|_{2}\leq\sqrt{d}\quad\forall{\bm{x}}_{z}\in Q_{ z},\forall{\bm{x}}_{z-1}\in Q_{z-1}\quad\textit{s}.\textit{t}.\quad\bm{\beta}_ {d}=0.$ (14) $$

Informally, in this $\delta^{-1}$-base representation, the least significant digit corresponds to the index of the $d$-th dimension, $\bm{\beta}_{d}$. When incrementing the token ID one by one, the index in the $d$-th dimension increases sequentially $0,1,2,\dots,\delta^{-1}-1$, without changing the other digits; thus, consecutive IDs imply that the tokens are “similar” to each other in $d$-dimensional space, with a distance of at most $\delta\sqrt{d}$. However, when a carry occurs, the higher-order digits may change significantly, causing tokens that were not originally adjacent to appear next to each other in terms of their indices, as we are essentially projecting a $d$-dimensional space onto a $1$-dimensional space. In this case, the distance is only bounded by $\sqrt{d}$.

#### Step 2. Contextual Mapping.

The networks, with $N$-loops, map the list of $N$ token IDs, denoted by ${\bm{z}}\in\{0,1,\dotsc,\delta^{-d}-1\}^{N}$, to sequence IDs bijectively. Furthermore, this mapping is not only bijective; it also requires the inclusion of additional *dummy indices*.

Note: We consider only the case where all $N$ input tokens are distinct, disregarding other cases, which can be treated as negligible when $\delta$ is small. The number of subsets where one of the $N$ tokens is duplicated is

$$ ${(\delta^{-d})}^{N}-\delta^{-d}\cdot(\delta^{-d}-1)\dotsc(\delta^{-d}-N-1)<C \delta^{-(N-1)d},$ $$

where $C$ is a constant. The volume of these subsets is $C\delta^{-(N-1)d}/\delta^{-Nd}=C\delta^{d}$, so the error with respect to the $L^{p}$ norm is $\mathcal{O}(\delta^{d})$.

Let ${\mathbb{L}}_{\delta}$ denote the set, composed of distinct $N$ tokens, i.e.

$$ ${\mathbb{L}}_{\delta}\coloneq\{{\bm{z}}\in\{0,1,\dotsc,\delta^{-d}-1\}^{N}\mid {\bm{z}}_{i}\neq{\bm{z}}_{j}\text{ for all }i\neq j\}.$ $$

Due to permutation equivalence, we can assume without loss of generality that elements of ${\bm{z}}\in{\mathbb{L}}_{\delta}$ is ordered, i.e., ${\bm{z}}_{1}>{\bm{z}}_{2}>\cdots>{\bm{z}}_{N}$.
Define ${\bm{u}}_{(\delta^{-d})}\coloneqq(\delta^{-(N-1)d},\dotsc,\delta^{-d},1)^{\top}$, which satisfy

$$ $\left\lVert{\bm{u}}_{(\delta^{-d})}^{\top}{\bm{z}}-{\bm{u}}_{(\delta^{-d})}^{ \top}{\bm{z}}^{\prime}\right\rVert>1,\quad\text{if }{\bm{z}}\neq{\bm{z}}^{ \prime}\text{ for any }{\bm{z}},{\bm{z}}^{\prime}\in{\mathbb{L}}_{\delta}.$ $$

This mapping, ${\bm{u}}_{(\delta^{-d})}^{\top}{\bm{z}}$, represents ${\bm{z}}$ in a $\delta^{-d}$-base system.
Then, we define sequence ID of ${\bm{z}}\in{\mathbb{L}}_{\delta}$ as:

$$ $\mathrm{s}({\bm{z}})\coloneqq{\bm{u}}_{(\delta^{-d})}^{\top}{\bm{z}}=\sum_{n=1 }^{N}{\bm{z}}_{n}\delta^{-(N-n)d}.$ (15) $$

By Lemma A.5, there exist a Transformer block $\mathrm{TF^{\prime(2)}}:\mathbb{R}^{3\times N}\to\mathbb{R}^{3\times N}$ with single head, head size $s=1$, and width size $q=3$, and two affine linear maps $\mathcal{L}^{\prime(2)}_{1}:\mathbb{R}\to\mathbb{R}^{3}$ and $\mathcal{L}^{\prime(2)}_{2}:\mathbb{R}^{3}\to\mathbb{R}$ such that

$$ $\bm{\mathcal{L}}^{\prime(2)}_{2}\circ{{\mathrm{TF}}^{\prime(2)}}^{\circ N} \circ\bm{\mathcal{L}}^{\prime(2)}_{1}({\bm{z}}^{\top})=\mathrm{s}({\bm{z}}) \cdot\mathbf{1}_{N}^{\top},$ $$

where ${\bm{z}}^{\top}\to{\bm{u}}_{(\delta^{-d})}^{\top}{\bm{z}}$ is a contextual mapping.

Furthermore, we have to add *dummy indices* to alleviate the approximation error caused by the looped architecture in Step 3.
Recall that ${\bm{\mathcal{B}}\in\{0,1,\dotsc,\delta^{-1}-1}^{d\times N}\}$ represents the coordinates of the inputs. Let ${\bm{Z}}\in\{0,1,\dotsc,\delta^{-1}-1\}^{d\times N}$ denote the coordinates where the tokens are ordered by their token IDs, i.e., ${{\bm{u}}_{(\delta^{-1})}^{\top}{\bm{Z}}}_{1}>{{\bm{u}}_{(\delta^{-1})}^{\top}
{\bm{Z}}}_{2}>\cdots>{{\bm{u}}_{(\delta^{-1})}^{\top}{\bm{Z}}}_{N}$, i.e., ${\bm{z}}={\bm{u}}_{(\delta^{-1})}^{\top}{\bm{Z}}$.
Thus, by redefining the sequence ID of Eq. 15 for ${\bm{Z}}$ instead of ${\bm{z}}$, sequence IDs in $\delta^{-d}$-base can be rewritten in the $\delta^{-1}$-base system as follows:

$$ $\displaystyle\mathrm{s}({\bm{Z}})$ $\displaystyle={\bm{u}}_{(\delta^{-d})}^{\top}({\bm{u}}_{(\delta^{-1})}^{\top}{ \bm{Z}})$ $\displaystyle=\sum_{i=1}^{d}\sum_{n=1}^{N}{\bm{Z}}_{i,n}\delta^{-\big{(}(N-n)d +(d-i)\big{)}}.$ $$

We also define the sequence IDs for each $\bm{\mathcal{B}}$ as:

$$ $\displaystyle\mathrm{s}(\bm{\mathcal{B}})=\sum_{i=1}^{d}\sum_{n=1}^{N}\bm{ \mathcal{B}}_{i,n}\delta^{-\big{(}(N-n)d+(d-i)\big{)}}.$ $$

Then, we define *extended* sequence IDs as:

$$ $\displaystyle\mathrm{s}_{\text{extend}}({\bm{Z}})$ $\displaystyle\coloneq 2\mathrm{s}({\bm{Z}})-{\bm{Z}}_{d,N}$ $\displaystyle={\bm{Z}}_{d,N}+\sum_{\begin{subarray}{c}i=1\\ (i,n)\neq(d,N)\end{subarray}}^{d}\sum_{n=1}^{N}2{\bm{Z}}_{i,n}\delta^{-\big{(} (N-n)d+(d-i)\big{)}},$ $\displaystyle\mathrm{s}_{\text{extend}}(\bm{\mathcal{B}})$ $\displaystyle\coloneq\bm{\mathcal{B}}_{d,N}+\sum_{\begin{subarray}{c}i=1\\ (i,n)\neq(d,N)\end{subarray}}^{d}\sum_{n=1}^{N}2\bm{\mathcal{B}}_{i,n}\delta^{ -\big{(}(N-n)d+(d-i)\big{)}}.$ $$

Recall that $\bm{\mathcal{B}}_{d,N}\in\{0,1,\dotsc,\delta^{-1}-1\}$. We define the *dummy indices* as:

$$ $\displaystyle\mathrm{s}^{b}_{\text{dummy}}(\bm{\mathcal{B}})$ $\displaystyle\coloneq b+\sum_{\begin{subarray}{c}i=1\\ (i,n)\neq(d,N)\end{subarray}}^{d}\sum_{n=1}^{N}2\bm{\mathcal{B}}_{i,n}\delta^{ -\big{(}(N-n)d+(d-i)\big{)}},\quad\text{for }b=\delta^{-1},\delta^{-1}+1, \dotsc,2\delta^{-1}-1.$ $$

We define each set as follows:

$$ $\displaystyle{\mathcal{S}}_{\text{distinct}}\coloneqq\left\{\mathrm{s}_{\text{ extend}}({\bm{Z}})\right\},\quad{\mathcal{S}}_{\text{distinct+duplicate}} \coloneqq\left\{\mathrm{s}_{\text{extend}}(\bm{\mathcal{B}})\right\},$ $\displaystyle{\mathcal{S}}_{\text{dummy}}\coloneqq\left\{\mathrm{s}^{b}_{\text {dummy}}(\bm{\mathcal{B}})\mid b=\delta^{-1},\delta^{-1}+1,\dotsc,2\delta^{-1} -1\right\}.$ $$

Then, we can see that

$$ $\displaystyle{\mathcal{S}}_{\text{distinct}}\subset{\mathcal{S}}_{\text{ distinct+duplicate}},\quad{\mathcal{S}}_{\text{distinct+duplicate}}\cap{ \mathcal{S}}_{\text{dummy}}=\emptyset,$ $\displaystyle{\mathcal{S}}_{\text{distinct+duplicate}}\cup{\mathcal{S}}_{\text {dummy}}=\{0,1,\dotsc,2\delta^{-Nd}-1\}.$ (16) $$

We define the input cubes for each token assigned to extended sequence ID $s$ as follows:

$$ $\bm{Q}_{s}\coloneqq\Big{\{}{\mathbf{X}}\in[0,1]^{d\times N}:{\mathbf{X}}_{i,n} \in\big{[}\bm{\mathcal{B}}_{i,n}\delta,\bm{\mathcal{B}}_{i,n}\delta+1\big{)} \quad\text{for }i=1,2,\dotsc,Nd\quad\textit{s}.\textit{t}.\quad s=\mathrm{s}_{ \text{extend}}(\bm{\mathcal{B}})\Big{\}}.$ $$

We can confirm for that

$$ $\|{\bm{X}}_{s}-{\bm{X}}_{s-1}\|_{2}\leq\delta\sqrt{Nd}\quad\forall{\bm{X}}_{s} \in Q_{s},\forall{\bm{X}}_{s-1}\in Q_{s-1}\quad\textit{s}.\textit{t}.\quad\bm{ \mathcal{B}}_{d,N}=1,2,\dotsc,\delta^{-1}-1.$ (17) $$

While, we have

$$ $\|{\bm{X}}_{s}-{\bm{X}}_{s-1}\|_{2}\leq\sqrt{Nd}\quad\forall{\bm{X}}_{s}\in Q_ {s},\forall{\bm{X}}_{s-1}\in Q_{s-1}\quad\textit{s}.\textit{t}.\quad\bm{ \mathcal{B}}_{d,N}=0.$ (18) $$

Informally, this set represents a collection of consecutive sequence IDs, where the first $N-1$ tokens are identical, and the $N$-th token is “similar”.
For example: (1) I write ‘a paper’, and (2) I write ‘papers’.
As explained in token IDs, if only the $d$-th dimension of the $N$-th token differs, then the tokens remain adjacent in the $d$-dimensional space. As a result, the total variation in the input sequence within the set is bounded by $\delta\sqrt{d}$.

To implement this, we slightly modified ${\mathrm{TF}}^{\prime}_{2}$. By Corollary A.6, there exist a Transformer block $\mathrm{TF}^{(2)}:\mathbb{R}^{5\times N}\to\mathbb{R}^{5\times N}$ with single head, head size $s=1$, and width size $q=4$, and two affine linear maps $\mathcal{L}^{(2)}_{1}:\mathbb{R}^{2}\to\mathbb{R}^{5}$ and $\mathcal{L}^{(2)}_{2}:\mathbb{R}^{5}\to\mathbb{R}$ such that

$$ $\bm{\mathcal{L}}^{(2)}_{2}\circ{{\mathrm{TF}}^{(2)}}^{\circ N}\circ\bm{ \mathcal{L}}^{(2)}_{1}(\left(\left[\begin{array}[]{ccc}{{\bm{z}}}^{\top}\\ {{\bm{Z}}_{d,:}}\\ \end{array}\right]\right))=\mathrm{s}_{\text{extend}}({\bm{z}})\cdot\mathbf{1} _{N}^{\top}\quad\textnormal{for any ${\bm{z}}\in{\mathbb{L}}_{\delta}$.}$ $$

#### Step 3. Token-wise Mapping.

From Step 1 and Step 2, the network takes a token ID and sequence ID as input for each token, which together form a *contextual token ID*. The network with $\delta^{-2(N+1)d}-1$ loops approximately maps these *contextual token IDs* to the output token embeddings of the target function.

To construct contextual token IDs, we define a bijective mapping $\mathcal{L}^{(3)}_{0}:\mathbb{N}^{2}\to\mathbb{N}$ as follows:

$$ $\mathcal{L}^{(3)}_{0}(z,s)\coloneq 2\delta^{-Nd}z+s,$ $$

where $z$ represents a token ID and $s$ represents a extended sequence ID.
Note that sequence IDs are less than $2\delta^{-Nd}$, so informally, it’s as if we are adding another digit, ${\bm{z}}_{n}$, as the most significant digit in a $\delta^{-d}$-based system.
Define the set of contextual token IDs for distinct $N$ token IDs as:

$$ ${\mathcal{K}}_{\text{distinct}}\coloneqq\bigg{\{}\mathcal{L}_{3}\big{(}{\bm{z} }_{n},\mathrm{s}_{\text{extend}}({\bm{Z}})\big{)}:n=1,2,\dotsc,N\bigg{\}}.$ $$

We also define the set for all inputs, including cases where some token IDs are duplicated, as

$$ $\displaystyle{\mathcal{K}}_{\text{distinct+duplicate}}$ $\displaystyle\coloneqq\bigg{\{}\mathcal{L}_{3}\big{(}\bm{\beta}_{n},\mathrm{s} _{\text{extend}}(\bm{\mathcal{B}})\big{)}:n=1,2,\dotsc,N\bigg{\}}$ $\displaystyle=\bigg{\{}s+2\delta^{-Nd}z:s\in{\mathcal{S}}_{\text{distinct+ duplicate}},z=0,1,\cdots,\delta^{-d}\bigg{\}}$ $\displaystyle=\bigg{\{}\bm{\mathcal{B}}_{d,N}+\sum_{\begin{subarray}{c}i=1\\ (i,n)\neq(d,N)\end{subarray}}^{d}\sum_{n=1}^{N}2\bm{\mathcal{B}}_{i,n}\delta^{ -\big{(}(N-n)d+(d-i)\big{)}}+2\delta^{-Nd}z\bigg{\}},$ $$

and dummy contextual token ID as

$$ ${\mathcal{K}}_{\text{dummy}}\coloneqq\bigg{\{}s+2\delta^{-Nd}z:s\in{\mathcal{S }}_{\text{dummy}},z=0,1,\cdots,\delta^{-d}\bigg{\}}.$ $$

We can confirm, from Eq. 16,

$$ $\displaystyle{\mathcal{K}}_{\text{distinct}}\subset{\mathcal{K}}_{\text{ distinct+duplicate}},\quad{\mathcal{K}}_{\text{distinct+duplicate}}\cap{ \mathcal{K}}_{\text{dummy}}=\emptyset,$ $\displaystyle{\mathcal{K}}_{\text{distinct+duplicate}}\cup{\mathcal{K}}_{\text {dummy}}=\{0,1,\dotsc,2\delta^{-(N+1)d}-1\}.$ $$

We now define the target output token embedding for each extended contextual mapping ID. Define $g:{\mathcal{K}}\to\mathbb{R}^{d}$, the mapping from extended contextual IDs to their token embeddings, as follows:

$$ $g(k)=\begin{cases}\bar{f}({\bm{X}})_{:,n}\quad\textit{s}.\textit{t}.\quad \mathcal{L}_{3}\big{(}\bm{\beta}_{n},\mathrm{s}_{\text{extend}}(\bm{\mathcal{B }})\big{)}=k&\text{for }k\in{\mathcal{K}}_{\text{distinct+duplicate}},\\ \operatorname{lin_interp}\left(g(\operatorname{near^{-}}(k)),g(\operatorname{ near^{+}}(k)),k-\operatorname{near^{-}}(k),\delta^{-1}\right)&\text{for }k\in{ \mathcal{K}}_{\text{dummy}}.\end{cases}$ $$

When $\bm{\mathcal{B}}_{d,n}\neq 0$, from Eq. 17, we have

$$ $g(k)-g(k-1)\leq\omega^{cont}_{f}(\delta\sqrt{d}),$ $$

since the token ${\bm{z}}_{n}$ remains the same for both $k$ and $k-1$.
While $\bm{\mathcal{B}}_{d,n}=0$, a carryover occurs, and we only bound the input sequences as in Eq.18, with the possibility that the token ID differs, the variation is bounded as Eq.14.
Since we have dummy token IDs, there are redundant indices before the carryover occurs. As a result, the variation in the next index is bounded by a factor of $\delta$, i.e.,

$$ $\|g(k)-g(k-1)\|_{p}\leq\delta\big{(}\omega^{cont}_{f}(\sqrt{Nd})+\omega^{tok}_ {f}(\sqrt{d})\big{)}.$ $$

Since that $\omega_{f}^{cont,tok}(n\cdot t)\leq n\cdot\omega_{f}^{cont,tok}(t)$ for any $n\in\mathbb{N}$ and $t\in[0,\infty)$ with $\delta<1$, we have

$$ $\|g(k)-g(k-1)\|_{p}\leq\omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}( \delta\sqrt{Nd}).$ $$

Thus, we have

$$ $|(g(k)_{i}-g(k-1)_{i}|\leq\omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}( \delta\sqrt{Nd})\quad\textnormal{for $i=1,2,\dotsc,d$}.$ $$

By Lemma 4.1,
there exist feed-forward layer $\mathrm{FF}^{(3)}:\mathbb{R}^{14d}\to\mathbb{R}^{14d}$ of width size $20d$ and two affine linear maps ${\mathcal{L}}^{(3)}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{14d}$ and ${\mathcal{L}}^{(3)}_{2}:\mathbb{R}^{14d}\to\mathbb{R}^{d}$ such that

$$ $|\big{(}{\mathcal{L}}^{(3)}_{2}\circ{(\mathrm{id}+{\mathrm{FF}}^{(3)})}^{(2 \delta^{-(N+1)d}-1)}\circ{\mathcal{L}}^{(3)}_{1}(k)-g(k)\big{)}_{i}|\leq\omega ^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}(\delta\sqrt{Nd}),$ $$

for any $i=1,2,\dotsc,d$ and $k=0,1,\dotsc,2\delta^{-(N+1)d}-1$.

#### Consolidation into Single Looped Transformer.

Lastly, we demonstrate $\tilde{f}$ can be represented in Looped Transformers. Let

$$ ${\bm{X}}^{(0)}\in\mathbb{R}^{1\times N},\quad{\bm{X}}^{(1)}\in\mathbb{R}^{5d \times N},\quad{\bm{X}}^{(2)}\in\mathbb{R}^{5d\times N},\quad{\bm{X}}^{(3)}\in \mathbb{R}^{14d\times N}$ $$

denote the divided input space.
Define single head ${\mathrm{Attn}}:\mathbb{R}^{(24d+1)\times N}\to\mathbb{R}^{(24d+1)\times N}$ with head size $1$ as:

$$ ${\mathrm{Attn}}\left(\left[\begin{array}[]{c}{\bm{X}}^{(0)}\\ {\bm{X}}^{(1)}\\ {\bm{X}}^{(2)}\\ {\bm{X}}^{(3)}\end{array}\right]\right)=\left[\begin{array}[]{c}\bm{0}_{5d+1 \times N}\\ {\mathrm{Attn}}^{(2)}({\bm{X}}^{(2)})\\ \bm{0}_{14d\times N}\end{array}\right],$ $$

where ${\mathrm{Attn}}^{(2)}$ denote the self-attention layer of ${\mathrm{TF}}^{(2)}$.
Let

$$ ${\bm{x}}_{0}\in\mathbb{R},\quad{\bm{x}}_{1}\in\mathbb{R}^{5d},\quad{\bm{x}}_{2 }\in\mathbb{R}^{5d},\quad{\bm{x}}_{3}\in\mathbb{R}^{14d}$ $$

denote the token-wise input space.
Define ${\mathrm{FF}}:\mathbb{R}^{24d+1}\to\mathbb{R}^{24d+1}$, with impulse function defined in Proposition A.3, as:

$$ $\displaystyle{\mathrm{FF}}\left(\left[\begin{array}[]{c}{\bm{x}}_{0}\\ {\bm{x}}_{1}\\ {\bm{x}}_{2}\\ {\bm{x}}_{3}\\ \end{array}\right]\right)=\left(\left[\begin{array}[]{c}{\bm{x}}_{0}+1\\ {\mathrm{FF}}^{(1)}({\bm{x}}_{1})\\ {\mathrm{FF}}^{(2)}({\bm{x}}_{2})+\bm{\mathrm{impulse}}_{(\delta^{-1}-1)}\big{ (}\mathcal{L}^{(2)}_{1}\circ\mathcal{L}^{(1)}_{2}({\bm{x}}_{1},({\bm{x}}_{1})_ {d})\big{)}\\ {\mathrm{FF}}^{(3)}({\bm{x}}_{3})+\bm{\mathrm{impulse}}_{(\delta^{-1}+N)}\big{ (}\mathcal{L}^{(3)}_{1}\circ\mathcal{L}^{(3)}_{0}({\bm{x}}_{1},{\bm{x}}_{2}) \big{)}\\ \end{array}\right]\right)$ $$

where ${\mathrm{FF}}^{(2)}$ denotes the feed-forward layer of ${\mathrm{TF}}^{(2)}$, and $\bm{\mathrm{impulse}}$ refers to the dimension-wise application of $\mathrm{impulse}$.
Note that ${\bm{x}}_{0}$ serves the role of a counter.
As shown in Proposition A.3, the impulse requires 4 ReLU functions per dimension. With $19d$ dimensions, this results in an additional width of $72d+4$.
The total width consists of $7d$ for ${\mathrm{FF}}^{(1)}$, $4$ for ${\mathrm{TF}}^{(2)}$, and $20d$ for ${\mathrm{FF}}^{(3)}$, resulting in a total width of $99d+8$.
Define two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{24d+1}$ and $\mathcal{L}_{2}:\mathbb{R}^{24d+1}\to\mathbb{R}^{d}$ such that

$$ $\mathcal{L}_{1}({\bm{x}})={(0,\mathcal{L}^{(1)}_{1}({\bm{x}}),{\bm{0}}_{19d})} ^{\top},\quad\mathcal{L}_{2}\big{(}{({\bm{x}}_{0},{\bm{x}}_{1},{\bm{x}}_{2},{ \bm{x}}_{3})}^{\top}\big{)}=\mathcal{L}^{(3)}_{2}({\bm{x}}_{3}).$ $$

Thus, we have

$$ $\displaystyle\tilde{f}({\bm{X}})=\bm{\mathcal{L}}_{2}\circ\mathrm{TF}^{\circ( \delta^{-1}+N+2\delta^{-(N+1)d})}\circ\bm{\mathcal{L}}_{1}({\bm{X}})\quad \textit{s}.\textit{t}.\quad\|\tilde{f}({\bm{X}})-\bar{f}({\bm{X}})\|_{p}\leq \omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}(\delta\sqrt{Nd})+\mathcal{O }(\delta^{d}).$ $$

#### Deriving Approximation Error.

Generally, the following inequality holds.

$$ $\sum_{i=1}^{n}x_{i}^{p}\leq\left(\sum_{i=1}^{n}x_{i}\right)^{p}\quad\text{for }x_{i}\geq 0\text{ and }p\geq 1.$ $$

Substituting $x_{i}=\|\bar{f}({\bm{X}})-f({\bm{X}})\|_{p}$ into the above inequality results in

$$ $\|\bar{f}-f\|_{L^{p}([0,1]^{d\times N})}=\Big{(}\int\left\lVert\bar{f}({\bm{X} })-f({\bm{X}})\right\rVert_{p}^{p}d{\bm{X}}\Big{)}^{1/p}\leq\int\|\bar{f}({\bm {X}})-f({\bm{X}})\|_{p}d{\bm{X}}.$ (19) $$

Also we can bound the entire norm with the token-wise norm as:

$$ $\displaystyle\|\tilde{f}({\bm{X}}_{k})-\bar{f}({\bm{X}}_{k})\|_{p}\leq\max_{n} \|\tilde{f}({\bm{X}}_{k})_{:,n}-\bar{f}({\bm{X}}_{k})_{:,n}\|_{p}\leq\max_{n,k ^{\prime}\in{\mathcal{K}}}\|\bar{f}({\bm{X}}_{k^{\prime}})_{:,n}-\bar{f}({\bm{ X}}_{k^{\prime}-1})_{:,n}\|_{p}.$ $$

With the triangle inequality, we have approximation error as

$$ $\displaystyle\|\tilde{f}-f\|_{L^{p}([0,1]^{d\times N})}$ $\displaystyle\leq\|\tilde{f}({\bm{X}})-f({\bm{X}})\|_{p}\cdot 1$ $\displaystyle\leq\|\tilde{f}({\bm{X}})-\bar{f}({\bm{X}})\|_{p}+\|\bar{f}({\bm{ X}})-f({\bm{X}})\|_{p}$ $\displaystyle\leq\max_{n,k^{\prime}\in{\mathcal{K}}}\|\bar{f}({\bm{X}}_{k^{ \prime}})_{:,n}-\bar{f}({\bm{X}}_{k^{\prime}-1})_{:,n}\|_{p}+\|\bar{f}({\bm{X} })-f({\bm{X}})\|_{p}+\mathcal{O}(\delta^{d})$ $\displaystyle\leq\omega^{tok}_{f}(\delta\sqrt{d})+\omega^{cont}_{f}(\delta \sqrt{Nd})+\omega_{f}(\delta\sqrt{Nd})+\mathcal{O}(\delta^{d}).$ $$

Then, $\delta$ must be expressed in terms of the number of loops $r$ to determine the approximation rate.

$$ $\displaystyle r=\delta^{-1}+N+2\delta^{-(N+1)d}$ $\displaystyle\Leftrightarrow\delta^{-1}+2\delta^{-(N+1)d}=r-N$ $\displaystyle\Leftrightarrow\delta^{-1}\cdot 2\delta^{-(N+1)d}\geq r-N$ $\displaystyle\Leftrightarrow 2\delta^{-(N+1)d-1}\geq r-N$ $\displaystyle\Leftrightarrow\delta\leq{\Big{(}\frac{r-N}{2}\Big{)}}^{-1/\big{( }(N+1)d+1\big{)}}.$ (20) $$

Thus, we can derive Theorem 3.7.
∎

### A.3 Piecewise Linear functions

Here, we define three functions implemented with ReLU functions.

###### Proposition A.1 (Rectangular function) .

Given $t\in\mathbb{R}$, define $\mathrm{rect}_{t}:\mathbb{R}\to\mathbb{R}$ as:

$$ $\mathrm{rect}_{t}(x)=\begin{cases}1\quad&\text{if }x\in[t,t+1),\\ 0\quad&\textnormal{otherwise}.\end{cases}$ $$

Four $\mathrm{ReLU}$ functions can approximate $\mathrm{rect}_{t}$ with arbitrarily small error via

$$ $\phi_{t}(x)\coloneqq{\sigma_{R}}\big{(}\tfrac{x-t+\epsilon}{\epsilon}\big{)}-{ \sigma_{R}}\big{(}\tfrac{x-t}{\epsilon}\big{)}+{\sigma_{R}}\big{(}\tfrac{-x+t+ 1}{\epsilon}\big{)}-{\sigma_{R}}\big{(}\tfrac{-x+t+1-\epsilon}{\epsilon}\big{) }-1,$ $$

where $\lim_{\epsilon\to 0}\phi_{t}(x)=\mathrm{rect}_{t}(x)$.

###### Proposition A.2 (Step function) .

Define $\mathrm{step}:\mathbb{R}\to\mathbb{R}$ as:

$$ $\mathrm{step}(x)=\begin{cases}0\quad&\text{if }x<0,\\ 1\quad&\text{if }x\geq 0.\end{cases}$ $$

Two $\mathrm{ReLU}$ functions can approximate $\mathrm{step}$ with arbitrarily small error via

$$ $\phi(x)\coloneqq{\sigma_{R}}\big{(}\tfrac{x}{\epsilon}+1\big{)}-{\sigma_{R}} \big{(}\tfrac{x}{\epsilon}\big{)},$ $$

where $\lim_{\epsilon\to 0}\phi(x)=\mathrm{step}(x)$.

###### Proposition A.3 (Impulse function) .

Given $\theta\in\mathbb{N}$, define $\mathrm{impulse}_{\theta}:\mathbb{R}\to\mathbb{R}$ for $x\in[-M,M]$ and $t\in\mathbb{N}$ as:

$$ $\mathrm{impulse}_{\theta}(x,t)=\begin{cases}x\quad&\textnormal{if}\ t=\theta, \\ 0\quad&\textnormal{otherwise}.\end{cases}$ $$

Four $\mathrm{ReLU}$ functions can approximate $\mathrm{impulse}_{\theta}$ with arbitrarily small error via

$$ $\displaystyle\mathrm{impulse}_{\theta}(x,t)\coloneqq$ $\displaystyle{\sigma_{R}}\big{(}x+2M(t-\theta+1/2)\big{)}-2M{\sigma_{R}}(t- \theta+1/2)$ $\displaystyle-{\sigma_{R}}\big{(}x+2M(t-\theta-1/2)\big{)}+2M{\sigma_{R}}(t- \theta-1/2).$ $$

###### Proposition A.1 (Rectangular function) .

###### Proposition A.2 (Step function) .

###### Proposition A.3 (Impulse function) .

### A.4 Step 1. Token-wise Quantization

Figure: Figure 4: An illustration of $h_{k}(x)$.
Refer to caption: extracted/5911203/figures/hk.png

We aim to construct quantization function $\mathbf{g}:[0,1]^{d}\to\{0,1,\dotsc,\delta^{-1}\}^{d}$ for each dimension as

$$ $\displaystyle\mathbf{g}({\bm{x}})$ $\displaystyle={(g({\bm{x}}_{1}),g({\bm{x}}_{2}),\dotsc,g({\bm{x}}_{d}))}^{\top},$ $\displaystyle\text{where}\quad g(x)$ $\displaystyle=\begin{cases}k&\textnormal{if }x\in[k{\delta},(k+1){\delta}),\\ 0&\quad\textnormal{otherwise}.\end{cases}$ $$

This function $g:\mathbb{R}\to\mathbb{R}$ can be expressed as

$$ $g(x)=\sum_{i=0}^{n-1}i\cdot\mathrm{rect}_{i}(x)$ $$

for any $n\in\mathbb{N}$ and $x\in\mathbb{R}$. The illustration of $h_{k}(x)\coloneq k\cdot\mathrm{rect}_{k}(x)$ is shown in Fig 4.
The key idea is that $h_{k}(x)$ can be represented with a single function $h$ in the form of $h\big{(}kx,k^{2},k\big{)}$.
Lemma A.4 implement $h\big{(}kx,k^{2},k\big{)}$ with a feed-forward layer and perform the summation through a skip connection.

###### Lemma A.4 .

Given any $\delta^{-1}\in\mathbb{N}$ and ${\bm{x}}\in\mathbb{R}^{d}$, there exist a feed-forward layer ${\mathrm{FF}}:\mathbb{R}^{5d}\to\mathbb{R}^{5d}$ of width size $q=7d$, and two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{d}\to\mathbb{R}^{5d}$ and $\mathcal{L}_{2}:\mathbb{R}^{5d}\to\mathbb{R}^{d}$ such that

$$ $\left(\mathcal{L}_{2}\circ\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}^{\circ( \delta^{-1}-1)}\circ\mathcal{L}_{1}({\bm{x}})\right)_{i}=\begin{cases}k&\quad \textnormal{if }{\bm{x}}_{i}\in[k{\delta},(k+1){\delta}),~{}~{}k=0,\dotsc, \delta^{-1}-1,\\ 0&\quad\textnormal{otherwise},\end{cases}$ $$

for any $i=1,2,\dotsc,d$.

###### Proof.

On the basis of proposition A.1, define function $h_{k}(x)=k\cdot\mathrm{rect}_{k}(x)$ via

$$ $\displaystyle h_{k}(x)\coloneqq{\sigma_{R}}\big{(}\tfrac{k}{\epsilon}(x-\tfrac {k}{\delta}+\epsilon)\big{)}$ $\displaystyle-{\sigma_{R}}\big{(}\tfrac{k}{\epsilon}(x-\tfrac{k}{\delta})\big{ )}+{\sigma_{R}}\big{(}\tfrac{k}{\epsilon}(-x+\tfrac{k}{\delta}+1)\big{)}$ $\displaystyle-{\sigma_{R}}\big{(}\tfrac{k}{\epsilon}(-x+\tfrac{k}{\delta}+1- \epsilon)\big{)}-k,$ $$

which satisfies

$$ $h_{k}(x)=\begin{cases}k&\textnormal{if }x\in[k{\delta},k{\delta}+1),\\ 0&\quad\textnormal{otherwise}.\end{cases}$ $$

For any $x\in[k{\delta},k{\delta}+1)$ where $k=0,1,\dotsc,\delta^{-1}-1$, we have

$$ $\sum_{i=0}^{\delta^{-1}-1}h_{i}(x)=h_{k}(x)=k.$ $$

Define function $h:\mathbb{R}^{3}\to\mathbb{R}$ to represent $h_{k}$ via

$$ $\displaystyle h(kx,k^{2},k)\coloneqq{\sigma_{R}}\big{(}\tfrac{kx}{\epsilon}- \tfrac{k^{2}}{\epsilon}+k\big{)}$ $\displaystyle-{\sigma_{R}}\big{(}\tfrac{kx}{\epsilon}-\tfrac{k^{2}}{\epsilon} \big{)}+{\sigma_{R}}\big{(}-\tfrac{kx}{\epsilon}+\tfrac{k^{2}}{\epsilon}+ \tfrac{k}{\epsilon}\big{)}$ $\displaystyle-{\sigma_{R}}\big{(}-\tfrac{kx}{\epsilon}+\tfrac{k^{2}}{\epsilon} +k\tfrac{1-\epsilon}{\epsilon})-{\sigma_{R}}(k)=h_{k}(x).$ $$

Define $\bm{\xi}_{k}$ as

$$ $\bm{\xi}_{k}={\Big{(}kx,\ k^{2},\ k,\ x,\ \sum_{i=0}^{k-1}h_{i}(x)\Big{)}^{ \top}}.$ $$

Then, construct a feed-forward layer ${\mathrm{FF}}:\mathbb{R}^{5}\to\mathbb{R}^{5}$ with a skip connection such that

$$ $\displaystyle\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}(\bm{\xi}_{k})$ $\displaystyle=\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}\bigg{(}(kx,\ k^{2},\ k, \ x,\ \sum_{i=0}^{k-1}h_{i}(x))^{\top}\bigg{)}$ $\displaystyle=\Big{(}(k+1)x,\ (k+1)^{2},\ k+1,\ x,\ \sum_{i=0}^{k}h_{i}(x)\Big {)}^{\top}$ $\displaystyle=\bm{\xi}_{k+1}.$ $$

via

$$ $\displaystyle\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}\left(\left[\begin{array}[ ]{c}kx\\ k^{2}\\ k\\ x\\ \sum_{i=0}^{k-1}h_{i}(x)\\ \end{array}\right]\right)$ $\displaystyle=\left[\begin{array}[]{c}kx\\ k^{2}\\ k\\ x\\ \sum_{i=0}^{k-1}h_{i}(x)\\ \end{array}\right]+$ $\displaystyle\begin{bmatrix}1&0&0&0&0&0&0\\ 0&1&0&0&0&0&0\\ 0&0&0&0&0&0&0\\ 0&0&0&0&0&0&0\\ 0&0&1&-1&1&-1&-1\\ \end{bmatrix}$ $\displaystyle{\sigma_{R}}\left(\begin{bmatrix}0&0&0&1&0\\ 0&0&2&0&0\\ \tfrac{1}{\epsilon}&-\tfrac{1}{\epsilon}&1&0&0\\ \tfrac{1}{\epsilon}&-\tfrac{1}{\epsilon}&0&0&0\\ -\tfrac{1}{\epsilon}&\tfrac{1}{\epsilon}&\tfrac{1}{\epsilon}&0&0\\ -\tfrac{1}{\epsilon}&\tfrac{1}{\epsilon}&\tfrac{1-\epsilon}{\epsilon}&0&0\\ 0&0&1&0&0\\ \end{bmatrix}\left[\begin{array}[]{c}kx\\ k^{2}\\ k\\ x\\ \sum_{i=0}^{k-1}h_{i}(x)\\ \end{array}\right]\right)+\left[\begin{array}[]{c}0\\ 1\\ 1\\ 0\\ 0\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{c}kx\\ k^{2}\\ k\\ x\\ \sum_{i=0}^{k-1}h_{i}(x)\\ \end{array}\right]+\left[\begin{array}[]{c}x\\ 2k+1\\ 1\\ 0\\ h_{k}(x)\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{c}kx+x\\ k^{2}+2k+1\\ k+1\\ x\\ \sum_{i=0}^{k-1}h_{i}(x)+h_{k}(x)\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{c}(k+1)x\\ (k+1)^{2}\\ k+1\\ x\\ \sum_{i=0}^{k}h_{i}(x)\\ \end{array}\right].$ $$

Then, define two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{1}\to\mathbb{R}^{5}$ and $\mathcal{L}_{2}:\mathbb{R}^{5}\to\mathbb{R}^{1}$ by

$$ $\mathcal{L}_{1}(x)\coloneqq(0,0,0,x,0),\quad\mathcal{L}_{2}(x_{1},x_{2},x_{3}, x_{4},x_{5})\coloneqq x_{5}.$ $$

Thus, we have

$$ $\displaystyle\mathcal{L}_{2}\circ\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}^{ \circ(\delta^{-1}-1)}\circ\mathcal{L}_{1}(x)$ $\displaystyle=\mathcal{L}_{2}\circ\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}^{ \circ(\delta^{-1}-1)}(\bm{\xi}_{1})$ $\displaystyle=\mathcal{L}_{2}(\bm{\xi}_{\delta^{-1}})$ $\displaystyle=\sum_{i=0}^{\delta^{-1}-1}h_{i}(x).$ $$

For $d$-dimensional inputs, we need $d$-times parameters.
∎

###### Lemma A.4 .

###### Proof.

### A.5 Step 2. Contextual mapping

The network takes token IDs as inputs, denoted by ${\bm{z}}\in\{0,1,\dotsc,\delta^{-d}-1\}^{N}$.
We consider only the case where all token IDs are distinct, since this accounts for the majority when $\delta$ is small.
The network maps token IDs into a sequence ID using inner product with the vector ${\bm{u}}\in\mathbb{R}^{N}$ defined as ${\bm{u}}\coloneqq(\delta^{-(N-1)d},\delta^{-(N-2)d},\dotsc,\delta^{-d},1)^{\top}$ i.e.

$$ $\mathrm{CM}({\bm{z}})\coloneq{\bm{u}}^{\top}{\bm{z}}.$ $$

Due to permutation equivalence, we can assume without loss of generality that elements of ${\bm{z}}\in{\mathbb{L}}_{\delta}$ is ordered, i.e., ${\bm{z}}_{1}>{\bm{z}}_{2}>\cdots>{\bm{z}}_{N}$. Then the map $\mathrm{CM}$ satisfies

$$ $\left\lVert{\bm{u}}^{\top}{\bm{z}}-{\bm{u}}^{\top}{\bm{z}}^{\prime}\right \rVert>1,\quad\text{if }{\bm{z}}\neq{\bm{z}}^{\prime}.$ $$

In other words, $\mathrm{CM}$ represent ${\bm{z}}$ in $\delta^{-d}$-base system.
The network computes ${\bm{u}}^{\top}{\bm{z}}$ in the form of $\sum_{i=1}^{N}\delta^{-(N-i)d}{\bm{z}}_{i}$.
The network computes ${\bm{s}}^{(k)}\coloneq\sum_{i=1}^{k}\delta^{-(k-i)d}{\bm{z}}_{i}$ in each loop, and after $N$-loops, it outputs ${\bm{s}}^{(N)}={\bm{u}}^{\top}{\bm{z}}$.
To implement this, the self-attention layer select ${\bm{z}}_{k}$ in the $k$-th loop iteration. We design the key and value weights to select the maximum token ID. The feed-forward layer post-processes the token ID as if selected, then replaces it with negative value to prevent selection in subsequent iterations, i.e., the post-processed token IDs for the $k$-th loop is

$$ $\displaystyle{\bm{z}}^{(k)}_{i}$ $\displaystyle=z\quad\textit{s}.\textit{t}.\quad\begin{cases*}z<0\quad \textnormal{if}\quad i\leq k,\\ z={\bm{z}}_{i}\quad\textnormal{otherwise}.\end{cases*}$ $$

###### Lemma A.5 .

Consider the set, composed of distinct indices for $d$-dimension $\delta$-discretized regions of $N$ tokens, i.e.

$$ ${\mathbb{L}}_{\delta}\coloneq\{{\bm{z}}\in\{0,1,\dotsc,\delta^{-d}-1\}^{N}\mid {\bm{z}}_{i}\neq{\bm{z}}_{j}\text{ for all }i\neq j\}.$ $$

There exists a function $\mathrm{CM}:\mathbb{R}^{N}\to\mathbb{R}$ composed of Transformer block $\mathrm{TF}:\mathbb{R}^{3\times N}\to\mathbb{R}^{3\times N}$ with single head, head size $s=1$, and width size $q=3$, and two affine linear maps $\mathcal{L}_{1}:\mathbb{R}\to\mathbb{R}^{3}$ and $\mathcal{L}_{2}:\mathbb{R}^{3}\to\mathbb{R}$, such that

$$ $\bm{\mathcal{L}}_{2}\circ{{\mathrm{TF}}}^{\circ N}\circ\bm{\mathcal{L}}_{1}({ \bm{z}}^{\top})=\mathrm{CM}({\bm{z}}^{\top})\cdot\mathbf{1}_{N}^{\top},$ $$

for any ${\bm{z}}\in{\mathbb{L}}_{\delta}$, where $\bm{\mathcal{L}}_{1}$ and $\bm{\mathcal{L}}_{2}$ denote the token-wise applications of $\mathcal{L}_{1}$ and $\mathcal{L}_{2}$, respectively, and $\mathrm{CM}$ is a contextual mapping, which satisfies the following properties:

- 1.
For any ${\bm{z}}\in{\mathbb{L}}_{\delta}$, the $N$ entries in $\mathrm{CM}({\bm{z}}^{\top})$ are all distinct.
- 2.
For any ${\bm{z}},{\bm{z}}^{\prime}\in{\mathbb{L}}_{\delta}$, if ${\bm{z}}$ is not a permutation of ${\bm{z}}^{\prime}$, all entries of $\mathrm{CM}({\bm{z}}^{\top})$ and $\mathrm{CM}({\bm{z}}^{\prime\top})$ are distinct.

###### Proof.

Due to permutation equivalence, we can assume without loss of generality that elements of ${\bm{z}}\in{\mathbb{L}}_{\delta}$ is ordered, i.e., ${\bm{z}}_{1}>{\bm{z}}_{2}>\cdots>{\bm{z}}_{N}$.
Define ${\bm{u}}\in\mathbb{R}^{N}$ as ${\bm{u}}\coloneqq(\delta^{-(N-1)d},\dotsc,\delta^{-d},1)^{\top}$, which satisfy

$$ $\left\lVert{\bm{u}}^{\top}{\bm{z}}-{\bm{u}}^{\top}{\bm{z}}^{\prime}\right \rVert>1,\quad\text{if }{\bm{z}}\neq{\bm{z}}^{\prime}\text{ for any }{\bm{z}}, {\bm{z}}^{\prime}\in{\mathbb{L}}_{\delta}.$ $$

Then ${\bm{z}}\to{\bm{u}}^{\top}{\bm{z}}$ is a contextual mapping.
We show how to construct Transformer block ${\mathrm{TF}}:\mathbb{R}^{3\times N}\to\mathbb{R}^{3\times N}$ with single head and head size $s=1$ such that, for any ${\bm{z}}\in{\mathbb{L}}_{\delta}$,

$$ $\mathrm{TF}^{\circ N}\left(\left[\begin{array}[]{ccc}{\bm{z}}^{\top}\\ {\mathbf{0}}_{N}^{\top}\\ {\mathbf{0}}_{N}^{\top}\\ \end{array}\right]\right)=\left[\begin{array}[]{ccc}{\mathbf{0}}_{N}^{\top}\\ {\mathbf{0}}_{N}^{\top}\\ {\bm{u}}^{\top}{\bm{z}}\cdot\mathbf{1}_{N}^{\top}\\ \end{array}\right].$ $$

where ${\mathbf{0}}_{N}\in\mathbb{R}^{N}$ denote a zero vector.
For ${\bm{z}}\in{\mathbb{L}}_{\delta}$, we define two series ${\bm{z}}^{(k)}$ and ${\bm{s}}^{(k)}$ by

$$ $\displaystyle{\bm{z}}^{(k)}_{i}$ $\displaystyle=z\quad\textit{s}.\textit{t}.\quad\begin{cases*}z<0\quad \textnormal{if}\quad i\leq k,\\ z={\bm{z}}_{i}\quad\textnormal{otherwise},\end{cases*}\quad\text{for }i=1, \dotsc,d.$ $\displaystyle{\bm{s}}^{(k)}$ $\displaystyle=\sum_{i=0}^{k}\delta^{-id}{\bm{z}}_{i}.$ $$

While ${\bm{z}}^{(k)}$ is not uniquely determined, any vectors that satisfies the conditions is accepted as ${\bm{z}}^{(k)}$.
We can confirm that ${\bm{s}}^{(k)}$ satisfies

$$ $\begin{split}{\bm{s}}^{(k)}&=\sum_{i=1}^{k}\delta^{-(k-i)d}{\bm{z}}_{i}\\ &=\big{(}\sum_{i=1}^{k-1}\delta^{-(k-i)d}{\bm{z}}_{i}\big{)}+{\bm{z}}_{k}\\ &=\big{(}\sum_{i=0}^{k-1}\delta^{-d}\cdot\delta^{-(k-1-i)d}{\bm{z}}_{i}\big{)} +{\bm{z}}_{k}\\ &=\delta^{-d}\cdot{\bm{s}}^{(k-1)}+{\bm{z}}_{k}.\end{split}$ (21) $$

Recall that ${\bm{s}}^{(N)}={\bm{u}}^{\top}{\bm{z}}$.
Define a single-head self-attention ${\mathrm{Attn}}:\mathbb{R}^{3\times N}\to\mathbb{R}^{3\times N}$ such that

$$ ${\mathrm{Attn}}\left(\left[\begin{array}[]{ccc}{{\bm{z}}}^{\top}\\ {{\bm{y}}}^{\top}\\ {{\bm{s}}}^{\top}\\ \end{array}\right]\right)=\left[\begin{array}[]{ccc}0\\ \max_{j}{{\bm{z}}}_{j}\cdot\mathbf{1}_{N}^{\top}\\ 0\\ \end{array}\right],$ $$

via the weight parameters

$$ $\displaystyle{\bm{W}}_{O}=\left[\begin{array}[]{ccc}0\\ 1\\ 0\\ \end{array}\right],\quad{\bm{W}}_{V}={\bm{W}}_{K}={\bm{W}}_{Q}=\left[\begin{ array}[]{ccc}1&0&0\end{array}\right]$ $$

Define ${\mathrm{FF}}:\mathbb{R}^{3}\to\mathbb{R}^{3}$ of width size $q=3$ via:

$$ $\displaystyle{\mathrm{FF}}\left(\left[\begin{array}[]{ccc}x_{1}\\ x_{2}\\ x_{3}\\ \end{array}\right]\right)$ $\displaystyle=\begin{bmatrix}-M&0&0\\ 0&-1&0\\ 0&1&\delta^{-d}-1\\ \end{bmatrix}{\sigma_{R}}\left(\begin{bmatrix}1&-1&0\\ 0&1&0\\ 0&0&1\\ \end{bmatrix}\left[\begin{array}[]{ccc}x_{1}\\ x_{2}\\ x_{3}\\ \end{array}\right]+\left[\begin{array}[]{ccc}\epsilon\\ 0\\ 0\\ \end{array}\right]\right)$ $\displaystyle=\left[\begin{array}[]{ccc}-M\sigma_{R}(x_{1}-x_{2}+\epsilon)\\ -\sigma_{R}(x_{2})\\ (\delta^{-1}-1)\sigma_{R}(x_{3})+\sigma_{R}(x_{2})\\ \end{array}\right],$ $$

where $0<\epsilon<\delta^{-1}$ and $M>\epsilon\delta^{-1}$.
For $x_{1}\in\{0,1,\dotsc,\delta^{-d}\}$, we have

$$ $x_{1}-M\sigma_{R}(x_{1}-x_{2}+\epsilon)=z,\quad\textit{s}.\textit{t}.\quad \begin{cases*}z=x_{1}\quad\textnormal{if}\quad x_{1}<x_{2},\\ z<0\quad\textnormal{otherwise}.\end{cases*}$ $$

This post-processes the token ID as if selected, then replaces it with negative value. We have

$$ $\begin{split}{\bm{z}}^{(k)}_{i}-M\sigma_{R}({\bm{z}}^{(k)}_{i}-{\bm{z}}_{k}+ \epsilon)&=\begin{cases*}z<0\quad\textnormal{if}\quad i\leq k+1,\\ z={\bm{z}}_{i}\quad\textnormal{otherwise},\end{cases*}\\ &={\bm{z}}^{(k+1)}_{i}\quad\text{for }i=1,\dotsc,d.\end{split}$ (22) $$

We can confirm that Transformer block ${\mathrm{TF}}:\mathbb{R}^{3\times N}\to\mathbb{R}^{3\times N}$ composed of ${\mathrm{Attn}}$ and ${\mathrm{FF}}$ satisfies

$$ $\displaystyle{\mathrm{TF}}\left(\left[\begin{array}[]{ccc}({\bm{z}}^{(k)})^{ \top}\\ {\mathbf{0}}_{n}^{\top}\\ ({\bm{s}}^{(k)})^{\top}\\ \end{array}\right]\right)$ $\displaystyle=(\mathrm{id}+\bm{{\mathrm{FF}}})\circ(\mathrm{id}+{\mathrm{Attn} })\left(\left[\begin{array}[]{ccc}({\bm{z}}^{(k)})^{\top}\\ {\mathbf{0}}_{n}^{\top}\\ ({\bm{s}}^{(k)})^{\top}\\ \end{array}\right]\right)$ $\displaystyle=(\mathrm{id}+\bm{{\mathrm{FF}}})\left(\left[\begin{array}[]{ccc} ({\bm{z}}^{(k)})^{\top}\\ {\bm{z}}_{k}\cdot\mathbf{1}_{N}^{\top}\\ ({\bm{s}}^{(k)})^{\top}\\ \end{array}\right]\right)$ $\displaystyle=(\mathrm{id}+\bm{{\mathrm{FF}}})\left(\left[\begin{array}[]{ccc} ({\bm{z}}^{(k)})^{\top}\\ {\bm{z}}_{k}\cdot\mathbf{1}_{N}^{\top}\\ ({\bm{s}}^{(k)})^{\top}\\ \end{array}\right]\right)$ $\displaystyle=\left[\begin{array}[]{ccc}({\bm{z}}^{(k)})^{\top}\\ {\bm{z}}_{k}\cdot\mathbf{1}_{N}^{\top}\\ ({\bm{s}}^{(k)})^{\top}\\ \end{array}\right]+\left[\begin{array}[]{ccc}-M\sigma_{R}\big{(}({\bm{z}}^{(k) })^{\top}-{\bm{z}}_{k}\cdot\mathbf{1}_{N}^{\top}+\epsilon\mathbf{1}_{N}^{\top} \big{)}\\ -{\bm{z}}_{k}\cdot\mathbf{1}_{N}^{\top}\\ (\delta^{-1}-1)({\bm{s}}^{(k)})^{\top}+\sigma_{R}(x_{2})\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{ccc}({\bm{z}}^{(k+1)})^{\top}\\ {\mathbf{0}}_{n}^{\top}\\ \delta^{-1}({\bm{z}}^{(k)})^{\top}+{\bm{z}}_{k}\mathbf{1}_{N}^{\top}\\ \end{array}\right]\quad\text{because Eq.~{}\ref{eq:vxk}}$ $\displaystyle=\left[\begin{array}[]{ccc}({\bm{z}}^{(k+1)})^{\top}\\ {\mathbf{0}}_{n}^{\top}\\ ({\bm{s}}^{(k+1)})^{\top}\\ \end{array}\right]\quad\text{because Eq.~{}\ref{eq:vs}}.$ $$

We define two affine linear maps $\mathcal{L}_{1}:\mathbb{R}\to\mathbb{R}^{3}$ and $\mathcal{L}_{2}:\mathbb{R}^{3}\to\mathbb{R}$ as $\mathcal{L}_{1}(x)\coloneqq(x,0,0)$ and $\mathcal{L}_{2}(x_{1},x_{2},x_{3})\coloneqq x_{3}$. Thus, we have

$$ $\displaystyle\bm{\mathcal{L}}_{2}\circ{\mathrm{TF}}^{\circ N}\circ\bm{\mathcal {L}}_{1}(x^{\top})=\bm{\mathcal{L}}_{2}\circ{\mathrm{TF}}^{\circ N}\left(\left [\begin{array}[]{ccc}{\bm{z}}^{\top}\\ {\mathbf{0}}_{N}^{\top}\\ {\mathbf{0}}_{N}^{\top}\\ \end{array}\right]\right)=\bm{\mathcal{L}}_{2}\left(\left[\begin{array}[]{ccc} {\mathbf{0}}_{N}^{\top}\\ {\mathbf{0}}_{N}^{\top}\\ {\bm{s}}^{N}\cdot{\mathbf{1}}_{N}^{\top}\\ \end{array}\right]\right)=({\bm{u}}^{\top}{\bm{z}})\cdot{\mathbf{1}}_{N}^{\top}.$ $$

Recall that ${\bm{z}}\to{\bm{u}}^{\top}{\bm{z}}$ is a contextual mapping.
∎

###### Corollary A.6 .

There exist a Transformer block $\mathrm{TF_{2}}:\mathbb{R}^{5\times N}\to\mathbb{R}^{5\times N}$ with single head, head size $s=1$, and width size $q=4$, and two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{2}\to\mathbb{R}^{5}$ and $\mathcal{L}_{2}:\mathbb{R}^{5}\to\mathbb{R}$ such that

$$ $\bm{\mathcal{L}}_{2}\circ{{\mathrm{TF}}_{2}}^{\circ N}\circ\bm{\mathcal{L}}_{1 }\left(\left[\begin{array}[]{ccc}{{\bm{z}}}^{\top}\\ {{\bm{Z}}_{d,:}}\\ \end{array}\right]\right)=\big{(}2{\bm{u}}^{\top}{\bm{z}}-{\bm{Z}}_{d,N}\big{) }\cdot\mathbf{1}_{N}^{\top}\quad\textnormal{for any ${\bm{z}}\in{\mathbb{L}}_{ \delta}$.}$ $$

where ${\bm{u}}\coloneqq(\delta^{-(N-1)d},\dotsc,\delta^{-d},1)^{\top}$.

###### Proof.

Define a single-head self-attention ${\mathrm{Attn}}:\mathbb{R}^{5\times N}\to\mathbb{R}^{5\times N}$ such that

$$ ${\mathrm{Attn}}\left(\left[\begin{array}[]{ccc}{{\bm{z}}}^{\top}\\ {{\bm{Z}}_{d,:}}\\ {{\bm{y}}}^{\top}\\ {{\bm{s}}}^{\top}\\ {{\bm{q}}}^{\top}\\ \end{array}\right]\right)=\left[\begin{array}[]{ccc}0\\ 0\\ \max_{j}{{\bm{z}}}_{j}\cdot\mathbf{1}_{N}^{\top}\\ 0\\ {\bm{Z}}_{d,\operatorname*{arg\,max}_{j}{{\bm{z}}}_{j}}\cdot\mathbf{1}_{N}^{ \top}\\ \end{array}\right],$ $$

via the weight parameters

$$ $\displaystyle{\bm{W}}_{O}=\left[\begin{array}[]{ccc}0\\ 0\\ 1\\ 0\\ 1\\ \end{array}\right],\quad{\bm{W}}_{V}={\bm{W}}_{K}={\bm{W}}_{Q}=\left[\begin{ array}[]{ccccc}1&0&0&0&0\end{array}\right]$ $$

Define ${\mathrm{FF}}:\mathbb{R}^{4}\to\mathbb{R}^{4}$ of width size $q=4$ via:

$$ $\displaystyle{\mathrm{FF}}\left(\left[\begin{array}[]{ccc}x_{1}\\ x_{2}\\ x_{3}\\ x_{4}\\ x_{5}\\ \end{array}\right]\right)$ $\displaystyle=\begin{bmatrix}-M&0&0&0\\ 0&-1&0&0\\ 0&1&\delta^{-d}-1&0\\ 0&1&0&-1\\ \end{bmatrix}{\sigma_{R}}\left(\begin{bmatrix}1&0&-1&0&0\\ 0&0&1&0&0\\ 0&0&0&1&0\\ 0&0&0&0&1\\ \end{bmatrix}\left[\begin{array}[]{ccc}x_{1}\\ x_{2}\\ x_{3}\\ x_{4}\\ x_{5}\\ \end{array}\right]+\left[\begin{array}[]{ccc}\epsilon\\ 0\\ 0\\ 0\\ 0\\ \end{array}\right]\right)$ $\displaystyle=\left[\begin{array}[]{ccc}-M\sigma_{R}(x_{1}-x_{2}+\epsilon)\\ 0\\ -\sigma_{R}(x_{2})\\ (\delta^{-1}-1)\sigma_{R}(x_{3})+\sigma_{R}(x_{2})\\ -\sigma_{R}(x_{4})+\sigma_{R}(x_{2})\\ \end{array}\right],$ $$

where $0<\epsilon<\delta^{-1}$ and $M>\epsilon\delta^{-1}$.
Note that the fourth columns after $t$ loops are $z_{t}$, so at the end, we obtain $z_{N}$.
Then, we define two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{2}\to\mathbb{R}^{5}$ and $\mathcal{L}_{2}:\mathbb{R}^{5}\to\mathbb{R}$ as $\mathcal{L}_{1}(x1,x2)\coloneqq(x1,x2,0,0,0)$ and $\mathcal{L}_{2}(x_{1},x_{2},x_{3},x_{4},x_{5})\coloneqq 2x_{4}-x_{5}$.
∎

###### Lemma A.5 .

###### Proof.

###### Corollary A.6 .

###### Proof.

### A.6 Step 3. Function value mapping with bit extraction

We use a bit extraction technique (Bartlett et al., ) to approximately memorize the piecewise linear function.Consider $n\in\mathbb{N}$ input indices $k=0,1,\dotsc,n-1\in\mathbb{N}$ with associated values $y_{0},y_{1},\dotsc,y_{n-1}\in\mathbb{R}$.
The network approximately memorize the difference $y_{i}-y_{i-1}$ with base-2 representation.
Since binary representation limited to $\{0,1\}$, $y_{i}-y_{i-1}$ has to be re-scaled with $\epsilon\coloneq|y_{i}-y_{i-1}|$ as

$$ $a_{i}=\big{\lfloor}\frac{y_{i}}{\epsilon}\big{\rfloor},$ $$

where $\lfloor x\rfloor=\max\{n:n\leq x,\ n\in\mathbb{Z}\}$. Then, the difference $b_{i}=a_{i}-a_{i-1}$ satisfies $b_{i}\in\{-1,0,1\}$, $b_{i}$, and it can be represented using two binary values $c_{i},d_{i}\in\{0,1\}$ as follows:

$$ $b_{i}=c_{i}-d_{i},$ $$

and we have

$$ $a_{k}=a_{0}+\sum_{i=0}^{k}b_{i}=a_{0}+\sum_{i=0}^{k}d_{i}+\sum_{i=0}^{k}d_{i} \quad\text{for}\quad k=0,1,2.$ $$

Lemma A.8 and Lemma 4.1 show that $\sum_{i=0}^{k}c_{i}$ and $\sum_{i=0}^{k}d_{i}$ can be realized by composition of single feed-forward layer. Thus the networks can approximate $y_{i}$ with $\epsilon a_{i}$ denoted by $\tilde{y}_{i}$ with the following accuracy

$$ $|\tilde{y}_{i}-y_{i}|=|\epsilon\big{\lfloor}\frac{y_{i}}{\epsilon}\big{\rfloor }-\epsilon\frac{y_{i}}{\epsilon}|=\epsilon|\big{\lfloor}\frac{y_{i}}{\epsilon} \big{\rfloor}-\frac{y_{i}}{\epsilon}|\leq\epsilon.$ $$

For $d$-dimensional inputs-outputs pair, we construct the networks for each dimension i.e.

$$ $\bm{\tilde{{\bm{y}}}}=(\tilde{y}^{1},\tilde{y}^{2},\dotsc,\tilde{y}^{d})$ $$

The key idea of our lemma and proof follows Lemma D.1 from Zhang et al. () as shown in below; however, we cannot directly apply their result here, as it requires depth-2 networks.

###### Proposition A.7 (Lemma D.1 in Zhang et al. ( 2023 ) ) .

Given any $r\in\mathbb{N}^{+}$, there exist ${\mathrm{FF}}:\mathbb{R}^{3d}\to\mathbb{R}^{3d}$ of width size $8$ and depth $2$ with two affine linear maps ${\mathcal{L}}_{1}:\mathbb{R}^{2}\to\mathbb{R}^{5}$ and ${\mathcal{L}}_{2}:\mathbb{R}^{5}\to\mathbb{R}$
such that: For any $\theta_{1},\theta_{2},\dotsc,\theta_{r}\in\{0,1\}$, it holds that

$$ ${\mathcal{L}}_{2}\circ{\mathrm{FF}}^{\circ r}\circ{\mathcal{L}}_{1}\big{(}k,\ \textnormal{bin}\hskip 1.5pt0.\theta_{1}\theta_{2}\cdots\theta_{r}\big{)}=\sum _{\ell=1}^{k}\theta_{\ell}\quad\textnormal{for $k=0,1,\dotsc,r$,}$ $$

where $\textnormal{bin}\hskip 1.5pt0.\theta_{1}\theta_{2}\cdots\theta_{r}$ denote the binary representation of $\theta=\sum_{i=1}^{n}\theta_{i}2^{-i}$.

We found that loop unrolling allows us to reduce the number of layers from 2 to 1, replacing $x^{k+1}={\mathrm{ReLU}}({\mathrm{ReLU}}(x^{\prime k}))$ with $(x^{k+1},x^{\prime k})={\mathrm{ReLU}}(x^{\prime k},x^{k})$.
Although our method makes the weights dependent on $\theta_{1},\theta_{2},\dots,\theta_{r}\in{0,1}$, this does not present an issue for our construction in function approximation. Specifically, ${\theta_{1},\theta_{2},\dots,\theta_{r}}$ is fixed for each target function, and the role of the network is to learn the weights tailored to that single function.

###### Lemma A.8 .

Given $\theta_{1},\theta_{2},\dotsc,\theta_{r}\in\{0,1\}$ for $r\in\mathbb{N}^{+}$, there exist a feed-forward layer ${\mathrm{FF}}:\mathbb{R}^{7}\to\mathbb{R}^{7}$ of width size $10$ and two affine linear maps $\mathcal{L}_{1}:\mathbb{R}\to\mathbb{R}^{7}$ and $\mathcal{L}_{2}:\mathbb{R}^{7}\to\mathbb{R}$ s.t.

$$ $\mathcal{L}_{2}\circ{(\mathrm{id}+{\mathrm{FF}})}^{\circ r}\circ\mathcal{L}_{1 }\big{(}k)=\sum_{i=1}^{k}\theta_{i}\quad\textnormal{for $k=0,1,\dotsc,r$.}$ $$

###### Proof.

From proposition A.2, we have a function ${\mathrm{step}}(x)$ defined by

$$ ${\mathrm{step}}(x)\coloneqq{\sigma_{R}}\big{(}\tfrac{x}{\epsilon}+1\big{)}-{ \sigma_{R}}\big{(}\tfrac{x}{\epsilon}\big{)},$ $$

satisfies

$$ ${\mathrm{step}}(x)=\begin{cases}1&\text{if }x\geq 0,\\ 0&\text{if }x<0.\end{cases}$ $$

Define $\beta_{i}$ for $i=0,2,\dotsc,r$ as

$$ $\beta_{i}=\textnormal{bin}\hskip 1.5pt0.\theta_{i}\cdots\theta_{r},$ $$

where $\textnormal{bin}\hskip 1.5pt0.\theta_{1}\theta_{2}\cdots\theta_{r}$ denote the binary representation of $\theta=\sum_{i=1}^{n}\theta_{i}2^{-i}$ and $\theta_{0}\coloneq 0$. We have

$$ $\theta_{i}={\mathrm{step}}(\textnormal{bin}\hskip 1.5pt0.\theta_{i}\cdots \theta_{r}-\tfrac{1}{2})={\mathrm{step}}(\beta_{i}-\tfrac{1}{2}),$ $$

implying, for $i=1,2,\dotsc,r-1$,

$$ $\beta_{i+1}=2\beta_{i}-\theta_{i}=2\beta_{i}-{\mathrm{step}}\big{(}\beta_{i}- \tfrac{1}{2}\big{)}.$ $$

For given $k\in{0,1,\dotsc,r}$, we have

$$ $\begin{split}\sum_{i=1}^{k}\theta_{i}=\sum_{i=1}^{k}\theta_{i}+\sum_{i=k+1}^{r }0=\sum_{i=1}^{r}\theta_{i}\cdot{\mathrm{step}}(k-i)&=\sum_{i=1}^{r}{\sigma_{R }}\Big{(}\theta_{i}+{\mathrm{step}}(k-i)-1\Big{)}\\ &=\sum_{i=1}^{r}{\sigma_{R}}\Big{(}{\mathrm{step}}(\beta_{i}-\tfrac{1}{2})+{ \mathrm{step}}(k-i)-1\Big{)}.\end{split}$ (23) $$

To compute the right-hand side, we need two nested ${\mathrm{ReLU}}$ functions. By using loop unrolling, we precompute ${\mathrm{step}}(\beta_{i}-\frac{1}{2})$ and ${\mathrm{step}}(k-i)$ in the previous
loops, allowing us to require only a single layer.
Define $\bm{\xi}_{l}$ for $l=0,1,\dotsc,r$ as

$$ $\bm{\xi}_{l}={\Big{(}k-l,\ \beta_{l},\ \beta_{l+1},\ {\mathrm{step}}(\beta_{l} -\frac{1}{2}),\ {\mathrm{step}}(k-l),\ \mathrm{sum}(l)\Big{)}}^{\top},$ $$

where $\mathrm{sum}(l)\coloneq\sum_{i=1}^{l}{\sigma_{R}}\Big{(}{\mathrm{step}}(\beta_
{i}-\tfrac{1}{2})+{\mathrm{step}}(k-i)-1\Big{)}$. Note that we have $\beta_{l+1}$ in the $l$-th loop to precompute ${\mathrm{step}}(\beta_{l+1}-\frac{1}{2})$ and ${\mathrm{step}}\big{(}(k-(l+1)\big{)}$ for the $l+1$-th loop.

Define ${\mathrm{FF}}:\mathbb{R}^{7}\to\mathbb{R}^{7}$ such that

$$ $\displaystyle\big{(}\mathrm{id}+{\mathrm{FF}}\big{)}(\bm{\xi}_{l})=\big{(} \mathrm{id}+{\mathrm{FF}}\big{)}\left(\left[\begin{array}[]{c}k-l\\ \beta_{l}\\ \beta_{l+1}\\ {\mathrm{step}}(\beta_{l}-\frac{1}{2})\\ {\mathrm{step}}(k-l)\\ \mathrm{sum}(l)\\ \end{array}\right]\right)$ $\displaystyle=\left[\begin{array}[]{c}k-l\\ \beta_{l}\\ \beta_{l+1}\\ {\mathrm{step}}(\beta_{l}-\frac{1}{2})\\ {\mathrm{step}}(k-l)\\ \mathrm{sum}(l)\\ \end{array}\right]+\begin{bmatrix}0&0&0&0&0&0&0&0&0\\ 1&1&0&0&0&0&0&0&0\\ 0&0&1&-1&1&0&0&0&0\\ 0&0&0&1&-1&-1&0&0&0\\ 0&0&0&0&-1&1&-1&0&0\\ 0&0&0&0&0&0&0&0&1\\ \end{bmatrix}$ $\displaystyle{\sigma_{R}}\left(\begin{bmatrix}0&1&0&0&0&0\\ 0&0&0&-1&0&0\\ 0&0&1&0&0&0\\ 0&0&1/\epsilon&0&0&0\\ 0&0&1/\epsilon&0&0&0\\ 0&0&0&-1&0&0\\ 0&0&0&0&1&0\\ 1/\epsilon&0&0&0&0&0\\ 1/\epsilon&0&0&0&0&0\\ 0&0&0&1&1&0\\ \end{bmatrix}\left[\begin{array}[]{c}k-l\\ \beta_{l}\\ \beta_{l+1}\\ {\mathrm{step}}(\beta_{l}-\frac{1}{2})\\ {\mathrm{step}}(k-l)\\ \mathrm{sum}(l)\\ \end{array}\right]+\left[\begin{array}[]{c}0\\ 0\\ 0\\ -1/(2\epsilon)+1\\ -1/(2\epsilon)\\ 0\\ 0\\ -1/\epsilon+1\\ -1/\epsilon\\ -1\\ \end{array}\right]\right)+\left[\begin{array}[]{c}-1\\ 0\\ 0\\ 0\\ 0\\ 0\\ 0\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{c}k-l\\ \beta_{l}\\ \beta_{l+1}\\ {\mathrm{step}}(\beta_{l}-\frac{1}{2})\\ {\mathrm{step}}(k-l)\\ \mathrm{sum}(l)\\ \end{array}\right]+\left[\begin{array}[]{ccc}-1\\ {\sigma_{R}}(\beta_{l})-{\sigma_{R}}\big{(}{\mathrm{step}}(\beta_{l}-\frac{1}{ 2})\big{)}\\ {\sigma_{R}}(\beta_{l+1})-\big{(}{\sigma_{R}}(\tfrac{\beta_{l+1}-1/2}{\epsilon }+1)-{\sigma_{R}}(\tfrac{\beta_{l+1}-1/2}{\epsilon})\big{)}\\ -{\sigma_{R}}\big{(}{\mathrm{step}}(\beta_{l}-\frac{1}{2})\big{)}+{\sigma_{R}} (\tfrac{\beta_{l+1}-1/2}{\epsilon}+1)-{\sigma_{R}}(\tfrac{\beta_{l+1}-1/2}{ \epsilon})\\ -{\sigma_{R}}\big{(}{\mathrm{step}}(k-l)\big{)}+{\sigma_{R}}(\tfrac{k-(l+1)}{ \epsilon}+1)-{\sigma_{R}}(\tfrac{k-(l+1)}{\epsilon})\\ {\sigma_{R}}\Big{(}{\mathrm{step}}(k-l)+{\mathrm{step}}(\beta_{l}-\frac{1}{2}) -1\Big{)}\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{ccc}k-(l+1)\\ 2\beta_{l}-{\mathrm{step}}(\beta_{l}-\frac{1}{2})\\ 2\beta_{l+1}-{\mathrm{step}}(\beta_{l+1}-\frac{1}{2})\\ {\mathrm{step}}(\beta_{l+1}-\frac{1}{2})\\ {\mathrm{step}}\big{(}(k-(l+1)\big{)}\\ \mathrm{sum}(l+1)\\ \end{array}\right]=\left[\begin{array}[]{ccc}k-(l+1)\\ \beta_{l+1}\\ \beta_{l+2}\\ {\mathrm{step}}(\beta_{l+1}-\frac{1}{2})\\ {\mathrm{step}}\big{(}(k-(l+1)\big{)}\\ \mathrm{sum}(l+1)\\ \end{array}\right]=\bm{\xi}_{l+1},$ $$

Define $\mathcal{L}_{1}:\mathbb{R}^{2}\to\mathbb{R}^{3}$ and ${\mathcal{L}}_{2}:\mathbb{R}^{3}\to\mathbb{R}$ via

$$ $\mathcal{L}_{1}(k)\coloneqq{(k,\beta_{0},\beta_{1},0,0,0)}^{\top}=\bm{\xi}_{0} ,\quad\mathcal{L}_{2}(x_{1},x_{2},x_{3},x_{4},x_{5},x_{6},x_{7})\coloneqq x_{7},$ $$

respectively. Note $\beta_{1}$ is defined by given $\theta_{1},\theta_{2},\dotsc,\theta_{r}$
∎

We prove Lemma 4.1 with Lemma A.8.

See 4.1

###### Proof.

We prove this for the case where $d=1$, considering $y_{k}\in\mathbb{R}$ for $k=0,\dotsc,m$. Define

$$ $a_{i}=\big{\lfloor}\tfrac{y_{i}}{\varepsilon}\big{\rfloor}\quad\textnormal{for $i=0,1,\dotsc,m-1$},$ $$

where $\lfloor x\rfloor=\max\{n:n\leq x,\ n\in\mathbb{Z}\}$
and set

$$ $b_{i}=a_{i}-a_{i-1}\quad\textnormal{for $i=1,2,\dotsc,m-1$.}$ $$

Since $b_{i}\in\{-1,0,1\}$, there exist $c_{i}\in\{0,1\}$ and $d_{i}\in\{0,1\}$ such that

$$ $b_{i}=c_{i}-d_{i}\quad\textnormal{ for $i=1,2,\dotsc,m-1$.}$ $$

Thus, we have

$$ $\begin{split}a_{k}=a_{0}+\sum_{i=1}^{k}c_{i}-\sum_{i=1}^{k}d_{i}\quad \textnormal{for $k=0,1,\dotsc,m-1$.}\end{split}$ $$

From Lemma A.8, there exist ${\mathrm{FF}}^{(c)},{\mathrm{FF}}^{(d)}:\mathbb{R}^{7}\to\mathbb{R}^{7}$ and affine linear maps $\mathcal{L}^{\prime}_{2}:\mathbb{R}^{7}\to\mathbb{R}$ and $\mathcal{L}^{(c)}_{1},\mathcal{L}^{(d)}_{1}:\mathbb{R}\to\mathbb{R}^{7}$ such that

$$ $\displaystyle\mathcal{L}^{\prime}_{2}\circ{(\mathrm{id}+{\mathrm{FF}}^{(c)})}^ {\circ(m-1)}\circ\mathcal{L}^{(c)}_{1}\big{(}k)=\sum_{i=1}^{k}c_{i},\quad \mathcal{L}^{\prime}_{2}\circ{(\mathrm{id}+{\mathrm{FF}}^{(d)})}^{\circ(m-1)} \circ\mathcal{L}^{(d)}_{1}\big{(}k)=\sum_{i=1}^{k}d_{i},$ $$

for $k=0,1,\dotsc,m-1$. Then, define ${\mathrm{FF}}:\mathbb{R}^{14}\to\mathbb{R}^{14}$, for ${\bm{x}},{\bm{y}}\in\mathbb{R}^{7}$,

$$ ${\mathrm{FF}}({\bm{x}},{\bm{y}})\coloneqq({\mathrm{FF}}^{(c)}({\bm{x}}),{ \mathrm{FF}}^{(d)}({\bm{y}})).$ $$

Define $\mathcal{L}_{1}:\mathbb{R}\to\mathbb{R}^{14}$ and $\mathcal{L}_{2}:\mathbb{R}^{14}\to\mathbb{R}$ as

$$ $\mathcal{L}_{1}(x)\coloneqq\big{(}\mathcal{L}^{(c)}_{1}(x),\mathcal{L}^{(d)}_{ 1}(x)\big{)},\quad\mathcal{L}_{2}({\bm{x}},{\bm{y}})^{\top}\coloneqq{\epsilon} \big{(}a_{0}+\mathcal{L}^{\prime}_{2}({\bm{x}})-\mathcal{L}^{\prime}_{2}({\bm{ y}})\big{)}.$ $$

We can confirm that

$$ $\displaystyle{\mathcal{L}}_{2}\circ(\mathrm{id}+\mathrm{FF})^{\circ(m-1)}\circ {\mathcal{L}}_{1}(k)$ $\displaystyle={\mathcal{L}}_{2}\circ(\mathrm{id}+\mathrm{FF})^{\circ(m-1)} \circ\big{(}\mathcal{L}^{(c)}_{1}(k),\ \mathcal{L}^{(d)}_{1}(k)\big{)}$ $\displaystyle={\mathcal{L}}_{2}\circ\big{(}(\mathrm{id}+\mathrm{FF}^{(c)})^{ \circ(m-1)}\circ\mathcal{L}^{(c)}_{1}(k),\ (\mathrm{id}+\mathrm{FF}^{(d)})^{ \circ(m-1)}\circ\mathcal{L}^{(d)}_{1}(k)\big{)}$ $\displaystyle={\epsilon}\Big{(}a_{0}+\mathcal{L}^{\prime}_{2}\circ(\mathrm{id} +\mathrm{FF}^{(c)})^{\circ(m-1)}\circ\mathcal{L}^{(c)}_{1}(k)-\mathcal{L}^{ \prime}_{2}\circ(\mathrm{id}+\mathrm{FF}^{(d)})^{\circ(m-1)}\circ\mathcal{L}^{ (d)}_{1}(k)\Big{)}$ $\displaystyle={\epsilon}\big{(}a_{0}+\sum_{i=1}^{k}c_{i}-\sum_{i=1}^{k}d_{i} \big{)}={\epsilon}a_{k}.$ $$

Thus we have

$$ $|\big{(}{\mathcal{L}}_{2}\circ(\mathrm{id}+\mathrm{FF})^{\circ(m-1)}\circ{ \mathcal{L}}_{1}(k)-{\bm{y}}_{k}\big{)}_{i}|=|{\epsilon}a_{k}-y_{k}|\leq\varepsilon.$ $$

We can extend this for $d$-dimensional input.
∎

###### Proposition A.7 (Lemma D.1 in Zhang et al. ( 2023 ) ) .

###### Lemma A.8 .

###### Proof.

###### Proof.

## Appendix B Role of Time dependent scaling parameters

We demonstrate that time-dependent scaling parameters overcome the limitations inherent to the looped architecture and eliminate the dependence of the approximation rate on the modulus of continuity. We use the architecture defined in Section 4 as:

$$ ${\mathrm{FF}}({\bm{x}})\to\bm{\eta}(t)\odot{\mathrm{FF}}({\bm{x}})\quad\text{ for the $t$-th loops},$ $$

Following lemma demonstrate that time dependent scaling parameters can exactly map indices to output vectors.

See 4.2

The key idea of the proof is that we use the impulse function defined as

$$ $\displaystyle\mathrm{impulse}_{0}(({\bm{y}}_{l})_{i},k-l)=\begin{cases}({\bm{y }}_{l})_{i}\quad&\text{if}\ k=l,\\ 0\quad&\text{otherwise},\end{cases}$ $$

for $i=1,2,\dotsc,d$ and $l=0,2,\dotsc,m-1$, which extracts the corresponding $({\bm{y}}_{l})_{i}$ in the $l$-th loop if the index matches $k$.

###### Proof.

We consider the case when $d=1$, where $y_{k}\in\mathbb{R}$ for $k=0,\dotsc,m-1$. We update $y_{k}$ as follows:

$$ $y_{k}\to y_{k}+\epsilon,$ $$

where $\epsilon$ is chosen such that none of the $y_{l}$ values are zero.

Next, we define $\eta(l)$ as:

$$ $\eta(l)=(0,1,\frac{y_{l}}{y_{l-1}}-1,1)^{\top}\quad\text{for }l=1,2,\dotsc,m-1.$ $$

By Proposition A.3, we have

$$ $\displaystyle\mathrm{impulse}_{0}(x,t)$ $\displaystyle\coloneqq{\sigma_{R}}\big{(}x+2M(t+1/2)\big{)}-2M{\sigma_{R}}(t+1 /2)$ $\displaystyle-{\sigma_{R}}\big{(}x+2M(t-1/2)\big{)}+2M{\sigma_{R}}(t-1/2)$ $\displaystyle=\begin{cases}x\quad&\textnormal{if}\ t=0,\\ 0\quad&\textnormal{otherwise},\end{cases}$ $$

where $M>\max y_{k}$.
Define

$$ $s(l)\coloneq\sum_{i=0}^{l}\mathrm{impulse}_{0}\big{(}y_{(i-1)},k-(i-1)\big{)},$ $$

for $l=1,2,\dotsc,m-1$. This satisfies

$$ $s(m-1)=\sum_{i=0}^{m-1}\mathrm{impulse}_{0}\big{(}y_{(i-1)},k-(i-1)\big{)}=y_{ k}.$ $$

Define $\bm{\xi}_{l}$ as

$$ $\bm{\xi}_{l}=\Big{(}k,\ k-l,\ y_{l},\ s(l))\Big{)}^{\top}.$ $$

for $l=0,1,2,\dotsc,m-1$.

Then, define ${\mathrm{FF}}:\mathbb{R}^{4}\to\mathbb{R}^{4}$ of width size $q=6$ via:

$$ $\displaystyle(\mathrm{id}+\eta(l)\odot\mathrm{FF})(\bm{\xi}_{l-1})=\bm{\xi}_{l -1}+$ $\displaystyle\eta(l)\odot\left(\begin{bmatrix}0&0&0&0&0&0\\ 0&0&0&0&0&0\\ 1&-1&0&0&0&0\\ 0&0&1&-1&-2M&2M\\ \end{bmatrix}{\sigma_{R}}\left(\begin{bmatrix}0&0&1&0\\ 0&0&-1&0\\ 0&2M&1&0\\ 0&2M&1&0\\ 0&1&0&0\\ 0&1&0&0\\ \end{bmatrix}\bm{\xi}_{l-1}+\left[\begin{array}[]{ccc}0\\ -1\\ M\\ -M\\ 1/2\\ -1/2\\ \end{array}\right]\right)+\left[\begin{array}[]{ccc}0\\ -1\\ 0\\ 0\\ \end{array}\right]\right)$ $\displaystyle=\bm{\xi}_{l-1}+\left[\begin{array}[]{ccc}0\\ 1\\ \tfrac{y_{l}}{y_{l-1}}-1\\ 1\\ \end{array}\right]\odot\left[\begin{array}[]{ccc}0\\ -1\\ {\sigma_{R}}(y_{l-1})-{\sigma_{R}}(-y_{l-1})\\ \Bigg{(}{\sigma_{R}}\big{(}y_{l-1}+2M((k-(l-1))+1/2)\big{)}\\ -2M{\sigma_{R}}((k-(l-1))+1/2)\\ -{\sigma_{R}}\big{(}y_{l-1}+2M(k-(l-1)-1/2)\big{)}+2M{\sigma_{R}}(k-(l-1)-1/2) \Bigg{)}\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{ccc}k\\ k-(l-1)\\ y_{l-1}\\ s(l-1)\\ \end{array}\right]+\left[\begin{array}[]{ccc}0\\ 1\\ \tfrac{y_{l}}{y_{l-1}}-1\\ 1\\ \end{array}\right]\odot\left[\begin{array}[]{ccc}0\\ -1\\ y_{l-1}\\ \mathrm{impulse}_{0}\big{(}y_{(l-1)},k-(l-1)\big{)}\\ \end{array}\right]$ $\displaystyle=\left[\begin{array}[]{ccc}k\\ k-l\\ y_{l}\\ s(l)\\ \end{array}\right]=\bm{\xi}_{l}.$ $$

for $l=1,2,\dotsc,m-1$. Thus we have

$$ $(\mathrm{id}+\eta(m-1)\cdot\mathrm{FF})\circ\cdots\circ(\mathrm{id}+\eta(1) \cdot\mathrm{FF})(\bm{\xi}_{0})=\bm{\xi}_{m-1}$ $$

Then, define two affine linear maps $\mathcal{L}_{1}:\mathbb{R}^{1}\to\mathbb{R}^{4}$ and $\mathcal{L}_{2}:\mathbb{R}^{4}\to\mathbb{R}^{1}$ by

$$ $\mathcal{L}_{1}(x)\coloneqq(k,k,y_{0},0),\quad\mathcal{L}_{2}(x_{1},x_{2},x_{3 },x_{4})\coloneqq x_{4}-\epsilon.$ $$

We can extend this to $d$-dimensional input by using $d$ time parameters, by applying above to each dimension.
∎

###### Proof.

## Appendix C Supplement for Experiments

We generate problem instance for edit distance problem as shown in Algorithm 1.
The first string is randomly selected, and we use two methods to generate the second string. In the first method, we create random results that yield a large edit distance cost. In the second case, we copy the first string with corruption, resulting in a small edit distance. The two strings are concatenated with a separator.

Figure: Algorithm 1 ED Data Generation from Feng et al. ()

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
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).
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 3
questions: See weakness
soundness: 2
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.
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.
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.

Meta-review: **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.

Decision: Reject