Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators?

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

Q1. What are the size of FFNs and the number of model parameters in the construction?

A1. The implementation [6] of the FF layer we adopted in Corollary 1 and 2 requires the number of parameters corresponding to the possible number of context ids and embedding dimension $d$. So, in the worst case, the number of parameters of the FF layer is $2n\binom{|\mathcal{V}|}{n} + d$. As for the self-attention layer, weight matrices $\boldsymbol{W}^{(O)}, \boldsymbol{W}^{(V)}, \boldsymbol{W}^{(K)}$ and $\boldsymbol{W}^{(Q)}$ all require $s + d$ parameters each. Thus, the number of parameters of the self-attention layer is $4(s+d)$.

It is true that the number of parameters for the FF layer increases as the number of layers in an FF layer in Transformer models is fixed at $2$ layers. However, if we try to perform memorization in one-layer Transformers, it is necessary to perform contextual mapping in the self-attention layer, and therefore, we must necessarily construct a two-layer FF block that associates each context id with a label. If the depth of the FF layer is not limited to $2$ layers, then the implementation of FFNs with the optimal memorization capacity [7] can be used as in the previous study [8].

Again, we would like to stress that our primary contribution lies in the assertion that a single-layer, single-head self-attention can already be a contextual mapping, and this conclusion cannot be derived from previous studies which treated the softmax function as an approximation of the hardmax function.

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.

[1] Approximation theory of transformer networks for sequence modeling.

[6] Understanding Deep Learning Requires Rethinking Generalization.

[7] On the Optimal Memorization Power of ReLU Neural Networks.

[8] Provable Memorization Capacity of Transformers.

[9] Dimension of metric spaces and Hilbert’s problem 13.

As for Q5, it is also worth mentioning that the linear dependency can be mitigated to the sub-linear $\tilde{O}(\sqrt{nN})$ when allowing for the deeper FF block, under the condition that the size $|\mathcal{V}|$ of the vocabulary set is independent of the input length $n$ and the number $N$ of input sequences. In such a case, the intervals $\delta$ between context ids, that is, the right-hand of equation 8 satisfies $\log \delta^{-1} = \tilde{O}(1)$, and thus the implementation of the deep FF layer from [3] can be applied.

We have added this line of discussion to the Remark 2 after the Corollary 1.

[3] On the Optimal Memorization Power of ReLU Neural Networks.

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

Q1. 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?

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

Q2. The role of the number of attention heads in determining the memorization capacity remains unclear. If it does have an impact, are there any quantitative metrics provided to elucidate its influence?

A2. The impact of increasing attention heads is also a highly interesting question, and we speculate that the effects may be particularly significant in real-world datasets, where inputs possess high-dimensionality.

In the proof of Theorem 2, input tokens are projected into one-dimensional space using key and query matrices, and then put together to generate a context id through the softmax function. In cases where the data is scattered across multiple dimensions, however, it may be possible to focus on different dimensions with each attention head and subsequently integrate the information, like when deepening Transformers mentioned above. By making certain assumptions about the multidimensionality of the input distribution, it might be possible to theoretically analyze the effects of multiple attention heads. We think that this direction, too, is a very interesting future work.

Thank you for taking the time to review our paper and sharing your insightful comments. We will address your questions below.

Q1. What is a sequence id, similarly what is a $\boldsymbol{v}$-dependent sequence id?

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

Q2. Is there any implication regarding the value of delta in the contextual mapping? In particular, how are the following approximation results affected if delta decreases?

A2. The term $\delta$ does not have an effect on the number of parameters in the context of universal approximation. The number of weights of the feed-forward layer we adopt as an implementation [2] depends solely on the number of possible context ids and the embedding dimension $d$. It is particularly independent of the value of $\delta$.

Q3. How does the memorization capacity depend on the vocabulary size?

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

Q4. As is mentioned above, can we value of the rank-$1$ projection be computed as implied by the theory and kept at that value in the experiments to see how well that performs?

A4. The projection vectors obtained in Lemma 3 originates from Lemma 2, or more precisely, Lemma 13 in [3]. Unfortunately, Lemma 13 in [3] was demonstrated through a probabilistic proof, making it challenging to prepare vectors that satisfy the desired property, i.e., equation (73) before training.

Q5. Is there any difficulty with extending Prop. 1 using positional encoding to consider functions that are not permutation equivariant?

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

[1] Provable Memorization Capacity of Transformers.

[2] Understanding Deep Learning Requires Rethinking Generalization.

[3] Provable Memorization via Deep Neural Networks Using Sub-linear Parameters.

[4] Are Transformers Universal Approximators of Sequence-to-Sequence Functions?