Attention-Only Transformers via Unrolled Subspace Denoising

Authors: Peng Wang, Yifu Lu, Yaodong Yu, Druv Pai, Qing Qu, Yi Ma
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: transformers, subspace denoising, algorithm unrolling, mixture of low-rank Gaussians

Abstract:
Despite the great success of transformers in practice, their architectures have been empirically designed, hence lack of mathematical justification and interpretability. Moreover, many empirical studies have indicated that some components of the transformer architectures may be redundant and can be removed or replaced without compromising overall performance. Hence to derive a compact and interpretable transformer architecture, we contend that the goal of representation learning is to compress a set of noisy initial token representations towards a mixture of low-dimensional subspaces. Based on the existing literature, the associated denoising operation naturally takes the form of a multi-subspace self-attention (MSSA). By unrolling such iterative denoising operations as a deep network, we arrive at a highly compact architecture that consists of only an MSSA operator with skip connections at each layer, without MLP. We rigorously prove that each layer of the proposed transformer performs so highly efficient denoising that it improves the signal-to-noise ratio of token representations {\em at a linear rate} with respect to the number of layers. Despite its simplicity, extensive experiments on language and vision tasks demonstrate that such a minimalistic attention-only transformer can achieve performance close to conventional transformers, such as GPT-2 and CRATE.

Decision: reject
Reviewer scores: 6, 5, 3, 6

Reviews and discussion:
<thread>
creation date:
1730566241196

thread_type:
official_review

<thread_object>
creator: reviewer_t5de

<content>
summary:
This paper proposes a multi-subspace self-attention (MSSA) to replace Transformers. By unrolling such iterative denoising operations as a deep network, the proposed technique arrives at a highly compact architecture that consists of only an MSSA operator with skip connections at each layer, without MLP. They also rigorously prove that each layer of the proposed transformer performs so highly efficient denoising that it improves the signal-to-noise ratio of token representations at a linear rate with respect to the number of layers.

strengths:
1. The motivation is based on the existing linear representation hypothesis and the superposition hypothesis, which states that the token representations lie on a union of (possibly many) low-dimensional subspaces. 

2. The empirical performance is good. On zero-shot language tasks and fine-tuning ViT, the method achieves comparable results.

weaknesses:
1. 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].

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

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

>[1] Sparse Autoencoders Find Highly Interpretable Features in Language Models. ICLR 2024.

questions:
Please see the weaknesses.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730665817662

thread_type:
official_review

<thread_object>
creator: reviewer_q2fd

<content>
summary:
This paper introduces a class of interpretable attention-only transformers that compresses noisy initial token representations into a mixture of low-dimensional subspaces, primarily building on the multi-head subspace self-attention introduced in [1]. The authors provide theoretical guarantees on the model's denoising performance, demonstrating that the signal-to-noise ratio improves with each layer. Experimental results in both language and vision tasks illustrate the advantages of the proposed transformer.

strengths:
1. The paper provides a good review of the multi-head subspace self-attention from [1].

2. The paper is well-written with illustrative figures.

weaknesses:
1. 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.
 
2. 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.

3. The experimental results are unconvincing. 

* First, in Table 1, test perplexity should be used to evaluate the models rather than just validation loss and zero-shot accuracy. Also, in Table 1, the proposed models need many more parameters compared to the baseline GPT-2 while still worse than the baseline GPT-2 in most cases. 

* Second, in Table 2, the performance of the proposed model, AoT, on ImageNet-1K is much worse than the baseline CRATE-\alpha model. For example, AoT-Huge has more parameters (86M) than the baseline CRATE-\alpha-B/16 but is 2% worse than the baseline in top-1 accuracy. 

* Third, 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.

4. I highly doubt the efficiency of the proposed method. From Table 1 and Table 2, it is clear that the proposed model requires more parameters than the baseline transformer models. Also, the multi-head subspace self-attention in Eqn. 3 requires more computations compared to the baseline self-attention.

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

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

**References**

