Paper: Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR-2024)

Reviewer comment (Theory):
There is no characterization on the total number of model parameters in the memorization capacity result. It is not surprising that Transformers can memorize. It is the upper/lower bound on the required number of parameters for memorization that makes the result interesting.

Author response:
> Q4. The statements of the memorization capacity results are still not satisfactory. Standard memorization capacity results should characterize the required number of parameters for memorizing $N$ data samples [1,2]. The authors need to explicitly discuss this and update the statements of the Corollaries.

A4. We have added an evaluation of the required number of  parameters $4(s + d) + d(2nN + d)$ to the Corollaries.

We highly appreciate your suggestive comment to help us improve our Corollaries further.

> Q5. Besides, I have to point out the dependency on $\binom{|V|}{n}$ is very bad (and vacuous). Can this be replaced with a linear or even sub-linear dependency on $N$? I believe the results in [3] (and other related papers) may be useful for the authors and I encourage the authors to elaborate on this.

A5. The required number of parameters for the memorization of one-layer Transformers can be evaluated as $O(N)$.

In fact, when we focus on the order of parameters in terms of $N$, the possible number of context ids is at most $nN$. Thus, the construction [4] of the FF layer we adopt requires $2nN+d$ parameters (to be precise, $d \times (2nN + d)$ because of output dimension $d$). By adding the number of parameters in the self-attention layer, which is $4(s + d)$, the total number of parameters necessary for one-layer Transformers to memorize $N$ input-output pairs of sequences with sequence length $n$ and embedding dimension $d$ is $4(s + d) + d(2nN + d)$, which is linear dependency on $N$.

Note that this linear dependency is optimal up to logarithmic factors for the one-hidden-layer FF block [5]. We have added a remark on the parameter efficiency after Corollary 1.

It seems non-trivial as to whether the required parameters become sub-linear when allowing the FF layer to be deep.

This is mainly because prior works assume that the separatedness of inputs are almost independent of the number $m$ of data. More precisely, for $(r,\delta)$-separated $m$ inputs to memorize, it is necessary that $\log (r\delta^{-1}) = \tilde{O}(1)$ [3] or $\log (r\delta^{-1}) = O(m^{2/3})$ [6].
In contrast, when attempting contextual mapping with a one-layer self-attention, context ids become $(r, \delta)$-separated with $\log (r\delta^{-1}) = O(n^4N^4)$ (equation 8), a scenario not considered in prior works [3,6].

The question of whether there exists a deep FFN capable of achieving memorization with sub-linear parameters for datasets with such a narrow input intervals is highly interesting in itself, and could lead to a new line of research. Again, we deeply appreciate your insightful comments.

Reviewer comment (Novelty):
This paper is _not_ the first universal approximation theorem for two-layer Transformers with a self-attention of realistic size, and Section 4.2 _overclaims_ the contribution because of misinterpretation of [1]. The universal approximation result (Theorem 4.1) in [1] _is_ constructive and the construction is even simpler than in this paper. The authors also claim that [1] makes a particular assumption on the domain of functions. But those assumptions are only for the approximation rate result (Theorem 4.2), which is orthogonal to this paper.

Author response:
> Q2. Can you clarify the discussion on [1] in the paper?

A2. We refer to [1] as a non-constructive proof in the sense that the proof [9] of the Kolmogorov representation theorem for compact metric spaces the authors cited in [1] is non-constructive, and as a result, the inner and outer functions approximated by the FF layers in their construction are black-boxes.
Of course, we agree that non-constructive proofs also have their advantages, such as making proofs clearer.

Furthermore, we would like to point out that the construction in [1] includes an additional linear transformation layer before the 2-layer Transformer, in which input tokens are mapped to a high-dimensional space of dimension $4n^2d + 2n$.

