Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few
Authors: Qishuai Wen
Zhiyuan Huang
Chun-Guang Li
Venue: NeurIPS 2025 spotlight ()
Keywords: Attention Mechanisms
Interpretability
Efficiency
Primary area: deep_learning
TL;DR: We propose a unified optimization objective and thus derive a novel interpretable and efficient self-attention mechanism, named Contract-and-Broadcast Self-Attention (CBSA).

Abstract: Attention mechanisms have achieved significant empirical success in multiple fields, but their underlying optimization objectives remain unclear yet. Moreover, the quadratic complexity of self-attention has become increasingly prohibitive. Although interpretability and efficiency are two mutually reinforcing pursuits, prior work typically investigates them separately. In this paper, we propose a unified optimization objective that derives inherently interpretable and efficient attention mechanisms through algorithm unrolling. Precisely, we construct a gradient step of the proposed objective with a set of forward-pass operations of our \emph{Contract-and-Broadcast Self-Attention} (CBSA), which compresses input tokens towards low-dimensional structures by contracting a few representatives of them. This novel mechanism can not only scale linearly by fixing the number of representatives, but also covers the instantiations of varied attention mechanisms when using different sets of representatives. We conduct extensive experiments to demonstrate comparable performance and superior advantages over black-box attention mechanisms on visual tasks. Our work sheds light on the integration of interpretability and efficiency, as well as the unified formula of attention mechanisms. Code is available at \href{https://github.com/QishuaiWen/CBSA}{this https URL}.

Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few
Qishuai Wen, Zhiyuan Huang, and Chun-Guang Li
School of Artificial Intelligence,
Beijing University of Posts and Telecommunications, Beijing 100876, P.R. China {uqs, huangzhiyuan, lichunguang; hbw.hedd...up.uu...u...... edw.. edu. cn
Abstract
Attention mechanisms have achieved significant empirical success multiple fields, but their underlying optimization objectives remain unclear yet. Moreover, the quadratic complexity of self-attention has become increasingly prohibitive. Although interpretability and efficiency two mutually reinforeing pursuits, prior work typically investigates them separately. In this paper. we propose a unified optimization objective that derives inherently interpretable and efficient attention mechanisms through algorithm unrolling. Precisely, we construct gradient step of the proposed objective with set of forward-pass operations of Contract- and-Broadcat Self-Artention (CBSA), which compresses input tokens towards low-dimensional structures by contracting few representatives of them. This novel mechanism can not only scale linearly by fixing the number of representa- tives, but also covers the instantiations of varied attention mechanisms when using different sets of representatives. We conduct extensive experiments to demonstrate comparable performance and superior advantages over black-bon attention mecha- nisms on visual tasks. Our work sheds light on the integration of interpretability and efficiency, as well the unified formula of attention mechanisms. Code available at this https URL
Introduction
Attention mechanisms have been widely applied across diverse areas, including computer vision I 2),
natural language processing [ 4, and scientific discovery [5). Nonetheless, series of puzzling phenomena- such as emergent segmentation properties [6]. in-context learning ability [7], attention collapse [ 9 and extreme- token phenomena (10) have been uncovered in them, hindering the principled and trustworthy development. At the same time, the quadratic computational and memory complexity of self-attention with respect the sequence length impedes its broader applications in real-time systems as well as the processing of long documents 112 and high-resolution
images 13
In light of these challenges, has been more crucial to mathematically demystify attention mech- anisms, which offers deeper insights into their simplification and acceleration. Over the past few years, remarkable advances have been made addressing the interpretability or efficiency issue separately. On the one hand, in ante-hoc manner, attention mechanisms be interpreted by optimization objectives grounded clustering [14], denoising [3], energy minimization [16], matrix decomposition (17), and contrastive learning [18). These inherently interpretable approaches more rigorous than post-hoc explanations [19]. On the other hand, numerous techniques have been developed alleviate the quadratic complexity of self-attention, including sparse attention (0) and
linear attention 21
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

However, the joint development of interpretability and efficiency in attention mechanisms remains largely unexplored area of research. This leaves the design of efficient attention mostly heuristic, and the interpretations and explanations for attention mechanisms less instructive. To bridge this gap, we formulate a unified optimization objective by mildly modifying a compression-dre objective called MCR? (22]. Indeed, this objective has been utilized for designing an interpretable softmax attention, MSSA [23], and a linear-time attention, TSSA 124]. But MSSA also scales quadratically, and TSSA is effectively a channel attention mechanism, which contrasts sharply with both softmax attention (token mixer) and linear attention (channel mixer) Therefore, instead of an isolated mechanism, in an interpretable way, revealing how they are fundamentally connected yet distinctly presented, as well as the trade-off between expressive capacity and efficiency.
optimization
aim to develop framework that unifies these varied attention mechanisms
In this paper, we adopt two ante-hoc interpretations to constitute our proposed optimization objective: a) input tokens are compressed towards low-dimensional structures for compact and structured representation; and b) the geometry and information-thhoretii essence of input tokens can be captured by small number of representatives [25] 266 of them. Since the former has been formulated as the
MCR² objective [22 23 (see Section 2). the remaining task is to leverage the representatives to optimize it, thereby efficiently compressing all by contracting few (see Section [.T]. By unrolling the resulting optimization objective, we derive our Contract-anu santract-and-rroaddee
Self- Belf-Attention (CBSA), which contracts the representatives and broadcasts the contractions back to input tokens (see Section 3.2).
Given a fixed number of representatives, the computational and memory complexity of CBSA scales linearly with the number of input tokens. Moreover, CBSA covers the instantiations of varied attention mechanisms, including softmax attention, linear attention, and channel attention, by taking different sets of representatives (see Section 3.3). As a result, CBSA serves as a unified formula for these attention mechanisms, and attributes their differences to their distinct information propagation (more precisely, compression) patterns induced by the different number and structure of representatives.
Detpet Tokeu
Auh stecttic Lesi, keer
Harwan---ee Pask Objeccives MCR
ISTA
fr
Breedoser
Layerhorm
Conmaciion
BC Liais
Centraci-eaa--ee eaal--aad--- Cettaad-eaal-Buaaaao
Self-Alleation
of
LepoNNam
Tepur Tokuns
Concinseaal Enbeeding
KO COSA
Softia ABUESI
Iapel Innac
Figure 1: Overview of CBSA. Left panel: Besides projecting tokens onto subspaces and back the ambient space, there generally two stages in CBSA: 1) representative initialization and extraction; 2) representative contraction and contraction broadcast. The former extracts representatives satisfying the inequality constraints in ( while the latter is a gradient step of the compression term in (4). Right panel: CBSA covers instantiations of varied attention mechanisms. Their compression patterns are distinct as illustrated above. Further analysis is elaborated in Section 3.3
Paper contributions
The contributions of the paper are summarized as follows.
1. We formulate an optimization objective that unifies the interpretability and efficiency of attention mechanisms through the idea of compressing all by contracting a few.
2. We derive an inherently interpretable and efficient attention mechanism, CBSA, which is a potential unified formula for different attention mechanisms.
3. We validate the interpretability and efficiency of CBSA through extensive experiments on visual tasks.
'Softmax attention calculates all possible pairwise similarities between tokens, while TSSA just scales feature channels according their second moments; see [13

---

2 Notations and preliminaries
Notations. Given a positive integer n, let [n] = [1, For s 2 n, let (s,,n) R denote the set of s matrices with orthonormal columns, and O(s) (ss,s) denote the set of s x orthogonal matrices. Let I, denote an identity matrix of size n, and O2 denote a zero square matrix of size n. Given a vector E R", let (iag(v)> E R=** be a diagonal matrix with the entries of v along its diagonal. Let 2 E d denote N input tokens represented in the ambient space I&d Specially, let Z' denote these tokens feeding into the f-th attention layer. The same holds for the representatives of input tokens, Q E __ where m can be much smaller than N.
Union of subspaces. Although the union of nonlinear manifolds provides a better approximation (77. adopt much simpler structure: the union of (ow-dimensiona) linear subspaces Specifically, it is parameterized as K incoherent p-dimensional subspaces spanned by orthonormal bases [[K] = A (d,, p) where pK d and U. U Op vi + j. We refer to the i-th basis vector of U as Uki. Similar to [23] we implement U [K] as learnable parameters each attention layer and thus denote the U [K] implemented by the ε-th layer as UÍK |k]
Coding rate. To quantify the compactness, i.c., the extent to which tokens are compressed towards subspaces, adopt the (lossy) coding rale (28). which measures how efficiently the token distribution can be covered by e-balls under a given quantization precision e >> 0 (as illustrated in Fig. Ta)). The coding rate of input tokens in the ambient space Fd defined as:
d
R =
log del IN
Z z Ne
MCR objective. The Maximal Coding Rate Reduction (MCR?) defined on the coding rate as follows:
22
objective adopted in
23 is
K RU 2) A| lo- (2)
- All max AR(Z) R(Z) - (Z (
R(Z)
k=1
expansion
compression
sparsity
The input tokens are compressed towards K subspaces by the compression term, while expanded in the ambient space by the expansion term to avoid collapse, yielding compact and structured representation [22] Yu et al. (23] have demonstrated that the approximated gradient step of the
compression term corresponds an interpretable softmax attention mechanism.
Methods
3.1
Compressing all by contracting few
Due to the existence of Gram matrix in (T). the attention mechanism derived from 2 which is called Multi-head Subspace Self-Attention (MSSA), inevitably scales quadratically with the number of input tokens. While previous work has bypassed this issue by replacing the Gram matrix with the covariance matrix
(99 and further introduced variational formulation (24), these strategies degenerate the token mixer into a channel mixer and channel attention, respectively.
In this paper, inspired by the concept of landmarks [10,3], propose simple but flexible approach to streamline the optimization MCR2: compressing all input tokens by contructing small number of representatives of them. Before demonstrating that this achieves linear complexity in N and prevents the aforementioned degeneration, we first formulate
new optimization objective.
An initial attempt is to impose a set of equality constraints on the coding rates as follows:
K
max R(Z) Z
Σ - s.t. RÚ Q) = R(U 2), 4k E [K],
(3)
1=1
"As this structure faces challenges in adapting to all modalities and tasks, we excluded natural language processing tasks from our experiments. But we argue that it serves a feasible starting point for finer mnainer structures; further discussion available our OpenReview page.
3

---

where representatives Q = g(Z) are extracted from input tokens Z by differential function (() : RdxN + Rdxm This new objective in (3) is equivalent to the original objective in (2) but more efficient to handle because the number of representatives (c.g. m = p d/k) can be far smaller.
Since that the equality constraints in is overly restrictive in practice, attempt to relax them by introducing a tolerance т, which uniformly bounds the absolute difference of the two coding rates within each subspace, i.e.. ? Q) R Z) 3 т. Therefore, contracting the representatives will correspondingly compress the imput tokens well up to the tolerance т.
Consequently, we have relaxed optimization problem for our subsequent derivations as follows:
K Σ R Q) >/ s.t. |RU - R 2) [kE]]
max 1(2) 2
ST
[K]. (4)
In this paper, we employ an arguably simplest way to extract Q in each subspace: Q - U A where A₁ E Nx is the coefficient matrix over dictionary U Z E RRxX, i.e., the projected representatives in each subspace are linear combinations of the projected input tokens.
3.2 Contract-and-Broaddar Self-Attention
We now are ready to derive an attention mechanism in an ante-hoc interpretable manner, by im- plementing a gradient step of the compression term in the proposed objective as its forward pass. This methodology dates back its origins to the pioneering work [32] and is referred to as algorithm unrolling or unfolding (33]
Representative initialization and extraction. Inspired by the fact that cross-attention can be interpreted approximating the coding rate (34), employ this idea extract representatives that satisfy the inequality constraints in (4), thus capturing the information information-theoretie essence of input tokens. Specifically, we take an initial guess of the representatives as the query, and the input tokens as the key and value matrices, i.e.,
UTQ =.QUUk U Z softmax (U (U Qmi)) vk E [K],
(5)
Ax
where the initial guess Q is treated as a constant with respect to Z such that its strategy (whether input-dependent or not) does not affect the subsequent derivation. To be more specific, following (35). we initialize Qini via an average pooling over the input tokens; see Appendix A for discussion. More importantly, the representatives extracted by this cross-attention operation are linear combinations of the input tokens, thus naturally leading to the form we desire. Therefore, the attention matrix in effectively is the coefficient matrix A
Representative contraction and contraction broadcast. To derive the attention mechanism, following 233, we focus on optimizing the compression term via gradient descent step. Having the extracted representatives satisfying the inequality constraints, we take a gradient descent step on the compression term of the objective function in with respect to input tokens as follows:
Z -- - CBSA(Z Z
UKK)
where
(6)
K
CBSA(Z U (KK
Σ
EUAIIQ (Im + P (UQQ) me?
A K
(7)
k=1
Contraction
Brendcast
in which the step size parameter is learnable in our implementation. One can verify that 7 is proportional to the gradient of the compression term in (4) with respect to input tokens, while the contraction term 7 proportional to the gradient with respect to representatives. Hence, we refer to this formula as a Contract-and-Be aaniraact-nn-BBradasss Self-Attention (CBSA), reflecting that: a) the contraction term gives the contracting directions of the representatives (abbreviated as contractions);
'The representatives are not necessarily a subset of the input tokens: they resemble the cluster centroids in k-means, which continuously extracted, rather than discretely selected.

---

b) the broadcast term, which reuses the attention matrix in (5) broadcasts the contractions back to all input tokens.
Contraction via self-attention. To avoid computing the expensive matrix inverse in 7 similar to approximate it by Gram matrix and softmax function (i.e., attention matrix)
23
CBSA(Z U()) =
k Q softmax (UT (U Q)
A
(8)
Contraction via self atteettio
Broadcast
Note that the contraction term now effectively constitutes a self-attention operation in which the linear projections for the query, key, and value all identical the subspace basis, i.e., query
Wkey =
Valuee
U By default, we implement CBSA via rather than ( in our experiments.
Overview of CBSA. The workflow of CBSA is illustrated in the left panel of
Fig. We also construct an inherently interpretable
Transformer (CBT) by stacking CBSA with the ISTA module [23]. which is also derived via algorithm unrolling We report the computational complex- ities of CBSA and its sub-operations in Table [ and compare the FLOPs of different attention mechanisms in Fig. 2 where d = 384, H 6, and a patch size of 16 x 16. Provided N > 2d/H = 2p (typically 128 in Transformers), the FLOPs of CBSA are lower than those of MSSA. Further comparisons with other modules, including MHSA and MLP, are provided in Fig.
201
12
201
Figure 2: Computation complexity.
Table 1: Computational complexities. By default,
p 4/H1, where H K denotes the number of attention heads (interpreted as subspaces in our case). The complexity of each sub-operation is computed by summing the costs across all heads. is worth noting that the projection operations, which are essential to almost all attention mechanisms, confine the overall complexity at least O(dd").
operations of CBSA
Q(MSSA)
((CBSA)
(extraccion) / (broadecs))) (econtraction)
(pprojoction)
2Nd +Ndd
2Nd + 3N 2m²d
Nmd
m d
Nd
A
Marual
Maneii
Matzaa pic Dass
Matual
Matural
2
Seftrow
Morrral
Softnas
1
Manal
Mitural
Marral
Matroa
Monred
Mattral
Manrrd
Marral
Marrral
Worr
Whe
Watee
Weorr
Wher
Wabe
Pooting
Peetze
(4) Sodimax Atteniis
(b) Linze Alication
(6 Channel Alleation
(6) Ageni Atteniius
CBSA (curs)
Figure 3: Different attention mechanisms. The mapping function and agent hias in the diagram of Agent stands for the matrix transpose.
Attention (35 are omitted for simplicity, and"
3.3 CBSA as a unified attention formula
In this subsection, explore the potential of CBSA to serve unified formula for different attention mechanisms. Unlike recent work 37 38). CBSA encompasses a broader spectrum of mechanisms (see Fig. B} in interpretable and mathematically grounded manner.
Our analysis reveals that, CBSA 7 can derive multiple variants corresponding to existing attention mechanisms by varying the choice of representatives. In these variants, the distinct initialization, extraction, contraction, and broadcast steps of CBSA may not be explicitly observed, as some of them are simplified or fused together due to the specific number and structure of representatives, or engineering concerns.
Softmax attention variant. Obviously, the input tokens themselves satisfy the constraints in 4 and thus can be directly used as the representatives, i.e., Q Z and m N. In this case, we call the
"The scaling factor and sign inversion 36) introduced by this approximation are absorbed into the learnable step size x for notational simplicity.
"Note that the ISTA module is designed to unroll the suitably-relaxed proximal gradient step to address the difference of the sparsity penalty and the expansion term A Z R(( Z).
5

---

input tokens are self-expressed [39], where data samples are linearly represented over a dictionary composed of themselves. Then we can take trivial solution for the regularization term where all coefficient matrices are identity matrices. Substituting them into (8) yields the following operator, known as MSSA in the white-box transformer 23
M U (K]
Σ U U
softmax (UZ2))
Wain vas softmax 2) (W 21)
Linear attention variant. To analyze the case of orthogonal representatives, we start with a canonical choice: the principal directions of input tokens. We thus perform singular value decomposition (SVD) within each subspace:
=
Ak [K],
(10)
where L3 E Ov)), RS E ((N, p). Σ is a x p diagonal matrix of singular values, and the columns of L are known as the principal directions. Then, by right multiplying both sides by Rg, we have:
U ZR, = L, = L, vk E [K].
(11)
By comparing 11 with the way to form the representatives, i.e., U U U ZA.. we let U. Q L& and Aş Rg. Substituting them into leads to the following operator:
K J (U?2) k=1
-1 L L (12)
Σ
L
I'm +
V.s. Wistee 20 Who 2)
where F is function defined the spectrum of a positive semi-definite matrix and applies f(A:)
+(()+A,) to each eigenvalues {>i]== of the covariance matrix This operator highly resembles the linear attention 21] due to that it also factorizes the N x N attention matrix and multiplies the key and value first to linearize the computational complexity In Appendix
B we prove that a similar
result holds for any set of orthogonal representatives.
Channel attention variant. Assuming that the basis vectors of U are the principal directions for any set of input tokens (which is impossible but simplifies the computation), the directions of the representatives can be fixed along these basis vectors, i.e., U = Lx, thereby being orthogonal and input-agnostic (fixed). Then, 12
is simplified to:
Σ UUD.UUT Z. where D = Diag (If( ()
(13)
which basically recovers TSSA 24 In 13 the feature channels are adaptively scaled according
to their second moments of token projections, while channel attention typical employs an MLP to predict the channel-wiss wise scaling factors 41
Agent attention variant. Agent Attention is basically a variant of CBSA with the contraction step removed which can be perceived in Fig. 3 Although this removal appears confine the token-mixing ability, it is compensated by the pooling-basee based initialization, which is also token mixer 42
To gain some intuition of the gap in expressive capacity among the aforementioned mechanisms, illustrate their compression patterns in the right panel of Fig. [ The channel attention variant is restricted to compressing input tokens along fixed axes parametenzee by U KK) whereas the linear attention variant compresses them along principal directions that dynamically determined by the input. We argue that such a dynamism is crucial for in-context learning 7 and for mitigating
As I is monotonically decreasing, the effect of 12 is to preserve the principle directions (i.e.. representa- tives) with large variance while suppressing the other directions with vanishing variance (29).
In fact, 12 is less expressive than black-hox lincar attention, hecause (U. (U : 2) is symmetric hut (W value Z) ( Waey Z) is generally not.
"Sine the query, key, and value are identical in CBSA, the broadcast term can be obtained directly from the extraction step, eliminating the separate computation branch in Agent Attention.

---

superposition 433 In contrast, softmax attention exhibits much greater flexibility, as it manipulates each token independently. Actually, its compression can be viewed as operating in an N-dimensional space, rather than in the d-dimensional feature space. Our proposed CBSA aims to approximate the behavior of softmax attention while significantly reducing computational cost.
The above findings can also be interpreted from dictionary learning perspective, where the rep- resentatives correspond the atoms of dictionary. When the representatives are orthogonal and fixed, they form a complete dictionary; when they are orthogonal yet input-dependent, they resemble a submatrix of an overcomplete dictionary as in compressed sensing 44 When the input tokens themselves serve representatives, they constitute a self-expressivee dictionary [39].
4 Experiments
In this section, we evaluate the interpretability and the efficiency of the proposed CBSA and the CBTs built upon CBSA. As natural images often lie on low-dimensional subspaces 44 46], we focus on classical visual tasks such as image classification and semantic segmentation, where higher
resolutions generally lead to better accuracy 47
Baseline and training configuration. We compare our CBSA to the vanilla softmax attention 44
D. and interpretable attention mechanisms based on MCR2, e.g., CRATE (23), ToST [24] and DEPICT [34]. Table 2 summarizes the baselines with brief descriptions. The results in gray are cited directly from the corresponding papers; whereas the others reproduced under varied settings
for fair comparisons. By default, the training configuration follows the baselines, with detailed information provided in Appendix
Implementation detail. The projection back to the ambient space, which should theoretically be a left multiplication by Uk- is over- over-parameterized with independently learnable matrix. This strategy is also adopted MSSA [23] and TSSA [24). and its effect has been analyzed in [36]. In short, although this relaxation compromises the theoretical rigour, it is crucial for achieving better accuracy. In addition, the step size in (6) is implemented as a learnable parameter without constraining its sign. This allows the model flexibly choose between compression and decompression. The PyTorch implementation is provided in Appendix D
Table 2: Summary of baselines. Note that these methods are not limited to the tasks listed here, our descriptions only indicate their usages in the experiments of this paper.
Attention Mechanism
Complexity
Interpretable
Tasks
VIT 0
MHSA (sofimax attention)
quadratic
X
image classification
CRATE
MSSA (softmax attention)
quadratic
image classification
ToST 24
TSSA (channel attention)
lincar
V
image classification
Agent Attention
Agent Attention attention)
linear
imuge classification
Segmenter 50
MHSA
quadratic
X
semantic segmentation
DEPICT
MSSA
quadratic
semantic segmentation
4.1 Advantages enabled by interpretability
In this subsection, we show superior advantages of CBSA black-box attention mechanisms. The most essential aspect of our CBSA is its interpretability. which induces other desirable properties such robustness and emergent segmentation.
(a) Raw
(b) 1st ser
(c) 256 iters
(d) 512 iters
(c) 640 sters
(f) 768 iters
(g) 896 ites
(h) 1024 hees
Figure 4: Compact and structured representation. There 10 classes indicated by colors. Points in each
class are generated by sampling from a one-dimensional subspace and are then perturbed by adding noise.
Compact and structured representation. Similar to MCR2 2 22 our optimization objective aims learn compact and structured representation by compressing input tokens towards low-dimensional subspaces. To confirm whether its iterative gradient steps actually achieve this goal, we iterates
7

---

the linear attention variant of CBSA 12 on synthetic data, where image tokens are modeled as the points siree-dimensional space R. Specifically, it conducted on each class in the ambient space (thus being parameter-free) with forward-only manner. As shown in Fig. 4 the representation ultimately admits union of well-separated one-dimensional subspaces after 1024 iterations.
Codling
(ste
Horaaizdd Coding rodelo
Harma ladd Cording tonal medeld
1000
15%
600
Figure 5: Evaluation on compression effect. We measure the compression term a function of layers.
Compressing all by contracting a few. To confirm our interpretation that CBSA compresses all input tokens by contracting few representatives, we check that: whether the input tokens are indeed compressed; if so, whether the compression is driven by contracting the representatives. We measure the compression term of (2 as well as its normalized varian P in Fig. 5 This normalized coding rate is invariant to in-place scaling and depends the angles between input tokens. We observe two pieces of supporting evidences: a) the more compact the ultimate representation is, i.e., the compression term measured in the last layer is lower, the better the model performs on ImageNet-IK (see Table 3 10 the latter half of the layers exhibit consecutive compression, in nearly all models Then, in Fig. measure the reduced coding rate of the input tokens and representatives, respectively, after they are processed by CBSA within each subspace. We observe that the two kinds of reduced coding rates show highly similar trends across most subspaces
Requree cntegrote ccouy layer heads
Aaduced coding rats acues byyen and baack
Figure 6: Comparison reduced coding rates. The reduction with respect to input tokens is calculated between input tokens and compressed tokens, while the reduction with respect to representatives is calculated between extracted representatives and contracted representatives. We measure the reduced coding rate with respect to the input tokens and the representatives, respectively, across all heads of a model. Results of CBT- T/S are presented here, those of CBT-B/L can be found in Fig. I0 Note that we set = 1 to exclude the de-compression observed in Fig.
CUTOCE
CRATE
CRATE CUT
Figure 7: Visualization of [CLS] attention map. We empirically estimate the full attention matrix by ALA E R NxN and visualize the attention maps of the [CLS] token from the early, middle, and late layers of CBT, CRATE, and their hybrid model, respectively.
Emergent segmentation properties. It has been reported that segmentation properties emerge in CRATE with merely standard supervised classification training owing to MSSA (51]. Compared to CRATE, our CBT attends to more semantically meaningful regions in the early layers, but the segmentation properties fail to persist in subsequent layers. shown in Fig. 7 To address this
9That is. the input tokens are normalized to unit vector before measuring their coding rate. Noot that this correlation does not hold for the normalized coding rate, indicating that compression in magnitude is a critical asppct.
We regard the decompression phenomenon in the early layers "known unknown" requiring further investigation.
We hypothesis that the higher coding rate of the representatives makes them act as "scalpels" in order to give "surgery" the input tokens.
8

---

phenomenon, we construct a hybrid model, termed CRATE+CBT, where the first half of the attention layers employ MSSA and the latter half employs CBSA. In this hybrid model, we observed that the segmentation properties not only emerged in the very first layer, but are also progressively enhanced in the following layers, rather than fading as in CBT. Qualitative results supporting this conclusion can be found in Appendix C
Semmnttic segrontaation
$3.2
43.0
CBT Izars) DCTCT
42.0
OIT lter
100
10-1
1241
GRLOGs (dessder-sns) stis
Molee lesl
No lwat
Figure 8: Semantic segmentation on ADE20K. All results evaluated on the ADE20K validation set. CBSA indicates that the CBSA layers are implemented rigorously, ie.. without the overparameterization trick. Middle and Right: Random Gaussian noises with zero mean and standard deviation are independently added each parameter of the attention layers in the decoder.
Robustness against parameter perturbation, As the attention heads of CBSA are modeled as low-dimensional subspaces, perturbing their projection matrices (i.e., subspace bases) with relatively small noise does not significantly alter the subspaces they span (34). Consequently, as shown in the right two panels of Fig. 8. CBSA is extremely robust against parameter perturbation, whereas the black-box method (i.e., Segmenter (50) collapses under the same perturbation.
Table 3: Classification on ImageNet-lk. Models are trained on images of resolution 224 x 224 with patch size of 16 16 where CBT and ViT use convolutional embedding layers but CRATE uses linear embeddings.
Datasets
CBT-T(iny)
CBT-S(mall)
CBT-B(ase)
CBT-L(arge)
CRATE-B
CRATE-L
VIT-S
# parameters FLOPS
18M
6.7M
25.7M
83.1M
22.8
77.6M
22.1M
L.G
4.0G
15.1G
47.3G
12.60
43.3G
9.80
ImageNet-IK
63.2
71.4
73.4
74.4
70.8
71.3
72.4
CIFARIO
94.8
96.3
96.7
97.3
96.8
97.2
97.2
CIFAR100
76.5
80.4
82.0
83.4
82.7
83.6
83.2
Oxford Flowers- 102
88.4
91.7
93.6
93.9
88.7
88.3
88.5
Oxford-IITT-Pts
86.8
91.6
92.6
92.9
85.3
87.4
88.6
Table 4: Fair comparisons on ImageNet-1K. All models employ convolutional embedding layers and ISTA feedforward blocks. The only difference lies in the attention mechanism, where A and Agent-S 3SS] implemented as CBSA but without contraction step (as derived Section 3.3).
ImageNet-1e IK
CBSA-T
CBSA-S
TSSA-T
ISSA-S
MSSA-T
MSSA-S
Agent-T
Agent-S
pairwise similarities
0.53M
1.M
0.45M
0.91M
1.4M
2.8M
0,52M
1.0M
Top-1 accuracy
63.2
71.4
61.2
68.5
64.7
72.1
63.8
71.8
mage resoiution: 224x224
Image resolution: 384x284
Acusic 264
Aucra 2601
stans
went
30
Losr Mirtin
Lrear Etrttin
Jan
325
1.000
4000
200
276
3,20
4006
203 O1LOPS
GPLOPS
Number J0horr Image reselation
Numb bolkenta imag: rcsolation
Figure 9: Adapting pre-trained VITs into CBSA style. We finetune both the adapted models and the original ViT's on ImageNet- 1k for 50 epochs, and report the top-1 accuracy on the validation set. CBSA" leverages the three distinct projection matrices inherited from the pretrained MHSA to calculate query, key, and value, instead of using single projection matrix in (8). CBSA refers to the CBSA" without the contraction step, which is essentially an Agent attention.
4.2 Evaluations on real-world visual tasks
We pretrain CBT models on the ImageNet-Ik 1k dataset, and finetune them several downstream datasets. The top-1 accuracy on validation sets is reported in Table particular, our CBT-Small achieves comparable t accuracy to ViT-S using only 30% of the parameters and 40% of the

---

FLOPs. Compared to CRATE models, our CBTs (with convolutional embedding layers) perform remarkably better while using fewer parameters and FLOPs.
We also conduct a set of fair comparisons across different attention mechanisms, which can be regarded as variants of CBSA, and report the results in Table 4 In this setting, CBT models remain competitive with CRATE while computing significantly fewer pairwise similarities. These results confirm that, from MSSA CBSA and then to TSSA, the number of pairwise similarity computations decreases at the cost of performance sacrifice. In addition, we present the throughput comparisons on nigh-resolution images (e.g., 512 x 512) Table 7 (Appendix C. showing that our methods consistently achieves superior training and inference efficiency.
To investigate the potential of applying CBSA to large-scatt pretraining, preliminarily finetune ViT models pretrained on ImageNet-211 by adapting their attention blocks into the CBSA A style. For comparison, also adapt them into linear attention. Experimental results are shown in Fig. [ Although CBSA deviates more from MHSA than linear attention and is consequently harder to adapt from pretrained ViTs, CBSA achieves comparable performance that of linear attention with nearly the same FLOPs. Interestingly, when the contraction step is removed, CBSA surpasses linear attention, which has been extensively investigated in 35
For semantic segmentation, following the design of DEPICT 34 we build CBT decoders by stack- ing CBSA layers without the feed-forward modules the top of VIT encoders. We evaluate the performance of them on the ADE20K dataset [53] and show the results in the left panel of Fig. [ Clearly, our CBT decoder consistently surpasses both white-box (DEPICT) and black-box (Seg- menter) counterparts that rely on softmax attention. In particular, the test-performing CBT decoder improves upon Segmenter by 1.5% mIoU while using merely 20% of the FLOPs and 0.06% of the pairwise similarities in the decoder.
Related work
Efficient attention mechanisms can be roughly divided into two categories: sparse attention and linear attention [54]. Approaches in sparse attention sparsify the attention matrix proactively by restricting the attention span to either random, or fixed [55] 20). or learnable [56] [57] patterns, or
their combinations. Approaches in linear attention [21 58) decompose the attention matrix into a product of two low-rank matrices and thus avoids its explicit computation via the associative property of multiplication. The idea of using representative tokens has been applied to both. Global tokens memory be introduced in sparse attention to maintain global connectivity, thereby further shrinking the attention span [59] 60]. Meanwhile, Agent tokens [35] and landmarks 30 can also be incorporated into linear attention from different perspectives.
Our CBSA distinguishes itself from previous efficient attention mechanisms by being inherently interpretable and derived from an optimization objective which efficiently compresses input tokens towards dimeeniiiiona structures. Moreover, it can not only be viewed as sparse attention or linear attention for leveraging the representatives, but also mathematically generalizes softmax attention, linear attention, and channel attention as its special cases.
6 Conclusion
We have proposed an optimization objective for deriving attention mechanism and unifying the investigation towards the interpretability and the efficiency. By unrolling the gradient optimization steps of this objective, we derived an inherently interpretable and efficient attention mechanism, called Contract-and- aud--rraddcstt Self- Atlf-Attention (CBSA). We found that our CBSA covers the instantiations of softmax attention, linear attention, and channel attention by changing the number and structure of representatives, thus revealing their fundamental connections. We validated the effectiveness of our CBSA through extensive experiments on visual tasks. We believe that the preliminary framework established in this work offers a promising direction for exploring a unified formula for existing attention mechanisms well as new attention mechanisms in an inherently interpretable way.
The checkpoint is obtained from the timm 52 library.
10

---

Acknowledgments and Disclosure of Funding
This work is supported by the National Natural Science Foundation of China under Grant Nos. 62576048 and 61876022. C.-G. Li is the corresponding author.
References
[1] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[2] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 19979,1598,, 2022.
[3] Jacob Devlin, Ming-We Wei Chang. Kenton Lee, and Kristina Toutanova. BERT: pre- 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 Language Technologies, pages 4171-4186, 2019.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[5] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figumov, Olaf Ronneherger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andrew J Ballard, Andrew Cowie, Bernardino camerr-Parrede,, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure 2021.
prediction with AlphaFold. Nature, 596(7873):583
[6] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self- supervised vision transformers. In JEEE/CVF International Conference on Computer Vision, pages 9630 9640, 2021.
[7] Tom B. Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever. and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
[8] Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang. Xianchen Lian, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer. arXiv preprint arXiv: 103. 1886, 2021.
[9] Yihe Dong, Wean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2803, 2021.
[10] Tianyu Guo, Druv Pai, Yu Bai, Jiantao Jiao, Michael I Jordan, and Song Mei. Active-dormant attention heads: Mechanistically demystifying extreme-token phenomena in llms. arXiv preprint arXiv:2410.1388 2024.
[11] Quoc- Vinh Lai-Dang. A survey of vision transformers in autonomous driving: Current trends and future directions. arXiv preprint 2735522403.07542, 2024,
[12] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- document transformer. arXiv preprint arXiv:2004.05 2020.
[13] Liu, Han Hu, Yutong Lin, Zhuliang Yao. Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo. Swin transformer V2: scaling up capacity and resolution. In JEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1
2022.
[14] Lemeng Wu, Xingchao Liu, and Qiang Liu. Centroid transformers: Learning to abstract with attention. arXiv preprint 202.::1002.08000 2021.

---

[15] Peng Wang, Yifu Lu, Yaodong Yu, Druv Pai. Qing Qu, and Yi Ma. Attention-only transformers via unrolled subspace denoising. In International Conference Machine Learning, 2025.
[16] Hubert Ramsauer, Bernhard Schäfl, Jobannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitmer, Thomas Adler, David P. Kreil, Michael K. Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations, 2021.
[17] Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is attention better than matrix decomposition? In International Conference on Learning Representations, 2021.
[18] Ruifeng Ren and Yong Liu. In-context learning with transformer is really equivalent to a contrastive learning pattern. arXiv preprint, 2023.
[19] Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature machine intelligence, 1(5):206-215 2019.
[20] Rewon Child, Scott Gray, Alee Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 105500
10509, 2019.
[21] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, volume 119, pages 5156-5165, 2020.
[22] 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. In Advances in Neurul Information Processing Systems, 2020.
[23] Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D. Haeffele. and Yi Ma. White-hox transformers via sparse rate reduction. In Advances in Neural Information Processing Systems, 2023.
[24] Ziyang Wu, Tianjiao Ding, Yifu Lu, Druv Pai, Jingyuan Zhang, Weida Wang. Yaodong Yu, Yi Ma, and Benjamin David Haeffele. Token statisties transformer: Linear-time attention via variational rate reduction. In International Conference on Learning Representations, 2025.
[25] Ehsan Elhamifar, Guillermo Sapiro, and René Vidal. See all by looking at a few: Sparse modeling for finding representative objects. In IEEE Conference on Computer Vision and Pattern Recognition, pages 1600-1607, 2012.
[26] Chong You, Chi Li, Daniel P Robinson, and Rene Vidal. Self-representation based unsupervised exemplar selection in union of subspaces. IEEE Transactions Pattern Analysis and Machine Intelligence, 44(5): 2698-2711, 2020.
[27] Tianjiao Ding, Shengbang Tong. Kwan Ho Ryan Chan, Xili Dai, Yi Ma, and Benjamin D. Haeffele. Unsupervised manifold linearizing and clustering. JEEE/CVF International Conference on Computer Vision, pages 5427-5438, 2023.
[28] Ma, Harm Derksen, Wei Hong, and John Wright. Segmentation multivariate mixed data via lossy data coding and compression. IEEE transactions on pattern analysis and machine intelligence, 29(9): 1546-1562, 2007.
[29] 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 (114):1-100 2022.
[30] 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. AAAI Conference on Artificial Intelligence, pages 14138-14145 2021.
[31] Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random- access infinite context length for transformers. arXiv preprint arXiv:2305. 16300, 2023.
[32] Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In International Conference on Machine Learning, pages 399-406, 2010.
[33] 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, 2021.
12

---

[34] Qishuai Wen and Chun-Guang Li. Rethinking decoders for transformer-based semantic segmentation: A compression perspective. In Advances in Neural Information Processing Systems, 2024.
[35] Dongchen Han, Tianzhu Ye, Yizeng Han, Zhuofan Xia, Siyuan Pan, Pengfei Wan, Shiji Song, and Gao Huang. Agent attention: On the integration of softmax and linear attention. In Europea conference on computer vision, pages 124-140, 2024.
[36] Yunzhe Hu, Difan Zou, and Dong Xu. An in-depth investigation of sparse rate reduction in transformer- like models. In Advances in Neural Information Processing Systems, 2024.
[37] Jusen Du, Jiaxi Hu, Tao Zhang, Weigao Sun, and Yu Cheng. Native hybrid attention for efficient sequence modeling, arXiv preprint arXiv:2510.07019 2025.
[38] Han Guo, Songlin Yang. Tarushii Goel, Eric Xing, Tri Dao, and Yoon Kim. Log-linear attention. arXiv preprint 2xiv:2506.04477 2025.
[39] Ehsan Elhamifar and René Vidal. Sparse subspace clustering. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pages 2790- 2700002277.220 2009.
[40] Jie Hu, Li Shen, and Gang Sun. Squeeze-and-eee eendexxittation networks. In CVPR, pages 7132-7141. Computer Vision Foundation LEEE Computer Society, 2018,
[41] Sanghyun Woo, Jongchan Park. Joon- Young Lee, and In So Kweon. CBAM: convolutional block attention module. ECCV (7), volume 11211 of Lecture Notes in Computer Science, pages 3-19. Springer, 2018.
[42] Weihao Yu, Mi Luo, Pan 7hou, Chenyang Si, Yiehen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. In JEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 10809-10819, 2022.
[43] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan. Shauna Kravec, Zac Hatfield-Dodde Robert Lasenby, Dawn Drain, Carol Chen. Roger B. Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of superposition. arXiv preprint 2209 10652, 2022.
[44] Emmanuel Candes and Terence Tao. Decoding by linear programming. IEEE transactions on information theory, (1(1))42203 4215, 2005.
[45] René Vidal, Yi Ma, and Shankar Sastry. Generalized Principal Component Analysis. Springer, 2016. ISBN 978-0-387-87810-2
[46] Phillip Pope. Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. arXiv preprint arXiv:2104.00s: 04.00894, 2021.
[47] Duy-Kien Nguyen, Mido Assran, Unnat Jain, Martin R. Oswald, Cees G. M. Snoek, and Xinlei Chen. An image is worth more than 16x16 patches: Exploring transformers on individual pixels. In International Conference on Learning Representations. 2025.
[48] Feng Wang. Yaodong Yu, Guoyizhe Wei, Wei Shao, Yuyin Zhou, Alan Yuille, and Cihang Xic. Scaling laws in patchification: image worth 50,176 tokens and more. arXiv preprint arXiv:2502.03733, 2025.
[49] 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, pages
2017.
[50] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In /EEF/CVF International Conference on Computer Vision, pages 7242-7252, 2021.
[51] Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, and Yi Ma. Emergence of segmentation with minimalistic white-box transformers. In Conference Parsimony and Learning. volume 234, pages 72-93, 2024.
[52] Ross Wightman. Pytorch image 2019.
[53] Bolci Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision, 322-33302-3 2019.
[54] Yutao Sun, Zhenyu Li, Yike Zhang, Tengyu Pan, Bowen Dong, Yuyi Guo, and Jianyong Wang. Efficient attention mechanisms for large language models: A survey, arXiv preprint 20X5v22507.19595 2025.
13

---

[55] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer. Alexander Ku, and Dustin Tran. Image transformer. International Conference on Machine Learning, volume 80, pages 4052-4061, 2018.
[56] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da- Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning. volume 119, pages 9438-9447, 9447, 2020.
[57] Aurko Roy. Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 53 68, 2021.
[58] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song. Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
[59] Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang, tra-frranformmr. In Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1315-1325, 2019.
[60] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. Advances Neural Information Processing Systems, 2020.
[61] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pages 213-229, 2020.
[62] Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In International Conference on Learning Representations, 2024.
[63] Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang. Flatten transformer. Vision transformer using focused linear attention. In IEEE/CVF International Conference on Computer Vision, pages 5938-5938845 2023.
[64] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang. Hicu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, Yifeng Lu, and Quoe V. Lc. Symbolic discovery of optimization algorithms. In Advances in Neural Information Processing Systems, 2023.
[65] Nly Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representattons, 2019.
[66] Mark Everingham and John Winn. The pascal visual object classes challenge 2012 (voc2012) development kit. Pottern Analysis, Statistical Modelling and Computational Learning, Tech. Rep, 8(5):255, 2011.
14

---

Appendix
A Representative initialization
Besides pooling the input tokens, another intuitive choice for representative initialization is to set them as learnable (i.c., trainable) parameters, similar to the object queries in [6]], the class embeddings in [50], and the registers in [62 However, in this case found that the attention matrix in the extraction (cross-attention) step 14 is low rank. For example, although its maximal rank can be m = p = 64 in CBT models, it typically remains around 5. This indicates that far fewer initialized representatives are being utilized than expected, leading to unsatisfactory performance of CBSA.
By contrast, when pooling- boosigg-based initialization is employed, the rank of the attention matrix always attains its maximum, i.c. m. Nonetheless, as pointed out in [3], this still constitutes a low-rank bottleneck: linear attention underperforms softmax attention, whose attention matrix has rank N > m. Although this issue can be alleviated by introducing depth-wies wise convolution (DWC), which boosts performance while preserving linear complexity, our work does not involve this modification since we focus on theoretical contributions rather than empirical improvements.
B Derivations
Here, we show how
12 is derived by substituting U L, As
R& into
K
(Im ER
m me²
(k,En) (LrEn)) R
(14)
-1
Σ
LBLEE En
Im + マ m Li Ls
R
(15)
k=1
K
U&L
Im + हeج 2
R
(16)
k=
I -
U
R
(17)
Σ URLi
K
Im + 1
-
-1
L U
(18)
k=1
K
Σ
U (L.ERL%)
U Uiz
(19)
k=
[ U.Z
(20)
Note that an arbitrary set of orthogonal representatives can be expressed as the principal directions under an orthogonal rotation and scaling, i.c.,
U Q= =
= ZR..A.
(21)
where Ps. E O(p) denotes an orthogonal rotation within the k-th subspace, and A₁ is diagonal matrix of size p. Generalizing the way to form the representatives in (4) from U Q U. ZA, to U.Q. Q U AR, where B E Rddd we have:
R (KK & LUBBUUIQ (Im mea (UIQ) -1
A.
(22)
me²
MThat is. rank (softmax (UU (UU+))).
15

---

Similarly, by substituting U.. = PA A = R&Ax and B₁ Pk into 22 we obtain
UULLAA Im Im+ + U
(23)
C More experimental results
Training setup. We train the CBT models in Table B and all models in Table 4 150 epochs with the Lion optimizer [64]. The learning rate is 2.0 x 10 the weight decay coefficient is 0.05, and the batch size 256. We also incorporate warm-up strategy over the first 20 epochs. For data augmentation, we adopt rather simple choice: just random cropping and random horizontal flipping. We apply label smoothing with a smoothing coefficient of 0.1. For fine-tuning, we use the Adam optimizer [65] learning rate of 5 x 10-5, weight decay of 0.01, and batch size 64. The settings above are largely inherited from 233
Quantitative segmentation properties. Following prior works 6 5
we evaluate the zero-shot
segmentation performance ImageNet-IK pretrained models on the PASCAL VOC12 validation set [66]. Specifically, we assess the best- pess-perforring attention maps of the [CLS] token, but in simplified setting that considers only three segmentation targets, instead of fine- tnnegrained semantic classes. In Table 5 we find that the hybrid model has consistently better segmentation performance. We also measure the segmentation performance of different layers and report the results Table [
Table 5: Zero-shot segmentation. We use the Jaceard similarity, which is also called intersection over union (IoU), to quantify the alignment between the [CLS] token's attention map and the segmentation ground truth. The best-performing results are highlighted in bold.
Segmentation target (CRATE+CBT)-S
CRATE-S
CBT-S
ToST-S
Foreground
0.68
0.65
0.51
0.42
Background
0.78
0.76
0.72
0.75
Boundary
0.18
0.17
0.15
0.12
Table Zero-shot segmentation across layers. The top-performing three layers for each model underlined.
Layers
L2
L3
L4
L5
L77
L8
L9
L10
L11
L12
(CRATE-CIBT)
CRATE-S
0.34
0.45
0.51
0.48
0.53
0.56
0.53
0.57
0.55
0.52
0.50
0.39
CBT-S
0.39
0.41
0.43
0.42
0.37
0.40
0.33
0.35
0.34
0.36
Table 7: Throughput comparisons. The results below are obtained from experiments dataset CIFAR-10 with an image resolution of 512 x 512.
Images / see
CBT-
ViT-T
CRATE-T
ToST-T
CBT-S
ViT-S
CRATE-S
ToST-S
Training
336
174
203
323
205
94
111
199
Inference
572
370
395
533
405
211
220
429
Figure 10: Additional results on CBT-B and CBT-L for the experiment in Fig
16

---

D PyTorch implementation
Algorithm 1: PyTorch implementation of CBSA 8
class CBSA (nn Module)
def
-_init__ (self, dim, heads, dim dim_head)
super .w..i....______.. inner_dim heads dim_head self heads heads
self. dilfheim_eeaddeeo dim_head
self. dim_head ** -0.5
self - mn. cn.Softmax (dim=-11
self
ww..AaatieeevgPooll..o (ouphutttootuuttueeseeeee 8))
# subspace bases
11
self. nn. (dim, inner_dim, bias=False)
12
# over-parametrrizaiio
13
self - mn Linear (inner_din, dim)
# step-sizes
15
nn Parameeer (torch, randn (heads, 1, 1)) self self.ss_rep nn. Parameter (torch. randn (heads, 1, 1))
16
17
def Ettention(self. query, key, value):
dots (query key, transpose (-1, -2)) self scale attn self attend out = attn 0 value return out, attn
20
21
22
24
def forvard x):
25
b, n, с shapee height = width int (n ** 0.5)
26
# projection onto subspaces
self (x)
29
# representative initialization
30
rep self.pool pool :-1, :] reshappee,,, height, width, c).
permute (0, 3, 1, 2)) reshhee (b, c, -1) -1) permutue 2, 1)
reshaape
self heads, self dim_head) permute (0, 2,
3)
32
rep rep. reshape 64, self heads, self headh permute (0,
2. 1. 3)
# representative extraction rep_delta, attn self attention (rep, v, w)
35
rep rep self.ss_rep rep_delta # representative contraction z_delta, = attenntooonion (rep, rep, rep)
36
37
# contraction broadcast
x_delta attn transpose (-1, -2) xedet
40
x_delta self x_delta
41
x_delta rearrange (x_delta, "b h n k -> b n (h k)
42
# projection back to the ambient space return self tot (x_delta)
E
Limitations
Our work aims to propose unified framework for both interpretable and efficient attention. Currently, have proposed unified optimization objective derive such an attention mechanism, CBSA, which is capable to accommodate previous interpretable attention mechanisms derived from MCR Nonetheless, more fundamental and mathematical connections between our CBSA and existing efficient attention mechanism remain unexplored.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Our three-fold contribution listed in Section is justified in Section P
and
Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the current limitations of work in Appendix
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
18

---

The assumptions can be found in Section 2 and Section 3 and complete proofs are in our Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Our designs are clearly formulated Section 3 and illustrated Fig. We also provide the pseudo code of our CBSA in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
19

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The link to our models and code can be found in our abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips //ipss.cc/pill wttpcccpppppppppeeeeeeeeeeeeeeee htapa//eguddeee/eeeeabiiissssss (oooeooomisssssseeeeeeee ssi.iiiiccyyyyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We present our experimental detail, such the training setup, in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We only report error bars in Fig. For other experiments, it would be too computationally expensive for us.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
20

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: One can find the details in our Github page.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Our research mainly focuses on the interpretability in AL.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
21

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: Our methods pose no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Only publicly available assets are used in this paper, and we followed their license to reference their authors.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
22

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: There is no human subjects are involved in our work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: There is no human subjects are involved in our work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
23

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs do not participate any part of our work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https://neure) /ttp://///////ooeeeeeee///////// for what should or should not be described.
24


Reviews:
[{"review_id": "Qm7zA5Y21P", "reviewer": "Reviewer_FLEn", "summary": "Attention in modern deep learning models has led to significant performance in many domains, but it suffers from a quadratic computational overhead. Additionally, there is still a lack of understanding of their underlying mathematical optimization objective. While prior works work on either of these problems, this paper claims to address both problems with their Contract-and-Broadcast Self-Attention (CBSA). CBSA compresses the input tokens by using a subset of them to create a mixture of low-dimensional subspaces that preserve the original coding rate of the entire input token space. CBSA is a self-attention implementation made interpretable by unfolding a gradient-based optimization on the coding rate preservation objective. Experimental results show competitive performance, computational efficiency, and kept emergent properties such as segmentation and parameter perturbation robustness of their method.", "questions": "- Figure 2 (right) is still a bit unclear to me. Is the y-axis suppose to be labeled $\\Delta R$ and thus measure Eq. 2 just between $R_c(Z | U_{[K]})$ (blue) and $R_c(Q | U_{[K]})$ (orange)?", "limitations": "- See weaknesses\n- CBSA seems more like an alternative to TSSA than an improved method with regards to performance and computational efficiency.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Paper Strengths:\n- The paper provides a novel method for lower the computational complexity of the attention mechanism made interpretable through algorithmic unrolling.\n\t- Shown in both classification and segmentation tasks.\n- The method is linked to existing methods of MSSA if full attention is kept by having the number of representative equal to the number of tokens but also reduces to channel attention as in TSSA if representatives are orthogonal and fixed.\n- The model's parameters are significantly more robust than other methods reported.\n- The paper is well written and organized.\n\nPaper Weaknesses:\n- Other than parameter robustness, this method isn't much more efficient or stronger in performance than TSSA (ToST)\n- The claim of segmentation properties is based on one visual analysis in Figure 4 without any quantitative experiments showing it's segmentation performance. While the visual analysis is enough to show sign of segmentation abilities, it's unclear how strong this emergent property is.", "was_revised": true, "final_justification": "Based on my initial review, and the rebuttal on my initial concerns on clarity, marginal performance improvements to related works and the strength of segmentation abilities being resolved, my score is 5.", "rebuttal": "Thank you for your valuable comments. \n\nTo address the weaknesses and limitations above, we present new experimental results demonstrating:\n+ **Improved Performance**: With an improved initialization strategy of the representatives, CBSA outperforms TSSA remarkably while maintaining its efficiency.\n+ **Quantitative Segmentation**: We measure the segmentation properties using Jaccard similarity and observe that CBSA preserves or even enhances the emergent segmentation properties.\n\nIn the following, we will address each of the Weaknesses/Questions/Limitations point-to-point and provide the corresponding experimental results.\n\n> **Weakness 1**: This method isn't much more efficient or stronger in performance than TSSA (ToST).\n\n**Pooling-based initialization breaks through the performance bottleneck.**\n\nGiven that CBSA enjoys much more expressiveness than TSSA in principle, it is also confused us for the marginal performance gain. We found that the rank of the attention matrix of the extraction (i.e., the cross-attention) step, i.e,\n$\\text{rank} \\left( \\text{softmax} \\left( Z^\\top Q_0 \\right) \\right)$, \nwas overly low. \nThis suggests that there are fewer representatives being utilized than expected. \nThis low-rank bottleneck stemmed from the learnable initialization of $Q$, and resulted in the unsatisfactory performance of CBSA. Now, we adopt a pooling-based initialization strategy, i.e.,\n$Q_0 = \\text{AvgPool}(Z)$. \nThis enables the attention matrix to achieve maximal rank, thereby improving model performance. \nThis intriguing finding aligns with the recent advances in linear attention mechanisms [62,63]. \nA detailed analysis will be presented in the final paper.\n\nConsequently, we revise the Tables as follows, highlighting new results **in bold**.\n\n**Revised  Table 1**: Evaluating CBT on ImageNet-1K.\n| Models | CBT-T | CBT-S  | CBT-B  | CBT-L | ViT-S |\n| :- | - | - | - | - | - |\n| FLOPs  | 1.1G         | 4.0G         | 15.1G        | 47.3G        | 9.8G  |\n|acc@1 (old) | 62.9 | 69.5 | 71.7 | 72.9 | - |\n| acc@1 | **63.2**  | **71.4**  | **73.4**  | **74.4**  | 72.4  |\n\n**Revised Table 2**: Fair comparisons between CBSA, TSSA, and MSSA on ImageNet-1K. \n| Models | CBT-T | CBT-S | ToST-T | ToST-S | CRATE-T | CRATE-S |\n| :- | - | - | - | - | - |-|\n| acc@1 (old)  | 62.9| 69.5 | - | - | - | - |\n| acc@1  | **63.2**  | **71.4** |  61.2 | 68.5| 64.7| 72.1|\n\n**Revised Table 3**: Evaluating CBT on ADE20K.\n| Models | CBT-B  | CBT-L  | DEPICT-L [29] | Seg-L-Mask [48] |\n| :- | -| - | - | - | \n| Decoder FLOPs   | 7.8G | 22.3G  | 35.0G | 113.2G |\n| mIoU (old)  | 48.7 | 52.0 | - | - |\n| mIoU   | **49.3** | **53.3** | 52.9 | 51.8 |\n\n**Key improvements**:\n+ **Comparable to ViT**: CBT-S performs comparably to ViT-S (71.4 v.s. 72.4) with only 40% of the FLOPs (see revised Table 1); \n+ **Superiority to TSSA**: The performance gain between small size models has been increased from 1% acc@1 to about 3% (see revised Table 2);\n+ **Better Segmenation**:\n For semantic segmentation, our CBSA outperforms MSSA and MHSA (see revised Table 3).\n\n**An Additional Experiment: CBSA matches or even exceeds Linear Attention**\n\nTo further validate the performance of CBSA, we utilize open-source ViT models pre-trained on ImageNet-21K and modify their forward passes to implement either CBSA or linear attention mechanisms. All models are then fine-tuned on ImageNet-1K.\n\n**Table 7 (new)**: Evaluating pre-trained ViTs adapted to linear complexity on ImageNet-1K.\n\n| Model Sizes | Tiny@224$^2$| Small@224$^2$| Base@224$^2$| Large@224$^2$| Tiny@384$^2$| Small@384$^2$| Base@384$^2$| Large@384$^2$|\n|:-|-|-|-|-|-|-|-|-|\n|MHSA |76.9 | 81.4 | 84.7| 85.9 |79.7| 83.7| 85.9| 87.1|\n|LinearAttn|73.2 | 78.6 | 79.8 | 80.7| 74.9 | 80.1 | 81.8 | 82.2|\n|CBSA |**72.2** | **78.3** | **79.2** | **79.8** | **72.3** | **79.3** | **80.8** | **81.2** |\n|CBSA$^\\vee$|**74.7** |**79.9** |**80.9**  |**83.8**| **77.3** | **81.4** | **83.1** |**85.0**|\n\nPooling-based initialization is adopted.\n$\\vee$: the contraction (self-attention) step in CBSA is removed.\nIn this case, CBSA effectively derives another existing attention mechanism known as Agent Attention [79].\nA detailed discussion of the practical advantages of removing this step will be provided in the final version of the full paper.\n\n> **Weakness 2**: While the visual analysis is enough to show sign of segmentation abilities, it's unclear how strong this emergent property is.\n\nWe thank the reviewer for this constructive suggestion.\n\nWe supplement Figure 4 with quantitative experiments evaluating zero-shot segmentation performance. \nFollowing prior works [8, 47], we assess the [CLS] token attention maps on the validation set of PASCAL VOC12 [64]. \nHowever, as neither [8, 47] provides open-source code, our implementation will differ (demo will be released via a Colab Notebook).\n\n**Jaccard similarity** is also called intersection over union (IoU). \nWe adopt this metric to quantify the alignment between the attention map and the segmentation ground truth. Higher Jaccard similarity indicates more pronounced segmentation properties emerged in the model.\n\n**Table 5 (new)**: The Jaccard similarity on the validation set of PASCAL VOC12.\n| Models | (CRATE+CBT)-S | CRATE-S | CBT-S | ToST-S |\n| :- | - | - | - | - | \n| Foreground |0.68 | 0.65 | 0.51  | 0.42 |\n| Background | 0.78 | 0.76 | 0.72 | 0.75 |\n| Boundaries | 0.18| 0.17 | 0.15 | 0.12 | \n\n**Table 6 (new)**: The Jaccard similarity across layers.\n| Layers   | L1   | L2       | L3       | L4       | L5   | L6       | L7       | L8       | L9       | L10      | L11      | L12  |\n| :-| - | - | -| - | - | - | - | - | - | - | - | - |\n| (CRATE+CBT)-S | 0.37 | 0.44     | 0.55     | 0.46     | 0.56 | **0.58** | **0.60** | 0.53     | **0.59** | 0.56     | 0.55     | 0.55 |\n| CRATE-S       | 0.34 | 0.45     | 0.51     | 0.48     | 0.53 | **0.56** | 0.53     | **0.57** | **0.55** | 0.52     | 0.50     | 0.39 |\n| CBT-S         | 0.39 | **0.41** | **0.43** | **0.42** | 0.37 | 0.40     | 0.38     | 0.33     | 0.35     | 0.34     | 0.36     | 0.33 |\n\nThe top-performing three layers for each model are highlighted **in bold**.\n\n**Key observations**:\n+ The segmentation properties are gradually reduced from MSSA to CBSA to TSSA.\n+ The segmentation properties are enhanced in the hybrid model, CRATE+CBT, surpassing pure MSSA.\n\n> **Question**: Figure 2 (right) is still a bit unclear to me. Is the y-axis suppose to be labeled and thus measure Eq. 2 just between (blue) and (orange)?\n\nWe sincerely appreciate this feedback and apologize for any confusion caused by the incomplete definition of $\\Delta R_c$ in our original submission. \n\n**Definition of $\\Delta R_c$.** \nThe compression term $R_c$ measures the compactness of input tokens or their representatives.\nWe measure the change of $R_c$, i.e., $\\Delta R_c$, produced by CBSA to quantify its effect in the compression procedure.\nNote that this is different from the $\\Delta R \\doteq R(Z) - R_c(Z \\mid U_{[K]})$ in Eq.(2).\nWe present formal definitions:\n+ $\\Delta R_c$ w.r.t $Z^\\ell$ (**the blue line**) is\n$\\Delta R_c \\left(Z^\\ell  \\mid U_{[K]} \\right) \\doteq R_c\\left(\\text{CBSA}\\left(Z^\\ell \\mid U_{[K]}\\right)  \\mid U_{[K]} \\right) -  R_c(Z^\\ell \\mid U_{[K]})$, \nwhich reflects the extent to which input tokens are compressed by the $\\ell$-th CBSA layer;\n+ $\\Delta R_c$ w.r.t $Q^\\ell$ (**the orange line**) is $\\Delta R_c(Q^\\ell) \\doteq R_c\\left(\\text{MSSA}\\left(Q^\\ell \\mid U_{[K]}\\right)\\right) -  R_c(Q^\\ell)$, where the MSSA over $Q$ corresponds the contraction step of CBSA\n(i.e., the contraction term of Eq. (16)). Therefore, this reflects the extent to which representatives are contracted by the $\\ell$-th CBSA layer.\n\n**Figure 2 (Right) Interpretation.** In theory, we expect both the blue and orange lines are consistently above $x=0$, so that the input tokens and their representatives are compressed and contracted in every attention layers. In Figure 2 (Right), we find only 2 of the 12 layers violates this pursuit.\n\n\n> **Limitation**: CBSA seems more like an alternative to TSSA than an improved method with regards to performance and computational efficiency.\n\nPlease refer to our response to Weakness 1. \n\nWe hope the experiments above could resolve your concerns on our work. Please let us know if any further clarification is needed.\n\n[62] Dongchen Han et al. \"Flatten transformer: Vision transformer using focused linear attention.\" ICCV 2023\n\n[63] Qihang Fan et al. \"Breaking the low-rank dilemma of linear attention.\" CVPR 2025\n\n[64] Mark Everingham et al. \"The pascal visual object classes challenge 2012 (voc2012)\"\n\n[79] Dongchen Han et al. \"Agent attention: On the integration of softmax and linear attention.\" ECCV 2024.\n\nReferences [1–61] are provided in the manuscript."}, {"review_id": "ucX2O3Nqr2", "reviewer": "Reviewer_2uTu", "summary": "This paper introduces Contract-and-Broadcast Self-Attention (CBSA), a novel attention mechanism designed to jointly address interpretability and efficiency in Transformers. CBSA is derived from a unified optimization objective inspired by Maximal Coding Rate Reduction, where a small set of representative tokens is contracted via self-attention and their effects are broadcast back to the full input. The method is interpretable by construction and scales linearly with input size. The authors demonstrate that CBSA generalizes several known attention mechanisms (e.g., full attention, channel attention) as special cases. Experiments on standard image classification and segmentation tasks show that CBSA performs competitively with reduced computational cost and improved robustness.", "questions": "1. It is not entirely clear to me how the initial guess Q0 is optimized during training. While the update rule for Q (e.g., Eq. 10–11) is continuous, Q is supposed to represent a small subset of the input tokens Z, which feels more like a discrete selection. Could the authors clarify how this representative set is extracted and optimized? Are there constraints or assumptions that guide its learning?\n\n2. The proposed CBSA approach seems conceptually similar to LoRA, in that both reduce the representation to a few low-rank basis components. Could the authors explain how CBSA differs from or improves upon LoRA?\n\n3. Does the CBSA framework improve efficiency during inference, or is it primarily beneficial during training? Specifically, do the additional steps—such as projecting into subspaces and broadcasting from representatives—introduce overhead at test time?", "limitations": "There isn't a limitation section in the paper.", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper introduces a principled objective that unifies interpretability and efficiency, which are typically treated separately in prior attention models.\n\n2. The proposed Contract-and-Broadcast Self-Attention (CBSA) is derived directly from a gradient-based optimization on a compression objective (MCR²). This provides a clear, interpretable foundation for the mechanism.\n\n3. The authors show that MSSA (full attention) and TSSA (channel attention) are special cases of CBSA, depending on the choice of representative tokens. This highlights CBSA as a general and flexible framework.\n\nWeaknesses:\n1. CBSA depends on projecting data into orthogonal low-dimensional subspaces, which assumes that meaningful structure lies there. This assumption may not hold for all modalities or tasks, potentially limiting generality.\n\n2. In the experiments, there is no comparison to other methods (e.g., MSSA, TSSA) to show whether CBSA compresses more effectively or more meaningfully. Without a baseline, it’s unclear whether the observed compression is a distinctive advantage or just typical of deep models.", "was_revised": true, "final_justification": "The authors provide a thorough response to all of my questions and concerns. Most of them have been addressed with improved clarity.", "rebuttal": "Thank you for your insightful comments. It is our pleasure to discuss these points.\n\n> **Weakness 1**: CBSA depends on orthogonal low-dimensional subspaces, which assumes that meaningful structure lies there. This assumption may not hold for all modalities or tasks, potentially limiting generality.\n\nWe acknowledge that the structure of **U**nion **o**f **O**rthogonal **S**ubspaces (**UoOS**) faces several challenges in achieving generality, including its adaptation to all modalities and tasks as noted by the reviewer.\nIn light of this theoretical gap, we excluded natural language processing tasks from our experiments.\n\nHowever, we contend that adopting the structure of UoOS is not prohibitively restrictive in practice, though specific theoretical analyses in other modalities and tasks remain future work.\nUnderpinned by the following facts:\n+ **subspaces are over-parameterized** as each attention layer parameterizes a distinct $U_{[K]}$, yielding $L$ sets of $K$ orthogonal subspaces;\n+ tokens can be compressed towards **finer-grained low-dimensional structures** within each of the subspaces,\n\nUoOS could be a feasible starting point or prototype towards the optimal structures, e.g., simplex ETF [65] or its generalized case [66]. \n\nRegarding the language modality, an input sentence or article can be projected onto one or more subspaces that acts as its principal subspace  or multiple principal subspaces in the sense of PCA or GPCA [67].\nThis process mirrors Latent Semantic Analysis (**LSA**), which is a classical NLP technique.\nSimilarly, the channel-mixer variant of our CBSA employs the principal directions as the representatives (see Eqs. (18)-(19)).\n**Therefore, the representatives introduced in our work constitute finer-grained structures than UoOS, dynamically encoding a broader spectrum of semantic concepts.**\n\n> **Weakness 2**: In the experiments, there is no comparison to other methods (e.g., MSSA, TSSA) to show whether CBSA compresses more effectively or more meaningfully. Without a baseline, it’s unclear whether the observed compression is a distinctive advantage or just typical of deep models.\n\nWe thank the reviewer for this constructive suggestion.\n\n**The Prevalence of Compression**\n\nCompression is not a phenomenon unique to Transformers built upon the MCR$^2$ objective (e.g., CRATE, ToST, and our CBT), but rather pervasive across deep learning models.\nIn a collection of network architectures for classification, their representations are progressively compressed as the layers go deeper and the training continues [68], and eventually converges to neural collapse at the terminal phase of training [65].\nMeanwhile, compression toward low-dimensional distributions can also be observed in generative tasks, e.g., in diffusion models [69]. \nAnd another research direction assesses model performance based on representation compactness [70, 71].\n\n**Compactness Comparison to Other Methods**\n\nWe evaluate representation compactness in the final attention layer across related methods, quantified by the compression term $R_c$ of Eq. (2).\n\nWe observe that our CBSA compresses more effectively than TSSA but is inferior to MSSA. \n\nMore importantly, we find that the final representation compactness reflects the model performance well, indicating that the coding rate can act as a novel evaluation metric.\n\n**Table 8 (new)**: Measuring representation compactness in the final attention layer.\n\n|Models|CRATE-T| CBT-T|ToST-T|CRATE-S|CBT-S|ToST-S|\n|-|-|-|-|-|-|-|\n|$R_c(\\boldsymbol{Z}^L \\mid \\boldsymbol{U}^L_{[K]})$ |447|**622**|665|987|**1317**|1558|\n| acc@1|64.7|**63.2** |61.2|72.1|**71.4** |68.5|\n\nResults of our method are highlighted **in bold**.\n\nNote that the performance results of CBSA have been updated by changing the representative initialization strategy\n(please refer to our response to the weakness 1 identified by Reviewer FLEn for further details).\n\n> **Question 1**: It is not entirely clear to me how the initial guess Q0 is optimized during training. While the update rule for Q (e.g., Eq. 10–11) is continuous, Q is supposed to represent a small subset of the input tokens Z, which feels more like a discrete selection. Could the authors clarify how this representative set is extracted and optimized? Are there constraints or assumptions that guide its learning?\n\nWe appreciate the opportunity to clarify this confusing point.\n\n**The representatives are extracted through a continuous process.** The representatives are extracted via cross-attention, where the query is the initial guess and the key and value are the input tokens, i.e., $Q \\doteq Q_0 + \\eta Z\\text{softmax}(Z^\\top Q_0)$. This cross-attention can be viewed as a gradient step towards solving the $K$-means centroids of the input tokens. Therefore, the extraction is continuous rather than discrete.\n\n**The extraction can be easily adapted as discrete selection.**\nIf we replace the softmax with argmax and remove the residual connection, the extraction becomes\n$Q \\doteq Z\\text{argmax}(Z^\\top Q_0)$.\nThis selects the $1$-nearest neighbors of the initial guess from the input tokens, resulting in a discrete selection.\n\n**Assumptions and constraints behind the extraction:**\n+ **Assumption 1** (most important): When $Q$ is contracted, $Z$ can be correspondingly compressed (**compressing all by contracting a few**);\n+ **Assumption 2**: The $Q$ extracted by cross-attention satisfies Assumption 1 (validated in Figure 2 (right)).\n+ **Constraint**: The sole constraint is the model architecture, i.e., the optimization objective underlying its forward pass.\n\n> **Question 2**: Could the authors explain how CBSA differs from or improves upon LoRA?\n\n**A Quick Review of LoRA and CBSA**\n\nLoRA approximates the full fine-tuning update of a pre-trained weight matrix, $\\Delta W \\in \\mathbb{R}^{d\\times d}$, via the multiplication of two distinct low-rank matrices, i.e., $\\Delta W \\approx W_a W_b^\\top$, where $W_a, W_b \\in \\mathbb{R}^{d \\times r}$ and $r < d$.\n\nOur CBSA approximates the full attention matrix $A_{\\text{full}}  \\in \\mathbb{R}^{N \\times N}$ by the Gram matrix of a low-rank matrix $\\overline{A} \\in \\mathbb{R}^{N \\times m}$, i.e., $A_{\\text{full}} \\approx \\overline{A}\\ \\overline{A}^\\top$, where $m <  N$ is number of representatives, and $\\overline{A}$ is the attention matrix computed by the extraction (cross-attention) step.\n\n**Similarities and Differences between LoRA and CBSA**\n+ **Approximation targets**: LoRA approximates the weight matrix update ($\\Delta W$) while CBSA approximates the full attention matrix ($A_{\\text{full}}$);\n+ **Approximation flexibility**: Both allow tunable rank ($r$ and $m$), unlike linear attention (fixed rank $p$);\n+ **Number of matrices**: LoRA uses two low-rank matrices ($W_a, W_b$), while CBSA employs only one ($\\overline{A}$).\n+ **Approximation hardness**: The intrinsic rank of $\\Delta{W}$ is typically small [Table 6, 72], but $A_{\\text{full}}$ is always full-rank [62]. Consequently, LoRA matches full fine-tuning performance even with small rank r, whereas CBSA never exactly attains identical performance unless in the self-expressive case (Section 3.3);\n+ **Acceleration effects**: LoRA reduces fine-tuning overhead but maintains inference costs, while CBSA accelerates both training and inference;\n+ **A unified perspective**:\nAccording to meta-learning interpretations of linear attention [73,74], $W$ corresponds to slow weights (updated only during training), while $A_{\\text{full}}$ corresponds to fast weights (input-dependent). Therefore, LoRA approximates the slow-weight update and CBSA directly approximates the fast weight. Both methods are weight low-rank approximation techniques.\n\n> **Question 3**: Does the CBSA framework improve efficiency during inference, or is it primarily beneficial during training? Specifically, do the additional steps—such as projecting into subspaces and broadcasting from representatives—introduce overhead at test time?\n\nWe thank the reviewer for this constructive suggestion.\n\nWe measure computational throughput in processed images per second (img/s) across Transformer architectures, using a standard $512 \\times 512$ resolution typical for semantic segmentation tasks.\n\n|Models|ViT-T|CRATE-T|CBT-T|ToST-T|ViT-S|CRATE-S|CBT-S|ToST-S|\n|-|-|-|-|-|-|-|-|-|\n|im/sec (training)|174|203|**336**|323|94|111.4| **205**|199|\n|im/sec (inference)|370|395|**572**|533|211| 220.3|**405**|429|\n\nThe results show these extra steps save more time than they cost.\n\n> **Limitation**: There isn't a limitation section in the paper.\n\nThank you for this feedback.\n\nWhile our limitation discussion was originally confined to Appendix D, we will now include a reference to it in the main text to ensure visibility for interested readers.\n\nWe hope the discussions above could resolve your concerns on our work. Please let us know if any further clarification is needed.\n\n[65] Vardan Papyan et al. \"Prevalence of neural collapse during the terminal phase of deep learning training\", PNAS 2020\n\n[66] Jiachen Jiang et al. \"Generalized neural collapse for a large number of classes\", 2023\n\n[67] Rene Vidal et al. \"Generalized principal component analysis (GPCA)\", TPAMI 2005\n\n[68] Hangfeng He et al. \"A law of data separation in deep learning\", PNAS 2023\n\n[69] Peng Wang et al. \"Diffusion models learn low-dimensional distributions via subspace clustering\", CPAL 2025\n\n[70] Lai Wei et al. \"Diff-erank: A novel rank-based metric for evaluating large language models\", NeurIPS 2024\n\n[71] Xianzhen Luo et al. \"Is Compression Really Linear with Code Intelligence?\" 2025\n\n[72] Edward J Hu et al. \"Lora: Low-rank adaptation of large language models\", ICLR 2022\n\n[73] Jürgen Schmidhuber. \"Learning to control fast-weight memories: An alternative to dynamic recurrent networks\", Neural Computation 1992\n\n[74] Imanol Schlag et al. \"Linear transformers are secretly fast weight programmers\", PMLR 2021\n\n[75] Lemeng Wu et al. \"Centroid transformers: Learning to abstract with attention\", 2021"}, {"review_id": "2tXcgTuy4r", "reviewer": "Reviewer_AnYY", "summary": "The paper introduces Contract‑and‑Broadcast Self‑Attention (CBSA), a new attention operator derived by unrolling a unified \"compress‑all‑by‑contracting‑a‑few\" objective. A small set of learnable representative tokens are first extracted by a cross‑attention‑like step, then contracted (self‑attention among the representatives) and finally broadcast back to all tokens via the learned coefficients. With a fixed number of representatives m, CBSA scales linearly in sequence length and subsumes several known mechanisms: it reduces to full self‑attention when every token is its own representative, to an input‑dependent channel mixer when the representatives are orthogonal, and to TSSA‑style channel attention when they are orthogonal & fixed.\n\nExperiments on ImageNet‑1k classification and ADE20K segmentation show that CBSA matches or exceeds the accuracy of MSSA/CRATE while computing 2–3× fewer pair‑wise similarities and retaining desirable interpretability traits such as emergent segmentation and robustness to weight perturbations.", "questions": "- Could the authors provide some visualizations of segmentations produced by different models? For example, results on ADE20K.\n\n- Why is “Robustness under parameter perturbation” a desirable property? Would this property be useful in certain practical settings?\n\n- Could the authors provide a detailed mathematical description of the proposed Contract-and-Broadcast Self-Attention (CBSA) in the appendix (i.e., a mathematical version of Algorithm 1: PyTorch implementation of a CBSA layer in the appendix)? This could help readers quickly understand the model architecture.", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- Unified Principle: Provides a single, well‑motivated objective linking interpretability (subspace compression) and efficiency, rather than ad‑hoc architectural tweaks.\n\n- Shows, in closed form, how full spatial attention, channel mixers, and pure channel attention are special cases of CBSA. This is intellectually appealing and may guide future designs.\n\n- For fixed $m$, compute and memory are $O(N)$, yet the operator retains spatial interactions, also with detailed complexity analysis (Table 1) and wall‑clock measurements.\n\n- Competitive Top‑1 on ImageNet and mIoU on ADE20K with fewer pairwise similarity computations; robustness and segmentation emergence are analysed with convincing visualizations.\n\n\n**Weaknesses**\n\n- [minor] Statistical Rigor: Only some results include error bars (Fig. 2); other tables (e.g., Tables 2–4) lack variance estimates.\n\n- [minor] Writing density – The derivations (Sec 3) are terse; readers unfamiliar with previous work may struggle. Some notation (e.g., gradient signs, scaling constants) could be clarified.", "was_revised": true, "final_justification": "Most of my questions and concerns have been addressed. Overall, I believe this submission makes interesting and valuable contributions to interpretable and efficient attention mechanisms. I would like to maintain my original score and vote for acceptance.", "rebuttal": "We sincerely appreciate your recognition of our work and the constructive suggestions.\n\n> **Weakness 1**: Statistical Rigor: Only some results include error bars (Fig. 2); other tables (e.g., Tables 2–4) lack variance estimates.\n\nAs suggested, we will present all results in Figure 2 with their error bars.\n\nHowever, due to limited resources, we are currently unable to estimate the variance in model performance.\n**We humbly suppose that this concern stems from the marginal performance improvement when comparing our CBSA to TSSA.** If so, we believe our new experimental results will address this.\n\nOur CBSA now substantially outperforms TSSA by simply adopting a pooling-based initialization strategy for the representatives. \nPreviously, the representatives are initialized via learnable parameters, which turns out to be the performance bottleneck of our method. Please refer to our response to the **Weakness 1** identified by Reviewer FLEn for more details and improved results.\n\n**Revised Table 2**: Fair comparisons between CBSA, TSSA, and MSSA on ImageNet-1K. The representation compactness is measured in the final attention layer.\n\n| Models | CRATE-T  | CBT-T | ToST-T  | CRATE-S | CBT-S   | ToST-S  |\n| - | - | - | - | -| - | -|\n| $R_c(\\boldsymbol{Z}^L \\mid \\boldsymbol{U}^L_{[K]})$ | 447$\\pm$55 | **622**$\\pm$**44** | 665$\\pm$31 | 987 $\\pm$137 | **1317**$\\pm$**82** | 1558 $\\pm$51 |\n|acc@1| 64.7| **63.2**| 61.2| 72.1 | **71.4** | 68.5|\n\nResults of our method are highlighted **in bold**. The standard deviations are computed over 1,000 images.\n\nWe also observe that models with more compact final representations (i.e., smaller $\\small R_c(\\boldsymbol{Z}^L \\mid \\boldsymbol{U}^L_{[K]})$) demonstrate superior performance. \nThis suggests that $\\small R_c(\\boldsymbol{Z}^L \\mid \\boldsymbol{U}^L_{[K]})$ may serve as **a more reliable evaluation metric**.\n\n> **Weakness 2**: Writing density – The derivations (Sec 3) are terse; readers unfamiliar with previous work may struggle. Some notation (e.g., gradient signs, scaling constants) could be clarified.\n\nThank you for this valuable feedback. \n\nWe will continue refining our manuscript to achieve better balance between derivations, conclusions and discussions.\n\n> **Question 1**: Could the authors provide some visualizations of segmentations produced by different models? For example, results on ADE20K.\n\nWe will visualize and compare the segmentation masks predicted by different models in the appendix of our final version.\n\nIn addition to these qualitative results, \nwe report that the best performance of our method on ADE20K has been improved **from 52.0 to 53.3**,\nowing to the modified representative initialization strategy.\nWe respectfully refer the reviewer to our response to Reviewer FLEn.\n\n> **Question 2**: Why is “Robustness under parameter perturbation” a desirable property? Would this property be useful in certain practical settings?\n\n**The robustness is desirable because the following aspects**:\n+ **It aligns with the derivation**.\nOur CBT theoretically models the attention heads as subspaces, and the projection vectors as basis vectors.\nTherefore, perturbations to these basis vectors do not significantly alter their spanned subspaces, thereby largely preserving the model peformance.\n\n+ **It calls for advanced interpretation**.\nThe seminal work [76] has already noticed that it is the spanned subspaces rather than the individual units that contains the semantic information in the deeper layers of neural networks. This also aligns with more recent interpretability advances where individual neurons are difficult to interpret [77] and interpretable attention heads can be identified [78].\n\n+ **It enables potential applications** in model quantization and severe industrial environments where perfect parameter stability cannot be guanranteed. However, as these fields fall beyond our expertise, these applications remain speculative.\n\n> **Question 3**: Could the authors provide a detailed mathematical description of the proposed Contract-and-Broadcast Self-Attention (CBSA) in the appendix (i.e., a mathematical version of Algorithm 1: PyTorch implementation of a CBSA layer in the appendix)? This could help readers quickly understand the model architecture.\n\nThank you for this valuable feedback.\n\nWe present the mathematical version of Algorithm 1 as follows.\n\n***\n**Algorithm 1: Contract-and-Broadcast Self-Attention**\n***\n**Input**: input tokens $Z \\in \\mathbb{R}^{d \\times N}$, number of representatives $m$\n\n**Parameters**:  subspace bases $U_{[K]} = \\lbrace U_k \\in \\mathbb{R}^{d \\times p} \\rbrace_{k=1}^K$, initialization of representatives $Q_0 \\in \\mathbb{R}^{d \\times m}$ (optional), step-size for input tokens $\\kappa$, step-size for representatives $\\eta$\n\n**parallel for** $k \\in \\lbrace 1, \\ldots, K \\rbrace$ **do**\n\n$\\ \\ \\ \\ $\\# project input tokens onto the $k$-th subspace\n\n$\\ \\ \\ \\ Z_k \\coloneqq U_k^\\top Z$ $\\ \\ $\\# token projections onto the $k$-th subspace, $Z_k \\in \\mathbb{R}^{p \\times N}$\n\n$\\ \\ \\ \\ $\\# Initialize the representatives of $Z_k$\n\n$\\ \\ \\ \\ $**If** $Q_0$ is parameterized **then**  $\\ \\ $\\# not recommended\n\n$\\ \\ \\ \\ \\ \\ \\ \\ Q_k \\coloneqq U_k^\\top Q_0$\n\n$\\ \\ \\ \\ $**else** $\\ \\ $\\# recommended\n\n$\\ \\ \\ \\ \\ \\ \\ \\ Q_k \\coloneqq  \\text{AvgPool}(Z_k, m)$ $\\ \\ $\\# $Q_k \\in \\mathbb{R}^{p \\times m}$, $\\text{AvgPool}$ denotes average pooling\n\n$\\ \\ \\ \\ $\\# Extract the representatives of $Z_k$\n\n$\\ \\ \\ \\  \\overline{A}_k \\coloneqq \\text{softmax}\\left(Z_k^\\top Q_k \\right)$ $\\ \\ $\\# the attention matrix $ \\overline{A}_k \\in \\mathbb{R}^{N \\times m}$ \n\n$\\ \\ \\ \\  Q_k \\coloneqq Q_k + \\eta Z_k \\overline{A}_k$  $\\ \\ $\\# the representatives of $Z_k$\n\n$\\ \\ \\ \\ $\\# Contract the representatives\n\n$\\ \\ \\ \\  Q_k \\coloneqq Q_k \\text{softmax}\\left(Q_k^\\top Q_k \\right)$ $\\ \\ $\\# the contractions of the representatives\n\n$\\ \\ \\ \\ $\\# Broadcast the contractions back to input tokens\n\n$\\ \\ \\ \\ Z_k \\coloneqq \\kappa Q_k \\overline{A}_k^\\top$\n\n**end parallel** \n\n**return** $\\sum_{k=1}^{K} U_k Z_k$\n***\n\nWe hope the discussions above could resolve your concerns on our work. Please let us know if any further clarification is needed.\n\n[76] Christian Szegedy et al. \"Intriguing properties of neural networks.\" ICLR 2014\n\n[77] Nelson Elhage et al. \"Toy models of superposition.\" 2022\n\n[78] Catherine Olsson et al. \"In-context learning and induction heads.\" 2022\n\nReferences [1–61] are provided in the manuscript, references [62–64, 79] are listed in our response to Reviewer FLEn, and references [65-75] are listed in our response to Reviewer  2uTu."}, {"review_id": "6siJCSWPQu", "reviewer": "Reviewer_PryS", "summary": "The paper proposes a new approach to the computation of the Attention mechanism, called CBSA, that runs in linear time. This approach comes from optimizing an information-theoretic objective, inferred from the maximal coding rate reduction principle. The basic idea is to choose some representatives of the input tokens and contract them to reduce redundancy, then broadcast them and use them to update the full token set. \nThe CSBA operator proposed is interpretable, it indeed comes from information theory and aims to project over some small dimensional orthonormal basis, it is efficient as it runs in O(n) and it also has a very strong theoretical backbone by being obtained from unrolling an algorithm via GD steps. \nExtensive experiments on image classification and segmentation tasks highlight CBSA to be effective and computationally efficient, outperforming channel attention and matching full attention in performance, but with smaller computations required.", "questions": "none", "limitations": "•\tThe authors suggest representative selection can be made via SVD, fixing or learning, but they do not provide comparisons among them\n•\tThere is no analysis on the step size specifying how results vary according to them", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "STRENGTHS\n•\tThe paper provides a rigorous derivation of the algorithm starting from a theoretically well-grounded objective\n•\tFinds a way to unify two very promising methods like channel and full attention, combining them into a very powerful unified algorithm\n•\tThe procedure is interpretable from an information point of view\n•\tThe algorithm is efficient since it runs linearly in the sequence length\n•\tThe theoretical properties are backed by the empirical results\nWEAKNESSES\n•\tThe representative tokens might have undesired behaviors depending on the way they are computed, and this hasn’t been explored\n•\tNo code released yet to see the implementation", "was_revised": true, "final_justification": "Authors have thoroughly addressed my concerns and the paper has improved in clarity. My final score is 5.", "rebuttal": "We sincerely appreciate your recognition of our work and the insightful comments.\n\n> **Weakness 1**: The representative tokens might have undesired behaviors depending on the way they are computed, and this hasn’t been explored.\n\nWe are particularly pleased to have the opportunity to address this farsighted concern.\n\nOur findings actually attribute the differences between attention mechanisms to the distinct choices of representatives, \nas **their number and structure** (e.g. orthogonality) induce divergent information propagation patterns (i.e., the behaviors).\nGenerally, three cases of behaviors can be analyzed:\n+ **Behavior 1**: interactions among the representatives;\n+ **Behavior 2**: interactions between input tokens and their representatives;\n+ **Behavior 3**: interactions among input tokens.\n\nUndesirable/Desirable behaviors of representatives has been reported in our Section 3.3:\n1. **Undesired Behavior 1**: orthogonal representatives **do not interact with each other**, resulting the channel-mixer variant Eq. (19);\n2. **Undesired Behaviors 1 and 2**: Orthogonal and fixed (i.e., input-agnostic) representatives **scale in-place along their fixed directions**, leading to the channel attention Eq. (20).  \n3. **Desired Behavior 3**: In contrast, when input tokens are representatives themselves, their behaviors become **highly flexible and expressive** at the cost of efficiency, deriving the full attention variant Eq. (17).\n\nAs the analyses above remain special cases, we propose three promising perspectives to examine their general behaviors:\n1. **Graph perspective.** We can view **Behaviors 1–3** as the edges (i.e., the information flow pathways) in a graph where nodes are input tokens or their representatives. The more densely connected the graph is, which is determined by the number and structure of the representatives, the more expressive its information propagation becomes, but the harder it is to maintain efficiently.\nWhen input tokens are themselves representatives, it is a  fully-connected digraph.\nWhen representatives are orthogonal, it reduces to a bipartite graph.\n\n2. **Dictionary learning perspective.** We can view the representatives as atoms of a dictionary and recognize that the special cases we studies in Section 3.3 resembles the advances in dictionary learning.\nFor example, when the representatives are orthogonal and fixed, they form a complete dictionary.\nWhen they are orthogonal but input-dependent, the structure resembles a submatrix of an overcomplete dictionary in compressed sensing [80].\nWhen input tokens are themselves representatives, it becomes a self-expressive dictionary [44].\nTherefore, we can interpret the way the representatives are computed as a form of dictionary learning, and their behaviors as the structure of the atoms in the learned dictionary.\n\n3. **Rank perspective.**\nWe can measure **Behavior 2** via the rank of the broadcast matrix (i.e., the attention matrix computed during the extraction step), which can be interpreted as encoding the edge weights between input tokens and their representatives.\nA significantly low rank suggests the presence of an undesired behavior where many nodes of the representatives become “dead” (i.e., failing to relay information).\nFor **Behaviors 1 and 3**, recent studies [62, 63] have shown that the effectiveness of softmax attention can be largely attributed to the attention matrix being consistently full-rank.\n \n**Leveraging the Rank Perspective to Identify Undesired Behavior and Improve CBSA Performance**  \nFrom the rank perspective discussed above, \nwe found that the rank of the attention matrix of the extraction (cross-attention) step, i.e,\n$\\text{rank} \\left( \\text{softmax} \\left( Z^\\top Q_0 \\right) \\right)$, \nwas overly small.\nThis suggests that fewer representatives than expected are effectively utilized to gather, process, and broadcast information from the input tokens, which is an **undesired Behavior 2**.\nThis low-rank bottleneck stemmed from the learnable initialization of $Q$, and resulted in the unsatisfactory performance of CBSA.\nNow, we adopt a **pooling**-based initialization strategy, i.e., $Q_0 = \\text{AvgPool}(Z)$. \nThis enables the attention matrix to achieve maximal rank, thereby improving model performance. \nA detailed analysis will be presented in the full paper.\n\n**Revised Table 2**: Fair comparisons between CBSA, TSSA, and MSSA on ImageNet-1K.\n| Models | CBT-T | CBT-S | ToST-T | ToST-S | CRATE-T | CRATE-S |\n| :- | - | - | - | - | - |-|\n| acc@1 (old)  | 62.9| 69.5 | - | - | - | - |\n| acc@1  | **63.2**  | **71.4** |  61.2 | 68.5| 64.7| 72.1|\n\nThe performance gain between small size models has been increased from 1% acc@1 to about 3%.\nPlease refer to our response to Reviewer FLEn for additional empirical evidences demonstrating the improvements achieved by addressing the undesirable interactions between input tokens and their representatives.\n\n> **Weakness 2**: No code released yet to see the implementation.\n\nThank you for your interest in the implementation details. \n\nThe PyTorch implementation of our CBSA was provided in Appendix: Algorithm 1. \nWe will open-source our code and demo upon publication. \nIn fact, our implementation builds upon the official PyTorch code from [27], with key modifications to the attention layer.\n\n> **Limitation**: The authors suggest representative selection can be made via SVD, fixing or learning, but they do not provide comparisons among them. There is no analysis on the step size specifying how results vary according to them.\n\nThe points raised by the reviewer are highly valuable, as they may further enhance the model’s interpretability and offer deeper insights into its behavior. For instance, one intriguing question is whether the cross-attention mechanism implicitly implements or approximates an SVD-like procedure. \nAnother question is to analyze whether adaptive step sizes outperform fixed ones. It is also worth exploring whether the step sizes across different layers or attention heads correlate with their functional importance. \n\nWhile these questions are highly valuable for better understanding the learning dynamics, a thorough investigation is beyond the scope of this work. We appreciate the reviewer’s feedback and will consider them as promising directions for future research. \n\nWe hope the discussions above could resolve your concerns on our work. Please let us know if any further clarification is needed.\n\n[80] Emmanuel J Candes and Terence Tao. \"Decoding by linear programming.\" IEEE Transactions on Information Theory 2005\n\nReferences [1–61] are provided in the manuscript, references [62–64, 79] are listed in our response to Reviewer FLEn, and references [65-75] are listed in our response to Reviewer  2uTu."}]

Author rebuttal:

Meta-review: This paper proposes Contract-and-Broadcast Self-Attention (CBSA), a unified mechanism addressing Transformer attention’s interpretability and efficiency via a "compress-all-by-contracting-a-few" objective. Reviewers uniformly praised its strengths: a well-motivated unified principle linking subspace compression and efficiency, theoretical derivations showing it subsumes full and channel attention as special cases, linear computational scaling, and competitive performance on ImageNet and ADE20K with fewer computations.
However, concerns were noted: limited variance estimates in results, terse derivations in Section 3, under-explored representative token behaviors, lack of code release, and insufficient quantitative analysis of segmentation properties.
Despite these, the work makes valuable contributions by unifying key attention mechanisms under a rigorous framework with strong empirical support. Its potential to guide future attention design and bridge efficiency-interpretability gaps is significant. Thus, the AC concurs this is a valuable contribution and recommends acceptance.
Decision: Accept (spotlight)