Attention-Only Transformers via Unrolled Subspace Denoising

Q1. A lot of linear hypothesis works can be used to do some interpretability tests of the LMMs. I am wondering if this new architecture might naturally fit the interpretability tests. I would suggest that the authors try GPT's auto-interpretability and might do some circuit analysis as done in [1].

A1. Thanks for pointing this out. We agree that the proposed architecture naturally aligns with interpretability tests, given its minimalistic and theoretically grounded design. As recommended, we will explore GPT’s auto-interpretability framework and conduct circuit analysis similar to the approach in [1]. Due to current limitations in computational resources, we are working to implement these analyses and aim to update the results before the rebuttal deadline. We plan to incorporate these findings into the revised manuscript to provide additional insights into the interpretability of the proposed architecture.

Q2. Can the author do some experiments about training ViTs from scratch? That would be more convincing to show the effectiveness.

A2. Yes, we are running the experiment and plan to add them to our revised manuscript.

Q3. Can authors visualize some of the attention masks in ViTs for a few input images? That might help if each U can correspond to different regions.

A3. Yes. We have trained the proposed AoT architecture on Imagenet21k and visualized the self-attention heatmaps between the classification token and other image patches. Then we refer the reviewer to Figure 8 on Page 20 in the updated manuscript for these experimental results. It is observed that different heads capture different parts of objects, and display different semantic meanings.

Q4. My main interest is to see if this architecture has better interpretability. The authors emphasize the interpretability in the intro but do not really validate it in the experiments. I am curious to see the results.

A4. We should clarify that interpretability in our work refers to understanding how the transformer architecture transforms token representations across layers. This notion differs from the traditional definition of interpretability in large language models (LLMs), which often emphasizes extracting concise, human-readable explanations for specific token representations. Our approach to interpretability is rooted in the mathematical and structural understanding of the transformer’s behavior, providing a clearer view of how information is processed and denoised at each layer.

We thank the reviewer for the careful review. In the following, we address the reviewer’s comments individually.

Q1. The proposed attention-only transformer is just a simple extension by unrolling Eqn. 3 in the paper, which is the multi-head subspace self-attention in [1]. Most heavy-lifting work has been done in [1], and unrolling Eqn. 3 is just an incremental contribution.

A1. We respectfully disagree with the comment that our contribution is incremental. While [1] introduces the MSSA mechanism, it mainly focuses on the derivation and conceptual introduction of the operation, without conducting a fine-grained theoretical analysis of its properties and behavior in deep architectures. Our work makes several significant advancements beyond [1]:

Rigorous Theoretical Analysis: In our work, we provide a theoretical framework that rigorously proves the efficiency of the MSSA operator when unrolled into a deep transformer architecture. Specifically, we demonstrate that each layer achieves a linear improvement in the signal-to-noise ratio of token representations.

Architectural Innovation: By unrolling Eqn. 3, we derive a highly compact attention-only transformer architecture that entirely omits the MLP layers, reducing redundancy while maintaining strong performance. This minimalistic design differs significantly from the architectures in [1], which do not explore such an approach.

Q2. The motivation behind the proposed unrolling transformers is to denoise the noisy initial token representations towards a mixture of low-dimensional subspaces. However, it is not clear why denoising token representations should help improve the performance of the model. Also, there is no empirical evidence showing that the tokens are drawn from a mixture of noisy low-rank Gaussian distributions as assumed in Definition 1.

A2. We would like to clarify the motivation and assumptions behind our approach:

Why denoising token representations should improve the model performance: The goal of representation learning is to find the underlying structure of data while effectively removing noise. By progressively denoising token representations toward a structured and compact form, the transformer architecture extracts cleaner and more informative features, leading to improved performance on downstream tasks. A mixture of low-dimensional subspaces serves as one specific example of such a structured and compact representation. However, it is worth noting that transformers can also adaptively transform token representations into other meaningful structures, depending on the nature of the data and the task at hand.

Empirical evidence showing that the tokens are drawn from a mixture of noisy low-rank Gaussian distributions: Recent empirical studies on language tasks have raised the “linear representation hypothesis”, which posits that token representations can be linearly encoded as one-dimensional feature vectors in the activation space of LLMs [1], and the “superposition hypothesis”, which further hypothesizes that token representations are a sparse linear combination of these feature vectors [2]. These empirical studies demonstrate that the token representations lie on a union of (possibly many) low-dimensional subspaces.

[1] Park, K., Choe, Y. J., & Veitch, V. (2023). The linear representation hypothesis and the geometry of large language models. arXiv preprint arXiv:2311.03658.

[2] Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., ... & Olah, C. (2022). Toy models of superposition. arXiv preprint arXiv:2209.10652.

Q3. Comparison of the model performance between the proposed transformer architecture and the state-of-the-art ones