What we mean by "particular assumption on the domain of functions" is the encoded input space $\mathcal{X}^{(E)}$ that appears in Theorem 4.1 in [1].
This space is defined by $\mathcal{X}^{(E)} = \\{  \boldsymbol{x}\ :\ x(x) \in \mathcal{I}^{(s)} \subset \mathbb{R}^d,$ where $\mathcal{I}^{(i)}, \mathcal{I}^{(j)}$ are disjoint, compact $\forall i,j \in \\{1,\dots,\tau\\} \\}$, and to our understanding, $\mathcal{I}^{(1)},\dots,\mathcal{I}^{(\tau)}$, which appear for the first time in the definition of the encoded input space, have to be given beforehand.
If we consider universal approximation theorems for two-layer Transformers with positional encoding and continuous functions on a compact domain, then this encoded input space suffices. 

In contrast, we showed that $2$-layer Transformer without positional encoding is a universal approximator for continuous permutation equivariant functions on a compact domain.
Our result can be readily extended for continuous but not necessarily permutation equivariant functions on a compact domain by using positional encoding, and at the same time is significant from the perspective of geometric deep learning.

To make clearer the contributions of [1] and the novelty of our results, we have updated the discussion after Prop.1 in our paper, and changed the phrases in the abstract and "Our Contributions" part as follows.

Before: universal approximation theorem for continuous functions.

After: universal approximation theorem for continuous "perumtation equivariant" functions.

Does our reply above answer your question? We have been given a very important point to make.
Let us know if our answer is still incomplete.

Reviewer comment (Theory):
**Proof issues.** Equation 10 seems buggy. The authors should specify $\mathbf X^{(i)}\neq \mathbf X^{(j)}$ or $\mathcal V^{(i)}\neq \mathcal V^{(j)}$. Furthermore, if it holds for some $i,j$ and $k,l$, then one can consider the same inequality for $j,i$ and $l,k$ and obtain a contradiction. This raises doubt about why the condition can be true.

Author response:
> Q3. Can you clarify the proof issue mentioned in weaknesses?

A3. As for the bugs in equation (10), you are absolutely right. To be precise, the left-hand side of equation (10) should be its absolute value, and also the condition $\mathcal{V}^{(i)} \neq \mathcal{V}^{(j)}$ should be added.
Thank you for reporting the bugs.

We also very much thank you for informing us of the presentation issues and the typo regarding the year of publication [2]. We have fixed these issues and updated our paper.
We have also added [3,4,5] to the related work section.

Reviewer comment (Novelty):
**Related works.** The year of [2] is marked as 2023 in the paper, but [2] was actually published in ICLR 2020. The authors should correct this to avoid misleading readers. In addition, they are recommended to also discuss [3,4,5].

Author response:
> Q3. Can you clarify the proof issue mentioned in weaknesses?

A3. As for the bugs in equation (10), you are absolutely right. To be precise, the left-hand side of equation (10) should be its absolute value, and also the condition $\mathcal{V}^{(i)} \neq \mathcal{V}^{(j)}$ should be added.
Thank you for reporting the bugs.

We also very much thank you for informing us of the presentation issues and the typo regarding the year of publication [2]. We have fixed these issues and updated our paper.
We have also added [3,4,5] to the related work section.

Reviewer comment (Writing):
**Presentation issues.** There are many typos and minor errors in the paper. For example:
- Page 3, "the Transformer block is represented as a combination of ..." → "the Transformer block is represented as a composition of ..."
- Page 5, "this theorem indicates that one-layer Transformer does not have a memorization capacity": Please emphasize one-layer Transformer _with hardmax attention_.
- Theorem 2, $\mathbf{W}^V$ → $\mathbf{W}^{(V)}$.
- The appendix uses eq. $i$ and equation $i$ interchangeably. Please make it consistent.

Author response:
> Q3. Can you clarify the proof issue mentioned in weaknesses?

A3. As for the bugs in equation (10), you are absolutely right. To be precise, the left-hand side of equation (10) should be its absolute value, and also the condition $\mathcal{V}^{(i)} \neq \mathcal{V}^{(j)}$ should be added.
Thank you for reporting the bugs.

