Attention-Only Transformers via Unrolled Subspace Denoising
Authors: Peng Wang
Yifu Lu
Yaodong Yu
Druv Pai
Qing Qu
Yi Ma
Venue: ICML 2025 poster (2025)
Keywords: transformer
attention
subspace denoising
token representation
Primary area: optimization->nonconvex
TL;DR: We propose a minimal transformer with only self-attention and skip connections.

Abstract: Despite the popularity of transformers in practice, their architectures are empirically designed and neither mathematically justified nor interpretable. Moreover, as indicated by many empirical studies, some components of transformer architectures may be redundant. To derive a fully interpretable transformer architecture with only necessary components, 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. To compress these noisy token representations, an associated denoising operation naturally takes the form of a multi-head (subspace) self-attention. By unrolling such iterative denoising operations into a deep network, we arrive at a highly compact architecture that consists of \textit{only} self-attention operators with skip connections at each layer. Moreover, we show that each layer performs highly efficient denoising: it improves the signal-to-noise ratio of token representations \textit{at a linear rate} with respect to the number of layers. Despite its simplicity, extensive experiments on vision and language tasks demonstrate that such a transformer achieves performance close to that of standard transformer architectures such as GPT-2 and CRATE.

Attention-Only Transformers via Unrolled Subspace Denoising
Peng Wang's Yifu Lu' Yaodong Yu² Druv Pai? Qing Qu Ma
Abstract
and reinforcement learning (Chen al., 2021). Transformer architectures are generally constructed by stacking mul- tiple identical layers designed to process and learn from data. Each layer is composed of several interacting com- ponents arranged in a specific sequence, including multi- head self-attention operators, layer normalization, multi- layer perceptron (MLP) networks, and skip connections. In practice, transformers, such as BERT (Devlin, 2018) and GPT-4 (Achiam et al., 2023), are highly deep, often with dozens even hundreds of layers, and significantly over- parameterized, containing millions or even billions of param- eters. This considerable depth and over parameterrzaaio endow transformers with impressive learning capabilities, allowing them to capture complex patterns and relationships.
Despite the popularity of transformers in practice, their architectures are empirically designed and neither mathematically justified nor interpretable. Moreover, as indicated by many empirical studies, some components of transformer architectures may be redundant. To derive a fully interpretable transformer architecture with only necessary com- ponents, we contend that the goal of representa- tion learning is to compress set of noisy initial token representations towards a mixture of low- dimensional subspaces. To compress these noisy token representations, an associated denoising op- eration naturally takes the form of multi-head (subspace) self-attention. By unrolling such iter- ative denoising operations into deep network, we arrive at highly compact architecture that consists of only self-attention operators with skip connections at each layer. Moreover, we show that each layer performs highly efficient denois- ing: it improves the signal-to-noiss ratio of token representations at linear rate with respect the number of layers. Despite its simplicity, ex- tensive experiments on vision and language tasks demonstrate that such a transformer achieves per- formance close to that of standard transformer architectures such as GPT-2 and CRATE.
Despite the remarkable success of transformers, their deep and over peeepparamttrieed architecture renders them "black boxes", hindering an understanding of their inner mecha- nism. To address this challenge, a common approach in- volves systematically removing or modifying certain com- ponents in transformers to simplify the architecture; see, e.g., Dong et al. (2021); Alcalde et al. (2024); Noci et al. (2024); Geshkovski al. (2023a); Geva al. (2020); Guo et al. (2024). For example, Alcalde et al. (2024) studied pure-attention hard-max transformers with skip connections and showed that the output converges to a clustered equi- librium as the number of layers goes to infinity. Noci et al. (2024) analyzed a modified softmax- based attention model with skip connections, demonstrating that the limiting distri- bution can be described by a stochastic differential equation. These studies indicate that the most basic components of transformers are self-attention layers and skip connections. Although existing studies have provided valuable insights into different components of transformers, few of them elu- cidate the underlying mechanisms by which transformers process and transform input into output across layers.
1. Introduction
Over the past decade, transformers (Vaswani et al., 2017) have achieved remarkable empirical success across various modern machine learning applications, including large lan- guage models (L.Ms) (Brown et al., 2020; Devlin, 2018), vi- sion generative models (Bao al., 2023; Chen et al., 2020),
Existing empirical studies suggest that some components of transformers may not be essential and can be removed or modified without compromising performance. For ex- ample, He & Hofmann (2024) empirically demonstrated that transformer architecture can be simplified by removing components such as skip connections, value matrix, and normalization layers without degrading performance. Addi- tionally. Sukhbaatar al. (2019) investigated the effects of removing MLP blocks from transformers and augmenting
Department of Electrical Engineering and Computer Seience, University of Michigan, Ann Arbor, USA 'Department of Electri- cal Engineering and Computer Science, University of California, Berkeley, USA IInstitute of Data Science & School of Computing and Data Science, University of Hong Kong, Hong Kong, China Correspondence to: Peng Wang <peng8wang @ gomil.com>.
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Attention-Only Transformers
the self-attention layers to play similar role MLP blocks, showing that performance can be preserved. Similarly, Pires et al. (2023) examined the potential for reducing the fre- quency of MLP layers in transformers. Other works also studied other simplifications of transformers, such as linear attentions (Katharopoulos et al., 2020) and shared-QK at- tentions (Kitaev et al., 2020). Based on these discussions, this work focuses on addressing the following question:
sparse linear combination of these feature vectors (Elhage et al., 2022; Yun et al., 2021). Building on these hypothe- ses, various approaches have been proposed to understand and utilize token representations. For example, Templeton (2024) employed sparse autoencoders to decompose the token representations of Claude Sonnet into more inter- pretable components. Luo et al. (2024) leveraged sparse dictionary learning to explore token representations, de- composing them into interpretable components based a concept dictionary. Recently, Engels et al. (2025) conjec- tured that token representations in LLMs are the sum of many sparse multi-dimensionall features. This conjecture is supported by their experiments on GPT-2 and Mistral 7B, where they used sparse autoencoders to identify multi- dimensional features. Notably, all of these empirical studies conclude that the token representations lie on union of (possibly many) low-dimensional subspaces.
Can we design minimalistic transformer architecture consisting of fully interpretable layers that achieves performance close to that of standard transformers?
1.1. Related Works
Existing studies of self- attention mechanisms. One main factor of the power of transformers is the self-attention layers, which enable the model to capture long-range dependen- cies and contextual relationships between tokens by weighing token relationships across the input sequence (Vaswani et al., 2017). To explore the mechanism behind self-attention, numerous studies have inves- tigated the performance of pure self-attention networks, often incorporating only one additional component to prevent rank collapse and maintain expressiveness; see, e.g., Dong et al. (2021); Geshkovski et al. (2023a;b); Wu et al. (2024). We refer the reader to Appendix A for more discussions.
Z (++1)
MSSA
1.2. Our Contributions
Based on the above discussions and motivated by referenced empirical findings, we propose a simple yet evocative model for the structure of token representations in trained trans- formers. Specifically, we model the underlying distribution of token representations as a mixture of low-rank Gaussians, each supported on a subspace and corrupted by noise (see Definition 2.1). Then, the goal of representation learning is to denoise a set of noisy initial token representations to- wards the corresponding subspaces. Our contributions summarized as follows:
LayerNorm
(1)
Figure I. Each layer of our proposed trans- former architecture.
Attention-only transformer via unrolled optimization. Under the mixture of low-rank Gaussian model, we inter- pret multi-head (subspace) self-attention as a denoising operator, which compresses noisy token representations into their corresponding supporting subspaces. By itera- tively unrolling the multi-head (subspace) self-attention operator, we construct a new transformer architecture with streamlined design, consisting of only aelf-attention lay- with skip connections (see Figure 1). This design is much simpler compared to standard transformers.
Network architecture design via unrolled optimization. Several lines of work have proposed that the success of modern deep networks largely stems from their ability to transform the raw data into compact and structured repre- sentations, which facilitates downstream tasks (Chan et al., 2022; Chen et al., 2023; Ma et al., 2022: Yu et al., 2024; Huh et al., 2024). A principled and interpretable approach to learning such representations is to construct an architecture that incrementally transforms tokens into these representa- tions via unrolling optimization steps as layers of a deep network (Chan et al., 2022; Monga et al., 2021; Wang et al., 2016; Yu al., 2024; Zhang& & Ghanem, 2018). Appendix A contains more discussion on this point.
Theoretical guarantees for the proposed transformer. To quantify the denoising performance of the proposed transformer, we introduce a -o-niieto-noise (SNR) metric (see Eq. (8)) for the token representations. We prove that each layer of the proposed transformer improves the SNR at a linear rate when the initial token representations are sampled from a noisy mixture of low-rank Gaussians (see Theorem 3.1). This indicates that the multi-head (subspace) self-attention operator highly effective in de- noising token representations towards their corresponding subspaces.
Linear representation and superposition hypotheses. Recent empirical studies of language tasks have raised the "lineur representation hypothesis", which posits that token representations can be linearly encoded as one-dimensional feature vectors in the activation space of LLMS (Jiang et al., 2024; Park et al., 2024b), and "superposition hypothesis", which further hypothesizes that token representations are
In practice, LayerNorm layers may be added to enhance per- formance.
2