A3. The main goal of our empirical studies is not to surpass the performance of state-of-the-art transformers. Instead, our objective is to demonstrate that the proposed interpretable and attention-only transformer can achieve performance close to that of standard transformers while maintaining the simplicity and interpretability of its design. This focus allows us to highlight the potential of our architecture as a theoretically grounded alternative to existing approaches, particularly in scenarios where interpretability is a priority. Moreover, we will use a grid search approach to optimize our hyperparameters and some other approaches to improve the model performance in our revised manuscript.

Q4. Even though the authors claim that their proposed transformer is interpretable in the abstract and introduction, there are no experiments to show the interpretability. Similarly, one of the contributions mentioned in the introduction is to understand the role of self-attention and MLP layers, this is not really empirically discussed in the paper.

A4. Interpretability of our proposed transformer: We should clarify that interpretability in our work refers to understanding how the transformer architecture transforms token representations across layers. This notion differs from the traditional definition of interpretability in large language models (LLMs), which often emphasizes extracting concise, human-readable explanations for specific token representations. Our approach to interpretability is rooted in the mathematical and structural understanding of the transformer’s behavior, providing a clearer view of how information is processed and denoised at each layer.

Discussion of the role of self-attention and MLP layers: We have proved the role of self-attention in Theorem 1 and discussed the role of MLP layers, such as Lines 481-482.

Q5. Why shared-QKV subspace self-attention mechanisms, as mentioned in line 299?

A5. Because our studied MSSA in Eq (5) can be viewed a special instance of standard self-attention formula by setting $Q=K=V=U_k$.

Q6. Please provide the efficiency analysis, i.e., real-time running and memory usage, of the proposed transformer vs. the baseline transformer.

A6. Thanks for pointing this out. We have conducted the experiments and reported the Gflops, GPU memory, and running time per pass in Table 3 in Appendix B.0.3.

Q1. The paper doesn't offer any practical analysis on the theory. The results on language and vision are good, but using the paper's own arguments, the training dynamic of such systems and modern frameworks can reach convergence outside of this papers contributions. The key piece missing is an approach to understanding how the proofs of the theory play out in actual training.

A1. Thank you for your insightful comment. To clarify, our main theoretical result does not pertain directly to the training process of transformers. Instead, it demonstrates that under the proposed network architecture, we show that each layer of the transformer denoises token representations at a linear rate.

If we understand correctly, your concern is whether the proven linear denoising performance of the studied transformers can be observed when they are trained on real-world datasets (please let us know if our understanding is wrong). This theoretical result holds when the token representations approximately satisfy a mixture of low-rank Gaussians. Notably, any distribution can be well approximated by multiple Gaussian distributions [1]. In this sense, our result provides a meaningful foundation for understanding the behavior of the architecture in practical settings.

[1] Borkar, V. S., Dwivedi, R., & Sahasrabudhe, N. (2016). Gaussian approximations in high dimensional estimation. Systems & Control Letters, 92, 42-45.

Q2. How does the reformulation in the paper for AoT converge? How can we observe better SNR in these models?

A2. We apologize that we are not entirely clear on the first question. It would be great if you could provide more details about the reformulation for AoT converge.

According to Eq (8) in Theorem 1, when the step size $\eta$ or the thresholding parameter $\tau$ is large, the SNR increases at a faster rate as the number of layers grows. Additionally, as the number of layers increases, the SNR at the final layer is further improved.

Q. Validation on state-of-the-art transformer architectures, such as Llama, and comparisons to the latest transformer variants such as mamba and linear transformers

A. Thank you for the suggestion. We would like to clarify that the main goal of our empirical studies is not to surpass the performance of state-of-the-art transformers. Instead, our objective is to develop a minimalistic transformer-like deep architecture consisting of fully interpretable and provably effective layers that can achieve performance close to that of standard transformers.

That said, we are currently conducting experiments to validate the performance of our transformer on state-of-the-art models such as Llama and also compare our architecture with leading transformer variants such as Mamba across various benchmark tasks, as suggested. Due to limited computing resources, we aim to update these results before the rebuttal deadline. Furthermore, we will include these results in our paper to demonstrate the performance of our proposed architecture.

Updated: We have compared the in-context learning performance of our proposed architecture with Llama architectures in Figure 7 in the revised manuscript. It is observed that the proposed architecture achieves comparable performance to Llama.

Dear Reviewers,

Thank you again for your thoughtful review and the time you have spent evaluating our work.

As the author-reviewer discussion period approaches its conclusion, we noticed that we have not yet received a follow-up response from you. We are eager to know whether our responses have addressed your concerns. Your feedback is highly valuable to us, and if you have any remaining questions or require further clarification, we are more than happy to provide additional information.

If our responses have resolved your concerns, we would sincerely appreciate it if you could consider updating your score.

Thank you for your time and consideration.

Authors.