We also very much thank you for informing us of the presentation issues and the typo regarding the year of publication [2]. We have fixed these issues and updated our paper.
We have also added [3,4,5] to the related work section.

Reviewer comment (Writing):
**Question**: Can you clarify the discussion on [1] in the paper?

Author response:
> Q2. Can you clarify the discussion on [1] in the paper?

A2. We refer to [1] as a non-constructive proof in the sense that the proof [9] of the Kolmogorov representation theorem for compact metric spaces the authors cited in [1] is non-constructive, and as a result, the inner and outer functions approximated by the FF layers in their construction are black-boxes.
Of course, we agree that non-constructive proofs also have their advantages, such as making proofs clearer.

Furthermore, we would like to point out that the construction in [1] includes an additional linear transformation layer before the 2-layer Transformer, in which input tokens are mapped to a high-dimensional space of dimension $4n^2d + 2n$.

What we mean by "particular assumption on the domain of functions" is the encoded input space $\mathcal{X}^{(E)}$ that appears in Theorem 4.1 in [1].
This space is defined by $\mathcal{X}^{(E)} = \\{  \boldsymbol{x}\ :\ x(x) \in \mathcal{I}^{(s)} \subset \mathbb{R}^d,$ where $\mathcal{I}^{(i)}, \mathcal{I}^{(j)}$ are disjoint, compact $\forall i,j \in \\{1,\dots,\tau\\} \\}$, and to our understanding, $\mathcal{I}^{(1)},\dots,\mathcal{I}^{(\tau)}$, which appear for the first time in the definition of the encoded input space, have to be given beforehand.
If we consider universal approximation theorems for two-layer Transformers with positional encoding and continuous functions on a compact domain, then this encoded input space suffices. 

In contrast, we showed that $2$-layer Transformer without positional encoding is a universal approximator for continuous permutation equivariant functions on a compact domain.
Our result can be readily extended for continuous but not necessarily permutation equivariant functions on a compact domain by using positional encoding, and at the same time is significant from the perspective of geometric deep learning.

To make clearer the contributions of [1] and the novelty of our results, we have updated the discussion after Prop.1 in our paper, and changed the phrases in the abstract and "Our Contributions" part as follows.

Before: universal approximation theorem for continuous functions.

After: universal approximation theorem for continuous "perumtation equivariant" functions.

Does our reply above answer your question? We have been given a very important point to make.
Let us know if our answer is still incomplete.

Reviewer comment (Theory):
**Question**: Can you clarify the proof issue mentioned in **weaknesses**?

Author response:
> Q3. Can you clarify the proof issue mentioned in weaknesses?

A3. As for the bugs in equation (10), you are absolutely right. To be precise, the left-hand side of equation (10) should be its absolute value, and also the condition $\mathcal{V}^{(i)} \neq \mathcal{V}^{(j)}$ should be added.
Thank you for reporting the bugs.

We also very much thank you for informing us of the presentation issues and the typo regarding the year of publication [2]. We have fixed these issues and updated our paper.
We have also added [3,4,5] to the related work section.

Reviewer comment (Theory):
While the study provides insights into one-layer Transformers, it raises the question of whether these findings can be extended to two-layer or even deeper architectures. How scalable is the presented theory?

Author response:
We consider that one of the primary benefits of increasing the depth of the layers is that self-attention mechanisms may no longer need to identify the input sequence all at once.
In other words, it becomes possible to capture the context by having self-attention mechanisms in an earlier layers recognize different parts of the input sequence, and then consolidating this information in the final layer.

This computation strategy could allow the construction of models capable of calculating more separated context ids, which in turn provides a potential proof that deeper Transformers are more easily trainable.
This paper focuses on the expressive capacity of Transformers, and while we leave this possibility on an optimization aspect of deep Transformers as a future work, we believe that delving into this research direction would be a compelling and intriguing next step.

