Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction
Authors: Ziyang Wu
Tianjiao Ding
Yifu Lu
Druv Pai
Jingyuan Zhang
Weida Wang
Yaodong Yu
Yi Ma
Benjamin David Haeffele
Venue: ICLR 2025 Spotlight (2025)
Keywords: white-box deep neural networks
representation learning
transformer
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We build a novel transformer-like architecture with linear-complexity attention by algorithmic unrolling of a compression-based learning objective.

Abstract: The attention operator is arguably the key distinguishing factor of transformer architectures, which have demonstrated state-of-the-art performance on a variety of tasks. However, transformer attention operators often impose a significant computational burden, with the computational complexity scaling quadratically with the number of tokens. In this work, we propose a novel transformer attention operator whose computational complexity scales linearly with the number of tokens. We derive our network architecture by extending prior work which has shown that a transformer style architecture naturally arises by "white-box" architecture design, where each layer of the network is designed to implement an incremental optimization step of a maximal coding rate reduction objective (MCR$^2$). Specifically, we derive a novel variational form of the MCR$^2$ objective and show that the architecture that results from unrolled gradient descent of this variational objective leads to a new attention module called Token Statistics Self-Attention ($\texttt{TSSA}$). $\texttt{TSSA}$ has $\textit{linear computational and memory complexity}$ and radically departs from the typical attention architecture that computes pairwise similarities between tokens. Experiments on vision, language, and long sequence tasks show that simply swapping $\texttt{TSSA}$ for standard self-attention, which we refer to as the Token Statistics Transformer ($\texttt{ToST}$), achieves competitive performance with conventional transformers while being significantly more computationally efficient and interpretable.  Our results also somewhat call into question the conventional wisdom that pairwise similarity style attention mechanisms are critical to the success of transformer architectures.

Published as a conference paper at ICLR 2025
TOKEN STATISTICS TRANSFORMER: LINEAR-TIME ATTENTION VIA VARIATIONAL RATE REDUCTION
Ziyang Wu UC Berkeley
Tianjiao Ding UPenn
Yifu Lu" UMich
Druv Pai UC Berkeley
Jingyuan Zhang THU & TranscFnngram
Weida Wang Tsinghua SIGS
Yaodong Yu UC Berkeley
Yi Ma
Benjamin D. Haeffele JHU & UPenn
UC Berkeley & HKU
ABSTRACT
The attention operator is arguably the key distinguishing factor of transformer architectures, which have demonstrated state-of-the-art performance ety of tasks. However, transformer attention operators often impose a significant computational burden, with the computational complexity scaling quadratically with the number of tokens. In this work, we propose a novel transformer atten- tion operator whose computational complexity scales linearly with the number of tokens. We derive our network architecture by extending prior work which has shown that a transformer style architecture naturally arises by "white-box" architecture design, where each layer of the network is designed to implement an incremental optimization step of maximal coding rate reduction objective (MCR"). Specifically, we derive novel variational form of the MCR" objec- tive and show that the architecture that results from unrolled gradient descent this variational objective leads to new attention module called Token Statis- tics Belf-Attention (TSSA). TSSA has linear computational and memory com- plexity and radically departs from the typical attention architecture that putes pairwise similarities between tokens. Experiments on vision, language, and long sequence tasks show that simply swapping TSSA for standard self-attention, which we refer to as the Token Statistics Transformer (ToST), achieves compet- itive performance with conventional transformers while being significantly more computationally efficient and interpretable. Our results also somewhat call into question the conventional wisdom that pairwise similarity style attention mecha- nisms are critical to the success of transformer architectures. Code is available at https://g thtbbb/./gtuubbuuoooo com wttp:obwwwnn2111111
vari-
MOTIVATION
Transformer architectures have led to state-of-the-art performance across many applications in ma- chine learning, computer vision, natural language processing, and elsewhere (Vaswani et al. 2017]
Devlin et al. 2019; Radford et al. 2018 Chen et al., 2020 Dosovitskiy et al. 2020). Arguably, the defining component of transformers the attention operator, which originally motivated to allow for handling long-range interactions among data tokens (e.g., image patches, words, video frames). Attention can come in variety of forms, but transformer architectures typically employ self-attention (Vaswani et al.] (2017). In particular, the core aspect of the self-attention operator is scaled dot product attention, which generates output tokens as weighted averages of input tokens, with weights computed as a "similarity" between pairs of input tokens, i.c., the scaled dot product between pairs of tokens after multiplication by the "key" and "query" parameter matrices.
While the success of transformers (and by extension self-attention) in a wide variety of applica- tions demonstrates the utility of the approach, comes with a potentially significant drawback: requires one to compute similarity between all pairs of input tokens, which results in its compu- tation and memory complexity sealing quadratically with the number input tokens. Indeed, the computational demands of self-attention operators are well-noted in the literature, and number of
"Equal contribution

---

