Attention-Only Transformers via Unrolled Subspace Denoising

Q1. More fair discussions should be given in the experiments. For instance, the results in Table 1 show that the proposed method is 9.3% lower than the compared baseline. Using "comparable performance" is unfair.

A1. Thanks for the comment. To address the concern, we will change "comparable performance" to ''worse performance'' and revise the manuscript to provide a more accurate description of the results.

Additionally, after a more careful examination, we found that the CRATE baseline we compared against uses a patch size of 8, while our implementation uses a patch size of 16, making the comparison unfair. A smaller patch size results in more image tokens, which usually leads to better performance. We are currently training models with a patch size of 8 and will report the result shortly.

Q2. The advantage of the proposed method might be strengthened by comparing it with more transformer-based methods.

A2. Thanks for the suggestion. As suggested, we first compare AoT-MHSA to Vision Transformer (ViT) [1] on the ImageNet. Due to the time limitation, we train the model on ImageNet-1K from scratch and report the results as follows:

Next, we compare AoT-MSSA to Llama [2] architecture on the in-context learning task (Section 4.2.2) and report the result at the following anonymous link: https://postimg.cc/gallery/d0Ypwvc.

Based on the above experimental results, it is observed that the proposed architecture demonstrates performance that is comparable to, or slightly lower than, that of state-of-the-art transformers.

[1] Dosovitskiy, Alexey, et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.

[2] Touvron, Hugo, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.

Q1. Inconsistent notation: $Z^{(l)}$, $f^l$ have layer superscript $l$ in the equation but $U_k$ never does, giving the impression that $U_k$ are the same for each layer. Figure 3 architecture has $U^l$ and so does Yu et al., 2023.

A1. Thanks for pointing this out. This confusion stems from a gap between the theoretical construction of transformers under Definition 2.1 and their practical implementation for real-world data. Specifically, our approach constructs a transformer in a forward manner by interpreting each layer as a subspace denoising operator, assuming that the initial token representations satisfy Definition 2.1. In this theoretical setting, it suffices to consider $U_k^{(l)} = U_k$ across layers. However, in practical scenarios with real-world data, these subspace matrices $\{U_k^{(l)}\}$ need to be learned gradually via backpropagation and may be different across layers. We will include the above discussion and the following equation in Section 3.1 to clarify this point:

$$\mathrm{Z}^{(l+1)} = \mathrm{Z}^{(l)} + \eta \sum_{k=1}^K \mathrm{U}_k^{(l)} \mathrm{U}_k^{(l)^T} \mathrm{Z}^{(l)}\varphi\left( \mathrm{Z}^{(l)^T}\mathrm{U}_k^{(l)}\mathrm{U}_k^{(l)^T}\mathrm{Z}^{(l)}\right)$$

Q2. Layer norm: Layer norm appears in Fig. 3 but does not get mentioned much in the model or the analysis.

A2. Yes, layer normalization indeed appears in Fig. 3 as part of the practical implementation of the transformer architecture. Our theoretical framework primarily focuses on the core components of transformers, namely self-attention and skip connections, as denoising token representations can be achieved without layer normalization when the initial token representations satisfy Definition 2.1. However, in practice, layer normalization is necessary to stabilize training and improve convergence. We will include a brief discussion on the role of layer norm in the revised version to clarify its inclusion in the architecture.

Q3. I would prefer a clearer presentation of the architecture and not having to rely on the previous papers for disambiguation.

A3. Thanks for the valuable comment. To address this, we will include a more detailed and explicit description of the architecture in the revised manuscript to reduce the dependency on previous papers such as Yu et al. (2023a, b).

Weakness 1. Inconsistency between theory and implementation: According to the parameter calculations, MSSA ultimately employs the same methodology as MHSA, where all projection matrices $U_o$, $U_q$, $U_k$, and $U_v$ are not shared...

A1. We should clarify that the theory and implementation are consistent. Specifically, we have implemented transformers using both MSSA (see AoT-MSSA in Section 4.1) and MHSA (see AoT-MHSA in Section 4.2). Our implementation of AoT-MSSA strictly follows the theoretical framework presented in the paper, where the projection matrices are designed according to the structured attention mechanism.

Weakness 2. Suboptimal experimental results: The experimental results in Table 1 indicate that AoT's performance is significantly suboptimal.

A2. Thanks for the comment. After a more careful examination, we found that the CRATE baseline we compared against uses a patch size of 8 in Table 1, while our implementation uses a patch size of 16, making the comparison unfair. A smaller patch size results in more image tokens, which usually leads to better performance. We are currently training models with a patch size of 8 and will report the result shortly.

Moreover, we compare AoT-MHSA to Vision Transformer (ViT) [1] on the ImageNet. We train the model on ImageNet-1K from scratch and report the results as follows:

We believe that with further tuning, we can narrow the performance gap even more.

[1] Dosovitskiy, Alexey, et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.

Weakness 3. Problematic claims on the experimental results: In nanogpt experiments, according to some previous top-conference papers, improvements of 0.02-0.03 in validation loss are typically considered meaningful...

A3. Thank you for the feedback. We have retrained AoT-MHSA with additional hyperparameter tuning and present the updated experimental results below:

We acknowledge that our method is still underperforming. However, compared to Table 2, the performance of the proposed transformer has shown improvement than before, achieving a smaller validation loss. Here, its full potential has not been fully explored due to limited computing resources. With further hyperparameter tuning, we believe that performance can be further improved. Additionally, the primary focus of this paper is on the theoretical contributions rather than empirical performance.

Weakness 4. Inconsistency in Eq. (5): There is a dimensional inconsistency in Equation 5. Since $\mathrm{head}_k$ is d_h by n, $\mathrm{head}_k^T$ should be n by d_h, but the matrix dimensions do not align properly in the equation.

A4. Thanks for catching this error. To fix it, we revise Eq. (5) into

$$\mathrm{MHSA}(Z) = W_O \begin{bmatrix} \mathrm{head}_1 \\\\ \dots \\\\ \mathrm{head}_K \end{bmatrix}.$$

Here, $Z \in \mathbb{R}^{d\times n}$, $W^O \in \mathbb{R}^{d\times Kd_h}$, and $\begin{bmatrix} \mathrm{head}_1 \\\\ \dots \\\\ \mathrm{head}_K \end{bmatrix} \in \mathbb{R}^{Kd_h \times n}$ due to each head $\mathrm{head}_k \in \mathbb{R}^{d_h\times n}$.

Weakness 5. Difficult mathematical theory: Upon brief examination of Lemma B.1, potential errors were identified. For example, when $t=0$, $\delta=1$, and $d=64$, the lemma yields an impossible probability value. This suggests that additional constraints on t or other parameters may be necessary for the lemma to hold true.

A5. Thanks for the comment. We will provide additional explanations and clarification to facilitate understanding in the revised manuscript. The result in Lemma B.1 is a standard concentration inequality for Gaussian random vectors; see Theorem 5.6 & Example 5.7 in Ref [2].

It is important to note that the choice of $t \ge 0$ is crucial for the validity of the inequality. Setting t=0 would result in a trivial statement, as the inequality is designed to quantify significant deviations from the mean. Therefore, only when the deviation is sufficiently large does the inequality yield a meaningful result. For example, when $\delta=1$ and we set $t=2\sqrt{\log d}$, then it holds with probability at least $1-2/d^2$ that $||x-\sqrt{d}|| \le \sqrt{2\log d}+2$.

[2] Boucheron et al. (2013). Concentration Inequalities: A Nonasymptotic Theory of Independence. Oxford University Press.