Reviewer comment (Writing):
I had some issues with a few keywords that were not introduced; for example, in the discussion of Theorem 2 I had no idea what a sequence id was, or how it is used to construct the context id (which, if I understood correctly, is the output of the attention map).

Author response:
The term "sequence id" refers to an identifier assigned to distinguish each input sequence.  
To recognize the context of a token, it is sufficient to obtain the sequence id and the token id, which is an id corresponding to the token.  
In this study, we calculate the context id using the linear combination of a sequence id and a token id. To compute the context id in a single-layer self-attention mechanism, we use each input token vector directly for the token id. The sequence id, calculated through the self-attention mechanism, is appropriately scaled and then combined with the token id through the skip connection to obtain the context id.  

A $\boldsymbol{v}$-dependent sequence id for some token $\boldsymbol{v} \in \mathcal{V}$, on the other hand, is a more intricate concept deeply involved in the proof.  
As mentioned earlier, the context id is generated from the linear combination of a token id and a sequence id, so in this sense, it can be regarded as a pair of the token id and the sequence id, i.e., (token id, sequence id).  
In the prior research [1], sequence ids are uniquely generated for each input sequence. However, by considering the context id as such a pair, it becomes apparent that sequence ids need not be unique. Specifically,  given some input token $\boldsymbol{v}$, it is sufficient for the identifier to distinguish each input sequence containing $\boldsymbol{v}$. We refer to this identifier as $\boldsymbol{v}$-dependent sequence id.

Reviewer comment (Writing):
What is a sequence id, and similarly, what is a $v$-dependent sequence id?

Author response:
The term "sequence id" refers to an identifier assigned to distinguish each input sequence.  
To recognize the context of a token, it is sufficient to obtain the sequence id and the token id, which is an id corresponding to the token.  
In this study, we calculate the context id using the linear combination of a sequence id and a token id. To compute the context id in a single-layer self-attention mechanism, we use each input token vector directly for the token id. The sequence id, calculated through the self-attention mechanism, is appropriately scaled and then combined with the token id through the skip connection to obtain the context id.  

A $\boldsymbol{v}$-dependent sequence id for some token $\boldsymbol{v} \in \mathcal{V}$, on the other hand, is a more intricate concept deeply involved in the proof.  
As mentioned earlier, the context id is generated from the linear combination of a token id and a sequence id, so in this sense, it can be regarded as a pair of the token id and the sequence id, i.e., (token id, sequence id).  
In the prior research [1], sequence ids are uniquely generated for each input sequence. However, by considering the context id as such a pair, it becomes apparent that sequence ids need not be unique. Specifically,  given some input token $\boldsymbol{v}$, it is sufficient for the identifier to distinguish each input sequence containing $\boldsymbol{v}$. We refer to this identifier as $\boldsymbol{v}$-dependent sequence id.

Reviewer comment (Experiments):
How does the memorization capacity depend on the vocabulary size?

Author response:
Unlike $\delta$, the vocabulary size has an impact on the number of parameters of the FF layer. Specifically, the maximum number of possible context ids is $n\binom{|\mathcal{V}|}{n}$. Therefore, when the implementation of the FF layer follows [2], $2n\binom{|\mathcal{V}|}{n} + d$ parameters are required for memorization.

Reviewer comment (Theory):
Is there any difficulty with extending Prop. 1 using positional encoding to consider functions that are not permutation equivariant?

Author response:
We believe it is relatively straightforward to extend Prop. 1 for continuous functions on a compact domain that are not permutation equivariant.  
In the original proof of universal approximation theorems for Transformers [4], the authors showed in Theorem 3 that Transformers with positional encoding can approximate an arbitrary continuous function on a compact domain that are not necessarily permutation equivariant with arbitrary precision. So basically the same proof technique should be applicable to our Prop. 1.