Published as a conference paper at ICLR 2025
techniques have been proposed to allow for more efficient computation Examples include parti- tioning the tokens into blocks and having each attention head only compute with one subset of the tokens (Qiu et al.) 2020), computing attention over various sliding windows of tokens (Beltagy et al. 2020
Liu et al. 2021), finding a suitable low-rank projection of the tokens (Wang et al. 20020). or
approximating the pairwise attention computation via the Nystrom extension (Xiong et al.) (02T). Likewise, [Ali et al. (2021) propose to calculate scaled dot products between feature channels rather than between tokens, noting relationships between the eigenspaces of Gram matrices and correlation matrices, and dub their method the cross-coariiace transformer (XCIT). In this paper, we will also propose an efficient attention operator, but we will derive our method from a radically different ap- proach which suggests that compuling (or approximating) pairwise similarities between lokens muy not be necessary al all. To build an intuition for this notion is useful to consider the core operation of the attention operator, which has been noted by multiple authors in the literature (Chaudhari et al., 2021 Zhang et al. 2021 Vidal 2022) as performing a form of kernel regression (Nadaraya 1964, Watson 1964) with close connections to classical denoising techniques such non-local means or block-matching and 3D filtering (BM3D) (Buades et al. 2005 Dabov et al.) 2006). Namely, output tokens are formed by taking weighted average of input tokens which are deemed to be "similar" (as measured by the learned scaled dot product metric). More abstractly, this suggests that the self- attention operator a particular case of a more general class of operators which produces outputs based on statistics computed from the input tokens. In other words, in the case of self-attention we compute weighted averages based on the learned scaled dot product metric, but one could envision producing outputs based on other statistics of the input tokens which are more efficiently computed.
Here, to derive our proposed method we will leverage "white-box" archi- tecture design (also referred to as al- gorithmic unrolling), which takes the
Peak GPU Menory (GR)
Infervace Time (1a8)
VIT
10.0
100.0
ToST
1.0
philosophy that the operations of the
layers of a neural network can be
0.1
interpreted as performing incremen- tal updates to optimize an objective function (Gregor & LeCun (2010); ef. Monga et al. (2021) and its ref- erences). In particular, the structure of the architecture is defined by the form of an update that specific opti- mization method (e.g.. gradient descent/uscent) employs, but the parameters of the resulting update are still left to be learned]] Notably, recent work (Yu et al. 2023 (024) has argued that transformer- style architectures arise naturally from white-box design when incrementally updating an objective function based on the principle of maximum coding rate reduction (MCR²). While this provides in- teresting interpretations of transformer architectures, the resulting white-box architecture is largely identical to transformer, so it still suffers from the same computational challenges described above. Nevertheless, here we draw inspiration from this prior work and make the following contributions:
OK
5K
10K
15K
0K
5K
10K 15K
Number of o'Tokens (s)
Number of Tooees (n)
Figure Our ToST architecture, from unrolling a novel varia- tional form of MCR is faster and uses less memory than stan- dard transformer architectures such ViT (note the log-scale y- axes) and is based dramatically different notion of attention.
1. We provide a novel variational form for the MCR² objective that additionally allows one to upper-bound certain spectral functions of large matrices, which may be of independent interest.
2. Using white-box architecture design to derive objective, we propose a novel attention operator which is highly efficient and only requires linear complexity to compute (see Figure Digure[]compaiinnn comparing our method with ViT (Dosovitskiy al. 2020 ).
incremental optimization of our variational
3. Our resulting architecture is radically different from standard transformers in that it does not compute/approximate pairwise similarities between input tokens. Instead, only computes a data-dependent low-rank projection based token features. This leads us to name it the Token Statistics Transformer (ToST) (see Figure 3).
empirical second moment statistic of the input
4. We validate our approach experimentally and show that our proposed architecture achieves com- parable performance to standard transformer architectures while having significantly reduced computation and memory footprint, particularly for large numbers of high-dimensional tokens.
'See App. 1.1 for discussion on related efficient models such as linear attention and state-space models. Note that the objective function used to derive the architecture and the training objective used to learn the model parameters need not he the same.
2

---

Published as a conference paper at ICLR 2025
Notably, replacing self-attention with attention operator typically at least maintains perfor- mance on benchmark tasks and often improves performance for tasks with large token lengths.
2 BACKGROUND AND PRELIMINARIES
Notation. We define few notations that will be used throughout the rest of the paper. For positive integer n, let [n] = 21.2...... nn. Let 0 (resp. 1) be the vector/matrix of all zeros (resp. all ones). Let I be the identity matrix. For O. 1, and I, the dimension is inferred unambiguously from context. For a matrix A E RR=x, we denote by a, E Rm the 7th column of A, and by Aiy E IR the (th entry of A. For a vector v E R". let (vag(v) E R=** of v along the diagonal. For matrix A E Rmx" we denote by A A E R=*" the element square of A. For a vector v E IR" and a function f: I + I. let /[v] E IR" be the element-wise application of f to the entries of v. For m 2 n, denote by Om,,n) C R the set of m matrices with orthonormal columns; consequently, we denote by ((n) = O(n, n) the set of n x n orthogonal matrices. We denote by PSD(n) S Rux" the set of n x n positive semi-definite (PSD) matrices. For M E PSD(n) and E [m]. we denote by (M) the 7th largest eigenvalue of M.
be a diagonal matrix with the entries
2.1 REPRESENTATION LEARNING VIA MAXIMAL CODING RATE REDUCTION
Before presenting the development of our attention operator, we first review core ideas that lead to its derivation. Real data, such as images, videos, and text, are often assumed to be realizations of some random variable X drawn from some high-dimensional probability distribution which has an underlying rich low wiw-eimensional structure (Fefferman et al. 2016; Wright& & Ma 2022). Commonly, preprocessing methods tokenize the data, converting each sample into list of vectors (i.e. tokens) which typically represent a small and local part of the data (e.g., a word or piece of a word, a patch of image). As a result, for applications involving transformers, X [x1,,..., xx,] RDxn is typically a matrix-valued variable, and the transformer seeks to find representation mapping, or mapping of the tokens c: Dxx= + _gxx such that the features Z = <(X) [zi,..., Zn] E Rdxx are appropriate for the task at hand
3 While many works focus primarily on how well Z can be used to solve particular task (e.g.. classification or segmentation), recent work (Yu et al. 2020
Chan et al., 2022 Ding et al. [023]] Chu et al. 2024) has proposed more task- approach to learning a suitable representation which, rather than focusing on a task-specific metric, instead seeks to find the underlying low-dimensional structure in the data via an objective based on the principle of maximal coding rate reduction (MCR²). More specifically, the tokens are assumed to belong to an underlying set of K groups, which may (for example) represent different modes within the population of tokens, and the MCR² objective seeks to find token features of each group which compressed within group while the set of all token features is expansive as possible. In particular, let II = = [T]],... Tx] E
R=n be a stochastic "group assignment" matrix (i.e., II1 = 1 and x 20 0, V(i, k) E [m] [K] ) that assigns a group membership probability to each token (i.e.. the طHh row of II is probability vector for group membership for the 7th token). Then, letting e >> 0 and ni (TT, 1) for each is E [K], the MCR² objective, denoted as AR(Z, II), is:
K
d ZZ [+
Tkk
I d 7DDia((Tx)2) (1)
AR II) = log det
log det n
2² n
22 n
R((Z)
(2.11
Intuitively, the R(Z) term measures the volume of all the features, so maximizing it w.r.t. Z would promote that the features spread out (hence, we call it the "cxpansion term").. The (Z, II) term measures the sum of volumes of features from each of the K groups encoded by II, so minimizing it (or maximizing - R..) w.r.t. Z would push features in each group to be close (hence we call it the "compression term")). (Yu et al. 2020) showed that learning a representation by maximizing equa- tion I under suitable normalization constraints leads to the features being compressed, linearized, and discriminative: features from the same group lie in oiw-dimensional subspace corresponding to the intrinsic -iwddmennionaaltty of the token data, and the subspaces they form are orthogonal to each other. See Figure[]te for a visualization.
Herre, with some abuse of notation, we will use Z to refer to the the token features at some intermediate layer (or point) within the network, not necessarily just the final features at the output of the network. We also pack the tokens into the columns of a matrix, rather than the rows as is often done in the transformer literature, believe leads cleaner presentation.
3

---

Published as a conference paper at ICLR 2025
data
high-dimensional tokens
structured representations
value position embedding
representation
map Φ
Figure 2: Tokenization and representation for image data. Data (images) are split up into tokens (patches) X. which share semantics with other tokens from the same or different samples. Tokens with similar semantics may belong to the same geometrie structures in the original space and be grouped together by II. A Icarned mapping @ converts these tokens into features which are compressed, lincarized, and discriminative.
2.2 BUILDING WHITE- -Bx DEEP NETWORKS VIA UNROLLED OPTIMIZATION
Having introduced the MCR objective, we now discuss the principle of designing network archi- tectures in a "white-box" manner (i.e., where the structure of the architecture is well motivated and interpretable) via algorithmic unrolling (Gregor LeCun (2010): Yu et al. (2024); see also the re- view MMnga et al. (202T) and references therein). Algorithmic unrolling constructs each layer of the network to (approximately) implement a step of an optimization algorithm for our represen- tation learning objective, where the network architecture is defined by the structure of an update step from an optimization algorithm (e.g., gradient ascent/descent) but the parameters of the update not fixed at their analytical values (e.g., the computed gradient for gradient step) but param- eterized and learned through model training. Notably, Yu et al. (2024) argue that by unrolling a slightly modified MCR objective, which additionally adds 29 sparsity penalty on Z promote axis-aligned solutions, and uses a modified compression term (Z), one arrives at an architecture which is almost identical to a standard transformer (with the primary difference being the query, key, and value parameter matrices are all equal in the self- attention operator). In particular, a gradient step on the compression term, i.e., Z Z VV Re (Z), approximates a multi-head self-attention operator with weight sharing, plus a residual connection. That is, incremental optimization of the compression term yields a self-attention operator. Further, a proximal gradient step on the remain- ing terms R(ZZ) Z lo gives a soft thresholding operation which has similar functional form to a two-layer perceptron with ReLU activation. Taken together, these resulting operations closely ap- proximate a layer of standard transformer (self-attention followed by two-layer perceptron). The transformer-like network constructed by concatenating these layers is named CRATE, and performs comparably to the empirically engineered ViT architecture (Dosovitskiy et al.) 2020). However, as the attention mechanism in CRATE is largely identical to standard self-attention, still has time and memory complexity that scales quadratically with n. the number of tokens. This quadratic complexity presents a large barrier to scaling up transformer-like models both within and without the CRATE paradigm. In the sequel, inspired by how CRATE derives a self attention like operation as gradient step on a compression term similar to R., we will construct new attention mechanism as a gradient step on a novel variational reformulation of R..
PROPOSED ATTENTION FORMULATION
We now describe our new proposed attention mechanism, which we will derive by unrolled opti- mization of novel variational form of the compression term within the MCR² objective.
3.1 A NEW VARIATIONAL FORM FOR CODING RATES
To begin, we consider a general form of MCR objectives based on concave functions of the spectrum of a matrix. Namely, for a given PSD matrix M E PSD( and any scalar 2 0 we have that log det(I+ + cM) = >i= log(1 (1 + cA, (M)) where (recall) (M) is the 7th largest eigenvalue of M. Further note that log [ + co a concave non-decreasing function of Thus, we describe our results in terms of more general form of MCR² based on general spectral functions of PSD matrices of the form F F(M) recall from our above discussion that the attention mechanism arises from unrolling the compression component of MCR², we consider more general MCR2-style compression function:
Σ (A;(M), where is concave and non-decreasing. In particular,
4

---

Published as a conference paper at ICLR 2025
K
5 (Z, II)
Tkk
F
ZDiag ( 2 kT)
(2)
Tik
For objectives of this form, we now note the following result:
Theorem 1. Let f: [0, xe) + I be non-decreasing, concave, and obey f(0) 0, and let F: PSD(d) - have the form F FMM) E(=(((A(M))) Then for each M PSD(d) and
Q E Od)), we have
F(M) & E ((Q
(3)
Further, the inequality in equation 3 is achieved with equality for any Q which diagonalizes M, and f is strictly concave then the inequality in equation 3 achieved with equality if and only if Q diagonalizes M.
Theorem T is straightforward corollary of Theorem which is stated and proved in Appendix A Theorem T allows us to upper-but bound function of large matrices via just computing a scalar function of the diagonals of matrix product, and we know the spectral structure of the input matrix then we can craft orthogonal matrices for which this upper bound is made tight. This is powerful tool and may be of independent interest.
just
3.2 EFFICIENT DEEP ARCHITECTURE VIA UNROLLING THE VARIATIONAL. FORM
Using the above result, can replace equation with an equivalent variational objective with form
R ( (U K
Σ
1 (UT ZDiag( (Tx) Ux))) Tikk
where the equivalence in the sense that for an optimal choice of [Ux O(d)] matrices as described in Theorem (i.e., orthogonal matrices which diagonalize each ZDiag achieve a tight bound with (Z, {U*)<n) = Re, IT). Note that in general, achieving this bound would require selecting, for each sampled instance of Z, new optimal set Us pa- rameter matrices which diagonalize each ZDiag (T )Z which is clearly impractical for network architecture. Instead, as an alternative viewpoint, rather than considering the data (Z) as fixed and instead take
will
trying optimize the U& parameters achieve the tight variational bound, the algorithmic unrolling design principle described above and design an operator to perturb Z to incrementally minimize Re {U+)<r) To make this point explicit, each variational bound becomes tight when the eigenspaces of ZDiag ) align with the columns of Ug. so by rotating the appropriate columns of Z (namely, those which correspond to large entries in πT.) align with U we can approach a tight variational bound. That is, instead of rotating U₁ to align with the data for each instance Z, we can instead rotate the token features in each Z align with U
Following this approach, we compute a gradient descent step on R w.r.t. Z. To begin this com- putation, first let π E R" be any element- wise non-negative vector. Then we have
1
Vz 2
f((ZDiag( ) )k) Diag V f ZDDaggg)
where f is the gradient of f. and (recall) [ applies f to each element of the vector in the bracket. In particular, for fxr = log(1 (1 + (d/c² 2 )x), ff(z) (d/2)(1 + (d/2)z))1 is simply non-linear activation. Also, (recall) n2 ( 1). Thus, the gradient of R w.r.t. Z is:
Vz (2, (U,)
U Diag
v (U
02 T
U ZDiag(Ta).
(6)
n
(T 1)
=D U+)
(Note that the 1/n constant arises from a (nz/n)- (1/ng) 1/n constant each term of the sum.)
If now consider gradient step w.r.t. the 1th token Z₁ arrive at our proposed incremental compression operator, i.e., our surrogate for self attention + residual operator:
K
2 Zj TV
R R II
{ = E U,D,Z ग. U,)U
(7)
for each E [m), where - >> is step size parameter for the incremental optimization.
5

---

Published as a conference paper at ICLR 2025
Output Tokeas
(05)
Diugonal Attention (Ilead
ML.P
K heads
Memmberhhip Aggregate
2
Layer/orma
(0%)
Diugenal Attention (Ilead K)
Token Statistics
Complexity
Atteation
O(p) Space O(pm) Time
Projected
Updated Tokeas
(Vf) gean
Diagonal Attention
Ispat Tokens
Figure 3: One layer of the proposed Token Statistics Transformer (ToST), Notably, the self-attention of ToST transforms tokens Z' efficiently to Z via multiplying each row of the projected token by only scalar. This leads to reduced complexity of the attention (ef. Table I it has (p) space and (pp)) time complexity, where p is the dimension of the projected tokens of each head, and is the number of tokens.
Model interpretation. Given the proposed attention operator in equation 7 first recall that the of II non-negative and sum 1, so our operator takes weighted average of K "attention head" operators and then adds a residual connection. Using that Σ 11 we can rewrite equation 7 as:
K
II jk
Zj
U,D(Z n U,)
T
(8)
action one atteation head
That is, can view each attention head as first projecting the token features onto the basis UR multiplying by UK, multiplying by the diagonal matrix D(Z, ग. Ux) (abbreviated as D&), projecting back into the standard basis via multiplying by Ug. and subtracting this from the original token features via the residual connection. The core aspect of our attention layer is the computation R forms probability distribution
of Dj. Namely, 7 2 O, T&/ ( ,
which tokens belong the ttth group. As result, (U. Z) / 1) estimates the second moment of U Z under the distribution given by T ( 1). Further, since f is a concave non-decreasing function, æ increases, so the entries of D (which have form (x(r)) achieve their maximum at x = and decay monotonically to 0 as x increases.
V [(())) monotonically decreases towards
From this, arrive at the interpretation of attention head residual operators [I Ur/n),,D&&UE1 Namely, this operator does an approximate low-rank data-dependent projection, where directions which have large amount of "power" after the projection U. Z (i.e., directions which have large second moment (U, 2) / 1)) are preserved, while directions which do 0 as the
not suppressed. To see this, recall that the entries of D₁ decrease monotonically second moment increases, so for directions with large second moments the attention + residual op- erator largely as the identity operator. Conversely, for directions with small second moment our operator subtracts a projection of the tokens along those directions, resulting in those directions being suppressed. Compared to the standard self-attention operator, our method clearly does not compute any pairwise similarities between tokens. Rather, the interactions between the tokens in Z impact the operator solely through their contribution to the second moment statistic used to construct the D s. Nevertheless, similar to the standard self-attention operator, our method still has clear interpretation as performing a form of compression towards a data-dependent low-rank structure, in the sense that it performs approximate low-rank projection, where the specific directions which are suppressed are those which are not strongly aligned with other tokens in the group.
3.3 PRACTICAL IMPLEMENTATION CONSIDERATIONS
Having introduced our proposed attention operator, we now discuss further practical considerations. First, until this point in the presentation we have avoided discussion of how tokens are "grouped" into various attention heads via the II matrix, but clearly a means constructing II is needed to implement method. Additionally, be square and orthogonal, but one would ideally like to use smaller matrices (i.c., reduce the number of columns in U) for efficiency as well drop the orthogonality constraints.
variational form in Theorem T requires the U matrices to

---

Published as a conference paper at ICLR 2025
Table 1: Time and space complexity for attention operators different transformer architectures: ViT
(Dosovitskiy et al. 2020), CRATE (Yu et al. 2024). XCIT (Ali et el. 2021), and the proposed ToST.
ViT
CRATE
XCIT
ToST (ours)
Compute time complexity Memory complexity
O (pn)))
O(pn²)
Oly2n)
Opn)
On
O(n²)
Olp2)
Opp)
In practice, we do not enforce the orthogonality constraints. To reduce the number of columns in the U matrices, we note that similar to CRATE (Yu et (024), if we assume the features Z within group are (approximately) clustered around a low-dimensional subspace say of dimension p then the within- -oohii--roop-k covariance ZDiag (T. /Z" is low-rank, where recall that Yu et al. (2020) shows that the optimal geometry of Z will be for each group to be a low-rank subspace, orthogonal to the other groups. We can thus explicitly find a low-dimensional orthonormal basis for the image of this covariance, i.e. the linear span of the data in group k. If the dimension is p S d, the basis can be represented by d x p orthogonal matrix U3 E O(d, p). In this case, efficiently upper-bound R. using these low-rank orthogonal basis matrices. To show this, we use a more general version of Theorem (see Theorem Qn Appendix (A] to yield the following corollary.
more
Corollary Let f: [0, DO) - be non-decreasing, concave, and obey S(0) F: (pD(p) - T.. have the form F(M) = Σ (()(M)). Let Z, II be fixed. Then, for all S image( Un).or for all k E [K], we have
O, and let
U....., UK E O(d, p) such that image(ZDiagg TT )
Re, (Z, II) S Rey ZZ, II (U+)
(9)
where Reg. is formally defined in equation 4 Equality holds if U2 diagonalizes ZDiag( x for each 1e E [K, and if strongly concave Then this equality condition becomes an "if and only if."
The final step to define our attention operator is to estimate the group membership II. For this we posit simple model of how each feature z; deviates from its supporting subspace and then find the optimal subspace assignment. Yu et af. (2024) show that if we independently model each 2j as belonging to low-dimensional Gaussian mixture model, where each Gaussian has a covariance matrix with identical spectrum and is supported on a subspace with orthonormal basis Uk, plus independent Gaussian noise with covariance 171, then the posterior probability that each token z belongs each subspace is given by the assignment matrix II II(Z] { as follows:
v(z1 (U+)
| 2 12
where v(zy (U+) softmax
Vje [m), (10)
2n
plzn (U+)
where 17 becomes a learnable temperature parameter. Thus, given an input feature Z, we estimate II using equation T0 then compute the attention operator. Combining the construction of II in equation 10 with equation 7
we obtain the Token Statistics Self-Attention (TSSA) operator:
TSSA(Z (U+)
U.D(Z Te
U ZDiag(Tk),
(11)
n
where Te are the columns of II II(Z equation We use this TSSA operator to construct an architecture, the Token Statistics Transformer (TOST). Namely, our focus to construct an alternative attention mechanism, so to obtain ToST replace the attention operator in XCIT (Ali al. 2021) with the TSSA operator (with f(z) log(1 + (d/e²)z) as in Section3.11 along with other changes detailed in Section 4 which serve to simplify the architecture. We visualize one layer of ToST in Figure 3 and provide pseudocode for the model in Appendix
{ defined in equation 10 and D defined in
Complexity analysis. Table gives a comparison of computational complexity between each at- tention head of the TSSA operator in our proposed ToST architecture and those of attention mech- anisms in other existing transformers. As we see, the TSSA attention head asymptotically less complex than alternatives in either or p for both compute time and memory requirements. This confirmed in real computation metrics shown in Figure [ for ToST and ViT (with 12 attention layers, d 384, and K 8); we present more computation metrics for all models in Appendix D.2
7

---

Published as a conference paper at ICLR 2025
Measure coding rate across layers
R27 .nh ITSSA block)
train
val
Layer index
(a)
(b)
Figure 4: (Left) The variational compression term J (Z', II') the TSSA outputs Z' and estimated memberships IIf at different layers é of the ToST-S model. (Right) Visualization of estimated II for several images. For an image with tokens (patches), we visualize each row membership matrix II in the TSSA layer by reshaping it into vN x /N matrix. Here visualize membership matrices II in ToST-S, estimated in layer 9, of cach input image.
4 EXPERIMENTS
In this section, we conduct experiments to verify and study the properties and performance of our proposed Token Statistics Transformer (ToST) on real-world datasets and tasks. As detailed in Section 3 we design our white-box, attention operator by following the principle of MCR? and un- rolling the compression objective defined in equation We thus adopt a simplistic implementation as close to our theoretical derivation as possible. Hence, it is not the goal of this work to show the current implementation of ToST can outperform existing highly engineered architectures. Instead, empirical studies aim to provide answers and evidence for the following questions:
1. Does our proposed TSSA attention operator optimize the compression objective in practice?
If simply replace standard self-attention with our TSSA attention operator, leaving the rest of the architecture largely unchanged, do we maintain (or improve) task performance?
We provide positive answers both questions. First, since our TSSA attention operator is mathemat- ically derived via algorithmic unrolling, it should optimize quantitative objective of the representa- tion, which we show occurs layer-by-layer. Second, we demonstrate that our proposed ToST model, despite its simplicity and efficiency, already performs competitively on real-world large datasets when compared to architectures which use conventional self-attention mechanisms.
Model architecture. Since our main focus is the proposed TSSA operator as explained in Sec- tion 3] we follow minimalist approach of architecture design by directly replacing attention layers in existing transformer architectures with the TSSA operator. Concretely, our realization of ToST inherits most other design choices from XCIT (Ali al. 2021 and GPT-2 (Radford et al. 2019) for vision and language tasks, respectively. We also test different model sizes by varying the number of attention heads K, the token dimension d, and number of layers L. For supervised classification, we also follow XCTT by inserting learnable [CLS] token that aggregates features from image patches via global class attention layer. More details on implementation and model configurations are provided in Appendix
Datasets and optimization. For vision experiments, pre-train the proposed ToST models on the ImageNet- k (Deng et al. (209) dataset. We also use these pre-trained networks as initialization and fine-tune them smaller datasets including CIFAR 10/100 (Krizhevsky al.) 2009). Oxford Flowers (Nilsback & Zisserman 2008), Oxford-IIT-Pets (Parkh et al. 20122) for transfer learning evaluations. We also adopt the Long-Range Arena (Tay et al. 2021 benchmark analyze the long sequence modeling capability of ToST. For causal language modeling, we train ToST autoregres- sively on the Open WebText dataset(Gokaslan et al.
2019) and test the trained model on its test split 2019). Details on training and optimization are as well as other datasets as shown in (Radford et al.
provided in Appendix
"To align with our theoretical formulation, remove the Local Patch Interaction (LPI) module from XCIT.
See details in Appendix C
8

---

Published as a conference paper at ICLR 2025
Table 2: Top I accuracy of TOST on various datasets with different model sizes when pre-trained on ImageNet. For ImageNevlImageNeRRaall we directly evaluate the top-1 accuracy. For other datasets, we pre- train the model on ImageNet and then evaluate the transfer learning performance via fine-tuning.
Datasets
ToST-Tiny)
TOST-S(mall)
ToST-Mtedium
XCTF-
XCTI-M
VIT-S
ViT-B(ase)
parameters
5.MM
22.6M
68.1M
24.9M
80.2M
22.1M
86.6 M
ImngeNett
67.3
77.9
80.3
80.5
81.5
79.8
81.8
ImageNet Real
72.2
84.1
85.6
85.6
85.9
85.6
86.7
CIFARI0O
78.3
82.7
84.5
86.1
87.6
88.8
89.3
Oxford Flowers- 102
88.6
92.8
94.2
93.9
94.0
94.0
95.7
Oxford-IIIT-Pet
Input
ToST
CIT
VIT
Figure 5: Comparison of [CLS] token attention map visualization. We take the last head in the penultimate global class attention layer for visualization from FOCITS. and ViT-S, respectively.
4.1 LAYER WISE ANALYSIS OF THE TSSA OPERATOR
In this sub-section, we aim to answer the first question by providing both qualitative and quantitative results that demonstrate each layer of the proposed ToST attention operator indeed optimizes its designed objective. This layer-wise analysis is possible due to the white-box design of the TSSA operator, similar to studies conducted for CRATE, but different from most black-box box models.
Do ToST attention layers achieve their design goals? As detailed in Section 3 the ToST atten- tion operator aims optimize the variational compression term RN (Z, II) defined Equation (4). In order understand whether the learned attention operator indeed performs such optimization, explicitly measure and plot the objective value R (Z', II²) the outputs Z' and group estimations
L I I {U,, J. k of each layer f. We measure this term both training and validation set of the ImageNet- 1k dataset by taking 500 samples from each. We plot our observations Figure 4a The compression term mostly decreases by layer, aligning with the design goal.
Visualizing membership assignment II in TOST. At each layer f, the estimated membership matrix II has the interpretation that each column गे N represents the estimated probabilities for each of the N tokens be assigned the tth subspace. Hence, the II matrix represents a soft clustering of the N tokens to K clusters, or subspaces, via the K attention heads. In Figure 4b) we visualize the clustering a selected layer, and observe that foreground image patches are clustered in the same subspace, confirming our interpretation. See Figure 9 for more results.
4.2 EVALUATION ON REAL- WORLD VISUAL DATASETS
We now provide evaluations and analysis of ToST models on vision classification tasks, along with comparison against other otate-of-the-art vision transformer architectures.
ImageNet-1k 1k performance and transfer learning results. In Table 2 we report the top-1 accu- racy of ToST on ImageNet-1k dataset as well as their transfer learning performance on several smaller datasets. We observe that the proposed ToST models achieve comparable ImageNet Ik and transfer learning performance as other empirically engineered architectures, while being sig- nificantly more efficient and interpretable with fewer parameters. Notably, identify promising scaling behavior, where our models consistently improve and the performance gap between ToST and other networks decreases as scale up the model sizes.

---

Published as a conference paper at ICLR 2025
Table Long-Range Arena (IRR) performance comparison.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Avg
Reformer
37.27
56.10
53.40
38.07
68.50
50.56
BigBird
36.05
64.02
59.29
40.83
74,87
54.17
LinFormer
16.13
65.90
53.09
42.34
75.30
50,46
Performer
18.01
65.40
53.82
42.77
77.05
51.18
Transformer
37.11
65.21
79.14
42,94
71.83
59.24
ToST (ours)
37.25
66.75
79.46
46.62
69.41
59.90
$4
59.60
86.82
90.90
88.65
94.20
84.03
Table 4: Causal language modeling performance. (Left) We report zero-shot cross entropy loss GPT2- Base compared with ToST of varying model sizes on the test set of the datasets (lower is better). (Right Relative time and memory usage for base models.
Model
# params
OWT
Lambada
Wikitext
PTB
Avg
Model
Length 4k
Length Sk
GPT2-Buse
124M
2.84
4.32
4.13
5.75
4.26
Time
Memory
Time
Memory
ToST Base
110M
3.20
4.98
4.77
6.39
4.84
ToST-Medium
304M
2.88
4.45
4.30
5.64
4.32
GPT2-Base
1x
1x
ToST-Laage
655M
2.72
4.32
3.99
5.03
4.02
ToST ToST-Base
0.60
0.41x
0.46%
0.24x
Visualizing attention maps of ToST We can gain additional understanding of the features learned by ToST based on visualization of the self-attention maps. For the global class attention layers, we extract and visualize the attention map between [CLS] token and image patch tokens for each head using the widely adopted approach from (Caron et al.) (02T). In Figure 5 we visualize the attention maps in the penultimate layer, for ToST-S, XCiT-S, and VIT-S models. Based on the results, ToST autonomously learns to perform segmentation and clustering without the need for complex self-supervised training techniques or iggmentation-relate annotations. Compared to the XCIT (Ali et al., 2021) and ViT TTuvrro al. 2021) models, ToST demonstrates more meaningful segmentation results. Figure 10 contains more attention map visualization for each model.
4.3 EVALUATION ON LANGUAGE AND LONG SEQUENCE TASKS
Long sequence modeling. To assess the performance of ToST on long sequence tasks, we the Long-Range Arena (LRA) benchmark TTy et al. 202T) and compare ToST against the transformer and its efficient variants (Kitaev al. 2020) Zaheer et al., 2020 Wang et al. 2020 Choromanski jet al., 2021). The results are reported in Table 3 These results demonstrate promising long-range modeling capability of ToST which outperforms other transformer-based methods on average. We note that this benchmark is currently dominated by state-space models (SSMs) rather than trans- former models, but among transformer-Be based models ToST has highly competitive performance, demonstrating the utility of our TSSA attention operator even for long sequence lengths. For com- pleteness, we also include performance results for S4 (Gu et al., 2022), a representative SSM.
Scalability of ToST for causal language modeling. We can naturally extend the ToST architec- ture perform next-token prediction by introducing simple causal variant of the TSSA operator, as detailed in Appendix B We present our results in Table ToST on various text datasets steadily improve we scale up the size of the model, demonstrating promising scaling behavior. Due to computational constraints, we did not systematically sweep hy- perparameters for training ToST for this task and employed the same set as GPT2-Base for training all our models. Still, we find that swapping our TSSA attention operator for self- olf-attention largely preserves model performance, despite self-attention largely being heralded as the key component of transformer performance for language modeling tasks. Figure shows that the causal version of ToST also faster and uses less memory than standard transformers for autoregressive prediction.
We observe that performance of
5 CONCLUSION
In this work, we propose novel attention operator based on white-box design by unrolled optimiza- tion of a novel variational form of the MCR² objective. Our resulting architecture (ToST) is unique among attention operators in that it does not require computing pairwise interactions between to- kens and instead is constructed from second moment statistic of projected token features. This results in our operator being significantly more efficient than standard attention operators, while still achieving similar performance to comparable transformers. We believe that this work provides an initial demonstration of the tremendous potential in designing novel and efficient deep architectures from mathematical principles.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was partially supported by NIH grant R21AI169363, ARO MURI W911NF-17-100304
NSF grant 2031985, the Northrop Grumman Mission Systems Research in Applications for Learn- ing Machines (REALM) initiative, and UC Berkeley College of Engineering Fellowship. Ziyang Wu and Yi Ma acknowledge support from the joint Simons Foundation-NSF DMS grant #203 I 1899, the ONR grant 100014-22-1-2100 the NSF grant #2402951, and Yi Ma also acknowledges partial support from TBSI, InnoHK, and the University of Hong Kong.
REFERENCES
Alaaeldin Ali, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek, et al. Xcit: Cross-covariance image transformers. Advances in neural information processing systems, 34:20014-20027, 2021.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/3 055500 2020.
Antoni Buades, Bartomeu Coll, and J-M Morel. A non-local algorithm for image denoising. In 2005 IEEE computer society conference on computer vision and pattern recognition (CVPR 005), volume 2, pp. 60-65, leee, 2005.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the JEEE/CVF international conference on computer vision, pp. 9660-9660, 2021.
Kwan Ho Ryan Chan, Yaodong Yu, Chong You, Haozhi Qi, John Wright, and Yi Ma. Redunet: A white-box deep network from the principle of maximizing rate reduction. Journal of machine learning research, 23(1 331114)1--103 2022.
Sneha Chaudhari, Varun Mithal, Gungor Polatkan, and Rohan Ramanath. An attentive survey of attention models. ACM Transactions on Intelligent Systems and Technology (TIST), 12(5):1-32, 2021.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Hya Sutskever. Generative pretraining from pixels. In Proceedings of the 37th international conference on ma- chine learning, volume 119 of Proceedings of machine learning research, pp. 1691-1703. PMLR, July 2020.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. International Conference on Learning Representations, 2021.
Tianzhe Chu, Shengbang Tong, Tianjiao Ding, Xili Dai, Benjamin David Haeffele, Rene Vidal, and Yi Ma. Image clustering via the principle of rate reduction in the age of pretrained models. In The twelfth international conference on learning representations, 2024.
Kostadin Dabov, Alessandro Foi, Vladimir Katkovnik, and Karen Egiazarian. Image denoising with block-matching and 3d filtering. In Image processing: algorithms and systems, neural networks, and machine learning, volume 6064, pp. 354-365. SPIE, 2006.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. leee, 2009.
Jacob Devlin, Ming- Wii Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: N0..8853//1/N// 1423.

---

Published as a conference paper at ICLR 2025
Tianjiao Ding, Shengbang Tong, Kwan Ho Ryan Chan, Xili Dai, Yi Ma, and Benjamin D Haeffele. Unsupervised Manifold Linearizing and Clustering. In International Conference on Computer Vision, pp. 5450-545045 54600-546,,10 October 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint WwXiv:2010.11929 2020.
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. Journal of the American Mathematical Society, 29(4):983-1049, 2016.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus.
http:
//Skylion007. htw.iioonoo..iittbbb io/ ht/ppenneebTextCorppaaa 2019.
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th international conference on international conference on machine learning, pp. 399 406, 2010.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In Infernational Conference on Learning Representations, 2022.
2022. Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations, 2020.
Andrej Karpathy. NanoGPT. http:s///ggttub...cokaaaaaaaaapp
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. Advances in neural infor- mation processing systems, 31, 2018.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEFF/CVF international conference on computer vision, pp. 10012-10022, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
Andrew Maas, Raymond E Daly, Peter Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings the 49th annual meeting of the association for computational linguistics: Human language technologies, pp. 142-150, 2011.
Vishal Monga, Yuelong Li, and Yonina C. Eldar. Algorithm Unrolling: Interpretable, Efficient
Deep Learning for Signal and Image Processing. IEEE Signal Processing Magazine, 38(2):18 44, March 2021. ISSN 1558-0792. doi: 10.1100/MSP222203306660 Signal Processing Magazine.
Conference Name: IEEE
Elizbar A Nadaraya. On estimating regression. Theory of Probability & Its Applications, 9(1): 141-142, 1964.
Nikita Nangia and Samuel Bowman. Listops: diagnostic dataset for latent tree learning. arXiv preprint arXiv: 1804.00024888 2018.
Maria-Eles Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. Indian Conference Computer Vision, Graphics and Image Processing, 2008
O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. V. Jawahar. Cats and dogs. In IEEE Conference on Computer Vision and Patter Recognition, 2012.
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Trevor Cohn, Yulan He, and Yang Liu (eds.), Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 2555-2565, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020. findings-emnlp.222.
12

---

Published as a conference paper at ICLR 2025
Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47:919 9744, 2013.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference Learning Representations, 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, volume 30. Curran Associates, Inc., 2017.
Rene Vidal. Attention: Self-expression is all you need, 2022.
Ke Alexander Wang, Jiaxin Shi, and Emily B Fox. Test-time regression: a unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501. 12352, 2025.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04770 04768, 2020.
Geoffrey S Watson. Smooth regression analysis. Sankhyā: The Indian Journal of Statistics, Series A. pp. 359-372, 1964.
John Wright and Yi Ma. High-dimensional data analysis with low-dimensional models: Principles, computation, and applications. Cambridge University Press, 2022.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138-14148, 2021.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- with the delta rule over sequence length. Advances in Neural Information Processing Systems, 17:115491-11552. 2025.
Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. Learning diverse and discriminative representations via the principle of maximal coding rate reduction. Advances in neural information processing systems, 33:9422-944 9434, 2020.
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Hao Bai, Yuexi- ang Zhai, Benjamin D. Haeffele, and Yi Ma. White-Box Transformers via Sparse Rate Reduction: Compression All There Is?, November 2023. arXiv:2311.13110 [cs].
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin Ha- effele, and Yi Ma. White-box transformers via sparse rate reduction. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañóó, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems, 2020.
Aston Zhang. Zachary C Lipton, Mu Li, and Alexander J Smola. Dive into deep learning. arXiv preprint arXiv:2106 106. / 1342, 2021.
13

---

Published as a conference paper at ICLR 2025
A PROOFS OF THE THEORETICAL RESULTS
Below we present a more refined version of Theorem I
which easily proves Theorem I by taking
p d, and also is sufficient to prove Corollary
Theorem 2. Let f: [0, oc) - IR be non-decreasing, concave, and obey S(0) = 0, and let F: PSD(d) R have the form
F(M)- = Σ ((M(MM) VM E PSD(d).
(12)
Then, for all M E (d)(d) and all Q E =(d.p) with p S d such that image(M) C image (Q), we have
F(M)= & MQ(QTMQ)
(13)
==1 The inequality in equation 13 is achieved with equality for any Q which diagonalizes M, i.e. Q MQ is diagonal. Moreover, if J is strictly concave, then the inequality in equation 13 is achieved with equality if and only if Q diagonalizes M.
Proof. First, we have that rank(M) < rank(Q) p. Thus, let M UAU be spectral decomposition of M, where A E IRd is diagonal with non-negative entries (MM for i E [d], and U E Odd). In particular, we can write
p
U= U U
(14)
0
where Ap RPXP a diagonal matrix with non-negative entries (M) for E [p]. Up O p) and U, E (d,, - p). such that image Q) = image(U, (Up). so that image(M)) C image(Up)) and
image(() Q) is orthogonal (in the sense of linear subspaces) to im)ge(U).
We begin by noting the simple fact that if diagonalizes M, then have TMQ PT P for some permutation matrix P E Op). This gives
f((Q MQ)
Σ f(P T P) Σ f((A))))= Σ f(M(M)
(15)
t=1
Σ f(M(M) F(M),
(16)
showing that the inequality equation 13 is met with equality for Q which diagonalizes M.
We now show that the inequality holds for a general Q O(d, p) such that image C image Q)-
Define W = UT Q E Ruup, Then we have
U
Q
[U
Q 0
W = U --UQ=
Q
(17)
Q
Now we have
W W = Q ----QQQQQQ--
(18)
so that W has orthonormal columns, i.e., W O(d In particular, writing this a different way,
U Q
-- w U,, 0]
- Q,u Q Q) T Q).
(19)
0
so that U. Q E Opp).
Now note that have
QQ MQ):: ((UW) ((U)) W.UUMUU) -WTAWW)
AW) (20)
d = >
(21)
1=1
14

---

Published as a conference paper at ICLR 2025
Since W has orthonormal columns, Σj=1 W), 1. Thus (Q'MQ); is is convex combination of the cigenvalues of M, including some zeros. Using the concavity of f. this yields
P
((QQ)M))))+ Es
W))
(22)
d W (M))
(23)
P
E (()( Σ
W?
Σ {(2 j=1
Σ
w?
(24)
Σ f(A((M) [( Q)
(25)
i=1
4 Σ (()(M))=F(MM)
Σ f(A((M)
(26)
j=1
In particular, [UTO) = because Q is an orthogonal matrix, thus having unit-norm rows and columns. This proves the inequality equation [3]
Regarding the equality case, we note that the only inequality used to derive equation 13 is the concavity condition f. If f is strictly concave, then this inequality is met with equality if and only if, for each i, there exists j such that Σ/= (M))
>y- (M). Namely, the inequality met with equality if and only if for each i, if W f. & and W * 0 then A (M) = A (M). For example, if M has distinet eigenvalues then W is truncated permutation matrix. In general, however, a W with these properties implies that Q = UW diagonalizes M. This shows that if f is strictly concave then the inequality in equation 13 is met with equality if and only if Q diagonalizes M.
B EXTENSION TO CAUSAL ToST ATTENTION
To enable auto-regressive training inference similar to GPT, the ToST attention mechanism illus- trated in the main text needs to be adapted to causal version. Namely, for the t-th token, it can only attend to tokens from earlier timesteps I11).
In conventional self-attention mechanism, such operation can be achieved by applying causal mask on the attention matrix A E IR Txx,, where is the total number of tokens. The causal mask M is a T T-by-T lower-triangular matrix of binary values (i.e. 0 or 1). The resulting causal attention matrix M @ A is therefore also lower-triangular and prevents any token from attending to future tokens.
In ToST attention, such operation however requires some tweaking. Let first revisit the formula- tion of the ToST attention mechanism:
z
= 2j
K
UgDag Vf
(U.
Z)
G2
1) (T 1)
TTA
Vjj E [n].
(27)
To enforce causality between tokens, the key is to make sure for generating token 2/. we should only rely statistics computed from tokens (zi...., 2 Assume we have already computed ITs the causal version of ToST becomes:
2 = Zj
T Σ j
Uea UpDag
VV
( ( 22 02
Li TT LinTT
2g. (
Vje e [n].
(28)
Note IT itself may he causal depending on how it is computed. For now, just assume we have computed IT and we will explain how handle II with more complex initialization scheme the next section.
15

---

Published as a conference paper at ICLR 2025
where we define
J,n = Diag ....,.,.........
(29)
Notice here we have effectively restricted the attention operator to only use the first j tokens, pre- venting influence from future tokens. This mechanism, named Causal-ToST, could be efficiently implemented in practice with cumulative sum (i.e. cumsum) operation. This is an intuitive way to encode causality into our attention mechanism, contrast the XCA opeartor in XCIT, which does not have a direct way of implementing causal attention. We will introduce more technical details on the implementation of this causal variant in the next section.
B. CONNECTION TO OTHER EFFICIENT ARCHITECTURES
It is worth noting that there exist other works on efficient alternatives to self-attention. In particular, recent works on deep state-space models (Gu et al. 2022) or variants of linear attention (Yang al.) 2025) have become increasingly popular and similarly enjoy linear computational complexity. In this subsection, we address connections and distinctions of our work to these methods. First, SSMs and linear attention variants turn out to be deeply related and they can be viewed under unified framework of test-time regression (Wang et al. 2025). From this perspective, these methods can be similarly viewed as algorithmic unrolling: these operators are obtained via optimization (e.g. gradient descent) from some learning objectives. The fundamental differences between ToST and these efficient alternatives thus primarily lie in the modeling assumption and objective functions. SSMs, for example, are primarily derived from the state-space representation modeled by some first-order (continuous) differential equations. This choice entails distinct design choices for SSM architectures such as discretization and efficient convolution operators. On the other hand, ToST is derived from compression- based objective on the space of discrete tokens. Similar self-attention, our resulting operator acts on tokens by aggregating token-level information (but without direct pairwise interaction). Hence, our work doesn't share many design choices of SSMs and is placed much closer to transformers. Consequently, we conduct empirical studies by directly swapping the attention module of standard transformers with our derived ToST attention. Nevertheless, we believe it is interesting direction explore how related these two formulations are, but this is beyond the scope of this paper and we leave a deeper exploration their connections to future work.
C EXPERIMENT DETAILS
Here present additional settings and results for our implementation and experiments.
C.1 IMPLEMENTATION DETAILS
For case of optimization and implementation, below we make a few modifications to the conceptual architecture presented in Section
Overparametterization in TSSA.
First notice we can express the TSSA term in Equation as:
D U,) (iagg(m))
TSSA (
U
Uk]
(30)
D(Z. कं
U U ZDiag(Tk)
where D defined equation 6
Following
et al.
(2024). we replace the term
(T/n) UU,..., Ux] with learnable matrix W E 1 TSSA operator becomes:
Expanding the definition of D, the
m Z) (m ZDiag(m)
TSSA(Z | (U+) -W
(31)
Diag ( k (TK. 1) ZDiag(mx))
Absorbing constant coefficients. We absorb the coefficients d/e² in vf(z) (d/e)) + (d/ in Equation (11) into the trainable parameter (U+)K]- Therefore, the non-linear acti- vation function becomes (x)x) = +)-
16

---

Published as a conference paper at ICLR 2025
Estimating membership II. When estimating II in Equation 10 perform Eg normaliza- tion on the projected tokens U Z R=x" along each feature dimension, for each head. That is, restrict the norm of each row of the projected tokens to be 1. This normalization scheme is only applied to tokens in the computation of and does not impact anything else in Equation (TT). We found this design substantially improve and stabilize the training of ToST models. To express the aforementioned normalization, denote the i-th row of U Z as (U, 2) Then define
1/11(U 1)al
Yk
u/. U Yk-
(32)
1 p)
Then,
I=
(33)
plzn {U+)
|
where
v(z, 1 (U softmax
Vj E [n].
(34)
27
In the causal variant of TSSA operator previously introduced (Appendix B this normalization scheme means that, each token has its own corresponding group assignment matrix II that only relies on earlier tokens for normalization. Denote II IRPX a group assignment matrix corresponding to the j-th token. Define
1/1((T,,n U 1.
v
where (recall)
J Diag( 1...., 1, O,....O). n-j
(35)
[1/1(1,, U ZZ 2..
Then,
L (2) (U+)
m
(36)
(zj (U+) T
U 2, ovl
where v (z, {U+) = softmax
Vj [n].
(37)
2n
U 3 I
Similar to analysis in Appendix B this normalization scheme can be implemented using the cumsum operation. However, one can notice that a naive implementation would require an entirely token arrives time step I, each column A <
different II to be computed every time f - corresponding to earlier time steps needs to be recomputed. This would lead to quadratic computational and memory complexity. practice, instead of recomputing earlier columns, we find that using zero-initialized learnable additive bias along both the head and the sequence dimensions before the softmax operator is sufficient. Denote the bias by for k-th head and j-th token. Then,
MiZi {U*]Ke))
-=
(38)
vizn (UK))
| b2 Gyill bj
where v(z, (U*)K = softmax
Vj E [n].
(39)
2m
| y 122 + by
Therefore, we still have Equation 28 as the causal variant TOST.
We provide PyTorch-style pseudocode in Appendix that reflects the architecture modifications introduced above. In particular, Algorithm T and Algonthm iimplement the TSSA attention layer and its causal variant, respectively.
17

---

Published as a conference paper at ICLR 2025
Table Configurations the transformer models in image classification experiments.
ToST-S
ToST-M
XCIT-S
XCTT-M
VIT-S
VIT-B
purameters
5.8M
22.6M
68.1M
24.9M
80.2M
22.1M
86.6
M
attention heads K
8
8
6
12
layers I
12
12
12
feature dimension
384
512
384
512
384
head dimension
48
48
48
64
64
Table Configurations for ToST models language modeling experiments.
ToST-Base
TOST-Medium
ToST-L.are
parameters
110M
304M
655M
atlention heads
layers I.
12
24
36
feature dimension
768
1280
head dimension
64
64
64
C.2 MODEL CONFIGRAATIONS
In the visual classification experiments conducted in Section 4 we focus on three types of models: the proposed ToST, XCIT, and ViT models. As detailed in the main text, our implementation of ToST inherits choices of other non-attention layers from XCIT. One important modification we make is the removal of Local Patch Interaction (LPI) module, which is designed to encourage explicit between-token interactions and does not align with theoretical derivation. Therefore, we prevent any form explicit interactions among tokens by removing the LPI block. We remove this module for XCTT well for fair comparison. Hence, reported this paper. One notable design choice of XCiT is the addition of two class attention layers after the last XCTT block. The goal of these layers is to aggregate features from the processed patch tokens and then write to learnable [CLS] token via conventional self-attention. These layers perform attention only between the [CLS] token and patch tokens, thus incurring O(n) computational complexity. For VIT models, we directly use the checkpoints and results reported in Touvron et al. (202T). Specifically, we use the ViT-S(mall) and ViT-B(ase) models, as they contain parameters of similar size XCiT-S(mall) and XCIT-M(edium), respectively. To summarize, present details about the configurations of these models in Table
train all ToST and XCIT models
In the causal language modeling experiments, we train the causal variant of ToST of three different sizes and closely follow GPT-2 (Radford et al. 2019) family of models. We present their config- urations in Table Please Appendix C33 for other training details. One thing to note is that we adopt the same hyperparameters as in (Karpathy) (022) for all sizes due to the limitation of computing resources. A more careful tuning certainly could lead to better performance for ToST.
C.3 TRAINING SETUP
Pre-training on ImageNet-lk. We train our models using the AdamW optimizer with a learning rate of 10- for 400 epochs throughout our pre-training experiments. We configure our batch size to be 2048 for all our training experiments. All images are reshaped into resolutions of 224 224 and we tokenize each patch of size 16 16 into patch embeddings. For the other hyperparameters
including data augmentation strategies, we adopt the exact same recipe as in (Ali et al. 2021). Please refer this prior work for details. We conduct all pre- pre-training experiments on 128 NVIDIA V100 GPUs.
Fine-tuning. We conduct transfer learning experiments by using the pretrained TOST, XCIT, and VIT models as initialization and fine-tuning them the following target datasets: CIFAR10/100 Oxford 10owers-102, Oxford-IIIT-Pet. We use batch size of 256 throughout our fine-tuning exper- iments. For other settings, adopt the same hyperparameter setup and pipeline in (Yu et al. 2024). Regarding computational resources, we conducted experiments on two NVIDIA RTX 4090 GPUs.
18

---

Published as a conference paper at ICLR 2025
Table 7: Configurations our ToST models experiments on the Long Range Arena benchmark.
layers
feature dim.
hends
pre-norm
batch size
learning rale
dropont
weight docay
warmup iler.
learning
rate docay
ListOps
6
512
8
2048
True
64
lc-4
0.1
le-2
90000
3000
inverse
1e:2
10000
Retrieval
128
512
True
le-4
0.1
4e-2
91960
10000
linear
Image
8
64
128
True
64
lc-4
0.0
2c-2
180000
9000
liogar
Pathfinder
Table Model ablation. Test accuracy on CIFARIO after 400 epochs for ToST-S base model with various
combinations of &g normalization and non non-attention blocks.
Normalization
Non-Attention Block
Acc.
Yes
MLP
91.2
Yes
ISTA
87.4
No
MLP
85.2
No
ISTA
80.6
Long-Range Arena (LRA) Benchmark. We implement ToST on top of an off-the-shelf PyTocc version of the LRA benchmark and assess ToST on tasks of long list operations (ListOps: Nan- gia & Bowman 2018) byte-level text classification (Text; Maas et al. 201T). document retrieval (Retrieval: Radev et al. 2013), image classification (Image; Krizhevsky et al. (009), Pathfinder (Linsley et 2018),, and its extremely long version (Path-X: Tay et al. 2021). Due to the reported issues with the Pathfinder task, we use different Pytorch implementation of LRA evaluate
ToST on the task of Pathfinder. The hyperparameters of ToST are summarized in Table 7
Causal Language Modeling. We autoregressively train ToST models of sizes corresponding to GPT2-Base, Medium and Large following Radford et al. (2019), by substituting the self-attention operator with proposed causal version of TOST, with the same pre-processing and post- processing steps, on the Wex dataset (Gokaslan et al. 2019). Our implementation is based on Karpathy (2022). We use context length of 1024, and optimize the models using AdamW opti- mizer (Loshchilov & Hutter 2019) for 600,000 iterations, with batch size 480. We set the learning rate 0.0006 with 2000 warm up iterations and cosine decay, weight decay to 0.1. For more details, we refer readers to Radford et al. (2019); Karpathy (2022).
D ADDITIONAL EXPERIMENTAL RESULTS
D.1 ABLATION STUDIES ON MODEL DESIGN CHOICES.
Recall that we have made few design choices for our implementation of ToST as detailed in Ap- pendix C.T Additionally, differently from CRATE (Yu et al.) (2024), we use a generic two-layer MLP instead of the pparsity-promotte ISTA block in our implementation. In particular, we find the following two components to have non-trivial impact on the performance of our proposed model:
1. & normalization when computing the membership matrix II;
2. MLP (as adopted in ViT, XCIT, etc) instead of ISTA (as adopted in CRATE).
We study how these choices affect our model performance and present our observations in Table Specifically, we vary these design choices and train each model variant on CIFAR10 for 400 epochs. Our results demonstrate the necessity of each component in our model design.
httaa://thhbb.ccc//cccoooeeerere /
19

---

Published as a conference paper at ICLR 2025
1968
10.0
Pot Nanay retn
Tea car Meracyy (C1D
1008
0.1
12.0
Namberdd Tokcma
Nuntarr Toleno
Pours Daavcition (4)
Figure 6: Memory usage and inference time of attention operators different transformer architec- tures: VIT (Dosovitslyy et al. 2020), CRATE (Yu et 2023), XCIT (Ali et al., 2021), and ToST (ours). (Left) Each model has K = heads and d = 381 feature dimensions at every layer with varying number of input tokens. (Right) Each model has 4096 tokens and K 8 heads every layer with varying feature dimension d. The y-axes are log-scale.
D.2 ADDITIONAL COMPARISON OF MEMORY AND TIME COMPLEXITY ATTENTION OPERATORS
Beyond the theoretical comparison of the memory and time complexity of attention operators of different architectures in Table we further conduct experiments to provide empirical evidence by evaluating real world memory cost and inference speed measure inference time and memory footprint for 12 attention layers for each of the tested archi- tectures with batch size of 1. For each configuration take warm-up forward pass and then 1000 forward passes to record the average inference time.
NVIDIA H800 GPU. Specifically,
Smaller models with K heads. Figure repoorss the memory and inference time of attention operators of different transformers, where the y-axes are in log scale for better readability. The two subplots the left are the setting shown in Figure where fix feature dimension d = 384 and vary the number n of tokens, while for the two subplots on the right, we fix n 4096 tokens and vary the feature dimension d. Notably, in all cases, ToST's attention operator has much lower computational and memory footprint than the attention of other transformers, which includes the common VIT transformers as well as the reduced-complesee tokens, Figure 6 (Left) shows that ToST is nearly 10 times faster than ViT for inference and uses roughly 100 times less GPU memory. This aligns well with our theoretical comparison in Table [
XCiT transformers. Indeed, for 10k
Larger models with K = 32 heads. We now move to larger models, where we fix K = 32 heads and d 4096 feature dimensions, and vary the number n of tokens. Figure presents the memory and inference time of attention operators of different transformers, where again the y-axes are in log scale. It can be seen that the same trend (as the smaller models case) shows up, and ToST uses lower memory and inference time than ViT, CRATE, and XCIT.
Causal-ToST vs GPT. Here we compare the memory and time complexity of different architec- tures on GPT-like auto-regressive generation (i.e. next token prediction) tasks. Following the config of GPT2-base, we fix K = 12 heads and = 768 feature dimensions, and vary the number n of tokens for both architectures. Figure g presents the memory and inference time of Causal-ToST and GPT-2. Similarly, ToST uses lower memory and inference time than GPT-2.
D.3 ADDITIONAL VISUALIZATION OF MEMBERSHIP MATRIX AND ATTENTION MAPS
We here visualize the membership matrix II learned in ToST-S different layers in Figure 9 as well the [CLS] attention maps from ToST-S, XCIT-S (Ali al. 2021). and ViT-S Touvron jet al., 202I) in Figure 10 for better interpretation of the semantic clustering property ToST, as discussed in Section
20

---

Published as a conference paper at ICLR 2025
VIT
1000.0
VIT
Peak GPU Memary (GB)
CRATE
CRATE
10.0
XCIT
Inference (ms)
XCI
ToST
TOST
100.0
0K
5K
10K
15K
0K
5K
10K
15K
Number Tokens (n)
Number of T'lecs (e)
Figure 7: Memory usage and inference time of attention operators different transformer architectures with varying numbers of input tokens. Each model every layer. The y-axes are in log-scale.
32 heads and 4096 feature dimension
Peak GPU Memory (GB)
GPT
800.0
20.0
Inference Time (ms)
TOST
600.0
15.0
400.0
10.0
5.0
200.0
0.0 OK
0.0
5K
10K
15K
OK
5K
10K
15K
Number of Tokens (n)
Number of Tokens (n)
Figure 8: Causal version of ToST architecture is also faster and uses less memory than standard trans- former architectures (c.g. GPT-2) on auto-regressive generation tasks. We adopt the Base model size for both architectures.
Figure 9: Visualization of membership matrices II from ToST-S, estimated in layer (left) and 9 (right) of each input image.
21

---

Published as a conference paper at ICLR 2025
Figure 10: Comparison of [CLS] token attention map visualization in the penultimate layer from ToST-S (left), XCIT-S (Ali al. 2021) (middle), and VIT-S (Touvron et al. 2021) (right). Each row is different image and each column a different head. Note that the VIT-S model has K 6 heads instead of K 8.
D.4 VISUALIZING LAYER-WIT WISE SUBSPACES IN TSSA:
do not, in practice, strictly enforce the orthogonality of the columns or require the direct sum of the subspaces defined by U|x| to cover the entire Rº space. Instead, we make these projection matrices learnable via backprop for some target tasks (e.g. image clas- sification). We now visualize the learned [K] matrices of different layers in the TSSA block of
As detailed in Section 3.3
trained ToST-S model on ImageNet-Ik k. Figure we first normalize the columns in each U then we visualize [U...... [ [UG.... [%] RPK xpK The ( block in each sub-figure corresponds to (U+) U for i,j [K] at particular layer l. For better visual clarity, we visualize each block by randomly picking 4 directions for each subspace U, We observe that the learned Ufxi are indeed approximately incoherent, which aligns well with our assumptions.
Laysr Covarance Natro
Layer Covertarce Satro
Layer 10 Covoriance Matro
Figure 11: Visualization of learned subspaces in TSSA blocks at different layers.
E
PSEUDOCODE OF TOST ARCHITECTURES
See Algorithms and the following pages.
22

---

Published as a conference paper at ICLR 2025
Algorithm 1: TSSA Attention Layer ToST in PyTorch
class Attention (nn. (
def _inii_ (self, dim, num_heads 8, gky_bias-False) self heads num heads
self.attend = nn. Softmax (dim 1)
self.qkv = nn. Linear (dim, dim, bias-qky_biad bias)
self.temp
www.aaaaaeeeeeeeme (terch. (num_heads, 1))
sclf.to_out nn. Sequential
na. Ln.Linear (dim, dim)
def fseward x)
w rearrange (sk "b (h d) -> b bbhnn b
+ self. heads)
b, h, N, d s.shape
noreeed torch.nn. www.tiinnanaaan..mmmmmmmmmm. www.aooooeeeeeeeeeeee www.iiiin.......mmleeeee... (w, dim=-2) Denoising against union Gaussians
Pi aelf.attend (torch (torch, sum (v_normed ** 2, dim=-1) self.
dcts = torch.matmen ( (Pi / (Pi. (dim=-1, keepdin-True)
le-8)) unsqueeze (-2),
2)
attn 1. (1 dots)
out torch. mul (w.mul (Pi. uisuuuqqqueeee11 (-1)), attn)
out rearrange (out, 'b h n d -> n (h ") return self out (out)
Algorithm 2: Causal Variant of TSSA Attention in PyTorch
class Attention (nn. Module)
def ___niiinsss_____________ (self, dim, num num_heads 8, 1024, gkv_bias-Palse)
self.heads num_heads
seli.attend nn Sn.Sottmax (dim 1)
self.gkv = nn. In.ein (dim, dim, bias=gkv_bian bias) self.temp nn Parameter (torch. ones (num heads, 1))
self.to_out nn. Sequential
Ln. (dim, dim)
self.bias nn. Taaamtterrtorde zeros (n head, max def forward(self, x):
w (ecrrange(se1f (k) n (h bbaddd h + self
b, h, N, d whshape
w_sq -
x_sq_normed (w_sq / terch. cumsum (w_ dim--2))
+ bel
Denoising against union of Gaussians
Pi self.attend (torch. (torch sum (w_ normed, dim=-1) self temp))
dcts torch. cumsum sq Pi unsqueeze (-1), din=-2)
" (Pi. cumsum (dim=-1)) unsqueeze (dim=-))..nsqeeeze attn = 1. / (1 dots)
out = torch. mull (w.mul (P1. uns.ueezeeuuueeeee (-1)), attn)
out rearrange (out, "b h d -> n (h ")
return self.to allf.t_ouut (out)
23

---

Published as a conference paper at ICLR 2025
Algorithm 3: MLP Layer of ToST in PyTorch
class MLP (nn. Module)
def __ini__sse1 in_ hidden_ hidden_featureen -ew.fffeturueoooee Nut_iteatures-Nnee layer-nnn :
super()._init___ init_
out features out features or in features hidden_features hidden_features or in_ self.fc1 mn. mn.ein ti (in_features, hidden_features) self.act act_layer() self.fc2 Linear mnidden_featuree out_features)
def (sel,, x)
self.fcl (x)
self
x self.fe2 (x)
return x
Algorithm 4: Transformer Block of TOST in PyTorch
class Transformer (nn. Module)
def init (self, dim, depth, heads, dim head)
super (i_iti___ii_)__
sclf.layers = nn ModulcList (())
sclf. heads
self.depth depth
self.dim dim
for in range (depth)
self. layers. append append (nn. Mcdulelist (I
PreNorm (dim, Attention (dim, heads heads, dim head
- dim head) )
PreNorm (dim, _ features-dim, (dim + 4.0),
act_layer-nn.ce GELU))
11)
def (srw,, x):
depth = for attn, ff in self. layersay grad attn (x) x x = ff grad_ x
return
24

---

Published as a conference paper at ICLR 2025
F LIMITATIONS AND FUTURE DIRECTIONS
Due to limitations in computational resources and time, we have only tested and compared the new architecture up medium sizes on the ImageNet IK dataset. Although our ToST architecture enjoys higher computational efficiency, which is increasingly crucial at larger scales, it remains to verify in the future whether its accuracy remains competitive against other architectures in very large scale upplications. Finally, the focus of this work is on improving the efficiency and scalability of the attention block of the transformer, so we have kept the MLP block unchanged. Preliminary ablation studies in Table for TOST. We leave the study of designing more effective, efficient, and white-box replacement for the MLP block for future.
that the vanilla ISTA block used in CRATE might not be optimal
25


Reviews:
[{"review_id": "Q34gXFwsLD", "reviewer": "Reviewer_McCu", "summary": "The paper introduces an efficient transformer variant underpinned by a novel linear complexity attention mechanism. The new transformer variant is dubbed the Token Statistics Transformer. The paper builds heavily on previous work, specifically on CRATE (Yu et al., '23), to derive this attention mechanism. \n\nThe paper follows a mathematically principled approach, adapted from CRATE, where a high-level data compression objective is optimized by the network's construction, with each network layer corresponding to an optimization step.\n\nThe high-level objective introduced in this paper is a novel variational variant of the maximal coding rate reduction objective from CRATE. Optimizing this objective analytically leads to the efficient linear-complexity attention mechanism entitled \"token statistics self-attention\". This is pretty neat. The authors also include empirical results validating the new transformer variant on relevant tasks in vision & language.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The paper is very well written; the presentation is top notch.\n- The fact that the method is actually supported mathematically is quite appealing -- this is becoming rare in the age of deep learning.\n- The paper includes experiments on both vision & language modelling tasks, showing that the ToST can be made to work (especially for vision tasks).", "weaknesses": "- Experiments on causal language modelling show that the architecture, in its current form, is not very competitive with standard transformer architectures, as indicated by the results in Table 4 -- approximately tripling the number of parameters yields similar performance to GPT2-Base. Looking at Table 4 (right) this would mean that the running time of the ToST is considerably higher than that of GPT-2 when matched for performance -- not exactly re-assuring. Perhaps follow up works can build upon this and find a competitive variant of ToST.\n\n- Experiments on the Long Range Arena are not exactly like-for-like, as the ToST hyper-params are tuned individually for each task -- to my complete surprise. See Table 7 in line 972 in appendix. So, I'm not sure what the takeaway is from Table 3 -- if the method needs tuning for each specific task that definitely takes quite a bit away from its appeal. Transformers are amazing (in part) because they require little tuning from task to task.\n\n- An in-depth section comparing the resulting architecture with state-space models is surprisingly missing. They have similar properties, in terms of compute complexity, but SSMs can attain much better performance in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "qpL3PbzeX2", "reviewer": "Reviewer_zZ2W", "summary": "The authors present a mathematically-grounded architecture with operators that resemble the components of the popular transformer objective. Building upon previous works on \"white-box\" architecture design, specifically one that interprets each network layer as an optimization step of a maximal coding rate reduction (MCR$^2$) objective, and derive a variational form of the objective leading to linear compute and memory requirements. The authors provide detailed experiments on vision, language and long sequence tasks, and also provide a pseudocode implementation and visualize attention maps and estimated memberships at different layers in their proposed network.", "questions": "See weaknesses above", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This is a fundamentally strong work that builds on the concept of white box transformers, specifically one that proposes a MCR$^2$ objective, leading to alternate operators that fundamentally act as attention and mlp blocks in a conventional transformer layer. This work is a great follow-up to it, for reasons listed below - \n\n- The derived variational form of the objective, backed up by sound theory, leads to a linear compute and memory requiring alternative to transformer.\n- The approach leads to strongest results among transformer architectures on long sequence tasks and has reasonable performance on vision and language tasks\n- The visualisation of attention maps and estimated memberships indicate that minimization of the proposed objective leads to intuitive behaviour of the proposed method.\n- The paper, along with its theory and experimental segments, is very well written, and is easy to follow", "weaknesses": "- An evident weakness of the method in its current form is the low results on various benchmarks. However, the theoretically sound nature of the network and intuitive visualisations are extremely promising, and follow-up works should be able to extend these ideas to construct more empirically strong networks.\n\n- The attention maps produced by the method seem more coarse and consist of high-norm artifacts, an issue also observed in vision transformers but to a lesser extent. In the case of ViTs, some works have suggested [1] having additional tokens as a remedy to this problem, as certain tokens start behaving as global tokens, leading to high-norm artifacts in the attention maps. It would be good to understand the authors' view on this issue in their approach.\n\n- The results on vision and language tasks do not mention the compute and memory savings, it would be useful to have those, when comparing against other baselines.\n\n[1] Vision Transformers Need Registers, Darcet et. al, https://arxiv.org/abs/2309.16588", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "7xJFYTWTHI", "reviewer": "Reviewer_SUwu", "summary": "This paper proposes a new attention mechanism that is derived from optimizing general MCR2-like objective. In particular, the derived architecture performs (approximate) low-rank sub-space projection in each attention head. Due to elimination of the classical pairwise attention, the proposed architecture has linear computational complexity. Aside from theoretical derivation, the authors have performed extensive empirical studies to evaluate the performance of the proposed architecture.", "questions": "Please see comments above.\n\nAdditionally, in order to achieve linear complexity, the proposed architecture gets rid of pairwise token interactions and instead creates indirect weak interactions via sample level statistics between tokens. This however comes at a cost of accuracy degradation. This makes me wonder how much such indirect interactions actually contributes to the final accuracy anyway. In particular, what happens if we simply replace $\\pi$ with uniform probability distribution? In other words, is there any significant accuracy benefit to indirect interaction vs no interaction at all?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1- This paper proposes an interesting optimization-based and white-box design perspectives in designing attention-based architectures. I can see the proposed methodology can be used to derive similar attention architectures beyond what has been proposed in this paper, which is an important contribution.\n\n2- The theoretical derivation is sound and quite clear.\n\n3- The empirical studies reported in the paper along with ablation studies are quite comprehensive. It's specially commendable that the authors have covered a variety of tasks and real-world datasets.\n\n4- The paper is very well-written and clear with its flow.", "weaknesses": "1- It seems the low-dimensionality of subspace within each head is the key to the performance benefit of the proposed framework. However, such gain may come at some costs that have not been properly discussed in the paper. In particular:\n\na) There is no guarantee (or soft encouragement) in the proposed framework that the direct sum of the subspaces in each layer should cover the entire (or most of) original space. This can potentially lead to generalization issues specially for out-of-sample test datapoints, where the test datapoints may have small components within the learned subspaces.\n\nb) Fine-tuning might be more challenging as it may require learning new supspaces that are orthogonal to the pre-trained subspaces, which is quite challenging.\n\nTo further understand the potential compromises here, I'd recommend authors report further empirical results comparing the zero-shot, out-of-sample performance of TOST with that of regular transformers as well as their convergence rates for fine-tuning. \n\n2- The idea of suppressing less pronounced directions within a set of token vectors may remove discriminating information that might be otherwise crucial to the end task. For instance, in some language tasks, the presence of certain (outlier) words in the input text can be quite decisive for the classification task, but since such token vectors are somewhat outliers, they can be suppressed by the proposed architecture; whereas, in classical pairwise attention mechanism, such intricacies can be organically captured by pairwise attention weights. Since this can pose a significant limitation of the proposed framework for tasks with such characteristics, I'd recommend further empirical studies on such tasks and comparing its results to that of the standard transformers. For instance, one such task can be in-context retrieval for long context sizes in language, e.g. the Needle in A Haystack benchmark. It'd be also quite insightful to visualize rare, yet informative tokens across different layers.\n\n3- Please report the model sizes for all model-dataset pairs in Table 3 (LRA results). It's not clear how much of the difference might be due to having different number of parameters, which wouldn't be a fair comparison.\n\n4- The current work has not been properly placed in related literature. In addition to novel attention operators, I believe the current work is highly related to various formulations of linear attention, but such comparisons are largely missing in the current manuscript. I'd recommend the authors to include a Related Work section as well as adding one or two linear attention baselines (of the same size) in other experiments than LRA.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "B6cdb7CXQj", "reviewer": "Reviewer_VpkN", "summary": "This paper introduces Token Statistics Self-Attention (TSSA), a novel attention mechanism derived from a variational form of the Maximal Coding Rate Reduction objective. TSSA exhibits linear computational and memory complexity. Furthermore TSSA significantly departs from traditional pairwise similarity-based attention. By replacing standard self-attention with TSSA in transformer models, the resulting transformer achieves competitive performance on various tasks in terms of expressibility, while maintaining linear complexity.", "questions": "- I'm curious about the results in Figure 5. Could you provide some intuition on why you attention does such a good job on isolating the attention on the main object? This isolation could result into improved robustness over the standard attention, which is worth investigating.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The authors propose a novel attention mechanism, which is well justified.\n- Their attention mechanism has a time and memory complexity of $O(N)$.\n- Based on the provided results, their attention mechanism seems to do a better job than standard attention on focusing on the correct pixels.", "weaknesses": "- Since the main motivation behind the design is improved time and memory complexity, the comparisons should've been made with FlashAttention. Given that FlashAttention is calculating the softmax-based attention (regular attention) in an efficient manner, and is prevalent (arguably more prevalent the standard attention), it is not fair to compare your time and memory with standard attention. (For reference, FlashAttention has a memory complexity of O(N), and has demonstrated a 2-4$\\times$ speedup for context lengths of 4-8K [1]. Therefore, a GPT2 model equipped with FlashAttention would probably outperform your design for the experiment outlined in Table 4-right.)\n\n[1] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Spotlight)