---

Attention-Only Transformers
z
Z
(L-1)
Z = Z (b)
22
2
U2 o
p
ft
U1
U
U10
Figure 2. Layers of transformer gradually denoise token representations z0 towards their corresponding subspaces.
Understanding the roles of self-attention and MLP layers. Notably, the proposed transformer is valuable model for understanding the mechanism of attention, as it ablates the effect of MLP layers. Moreover, comparing the proposed transformer to standard transformers provides insights into the specific role and empirical benefits of the MLP layers in different tasks, such as for images and texts (see experiments in Section 4).
transformers typically a sequence of tokens denoted as X [......, ØN| I2 DxN Then, the goal of transform- ers is to learn map f that transforms these tokens into structured and compact token representations that facilitate downstream tasks, such as classification (Dosovitskiy et al., 2021) and generation (Saharia et al., 2022), by capturing the underlying patterns in the data.
2.1. Unrolled Optimization for Token Representations
Finally, we conduct extensive experiments on both vision and language tasks, including supervised image classifica- tion, causal language modeling, and in-context learning, to complement our theory and demonstrate the potential of our proposed transformer architecture. We emphasize that the goal of our experiments is not strive for state-of-the-de performance for these tasks. Instead, they are intended to validate our theory about the components of the transformer.
In this subsection, we introduce how to learn token repre- sentations using the approach of unrolling optimization al- gorithms (Chan et al., 2022; Gregor & LeCun, 2010; Monga et al., 2021; Sun al., 2019; Wang et al., 2016; Yu et al., 2024; Zhang & Ghanem, 2018). Specifically, this approach constructs each layer of a neural network according to a step of an iterative optimization algorithm. That is, the network's architecture is designed to implement a specific optimization algorithm, where each layer corresponds to a single iterative step. By unrolling the algorithm, we construct a "white-box transformer architecture as a multi-layer neural network that incrementally transforms input tokens into structured and compact representations. This iterative process can be described as follows:
Notation. Given an integer n, we denote by [n]] the set .n.. Given a vector a, let ||a]| denote the Euclidean norm of and diag(a) (a) denote the diagonal matrix with as its diagonal. Given a matrix A, let || denote the spectral norm of A. ||A||] F denote the Frobenius norm, and ag denote the (i, element. For sequences of positive numbers }an} and {b,} we write an S bn or bn Z a, if there exists an absolute constant C > 0 such that an S Cbn- Given a constant T > 0, we define [(x > T) = 1 > and [( > r) otherwise. We use O to denote the set of all n d matrices that have orthonormal columns.
f:X
(0) 8
S
2
fh (L) =: Z,
where f DxN + RadxNN is a pre-processing mapping (e.g., positional encoding, token embedding) that transforms input tokens X E P to initial token representations 2 E RdxN, l RRdxN + RdxN denotes an iterative step or layer, and z denotes the token representations the I-th layer for each L E [L]. Then, key question is how to design the operator f' at each layer learn meaningful token representations in principled manner.
2. Technical Approach and Justification
In this section, we introduce the basic setup of transformers for learning representations from real-world data. Real- world data, such as images, videos, and text, are often mod- eled as random samples drawn from a high-dimensional probability distribution with low-dimensional intrinsic struc- tures (Wright& & Ma, 2022). Instead of directly inputting raw data samples into transformers, a common prepro- cessing step is to convert each sample into a sequence of tokens, where each token represents a localized seg- ment of the data, such as an image patch, a snippet of text, or a frame in a video. Consequently, the input to
2.2. A Model for Token Representations
Before design such an operator f', we model the struc- ture of token representations in pretrained L.MMs. Notably, extensive works (Templeton, 2024; Luo et al., 2024; Engels et al., 2025) have empirically demonstrated that token rep- resentations in trained LLMs usually approximately lie in a union of low-dimensional subspaces. These subspaces
3

---

Attention-Only Transformers
encode distinct semantic meanings, capturing various lin- guistic contextual features that contribute to the model's overall understanding and performance. This motivates us to model the token representations as follows:
expressed as sparse linear combinations of these subspace bases, capturing the essential semantic components of the token while ignoring irrelevant dimensions.
2.3. Denoising Operator for Token Representations
.a......,, Cx be partition of the index set [N] and UR E Odxps denote the orthonormal basis of the k-th subspace for each & E [K]. We say that the token representations { S IR are sampled from a mixture of noisy low-rank Gaussian distributions if for each l E [K].
Definition 2.1. Let Ci....,
Now, we introduce a denoising operator to compress token representations into the corresponding subspace when the initial token representations ( is generated according to Definition 2.1. To simplify our development, we assume that the subspaces in Definition 2.1 are orthogonal to each other, i.c., U, = 0 for all & * j. Note this assumption is not restrictive, as in nigh-dimensional spaces, random low- dimensional subspaces are incoherent each other with high probability, i.e U, (Wright & Ma, 2022).
K 2, = U.a + -VE U Vi Ok.
(1)
signal
jzL
noise
Multi-head subspace self-attention. Without loss of gen- erality, we rearrange the initial token representations Z such that those from the same subspace are concatenated
where a; 1.1. NO0, Ipx) and eijj i.dd ~ 6(0,6%I,) Ip,) for
all i E C2 and & E [K]. {ar} and }eir} are respectively mutually independent, and {a.) is independent of {eaz)-
together, i.e., (0) = [Z (0)
[ with
Before we proceed, let us make some remarks on this model.
(0) 2 U
Σ U, Ak [K]. J
An idealized model for token representations. This model serves an idealized framework for approximat- ing token representations in real-world pretrained LLMs. It assumes that the token representations are sampled from a mixture of multiple low-rank Gaussian distributions with noise. Under this model, the goal of representation learning to compress a set of noisy initial token presen- tations into the corresponding subspace. We should point out that in real-world applications, where token represen- tations exhibit more complicated structures, the goal of representation learning is to find compact and structured representation via compressing token sets, as argued in Yu et al. (2024). In addition, this model has been widely used in other machine learning problems, such as subspace clustering (Wang et al., 2022; Elhamifar & Vidal, 2013) and diffusion models (Wang al., 2024).
Here, the columns of g denote the token representations from the k-th subspace for each k E [K], the columns of A& RP X. consists of {a, (act. and the columns of EE E RR, consists of {eij)eee for each k E [K] with N₂ = |Cpl for each k E [K]. Obviously, projecting token representations onto their corresponding subspace helps to separate the signal from the noise components using U, = 0 for all le + j. i.e.
URAA Ug
if E = k, if & k.
U.U 0
(2)
To denoise the token representations from k-th subspace, we compute the similarity of projected token representations via (U 2)² (U Z) and verify that the similarity between projected token representations from the k-th subspace is high, while the similarity between other pairs of projected token representations is low when 5 < 1. Then, we convert it to distribution of membership with function yo, such as worddthresholding or soft-max functions, and denoise the token representations towards to the corresponding subspace using this membership. Now, we formalize the considered operator as follows: for each I = 0, 1, 1....,44 L 1,
Connections to hypotheses on token representations. This model aligns well with two well-established hy- potheses about the structure of token representations in pretrained LLMS: the "linear representation hypothesis" (Jiang et al., 2024; Park et al., 2024b) and the "superpo- sition hypothesis" (Elhage et al., 2022; Yun et al., 2021; Arora et al., 2018). The linear representation hypothe- sis posits that token representations in LLMs lie low- dimensional linear subspaces that encode semantic fea- tures. Similarly, the superposition hypothesis suggests that these representations can be approximately expressed sparse linear combination of these feature vectors. In the context of our model, each basis U₂ of the subspaces can be interpreted as a set of semantic features, where each feature corresponds to a specific aspect of the token's meaning. Token representations are then approximately
z (++)) Z ( + MSSA(Z (Z)) )))
(3)
where 7j > 0 is the denoising step size, vp.)) : BaxN + RdxN is an operator applied column-wise, and
K
MSSA(Z)
Σ U.U Z (ZTU,
(4)
k=
'One may straightforwardly generalize our results to non- orthogonal subspaces, with slightly more sophisticated analysis.
4

---

Attention- Only Transformers
Z (1+1)
AttentionHead
U LN(2
(-))
MSSA
MSSA
LayerNorm
Attentionnlead/.. (
LN(
D
Z (I)
10
AllentionHoull...
Figure 3. The attention-only transformer (AoT) architecture. Each layer consists of the MSSA operator and skip connection. Additionally, LayerNorm can be incorporated to enhance performance. In practice, backpropagation is applied to train the model parameters.
Notably, the operator in (4), referred to as the multi-head subspace self-attention (MSSA), is first proposed by (Yu al., 2024; 2023) to approximately optimize the compres- sion term of the sparse rate reduction objective for construct- ing transformer-lie like architecture. It is worth noting that Yu et al. (2023); Pai et al. (2023) showed that the negative compression gradient of the objective points from the to- ken representation to the corresponding subspace. However, they did not study the quantitative denoising efficiency of the MSSA operator (4).
better way to denoise these complex token representations. sum up, while state-of-the-art models necessitate the use of more advanced mechanisms like MHSA effectively denoise and optimize token representations, the model in Definition 2.1 offers useful framework for understanding token representations in an idealized yet evocative setting.
3. Main Results
In this section, we formally present an attention-only trans- former architecture using unrolled optimization and provide theoretical guarantee its denoising performance.
Connections to multi-head self-attention. Notably, the denoising operator (4) is essentially a special instance of multi-head self-attention (MHSA) implemented with a skip connection in transformers. Specially, the multi-head self- attention is of the following form:
3.1. Attention-Only Transformer
Armed with the setup in Section 2, we formally introduce the proposed attention-only transformer architecture. Specif- ically, by unrolling the iterative optimization steps (3) as layers of deep network, we construct a transformer archi- tecture in Figure 3. Each layer of the proposed architecture consists only of the MSSA operator and a skip connection. To enhance the model's performance, we may addition- ally incorporate LayerNorm before the MSSA operator to improve performance in practice. The complete architec- ture is built by stacking such layers, along with essential task-specife pre-processing and post-processing steps, such as positional encoding, token embedding, and final task- specific head to adapt to different applications. Notably, we apply the same procedure to (5), we obtain an attention- only transformer that only consists of the MHSA operator.
head
MHSA(Z) wo
(5)
head
where W W.K. V. are learnable weight matrices for queries, keys, and values for head k, wo another learn- able weight matrix, and
headę (T Z 1 (TK)TZ)
Comparing the MHSA operator with the MSSA opera- tor in (4), by setting w w.K = WY = Uk and WO = -U,,.... Un] in (5), we can obtain the MSSA op- erator in (4). In this special case, MHSA can be interpreted as a denoising operation onto different subspaces. How- ever, token representations in state-of-the-an large models are inherently more complex than the simplified structures assumed in Definition 2.1. In practice, these token repre- sentations are subject to variety of factors such as noise, context dependence, and intricate dependencies that make their structure more dynamic and multifaceted. In this con- text, using the more flexible MHSA mechanism may provide
Comparison standard transformers. Generally speak- ing, the standard decoder-only transformer architecture is composed of the following key components (Brown et al., 2020; Radford et al., 2019): (1) positional encoding, (2) multi-hee QKV self-attention mechanisms, (3) feed- forward MLP networks, (4) layer normalization, and (5) residual connections. In contrast, our proposed transformer
5

---

Attention- Only Transformers
Noiie lovel O 0.2
Nosc lovel 5 0.5
SNR
SNR
2.6
Fitted line
Fitted line
2.4
SNR
1.8
1.-
1.6
1.5
num. of layers
num. of layers
Figure 4. Denoising performance of the attention-only transformer. Here. sample initial token representations from a mixture of low-rank Gaussians Definition 2.1. Then, we apply (4) to update token representations and report SNR each layer. Left: noise level 8 = 0.2. Right: noise level 8 = 0.5.
architecture adopts a streamlined design by incorporating several key simplications. Specifically, it employs shared- QKV subspace self-attention mechanisms, excludes MLP layers, and reduces the frequency of LayerNorm.
3.2. Denoising via Attention-Only Transformer
In this subsection, we study the denoising performance of the proposed transformer when the initial token represen- tations are sampled from mixture of low-rank Gaussians as introduced in Definition 2.1. To quantify the denoising performance, we define the signal toall-to--oo to-noise ratio (SNR) for each cluster of the token representations at the I-th layer as
Differences from previous works attention-only trans- formers. In the literature, some theoretical works have studied attention-only transformers. For example, Dong et al. (2021); ; Wu et al. (2024) showed that pure-attention transformers with skip connections or LayerNorm can pre- vent rank collapse. Additionally, Alcalde et al. (2024) stud- ied the clustering behavior of attention-only hardmax trans- formers. While these studies contribute significantly to our understanding of the role of self-attention in transformers, they lack empirical validation and practical implications. In contrast to these works, we not only show that each layer of the proposed attention-only transformer can denoise token representations but also conduct experiments on real-world language and vision tasks demonstrate the potential.
SNR ((4) = URU 2 - - z
Wk E [K].
(6)
To simplify our analysis, we assume that p Pt PK. N₁
NK N/K, and
U
Ukl OdxKp.
With the above setup, we now prove the following theorem. Theorem 3.1. Let ( be generated according to Defini- tion 2. and Z be generated according to (3) for each I E [L]. Here, sp(x) h (x))), o E R is the soft-max function, and h RR. R is an element-wise thresholding function with h(x) rI T} for each E [N]. Suppose that Z log N, 8 S vlog N/ /p, and
The role of backward propagation. Notably, our ap- proach constructs transformer architecture in the for- ward pass by interpreting each layer as a denoising op- erator, conditioned on the assumption that the subspace K bases { are known. However, in practice, these subspace bases are unknown and need to be learned grad- ually via backpropagation. Hence, the forward denoising operator (4) at the 1-th layer becomes as follows: For each 1= 0.1......
( N exp(-9pp330 0p/32)
1
For sufficiently large N, it holds with probability at least
1 - KLN (t(1) that for each i E [L [L-1]
K Σ
SNR (((11)) +1+
7 ( Vk [K]. (8)
(++1)_ 7
k=1
The proof is deferred to Appendix B. Here we comment on the significance of this theorem:
th U u( ( z²).
Linear denoising performance of the attention-only transformer. When the initial token representations are sampled from a mixture of low-rank Gaussian distribu- tions with noise level O /P), we show that
Now, the parameters (T()) depend on the layer index and may be different across layers. These matrices are learned through end-to-end training via backpropagation.
6

---

Attention-Only Transformers
each layer of the proposed transformer denoises token representations at linear rate. This indicates the MSSA operator's s efficiency in reducing noise across layers. No- tably, our theoretical results are well-supported by experi- mental observations in Figure 4, which further validate the practical denoising capability of the proposed transformer.
Table 2. Top-I accuracy on ImageNet: Evaluation of AoT-MHSA- V and comparison to VIT.
Models
Accuracy
# of Parameters
AOT-MIISA-V
69.5%
1SM 22M
ViT
72.4 %
Difficulties in analyzing the dynamics of the update (3). Note that the update (3) is highly nonlinear and compli- cated. These characteristics lead to intricate interactions among consecutive updates that complicate the analysis of the learning dynamics. Compared to the existing works (Ahn al., 2023; Zhang al., 2024; Schlag et al., 2021) that mainly focus on linear self-attention with (p()) being the identify function, our analysis provides more pertinent results for understanding the denoising performance and learning dynamics of attention mechanisms, capturing the nonlinear interactions and transformations across the layers of modern transformer architectures.
(Chen et al., 2024) to pre-train the AOT-MSSA-V trans- former on ImageNet-21K for 90 epochs and to fine-tune it on ImageNet-IK (Deng et al., 2009) for 50 epochs by minimizing the cross-entropy (CE) loss. We use different hyperparameters for pre-training and fine-tuning. During pre- training,,
learning rate of 2 x 10 weight decay of 0.7, label smoothing with parameter of 0.2, and batch size of 4096. For fine-tuning, the corresponding val- ues are 5 x 10-2, 0.3, 0.1, and 2048, respectively. Standard data augmentation techniques, including random cropping, random horizontal flipping, and random augmentation, are used in our implementation, following the same setup as in Yu et al. (2023).
4. Experimental Verification
In this section, we evaluate our proposed allention-only transformer (AoT) architecture using the MSSA (denoted by AoT-MSSA) and MHSA (denoted by AoT-MHSA) op- erators on both vision and language tasks. Since the model configurations on vision and language tasks are different, we use AOT-MSSA-V and AOT-MHSA-V to denote the models applied to vision tasks, and AoT-MSSA-L and AoT- MHSA-I for those applied to language tasks. Due to lim- ited computing resources, the goal of our experiments is not to outperform state-of of-the transformers but to verify that AoT can achieve comparable performance on both lan- guage and vision tasks. In our implementations, we set the operator (p-) in Eq. (4) to be the softmax function.
Comparison between MHSA and VIT. We next train AOT-MHSA-V from scratch on ImageNet-IK for 150 epochs and compare its performance with ViT (Dosovitskiy et al., 2021). ) The training setup follows the same configu- ration as above: we use the Lion optimizer with a learning rate of 5 x 10-4, weight decay of 0.1, label smoothing with smoothing parameter of 0.1, batch size of 2048, and identical data augmentation strategies.
Based on the above experimental setup, we report the top-1 accuracy of AoT-MSSA-V V and CRATE in Table 1, and that of AoT-MHSA-V V and VIT in Table 2. Due to the absence of MLP layers in AoT, AoT-based models achieve slightly worse performance comparable to CRATE and ViT while using only nearly half the number of parameters. This result shows the effectiveness of the attention-only architecture. We provide visualization for the self-attention heatmaps of AoT-MSSA-V trained on ImageNet- IK in Figure 7 in Appendix C.3. We observe that each head captures similar semantic meanings across different images, demonstrating the interpretability of our proposed architecture in practice.
Table 1. Top-1 accuracy on ImageNet: Evaluation of AoT-MSSA- and comparison CRATE.
Models
Accuracy
# of Parameters
AoT-MSSA-V
71.7%
22M 39M
CRATE
79.5%
4.1. Vision Transformers for Image Classification
4.2. Decoder-Only Transformers for Language Tasks
In this subsection, we evaluate the performance of AoT as a backbone architecture for supervised image classification on ImageNet and compare it against several state-of the- art models. To construct the AoT-based model, we adopt the same preprocessing pipeline and classification head defined in (Yu et al., 2024, Section 4.1.1).
To study the performance of our architecture on language tasks, we consider the widely used Generative Pre- Tre-Trannin (GPT) task (Radford et al., 2019). In the context of causal language modeling, the goal is to predit the next token in sequence based on the preceding context. To adapt to this task, we modify the AoT architecture by changing the MSSA (resp., MHSA) operator to a causally masked MSSA (resp., MHSA) operator. We follow the same pre- processing and post-processing steps in (Yu et al., 2024,
Comparison between MSSA and CRATE. We consider the A model and compare it against the CRATE model in Yu et al. (2024). We employ Lion optimizer
7

---

Attention-Only Transformers
Convergence of Training Loss
Convergence of Validation Loss
5.0
5.0
NOT MSSA L. Basc:
AoT MESA Basc:
Act MSSALL Modium: 182M
AOTEMSA4 Modium: 182M
4.5
MEEAL Buse: 122M GPT2
4.5
AOEMESALL Base: 122M
Train Loss
Val Loss
4.0
3.5
3.5
3.0
3.0
2
5
3
6
Iteration
le5
Iteration
1e5
Figure 5. Evaluating models on language tasks. We plot the training loss (left) and validation loss (right) of the AoT and GPT-2 models pretrained on Open WebText.
Table Zero-shot results on several language benchmark datasets and tasks: Evaluation of different sizes of AoT with the MSSA and
MHSA operators and comparison to the GPT2 model.
Models
LAMBADA (val loss) +
PTB (val loss) |
WikiText (val loss) +
LAMBADA
CBT CN
CBT NE (acc)
# of parameters
(acc) f
(acc) 1
AoT-MSSA Base (102M)
4.70
6.03
4.65
0.25
0.80
0.74
AoT-MSSA L Medium (182M)
4,47
5.08
4.22
0.29
0.84
0.77
AoT MHSA- Base (122M)
4.42
5.52
4.19
0.38
0.86
0.82
GPT-2 Base (124M)
4.32
5.75
4.13
0.40
0.87
0.84
Section 4. 4.1. Our implementation of the GPT-2 type trans- formers and training pipeline is based on the framework outlined in (Karpathy, 2022).
BADA (Paperno et al., 2016)', and PTB (Marcus et al., 1993) in Table 3. In addition, we report zero-shot accu- racy in Table 3 on LAMBADA for predicting the final word of sentences, well (CBT) (Hill et al., 2015), where the task is choose either common nouns (CN) or named entities (NE) from 10 possi- ble options for an omitted word in paragraph. We observe that the AoT models with medium and large parameter sizes can achieve comparable performance to the GPT-2 base model. Moreover, we found that adding MLP layers to AoT does not improve the zero-shot performance. These results highlight the potential of attention-only models achieve competitive results while maintaining interpretability.
the Children's Book Test
4.2. LANGUAGE MODELING
Pre-training language models. We pre-train the AoT- MSSA-L and AoT-MHSA-L models of different sizes, along with GPT-2 (see Table for model sizes), on Open Web- Text (Gokaslan & Cohen, 2019). We defer the details of the model architectures to Table 4. Here, we train these models over a 1024-token context using the AdamW opti- mizer (Loshchilov & Hutter, 2019). We plot the training loss and validation loss against the number of training it- erations in Figure 5(a) and (b), respectively. We observe that medium- and large-sized AoT-based models achieve training and validation losses comparable to those of the GPT-2 base model. In addition, compared to the GPT-2 base model, the AOT-MHSA-L model is identical the GPT-2 base model, except for the absence of MLP layers in the architecture. As shown in Figure 5, incorporating MLP layers can accelerate the training process.
4.2.2. IN-CONTEXT LEARNING
In-context learning (ICL. refers to the ability of modern language models to perform tasks by using examples pro- vided in the input prompt, along with a new query input, generating outputs without updating the parameters (Brown et al., 2020; Garg al., 2022; Park et al., 2024a). We eval- uate the ICL capabilities of our AoT models and compare their performance with that of GPT-2 (Radford et al., 2019). Each model is trained from scratch on specific tasks, in- cluding linear and sparse linear regressions. We mainly follow the setup (Garg et al., 2022) to train models to
Zero-shot evaluation. Using the above pre-trained mod- els, we compute the cross-entropy validation loss without training on datasets WikiTexi (Merity et al., 2017)S, LAM-
to as WikiText.
"To obtain the accuracy on LAMBADA dataset, we use greedy decoding.
*For Wiki Text2 and Wi (Merity et al., 2017), the test splits are same, so we merge them as a single dataset referred
8

---

Attention-Onl Transformers
Transformes AREMSSAL AoT-MHSAALL:
AUT-MSSA-L
MIISI
Squared Error
Squared Error
Least Siqueres
Least Squaree
Lesso (eigba-0.011
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
10
20
25
30
35
10
20
25
30
35
In-context Examples
In-context Examples
Figure Evaluating models on in-context learning tasks. We plot normalized squared error as function of numher in-context examples for linear regression (left) and sparse linear regression (right) tasks.
learn linear functions in context. Specifically, for a spe- cific function class G, we generate random prompts by sam- pling function g G from distribution D3 over func- tions random imputs 2i..... &N E E ii.d. from Dx over inputs. To evaluate the inputs g. we create prompt P (Z1.g(æ1), IN .NN.NNN. 9(xN)). We train the model fo(-) to minimize the expected loss over all prompts pre- fixes:
denoising operator to token representations, where these rep- resentations are assumed be sampled from a mixture of low-rank Gaussians. Remarkably, this simple architecture consists of multi-head (subspace) self-attentin and skip connections at each layer, without MLP layers all. We have rigorously proven that each such layer improves the signal-to-noise ratio token representations at a linear rate with respect to the number of layers. Extensive experiments on both language and vision tasks demonstrate that this simplified architecture achieves performance comparable to that of standard transformers. Our theoretical and empirical findings suggest that subspace denoising via attention heads is the core mechanism underlying transformer effectiveness, with MI.P layers contributing only marginal performance gains. We believe this work lays a foundation for future exploration of more efficient and principled architectural designs.
min Ep
Σ N
(fo(P') 9(x)))
(9)
where pi is the prompt prefix up to the input i-th in-context example P (x1.g(x)),.... x.).
Tasks. We consider both linear functions and sparse linear functions with dimension d 20. The in-context examples 2; are sampled from the isotropic Gaussian distribution. For linear functions, we define g = {g g(x) = x))x), where is sampled from isotropic Gaussian distribution as well. For sparse linear functions, the setup is similar, but with modification: only coordinates in the vector w are set non-zero, while the remaining
Impact Statement
This paper presents work whose goal is to advance the field of machine learning. Both theoretical analysis and experi- mental evaluation presented in this paper aim to help people better understand the attention layer in standard trans- former architecture. Except for values academics, do not anticipate any social ethical implications.
are set to be zero.
Training and evaluation. For all experiments, we set the number of heads and the embedding size 128. We present the model configurations Table in Appendix C. To train the model, we sample batch of random prompts with size 64 and train the models for 50,000 iterations using Adam optimizer (Kingma& & Ba, 2014). We evaluate models using same Dg and Dx to sample 1280 prompts. We refer the reader to (Park et al., 2024a) for more details. We plot the estimation error against the in-context samples in Fig- ure We observe that our AoT architecture can in-context learn linear functions and sparse linear functions, achieving performance close to that of the GPT-2 transformer.
A Acknowledgment
Peng Wang and Qing Qu would like to acknowledge support from the NSF grant #2402950. Druv Pai would like to acknowledge support from the UC Berkeley College of Engineering Fellowship. Yi Ma would like acknowledge support from the joint Simons Foundation-NSF DMS grant the ONR grant N00014-22-1-2102 the NSF grant #2402951, and also support from and the HKU startup, the Hong Kong Center for Construction Robotics Limited (HKCRC) Award 052245, and JC Club Hong Kong.
5. Conclusion
In this work, we proposed new and minimalistic trans- former architecture by interpreting each layer as a subspace

---

Attention-Only Transformers
References
In International Conference on Machine Learning, pp. 1691-1703. PMLR, 2020.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I,, Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. GPT-4 technical report. arXiv preprint arXiv:2303.08774. 2023.
Chen, X., Liang, C., Huang. D., Real, E.,, Wang, K., Pham, H., Dong. X., Luong, T., Hsieh, C.-J., Lu, Y. et al. Sym- bolic discovery of optimization algorithms. In Advances in Neural Information Processing Systems, volume 36, 2024.
Ahn, K., Cheng, X.. Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. In Advances in Neural Information Processing Systems, volume 36, pp. 45614- 45650, 2023.
Chen, Y., Yun, Z.,, Ma, Y., Olshausen, B., and LeCun, Minimalistic unsupervised representation learning with the sparse manifold transform. In The Eleventh Interna- tional Conference Learning Representations. 2023.
Alcalde. A., Fantuzzi, G., and Zuazua, E. Clustering in pure- attention hardmax transformers and its role in sentiment analysis. arXiv preprint 24Xiv:407.01600 2024.
Deng, J., Dong. W., Socher, R., Li, L..J., Li, K., and Fei-Fei, ImageNet: large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pat- tern Recognition, pp. 248 245.-255..02 IEEE, 2009.
Arora, S., Li, Y., Liang, Y., Ma. T., and Risteski, A. Linear algebraic structure of word senses, with applications to polysemy. Transactions of the Association for Computa- tional Linguistics, 6:483-495, 2018.
Devlin, J. Bert: Pre-training of deep bidirectional trans- formers for language understanding. arXiv preprint arXiv: urXiv:1810.044 1810.04805, 2018.
Bao, F. Nie, S., Xue, K., Cao, Y., Li, C., Su, H., and Zhu, J. All are worth words: A vit backbone for diffusion models. In Proceedings of the JEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 22669- 22679, 2023.
Dong, Y., Cordonnier, J.-B., and Loukas, A. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. International Conference on Machine Learning, pp. 2793-2803, PMLR, 2021.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D.,
Dosovitskiy. A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Confer- ence Learning Representations, 2021. arXiv preprint arXiv:2010.11929
Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Askell, A., Agarwal, S., Herbert- Vosss, A., Krueger, G.,
Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877 1901, 2020.
Elhage, N., Hume, T., Olsson, C.,, Schiefer, N., Henighan, T., Kravec, S., Hatfield-Doddss Z.,, Lasenby, R., Drain, D., Chen, C., al. Toy models of superposition. arXiv preprint ar arXiv:2209. 10652, 2022.
Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9650-9660, 2021.
Elhamifar, E. and Vidal, Sparse subspace clustering: Algorithm, theory, and applications. IEEE Transactions Pattern Analysis and Machine Intelligence, 35(11): 2765-2781, 2013.
Chan, K. R.R Yu, Y., You, C., Qi, H., Wright, J., and Ma, Y. Redunet: A white-box deep network from the principle of maximizing rate reduction. Journal of Machine Learning Research, 23(1 (6(114)11-103, 2022.
Engels, J., Michaud, E. J., Liao, I. Gurnee, W., and Tegmark, M. Not all language model features are one- dimensionally linear. In International Conference on Learning Representations, 2025.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, volume 34, pp. 15084-15097, 2021.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information Processing Systems, volume 35, pp. 30583-30598, 2022.
Chen, M., Radford, A., Child, R., Wu, J., Jun, H., Luan, D., and Sutskever, I. Generative pretraining from pixels.
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. The emergence of clusters in self-attention dynamics.
10

---

Attention-Only Transformers
In Advances in Neural Information Processing Systems, volume 36, pp. 57026-57037, 2023a.
Loshchilov, 1. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference Learning Representations, 2019.
Geshkovski, B., Letrouit, C.,, Polyanskiy, Y., and Rigollet, P. A mathematical perspective on transformers. arXiv preprint arXiv:2312. 2023b.
Luo, J., Ding, T., Chan, K. H. R.,, Thaker, D., Chattopadhyay, A., Callison-BBrch, C., and Vidal, R. Pace: Parsimonious concept engineering for large language models. In Ad- vances in Neural Information Processing Systems, 2024.
Geva, M., Schuster, R., Berant, J., and Levy, O. Transformer feed-forward layers are key-value memories. arXiv preprint arXiv:2012.14913 2020.
Ma, Y., Tsao, D., and Shum, H.-Y. On the principles of par- simony and self-consistency for the emergence of intelli- gence. Frontiers of Information Technology Electronic Engineering, 23(9): 1298 1328, 2022.
Gokaslan, A. and Cohen, V. Openwebtext cor- pus. OpenWebTextCorppl 2019.
http:://kee i ooo00.0000000 www.ooogithub...... io/
Gregor, K. and LeCun, Y. Learning fast approximations of sparse coding. In International Conference on Machine Learning, pp. 399-406, 2010.
Marcus, M., Santorini, B., and Mareinkiewicz, M. A. Build- ing a large annotated corpus of english: The penn tree- bank. Computational Linguistics, 19(2)313-330, 1993.
Guo, J., Chen, X., Tang, Y., and Wang, Y. Slab: Effi- cient transformers with simplified linear attention and progressive re- p--aammettrrzed batch normalization. In- ternational Conference on Machine Learning, pp. 16802- 16812. PMLR, 2024.
Merity, S., Xiong. C., Bradbury, J., and Socher, Pointer sentinel mixture models. In International Conference on Learning Representations, 2017.
Monga, V., Li, Y., and Eldar, Y. C. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2): 18-44, 2021.
He, B. and Hofmann, T. Simplifying transformer blocks. In The Twelfth International Conference on Learning Representations, 2024.
Noci, L.. Li, C., Li, M., He, B., Hofmann, T., Maddison, C. J., and Roy, D. The shaped transformer: Attention models in the infinite depth-and-width limit. In Advances in Neural Information Processing Systems, volume 36, 2024.
Hill, F., Bordes, A., Chopra, S., and Weston, J. The goldilocks principle: Reading children's books with explicit memory representations. arXiv preprint arXiv:1511.0230/1 2015.
Huh, M., Cheung, B., Wang, T., and Isola, Position: The platonic representation hypothesis. International Con- ference on Machine Learning, volume 235, pp. 20617- 20642. PMLR, 21-27 Jul 2024.
Pai, D., Buchanan, S., Wu, Z., Yu, Y., and Ma, Y. Masked completion via structured diffusion with white-box trans- formers. In The Twelfth International Conference Learning Representations, 2023.
Jiang, Y., Rajendran, G., Ravikumar, K., Aragam, B., and Veitch, V. On the origins linear representations in large language models. International Conference Machine Learning. volume 235, pp. PMLR, 21-27 Jul 2024.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, N.-Q., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G.. and Fernández, R. The lambada dataset: Word prediction requiring a broad discourse context. Proceedings of the 54th Annual Meeting of the Association for Compu- tational Linguistics (Volume I: Lonng Papers), pp. 1525- 1534, 2016.
Karpathy, NanoGPT. https: com/ karpathy/ / /an.GPT, 2022.
Katharopoulos, A., Vyas, A., Pappas, N., Fleuret, Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on PMLR, 2020.
Park, J., Park, J., Xiong, Z., Lee, N., Cho, J., Oymak, S., Lee, K., and Papailiopoulos, D. Can Mamba learn how learn? comparative study on in-context learning tasks. In International Conference on Machine Learning, pp. 39793-39812, 2024a.
Machine Learning, pp.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.2 2014.
Park, K., Choe, Y. J., and Veitch, V. The linear representa- tion hypothesis and the geometry of large language mod- els. In International Conference on Machine Learning, pp. 39643-369999 39666. PMLR, 2024b.
Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.

---

Attention-Only Transformers
Pires, T., Lopes, A. V., Assogba, Y., and Setiawan, H. One wide feedforward all you need. In Proceedings of the Eighth Conference on Machine Translation, pp. 1031- 1044, 2023.
subspace clustering. arXiv preprini arXiv:2409.0242. 2024.
Wang. S., Fidler, S., and Urtasun, R. Proximal deep struc- tured models. Advances Neural Information Process- ing Systems, 29, 2016.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Wright, J. and Ma. Y. High-dimensional data analysis with low-dimensional models: Principles, computation, and applications. Cambridge University Press, 2022.
Saharia, C., Chan, W., Saxena, S., Li, I.., Whang. J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. Photorealistic text-to-image dif- fusion models with deep language understanding. Ad- vances in Neural Information Processing Systems, 35: 36479-36494, 2022.
Wu, X., Ajorlou, A., Wang, Y., Jegelka, S., and Jadbabaie, A. On the role of attention masks and layernorm in trans- formers. Advances in Neural Information Processing Systems, volume 37, 2024.
Yu, Y., Buchanan, S., Pai. D., Chu, T., Wu, Z., Tong, S., Haeffele, B., and Ma, Y. White-box transformers via sparse rate reduction. In Advances in Neural Information Prucessing Systems, volume 36, pp. 9422-9457, 2023.
Schlag, I., Irie, K., and Schmidhuber, J. Linear transform- secretly fast weight programmers. In Interna- tional Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Yu, Y., Buchanan, S., Pai, D., Chu, T., Wu, Z., Tong, S., Bai, H., Zhai, Y., Haeffele, B. D., and Ma, Y. White-box transformers via sparse rate reduction: compression is all there is? Journal of Machine Learning Research, 25 1000):1-128, 2024.
Sukhbaatar, S., Grave, E., Lample, G., Jegou, H., and Joulin, A. Augmenting self-attention with persistent memory. arXiv preprint arXiv: 1907.01470, 2019.
Sun, X., Nasrabadi, M., and Tran, T. D. Supervised deep sparse coding networks for image classification. IEEE Transactions Image Processing, 29:405-411 2019.
Yun, Z., Chen, Y., Olshausen, B. A., and LeCun, Y. Trans- former visualization via dictionary learning: contextual- ized embedding as a linear superposition of transformer factors. arXiv preprint arXiv:2103. 10..29999,999,,,,, 2021.
Templeton, A. Scaling monosemanticiry: Extracting in- terpretable features from Claude sonnet. Anthropic, 2024.
Zhang, J. and Ghanem, B. ISTA-Net: Interpretable optimization-inspreed deep network for image compres- sive sensing. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 1828- 1837, 2018
Tsai, Y.-H. H., Bai, S., Yamada. M., Morency, L.-P., and Salakhutdinov, R. Transformer dissection: An unified understanding for transformer's s attention via the lens of kernel. Proceedings of the 2019 Conference on Em- pirical Methods in Natural Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-IJCNLP), pp. 4344-4353, 2019.
Zhang, R., Frei, S., and Bartlett, L. Trained transform- ers learn linear models in-context. Journal of Machine Learning Research, 25(49): 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, I. Atten- tion is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
Vershynin, R. High-dimensional probability: An introduc- tion with applications in data science, volume 47. Cam- bridge university press, 2018.
Wang, P., Liu, H., So, A. M.-C., and Balzano, L. Con- vergence and recovery guarantees of the K-subspaces method for subspace clustering. In International Con- ference on Machine Learning, pp. 22884-22918. PMLR, 2022.
Wang. P., Zhang, H., Zhang, Z., Chen, S., Ma, Y., and Qu, Q. Diffusion models learn low-dimensional distributions via
12

---

Attention-Only Transformers
To simplify our development, we introduce some further notations. We use BikDiag( 1 diagonal matrix whose diagonal blocks are X1, x,....; XK-
kk) denote block
A. Related Literature
Existing studies on self-attention mechanisms. is widely believed that the power of transformers primarily stems from their self-attention layers, which enable the model to capture long-range dependencies and contextual relationships between tokens by dynamically weighing token relationships across the input sequence (Tsai al., 2019; Vaswani et al., 2017). To explore the mechanism behind self-attention, numerous studies have investigated the performance of pure self-attention networks, often incorporating only one additional component to prevent rank collapse and maintain expressiveness. For example, Dong et al. (2021) showed that in pure-attention transformers without skip connections and MLP layers, token representations collapse exponentially to a rank-1 I matrix across layers. They also showed that self-attention networks with skip connections prevent rank collapse. Geshkovski et al. (2023a;b) have studied the dynamics of multi- head self seffaattentions and characterized clustering behaviors of learned representations. Recently, Wu al. (2024) showed that pure self-attention networks with LayerNorm can prevent rank collapse. While these studies have advanced the theoretical understanding of self-attention mechanisms in simplified transformer architectures, they cannot provide any empirical validation on real-world vision language tasks, offering little insight into the role of self-attention in practice.
Network architecture design via unrolled optimization. networks largely stems from their ability to transform the raw data into compact and structured representations, which facilitates downstream tasks (Chan et al., 2022; Chen et al., 2023; Ma et al., 2022: Yu et al., 2024). A principled and interpretable approach to learning such representations with transformers is to construct an architecture that incrementally transforms tokens into these representations via unrolling optimization steps as layers of deep network (Chan et al., 2022; Monga et al., 2021; Wang et al., 2016; Yu et al., 2023; Zhang Ghanem, 2018). Notably, Monga et al. (2021) demonstrate that such unrolled networks are more interpretable, parameter-<fficient, and effective compared to generic networks. Using this approach, each iteration of algorithm for learning compact and structured representations is represented one layer of deep networks. For example, Gregor & LeCun (2010) have demonstrated that sparse coding algorithms, such as ISTA, can be used to construct MLPs. Recently, Chan et al. (2022) constructed "white-box" network based on an iterative gradient descent scheme to optimize the maximal coding rate reduction objective. More recently, Yu et al. (2024) designed a "white-box" transformer architecture by implementing an approximate alternating minimization to optimize the sparse rate reduction objective. The proposed transformer achieves performance comparable to some popular ones such as ViT (Dosovitskiy et al., 2021) and DINO (Caron et al., 2021).
commonly believed that the success of modern deep
B. Proof of Theorem 3.1
B.1. Preliminary Results
To prove Theorem 3.1, we first establish several probabilistic results about Gaussian random vectors. First, we present probabilistic bound on the deviation of the norm of Gaussian random vectors from its mean. This is an extension of Vershynin (2018, Theorem 3.1.1).
Lemma B.1. Let æ (0,0 (2)) ) be a Gaussian random vector. It holds with probability at least 1 - 2 2exp (-t2/252) that
||||||| 5vd st+ 25.
(10)
Based on the above lemma, we can respectively estimate the norm of coefficients in the signal and noise parts, the products between different pairs of Gaussian random vectors, and the bounds the soft-max values of these products.
Lemma B.2. Consider the setting in Definition 2.1 with p = Pı
Pk and N Nx N/K. Suppose that
p 2 16(vlog N + 1)2 and
log N
N 2 8 log N, d < 8
(1I)
The following statements hold:
13

---

Attention-Onl Transformers
(i) With probability at least
2KN- we have
||| vpl s 2 <V(Vlg (ollg.++1),,N. -1I, ( [N],
(12)
||e 8vpl s 20 (VlogN+ ,ii vi Chl [K].
(13)
(ii) With probability at least
4KN-2 we have
(a a,) s 3v Bolog Mooooolallllll. vj E [K],
(14)
((,, e,, - 3/llog Mell, E Ck.j E C₁ k & I E [K],
(15)
(eei, e,,,) - /lg N1 Ci,j E Cm,l, m +
(16)
(iii) With probability at least 2N-1, we have
max (ai; ej,a) iEC2
(17)
(iv) With probability at least
1 - 4KN-1 we have
exp apaaaa,, ezaa))
S S2 Vi Ck,je C,,
(18)
E Eiec, exp (xp, kj.k)) exp (xp Cjaa)
Vi * j E C.j Cm. I m # + k.
(19)
Ensived exp ((ev An ej,x))
2'
Proof. (i) Applying Lemma B.I to ~ (N(0, 1p) with t = 2/log N yields
P (1|| vpl < 2( log N + 1)) 21 1N-222
This, together with the union bound, yields that (12) holds for all i E [N] with probability at least 1 - -2 Using the same argument, we obtain that (13) holds for all i E Cs and * k E [K] with probability at least 1 - 2(K - 1)
Finally, applying the union bound yields that the probability is 1 - 2K
(ii) For each pair (i.j) with i * j E Cs and E [K]. conditioned on a, we have (a),a,) ~ (N [|a;]/)). According to the tail bound the Gaussian random variable, we have
P (Ka,, ال)l 2 3|la.Ily log N a.) S 2N
This, together with the union bound, implies that conditioned on a, it holds with probability at least 1 - 2N that (a, aللl < 2||a, Il Vlog N for all i * E C and & E [K]. Using the same argument, we obtain (15) and (16). Finally, applying the union bound yields the probability.
(iii) Conditioned on Ej we obtain that X; = (a,, (xxe))/)//l/slll ~N(0 1) for each i E Cj are i.i.d. standard normal random variables. Then, we have
max X, 2 log N =1 -1-(p(Xi VlggN) Ne
(20)
LEC&
Using the property of the standard Gaussian random variable, we have
P(i 2 2) 2 ( - 13
1
1
82 (-2)
exp
2
Taking = Vlog N, we obtain
IP X1 2 vlog.N]
1
(1- 11-
1
exp
log N
2
(21)
vlog N
log N
N log N'
14

---

Attention-Only Transformers
where the inequality follows from N 2 exp(2). Substituting this into (20) yields
max X; 2 V Vlog 2 (1 iECk
1
N/K
F
2 N N
VN
21-N
exp
2K v 2T log
where the second incquality uses - r < exp (-x) for all x > 0 and the last inequality follows from N 2 k log N. This, together with the definition of X, implies (17).
(iv) Conditioned on Cj,k, we have X; := ej,e) ~ (0, ||s) for each i E C. are i.i.d. normal random variables.
Suppose that (13) holds for all i E C,, k E [K], which happens with probability at least - 2(K - 1)N according (i). This implies for j Ck and [K],
3 || 6 (V++ + 2v log N +2) - 2óVD.
(22)
where the last inequality follows from p 2 16( 16l. N + 1)² due to (11). For ease of exposition, let
σ a= == = Σ
exp(X.).
(23)
IEC
Obviously, showing (18) is equivalent proving
2 ex) S
Σ
exp (X+) S. vi E Ck.
(24)
PEC
Note that X 1) for all i E Cz. Using the tail bound of the standard normal random variable, we have
Xi P ( 2 2o N S 22N Vi E Cz.
This, together with the union bound, yields that holds with probability 2N that |X,]| S 20 20 N for all i [N]. Using this, (22), (23), and the union bound, we obtain with probability at least 2K
[X < 3 N, Vie (N].
Therefore, have
exp (-35vppl plog N) I <xx((X,) s exp 35 log N), Vi Vi [N].
(25)
Using this and (23), have
N K
exp (-38 log log.N
S
This, together with (25), implies that proving (24) sufficient proving
log N 2 66 668plop log N + log (2K).
which holds when N 2 max (16K', exp (6482p)) due (11). According to the union bound, (18) holds with probability least 1 2K Using the same argument, (19) holds with probability least 1 2KK -
B.2. Proof of Theorem 3.1
To simplify our development, let
15

---

Attention-Oni Onll Transformers
BA A1
PA E2,1
A ЕK,1 2 EK,1
DE
A
E
E2,1
M
ERNNN
(26)
K A
E E2,1
E EK,1_ ET EK2
E E12
0 A
8 E
A A2
A ΕK,2
M2 =
NNN
E E
BE A2
E E
E Ei E E2, E EIK
OE Ak 2 A
Eg
MK
ERNN
0A ELK
0 AK
where 0 2 1. Recall that
Z (1)
Z 0)
Z
U Ej41 UKAK + Ejtk U₁ U,EK]],
(27)
Lemma B.3, Consider the setting in Definition 2.1 with P1 Pk and Ni =..... NO -N/K. Let (p-)) be
<(x) h(([x))),
(28)
where N + RN is the soft-max function and h RN + R. is an element-wise thresholding function with (x) = 7Il {x > t} for each E [N]. Suppose that (11) holds. Suppose addition that 2 64( 64( (4(1/og.. N + 1)² and
L ( 2'1+ (exp(-9pp33) 9p/32)
1
TE
(29)
The following statements hold with probability least 1 KN- that
c BIkDiag(rI,O,.. .), (MM)) BlkDiag(O, OI,
(30)
Proof. Suppose that (12)-(19) hold, which happens with probability at least KN- according to Lemma B.2, (11), and the union bound. Now, we focus studying M1 defined in (26). For case of exposition, denote the i-th column of M1 by m; E R for each i E [N]. Moreover, recall that
(1
N OK
(K-1)M
1, (K IIN
G G=
1.2.......
1,2,
K + + ...
......... K
K
We now divide our proof into two cases. We first study the i-th column of Mj for each E C1. and then study the i-th column of M2 for each E C₁ with l: # 1.
Case 1. According to (26), we have for each i C1,
mij 8² (a;
aj) az), Vj E Ct. mij = 0(a,, ej.k), E Cr, #
For each pair (i,j) with E C1, compute
a.(m) (mm:))
(902p)
exp (mi mij) 2 exp (Allaill (0llaill 3 /lll log 2 exp
(31)
32
16

---

Attention-Only Transformers
where the first inequality follows from (14) and the second uses (12) and vp 8(ylog N + 1). Using the same argument, for each pair (1,5) with i Ci,j Ck. and k * obtain
(m(m)) (y(m))
'00pp
2 exp
32
This, together with Σj= of (m;) = 1, yields (1 + (N-1) 8 C1.
- 1) exp
(-902p/32)) oi(m)) 2 1. Therefore, we have for each
1
1
1
(m()) 2 1+ N pxp(-99pp/32)
2' a,(m.) 2' Vj +
(32)
where the last inequality follows from p 2 64( 64(vlog N + 1)2 This, together with the value of T in (29), yields for each i e C1.
< < <T<a)(m))) Vj i.
Using this and (28), we have for each i C1.
h (oy(m)) =T, (o,(())) = Vi +
Case 2. For each i Cj with & + 1, it follows from (26) that
maj (er,,, az)) Vj C1, may (e,,,, e,1), Vj E C,,
Consider a fixed E Cj with & * 1, it follows from (17) that there exists ji E C such that mij; 2 01les,1ll N. This implies
0 (m;) a;(m;)
exp (Omri - -ms) 2 exp (lea, (o (ovlog BeeeN-Nlessa -Iles, [le..1l)
380
25
2 exp
log
52pp
16
where the second inequality follows from (13). This, together with (m:) + Jj, (m.) < 1, implies
1
(m.) < + exp (350Vplog N/4 2586p/16) 1 exp (pg N/2)
(33)
2
where the second inequality uses 6.7 S Vlg N/8 due to (11). On the other hand, it follows from (18) and (19) that
(m,m)) Fj +i.
This, together with (33), 8 1/8, vp2 8(vlog N + 1), and the value of T by (29), yields for each E C2 with & o 1,
my(m)) < ST [N].
(34)
This directly implies
(a(m)) 0, Vi Ck k 1.
I
Then, we have oe
(Mi =
Applying the same argument to Ma.... MR. we obtain (30).
Armed with the above result, we are ready to prove Theorem 3.1.
17

---

Attention-Only Transformers
Proof of Thhorrem 3.1. For ease of exposition, let ~ == 2 U& ( for each E [K] and I [Z]. Suppose that (30) holds, which happens with probability at least KN 9(1) according to (11), and (29), Lemma B.3. We claim that for each I [L], we have
z
[1+ UIAI Ezgn UjEiá (1+ UKA + E5tKUEKKs
(35)
This, together with (6), yields for each
[K] and [T]]
|U 1/ 2 SNR (2 I1(I U² 2
(1 (++r))IAAs 1Efz U,,,,l
(1) 1|
which directly implies (8) for each & E [K] and E [L-1]/ 1].
According to the union bound, the probability is 1 - K -9(1),
The rest of the proof is devoted to proving the claim (35) using the induction method. First, we consider the base case 1. According to (27) and (7), we compute
U z
[UIAI U1E211
UK A A
A E2.1 E E2,1
AT E
E. E
E ÆK,1
0 ZU, 20
E A
K E2,1
E EK.1
Using the same argument, we can compute M for each &s E [K]. This, together with (30) for each & E [K], yields
k
Σ U z 10)
(M = [UU 7U2A2 UKAAK]
Using this, (27), and (4), we directly obtain that (35) holds for = 1. Next, we consider the case 2 2. Suppose that (35) holds for some 7 2 1. We compute
U.TTZ") + (T) U,E,
UKKKI],
(1 (1+ A
(1 + qr) p A E2,1
(1 + nT AT EK,1
(1 (1+
E E2,1
E EK.1
(I) E A
Ex E2,1
E EK.1
Using the same argument, compute M (1) for each A [K]. This, together with (30) for each k E [K], yields
K
2M [1+ (1+
(1 (1+
Using this, (27), and (4), we directly obtain that (35) holds for + Then, we prove the claim.
Table architectures GPT-2 and AoT models. GPT-2, each layer consists of attention operator and MLP, while each
layer only has attention operator in AoT.
Models
Num. of para.
Num. of layers
Embedding dim.
Num. of heads
Attention type
Has ML.P
BOT-MSSA-L Base
102M
24
1024
16
MSSA
No
AoT-MSSA-L Medium
182M
36
1280
MSSA
No
AoT-MHSA-L. Base
122M
24
896
MHSA
No
GPT-2 Base
124M
12
768
12
MHSA
Yes
18

---

Attention-Only Transformers
C. Additional Experiemental Details
C.1. Language Model Configuration
We provide details the language model architecture in Table
C.2. ICL Configuration
We study decoder-only transformer models in GPT-2 family (Radford et al., 2019) and its corresponding AoT variants. As in Park et al. (2024a), we perform the same grid search over learning rates in [10-4, x 10 5, 10 4,4 10 *] and clipping the gradient norm in (5.0, 10.0, 50.03.
Table detailed architectures transformer and AoT used ICL experiments. ensure fair comparison, AoT models are designed with a larger number of layers to match the size of the transformer.
Models
Num. of para.
Num. of layers
Embedding dim.
Num. of heads
Attention type
Has MLP
AOT-MSSA
7.5M
32
128
8
MSSA MHSA
No
AoT-MHSA-L
8.55M
32
128
8
No
Transformer
9.63M
16
128
MHSA
Yes
C.3. Emergence Semantic Meaning
The attention heads in our models have different semantic meanings, and indeed demonstrate the interpretability of our proposed architecture in practice. Figure 7, we train the AoT model with the MSSA operator on ImageNet-IK and visualize the self-attention heatmaps between the [CI.S] token and other image patches. Note that the [CI.S] token is the "class token", trainable model parameter inserted along with other image tokens to represent the class information. We select 5 attention heads by manual inspection and find that they capture different parts of objects, displaying different semantic meanings.
19

---

Attention-Oni Onll Transformers
face
head
body
torso
legs
Figure 7. Visualization of attention heads on ImtgeNet-1K. We feed trained AoT-MSSA mini-batch of images and extract the attention maps of different heads from the penultimate layer. We show that these heads capture certain semantic meanings across different images.
20


Reviews:
[{"review_id": "NwflzzQoCR", "reviewer": "Reviewer_KY7F", "summary": "This paper proposes an attention-only transformer (AoT) architecture that eliminates the feed-forward network (FFN) modules found in traditional transformers, including CRAFT's Multi-head Subspace Self-Attention (MSSA). The authors argue that representation learning should compress noisy token representations toward a mixture of low-dimensional subspaces, and that multi-head self-attention naturally implements this denoising operation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **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...\n\n**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. \n\n> **Weakness 2.** **Suboptimal experimental results**: The experimental results in Table 1 indicate that AoT's performance is significantly suboptimal.\n\n**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. \n\nMoreover, 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: \n\n| Model | # of Parameters | Accuracy |\n| -------- | -------- | -------- |\n| AoT-MHSA     |  15M    | 69.5%     |\n| ViT Small     | 22M     | 72.4%     |\n\nWe believe that with further tuning, we can narrow the performance gap even more. \n\n*[1] Dosovitskiy, Alexey, et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.*\n\n> **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...\n\n**A3.**  Thank you for the feedback. We have retrained AoT-MHSA with additional hyperparameter tuning and present the updated experimental results below:\n\n| Model (# of Parameters) | LAMBADA (val loss) | PTB  (val loss) | WikiText  (val loss) | LAMBADA (acc) | CBT CN (acc)| CBT NE (acc) | OWT (val loss) |\n| -------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- |\n| AoT-MHSA (122M)    | 4.42    | 5.52 | 4.19 | 0.38 | 0.86 | 0.82 |2.92     | \n| GPT2 (124M)     | 4.32    | 5.75  | 4.13 | 0.4 | 0.87 | 0.84 |  2.86     |\n\n\nWe 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.\n\n> **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. \n\n**A4.** Thanks for catching this error. To fix it, we revise Eq. (5) into \n$$\n\\mathrm{MHSA}(Z) = W_O \\begin{bmatrix}\n\\mathrm{head}_1 \\\\\\\\ \\dots \\\\\\\\  \\mathrm{head}_K\n\\end{bmatrix}.\n$$\nHere, $Z \\in \\mathbb{R}^{d\\times n}$, $W^O \\in \\mathbb{R}^{d\\times Kd_h}$, and $\\begin{bmatrix}\n\\mathrm{head}_1 \\\\\\\\ \\dots \\\\\\\\  \\mathrm{head}_K\n\\end{bmatrix} \\in \\mathbb{R}^{Kd_h \\times n}$ due to each head $ \\mathrm{head}_k \\in \\mathbb{R}^{d_h\\times n}$. \n\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.\n\n**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]. \n\nIt 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$. \n\n\n*[2] Boucheron et al. (2013). Concentration Inequalities: A Nonasymptotic Theory of Independence. Oxford University Press.*", "claims_and_evidence": "## Strengths:\n1. The paper is well-written. Figures are clear to understand the paper.\n2. The theoretical formulation is clear, particularly regarding the concept of a union of (possibly many) low-dimensional subspaces.\n3. The paper provides a solid mathematical foundation for its claims.\n\n## Weaknesses:\n\n1. There appears to be an 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. This contradicts the theoretical framework presented, suggesting a mismatch between the proposed theory and actual implementation. I am not sure about this point, I do not check the authors' code.\n\n2. The experimental results in Table 1 indicate that AoT's performance is significantly suboptimal. On ImageNet, improvements of 0.2 or more are typically considered statistically significant, with 20M-30M models usually achieving accuracy between 80-84%. However, AoT only reaches approximately 70%, making it difficult to consider this approach effective.\n\n3. The claims regarding experimental results are problematic. In nanogpt experiments, according some previous top-conference papers, improvements of 0.02-0.03 in validation loss are typically considered meaningful. While increasing network parameters generally leads to lower validation loss, the paper shows that under comparable parameter conditions, AoT's loss is approximately 0.1 worse than the baseline, yet the authors claim comparable performance, not fair claim. Even more concerning, AoT with 180M parameters performs worse than the baseline (124M), strongly suggesting that AoT is an ineffective or significantly underperforming approach.\n\n4. There is a dimensional inconsistency in Equation 5. Since head_k is d_h by n, head_k^T should be n by d_h, but the matrix dimensions do not align properly in the equation.\n\n5. The mathematical foundations of the paper are difficult to fully verify as they require extensive knowledge of high-dimensional probability theory and considerable time investment. 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.", "methods_and_evaluation_criteria": "Yes, it makes sense", "theoretical_claims": "I partly check its proofs, I mention some problems.", "experimental_designs_or_analyses": "yes, it is ok, but the experimental results are not good enough.", "relation_to_broader_scientific_literature": "It provides a new perspective of understanding self-attention.", "essential_references_not_discussed": "references are ok.", "other_strengths_and_weaknesses": "see above", "other_comments_or_suggestions": "see above", "questions_for_authors": "see above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "no ethical concerns"}}, {"review_id": "7ckR9pM4Ah", "reviewer": "Reviewer_5Af5", "summary": "The authors propose an attention-only architecture, using the multi-head subspace self-attention (MSSA), first proposed by Yu et al., NeurIPS 2023 ( and also JMLR 2024). They have a model in which the embeddings of the tokens come from a mixture of different subspaces, albeit with additional additive noise making each token slightly away from its natural subspace. The purpose of the network is to denoise these vectors at each layer, increasing the signal to noise ratio as one goes from layer to layer. They view their architecture as a special case of a multi-head attention model.\n\nThe authors provide some theoretical analysis assuming a particular model for $Z^{(l)}$. They try this architecture for vision transformer for ImageNet top-1 accuracy and then also for language models doing next token prediction. The performances are not meant to be competitive with the best engineered transformer models, but are still nontrivial. They also do some in-context learning for linear models. The number of parameters are only in tens of millions as opposed to billions.\n\n## update after rebuttal\n\nI am satisfied with the authors' clarifications. I am keeping the same score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **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. \n\n**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:\n$$\n \\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)\n$$\n\n> **Q2.** **Layer norm:**  Layer norm appears in Fig. 3 but does not get mentioned much in the model or the analysis.\n\n**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. \n\n\n> **Q3.** I would prefer a clearer presentation of the architecture and not having to rely on the previous papers for disambiguation.\n\n**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).", "claims_and_evidence": "I take that the authors' main claim is that transformers can act as denoisers which ultimately provide a compressed representation allowing easy training for downstream tasks. I think they have shown some evidence for this claim. They also claim the architecture is interpretable, and show that different heads have some semantic interpretability (Fig. 7, Appendix D).", "methods_and_evaluation_criteria": "Once more, the benchmarks and the datasets are less demanding that state-of-the-art transformers.", "theoretical_claims": "I have not checked the proofs in Appendix B in detail but got a sense of the techniques involved and the general flow of the argument.", "experimental_designs_or_analyses": "Experimental design and analysis seems sufficient.", "relation_to_broader_scientific_literature": "How transformers form such a powerful representation of sequence is something of a mystery. The authors are advancing a view of transformers as a signal denoiser via something like a subspace clustering. This MLP layer free transformer is easier to analyze a provide some credence for this view. Of course, its performance is going to be far from satisfactory. We have to see if explainability/interpretability compensates for the loss of performance.", "essential_references_not_discussed": "I am not aware of such references.", "other_strengths_and_weaknesses": "There are some notations in the paper that I think would confuse the first time reader. $Z^{(l)}, f^{l}$ have layer superscript $l$ in the equations 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. Unfortunately, Figure 2 does the conceptual explaining with $U_k$'s apparently the same for each layer. Also, layer norm appears in Fig. 3 but does not get mentioned much in the model or the analysis.", "other_comments_or_suggestions": "None", "questions_for_authors": "I would prefer a clearer presentation of the architecture, and not having to rely on the previous papers for disambiguation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rzfrjwNomz", "reviewer": "Reviewer_sksF", "summary": "This paper presents a new transformer architecture by interpreting each layer as a subspace denoising operator for the token representations. The work is interesting, and the paper is well written in general. The authors also provide theoretical analysis on the developed model. Experimental results show that the proposed method can yield comparable performance in different applications.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **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.\n\n**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. \n\nAdditionally, 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. \n\n> **Q2.** The advantage of the proposed method might be strengthened by comparing it with more transformer-based methods.\n\n**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: \n\n| Model | # of Parameters | Accuracy |\n| -------- | -------- | -------- |\n| AoT-MHSA     |  15M    | 69.5%     |\n| ViT Small     | 22M     | 72.4%     |\n\nNext, 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. \n\nBased 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.\n\n*[1] Dosovitskiy, Alexey, et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.*\n\n*[2] Touvron, Hugo, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.*", "claims_and_evidence": "The experiments support the claims in the paper.", "methods_and_evaluation_criteria": "The proposed methods and evalutions make sense in the paper.", "theoretical_claims": "The authors provide proofs to some of the theories.", "experimental_designs_or_analyses": "The experiments are sufficient to support the claims in the paper. However, it could further benefit from comparison with more recent SOTA transformer-based methods.", "relation_to_broader_scientific_literature": "The idea of interpreting each layer of transformer as a subspace denoising operator for token representations is interesting, which facilitates the understanding of the commonly used transformer network.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: 1) Very good presentation of the work and the motivation is clear. 2) interesting findings to interpret a layer of transformer as a subspace denoising operator for token representation. 3) the work is theoretically founded and empirically verified in different tasks. \n\nWeaknesses: 1) 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. \n2) The advantage of the proposed method might be strengthened by comparing with more transformer-based methods.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See the weaknesses above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper provides an interpretation of self attention layers as an unrolled iterative denoising process, and proposes a Transformer only architecture with a MSSA attention layer. They demonstrate that this architecture achieves non-trivial performance on a set of benchmarks, which lays a foundation for understanding Transformers in a simpler setting. Reviewers in general appreciate the angle the paper takes, and acknowledges the value of such study despite its lack of competitive performance with modern Transformers. Further clarifications and results in the rebuttal also helped resolve some of the concerns. The AC agrees with the reviewers, but also thinks that the paper can be further improved by better elaborating on the potential benefit of the denoising view of Transformers, or how that could help us better understand standard Transformers. Given these considerations, the AC leans towards accepting this work if there is space.
Decision: Accept (poster)