[1] Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Hao Bai, Yuex- iang Zhai, Benjamin D Haeffele, and Yi Ma. White-box transformers via sparse rate reduction: Compression is all there is? arXiv preprint arXiv:2311.13110, 2023a.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730679351998

thread_type:
official_review

<thread_object>
creator: reviewer_9zdx

<content>
summary:
Paper proposes an attention only transformer with guarantees that each subsequent layer improves the SNR of the signal. they provide theoretical insights in to how such a system is beneficial and also run some baseline experiments on how this holds up against traditional transformer networks.

strengths:
The approach is very good and sound. builds on strong pre-existing work and adds an important piece to it. The flow of the paper is good and adds theoretical justifications on attention only transformers.

weaknesses:
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.

questions:
How does the reformulation in the paper for AoT converge?
How can we observe better SNR in these models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706489605

thread_type:
official_review

<thread_object>
creator: reviewer_aka8

<content>
summary:
This paper presents a novel approach to transformer architecture by proposing an attention-only model that utilizes multi-subspace self-attention (MSSA) for efficient denoising of token representations. The authors demonstrate that their unrolled optimization framework simplifies transformer architectures significantly, achieving performance comparable to traditional models while enhancing interpretability. Rigorous theoretical analysis supports the denoising capabilities of their architecture, showing improvements in signal-to-noise ratio across layers. Extensive empirical evaluations across language and vision tasks validate the practical effectiveness of their proposed design.

strengths:
The use of unrolled optimization to design a simplified transformer architecture presents a fresh perspective on transformer design, focusing on interpretability and efficiency.

The paper rigorously proves the denoising effectiveness of the proposed architecture, providing a solid mathematical foundation for its claims.

The authors conduct comprehensive experiments across diverse tasks, supporting the practical viability of their approach and demonstrating competitive performance.

weaknesses:
One notable weakness of this study is the lack of validation on state-of-the-art transformer architectures, such as Llama, and benchmark tasks, including MMLU. It would strengthen the paper significantly to include evaluations on these platforms to better assess the proposed method's performance in comparison to leading models in the field.

The paper could benefit from more detailed comparisons to the latest transformer variants such as mamba and linear transformers, which may also claim reduced complexity or improved performance.

questions:
See the weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732326579744

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global response

comment:
We thank all the reviewers for their insightful comments and constructive feedback. To summarize, our work introduces an attention-only transformer architecture that is both interpretable and effective. We rigorously demonstrate that each layer of our proposed transformer denoises token representations at a linear rate, given initial token representations sampled from a mixture of low-rank Gaussians. Our contributions aim to bridge theoretical insights with empirical validation. Below, we address two main points raised in the reviews:

- **Comparison to the state-of-the-art transformers:** 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.

- **Difference of the interpretability between our work and the existing works:**  In our work, interpretability 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732525273230

thread_type:
official_comment

<thread_object>
creator: area_chair_sagl

<content>
title:
Author Reviewer Discussion

comment:
Dear Reviewers,

Thank you for your efforts in reviewing this paper. We highly encourage you to participate in interactive discussions with the authors before November 26, fostering a more dynamic exchange of ideas rather than a one-sided rebuttal.

Please feel free to share your thoughts and engage with the authors at your earliest convenience.

Thank you for your service for ICLR 2025.

Best regards,

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission proposes "attention-only" transformers, which have so-called multi-subspace self-attention (MSSA) layers. The resulting transformer architecture is very compact and does not require the usual MLP layers. The paper has receives critical and borderline reviews. The reviewers acknowledge the benefit of the simplified architecture and the solid mathematical background. One important criticism is the lack of comparison to state of the art transformer architectures (e.g. aKa8, Q2fD), the novelty of the work (e.g. Q2fD), interpretability (T5de) and the experimental validation in terms of practical analysis (9ZDx). An additional evaluation on llama was provided during the rebuttal but the concerns regarding novelty and insufficient experimental validation still persist after the rebuttal, leaving the paper with a borderline reject score.