Title: Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few
Venue: NeurIPS 2025
Accepted: yes
Keywords: Attention Mechanisms

Interpretability

Efficiency
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

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 Contractand-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 this https URL.

1 Introduction

Attention mechanisms have been widely applied across diverse areas, including computer vision [1, 2], natural language processing [3, 4], and scientific discovery [5]. Nonetheless, a series of puzzling phenomena-such as emergent segmentation properties [6], in-context learning ability [7], attention collapse [8, 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 to the sequence length impedes its broader applications in real-time systems [11], as well as the processing of long documents [12] and high-resolution images [13].

In light of these challenges, it has been more crucial to mathematically demystify attention mechanisms, which offers deeper insights into their simplification and acceleration. Over the past few years, remarkable advances have been made in addressing the interpretability or efficiency issue separately. On the one hand, in an ante-hoc manner, attention mechanisms can be interpreted by optimization objectives grounded in clustering [14], denoising [15], energy minimization [16], matrix decomposition [17], and contrastive learning [18]. These inherently interpretable approaches are more rigorous than post-hoc explanations [19]. On the other hand, numerous techniques have been developed to alleviate the quadratic complexity of self-attention, including sparse attention [20] and linear attention [21].

However, the joint development of interpretability and efficiency in attention mechanisms remains a 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-driven optimization objective called MCR 2 [22]. Indeed, this objective has been utilized for designing an interpretable softmax attention, MSSA [23], and a linear-time attention, TSSA [24]. 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). 1 Therefore, instead of an isolated mechanism, we aim to develop a framework that unifies these varied attention mechanisms in an interpretable way, revealing how they are fundamentally connected yet distinctly presented, as well as the trade-off between expressive capacity and efficiency.

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-theoretic essence of input tokens can be captured by a small number of representatives [25, 26] of them. Since the former has been formulated as the MCR 2 objective [22, 23] (see Section 2), the remaining task is to leverage the representatives to optimize it, thereby efficiently compressing all by contracting a few (see Section 3.1). By unrolling the resulting optimization objective, we derive our Contract-and-Broadcast Self-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.

LayerNorm

LayerNorm

ISTA

Forward-Pass Objective:  MCR

2

Contract-and-Broadcast

Self-Attention

（

CBSA

）

Backward-Pass Loss:  Cross-Entropy

minimize the compression term of Eqn. 4

Pooling

Cross-Attention

Projection

Query

Key

Value

Self-Attention

Query

Key

Value

Matmul

Projection

Contraction

Broadcast

Extraction

Initialization

a gradient step of obj. func.

to extract representatives

(c)

CBSA (ours)

(a) Channel Attention

(b) Linear Attention

orthogonal

representatives

orthogonal and fixed

representatives

self-expressed

representatives

(d) Softmax Attention

Contract-and-Broadcast Transformer

(

CBT

)

Distribution of

input tokens

Distribution of

compressed

tokens

Scaling along

orthogonal directions

Contraction of

representatives

Paper contributions The contributions of the paper are summarized as follows.

We formulate an optimization objective that unifies the interpretability and efficiency of attention mechanisms through the idea of compressing all by contracting a few .

We derive an inherently interpretable and efficient attention mechanism, CBSA, which is a potential unified formula for different attention mechanisms.

We validate the interpretability and efficiency of CBSA through extensive experiments on visual tasks.

2 Notations and preliminaries

Notations. Given a positive integer n , let [ n ] . = { 1 , 2 , . . . , n } . For s ≥ n , let O ( s, n ) ⊆ R s × n denote the set of s × n matrices with orthonormal columns, and O ( s ) . = O ( s, s ) denote the set of s × s orthogonal matrices. Let I n denote an identity matrix of size n , and O n denote a zero square matrix of size n . Given a vector v ∈ R n , let Diag( v ) ∈ R n × n be a diagonal matrix with the entries of v along its diagonal. Let Z ∈ R d × N denote N input tokens represented in the ambient space R d . Specially, let Z ℓ denote these tokens feeding into the ℓ -th attention layer. The same holds for the representatives of input tokens, Q ∈ R d × m , where m can be much smaller than N .

̸

Union of subspaces. Although the union of nonlinear manifolds provides a better approximation [27], we adopt a much simpler structure: the union of (low-dimensional) linear subspaces. 2 Specifically, it is parameterized as K incoherent p -dimensional subspaces spanned by orthonormal bases U [ K ] . = { U k ∈ O ( d, p ) } K k =1 , where pK = d and U ⊤ i U j = O p , ∀ i = j . We refer to the i -th basis vector of U k as u ki . Similar to [23], we implement U [ K ] as learnable parameters in each attention layer and thus denote the U [ K ] implemented by the ℓ -th layer as U ℓ [ K ] = { U ℓ k } K k =1 .

Coding rate. To quantify the compactness, i.e., the extent to which tokens are compressed towards subspaces, we adopt the (lossy) coding rate [28], which measures how efficiently the token distribution can be covered by ϵ -balls under a given quantization precision ϵ > 0 (as illustrated in Fig. 1(a)). The coding rate of input tokens in the ambient space R d is defined as:

MCR 2 objective. The Maximal Coding Rate Reduction (MCR 2 ) [22] objective adopted in [23] is defined on the coding rate as follows:

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 in (2) corresponds an interpretable softmax attention mechanism.

3 Methods

3.1 Compressing all by contracting a few

Due to the existence of Gram matrix in (1), 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 [29] and further introduced a 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 [30, 31], we propose a simple but flexible approach to streamline the optimization of MCR 2 : compressing all input tokens by contracting a small number of representatives of them . Before demonstrating that this achieves linear complexity in N and prevents the aforementioned degeneration, we first formulate it as a new optimization objective.

An initial attempt is to impose a set of equality constraints on the coding rates as follows:

where representatives Q . = q ( Z ) are extracted from input tokens Z by a differential function q ( · ) : R d × N → R d × m . 3 This new objective in (3) is equivalent to the original objective in (2) but is more efficient to handle because the number of representatives (e.g., m = p = d / K ) can be far smaller.

Since that the equality constraints in (3) is overly restrictive in practice, we attempt to relax them by introducing a tolerance τ , which uniformly bounds the absolute difference of the two coding rates within each subspace, i.e., | R ( U ⊤ k Q ) -R ( U ⊤ k Z ) | ≤ τ . Therefore, contracting the representatives will correspondingly compress the input tokens as well up to the tolerance τ .

Consequently, we have a relaxed optimization problem for our subsequent derivations as follows:

In this paper, we employ an arguably simplest way to extract Q in each subspace: U ⊤ k Q = U ⊤ k ZA k , where A k ∈ R N × m is the coefficient matrix over dictionary U ⊤ k Z ∈ R p × N , i.e., the projected representatives in each subspace are linear combinations of the projected input tokens.

3.2 Contract-and-Broadcast Self-Attention

We now are ready to derive an attention mechanism in an ante-hoc interpretable manner, by implementing 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 as approximating the coding rate [34], we employ this idea to extract representatives that satisfy the inequality constraints in (4), thus capturing the information-theoretic 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.,

where the initial guess Q ini 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 Q ini 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 (5) effectively is the coefficient matrix A k .

Representative contraction and contraction broadcast. To derive the attention mechanism, following [23], we focus on optimizing the compression term via a 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 (4) with respect to input tokens as follows:

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 in (7) is proportional to the gradient with respect to representatives. Hence, we refer to this formula as a Contract-and-Broadcast Self-Attention (CBSA), reflecting that: a) the contraction term gives the contracting directions of the representatives (abbreviated as contractions);

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 [23], we approximate it by a Gram matrix and a softmax function (i.e., an attention matrix): 4

Note that the contraction term now effectively constitutes a self-attention operation in which the linear projections for the query, key, and value are all identical to the subspace basis, i.e., W query = W key = W value = U ⊤ k . By default, we implement CBSA via (8) rather than (7) in our experiments.

Overview of CBSA. The workflow of CBSA is illustrated in the left panel of Fig. 1. We also construct an inherently interpretable Contract-and-Broadcast Transformer (CBT) by stacking CBSA with the ISTA module [23], which is also derived via algorithm unrolling. 5 We report the computational complexities of CBSA and its sub-operations in Table 1 and compare the FLOPs of different attention mechanisms in Fig. 2 where d = 384 , H = 6 , and a patch size of 16 × 16 . Provided N > 2 d / H = 2 p (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. 9.

128

196

224 2

576

384 2

1024

512 2

1600

640 2

4096

1024 2

Number of tokens / image resolution

0.25

0.5

1

2

4

8

16

GFLOPs

×2.0

×2.9

×4.1

×9.3

FLOPs of an attention layer in ViT-S

MSSA

CBSA (ours)

MHSA

Liner Attention

Ω(MSSA)

Ω(CBSA)

sub-operations of CBSA

Ω(

extraction

)

/

Ω(

broadcast

)

Ω(

contraction

)

Ω(

projection

)

2

Nd

2

+2

N

2

d

2

Nd

2

+3

Nmd

+2

m

2

d

Nmd

m

2

d

Nd

2

Matmul

Matmul

Matmul

Softmax

Matmul

Matmul

Matmul

Matmul

Matmul

Matmul

Matmul

Matmul

Matmul

Matmul

Softmax

Matmul

Softmax Attention

Pooling

Matmul

Softmax Attention

Pooling

Softmax Attention

(b) Linear Attention

(a)  Softmax Attention

(c) Channel Attention

(d) Agent Attention

(e) CBSA (ours)

Key

Value

Query

3.3 CBSA as a unified attention formula

In this subsection, we explore the potential of CBSA to serve as a unified formula for different attention mechanisms. Unlike recent work [37, 38], CBSA encompasses a broader spectrum of mechanisms (see Fig. 3) in an 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

input tokens are self-expressed [39], where data samples are linearly represented over a dictionary composed of themselves. Then we can take a 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]:

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:

where L k ∈ O ( p ) , R k ∈ O ( N,p ) , Σ k is a p × p diagonal matrix of singular values, and the columns of L k are known as the principal directions. Then, by right multiplying both sides by R k , we have:

By comparing (11) with the way to form the representatives, i.e., U ⊤ k Q = U ⊤ k ZA k , we let U ⊤ k Q = L k Σ k and A k = R k . Substituting them into (7) leads to the following operator:

where F is a function defined on the spectrum of a positive semi-definite matrix and applies f ( λ i ) = ϵ 2 / ( ϵ 2 + λ i ) to each eigenvalues { λ i } p i =1 of the covariance matrix. 6 This operator highly resembles the linear attention [21], due to that it also factorizes the N × N attention matrix and multiplies the key and value first to linearize the computational complexity. 7 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 k 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 k = L k , thereby being orthogonal and input-agnostic (fixed). Then, (12) is simplified to:

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-wise scaling factors [40, 41].

Agent attention variant. Agent Attention is basically a variant of CBSA with the contraction step removed 8 which can be perceived in Fig. 3. Although this removal appears to confine the token-mixing ability, it is compensated by the pooling-based initialization, which is also a token mixer [42].

To gain some intuition of the gap in expressive capacity among the aforementioned mechanisms, we illustrate their compression patterns in the right panel of Fig. 1. The channel attention variant is restricted to compressing input tokens along fixed axes parameterized by U [ K ] , whereas the linear attention variant compresses them along principal directions that are dynamically determined by the input. We argue that such a dynamism is crucial for in-context learning [7] and for mitigating

superposition [43]. 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 a dictionary learning perspective, where the representatives correspond to the atoms of a 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 as representatives, they constitute a self-expressive 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 [45, 46], we focus on classical visual tasks such as image classification and semantic segmentation, where higher resolutions generally lead to better accuracy [47, 48].

Baseline and training configuration. We compare our CBSA to the vanilla softmax attention [49, 1], and interpretable attention mechanisms based on MCR 2 , 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 are reproduced under varied settings for fair comparisons. By default, the training configuration follows the baselines, with detailed information provided in Appendix C.

Implementation detail. The projection back to the ambient space, which should theoretically be a left multiplication by U k , is over-parameterized with an independently learnable matrix. This strategy is also adopted in 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 to flexibly choose between compression and decompression. The PyTorch implementation is provided in Appendix D.

4.1 Advantages enabled by interpretability

In this subsection, we show superior advantages of CBSA over black-box attention mechanisms. The most essential aspect of our CBSA is its interpretability, which induces other desirable properties such as robustness and emergent segmentation.

(a) Raw

(b) 1st iter

(c) 256 iters

(d) 512 iters

(e) 640 iters

(f) 768 iters

(g) 896 iters

(h) 1024 iters

Compact and structured representation. Similar to MCR 2 [22], our optimization objective (4) aims to learn compact and structured representation by compressing input tokens towards low-dimensional subspaces. To confirm whether its iterative gradient steps can actually achieve this goal, we iterates

the linear attention variant of CBSA (12) on synthetic data, where image tokens are modeled as the points in a three-dimensional space R 3 . Specifically, it is conducted on each class in the ambient space (thus being parameter-free) with a forward-only manner. As shown in Fig. 4, the representation ultimately admits a union of well-separated one-dimensional subspaces after 1024 iterations.

1

2

3

4

5

6

7

8

9

10

11

12

Layer index -

300

400

500

600

700

800

Rc

(

Z

U

[

K

]

)

Coding rate (tiny models)

CRATE

TOST

CBT (ours)

1

2

3

4

5

6

7

8

9

10

11

12

Layer index -

0

500

1000

1500

2000

Rc

(

Z

U

[

K

]

)

Coding rate (small models)

CRATE (w/o conv)

1

2

3

4

5

6

7

8

9

10

11

12

Layer index -

220

240

260

280

300

320

Rc

(

Z

U

[

K

]

)

Normalized Coding rate (tiny models)

1

2

3

4

5

6

7

8

9

10

11

12

Layer index -

400

600

800

1000

Rc

(

Z

U

[

K

]

)

Normalized Coding rate (small models)

Compressing all by contracting a few. To confirm our interpretation that CBSA compresses all input tokens by contracting a 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 variant 9 in Fig. 5. This normalized coding rate is invariant to in-place scaling and depends on 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-1K (see Table 3); 10 b) the latter half of the layers exhibit consecutive compression, in nearly all models. 11 Then, in Fig. 6, we 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. 12

0

2

4

6

8

Reduction w.r.t. input tokens

w.r.t. input tokens

w.r.t. representatives

20

0

20

40

60

80

100

120

Reduction w.r.t. representatives

Reduced coding rate across all layers and heads

4

2

0

2

4

6

8

10

Reduction w.r.t. input tokens

w.r.t. input tokens

w.r.t. representatives

50

0

50

100

150

Reduction w.r.t. representatives

Reduced coding rate across all layers and heads

L1

L2

L3

L5

L6

L7

L8

L9

L10

L11

L12

L1

L2

L3

L5

L6

L7

L8

L9

L10

L11

L12

CBT (ours)

CRATE

CRATE+CBT

6

6

Z

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. as shown in Fig. 7. To address this

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.

0

20

40

60

80

100

GFLOPs (decoder only)

38.0

42.0

46.0

50.0

53.3

mIoU (%)

Semantic segmentation

CBT (ours)

DEPICT

CBT  (ours)

Segmenter

10

3

10

2

10

1

10 0

Noise level -

0

10

20

30

40

mIoU

Parameter robustness - tiny models

ToST

DEPICT

CBT  (ours)

Segmenter

10

3

10

2

10

1

10 0

Noise level -

0

10

20

30

40

mIoU

Parameter robustness - small models

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.

0

25

50

75

100

125

GFLOPs

72

76

80

84

Accuracy (%)

Image resolution: 224x224

MHSA

Linear Attention

CBSA *  (ours)

CBSA

0

100

200

300

400

GFLOPs

72

76

80

84

88

Accuracy (%)

Image resolution: 384x384

MHSA

Linear Attention

CBSA *  (ours)

CBSA

196

224 2

576

384 2

1600

640 2

4096

1024 2

Number of tokens / image resolution

0.1

0.2

0.4

0.8

1.6

3.2

6.4

12.8

GFLOPs

FLOPs of modules in ViT-T

MHSA

MLP

Liner Attention

CBSA *  (ours)

CBSA (ours)

196

224 2

576

384 2

1600

640 2

4096

1024 2

Number of tokens / image resolution

0.5

1

2

4

8

16

32

64

GFLOPs

FLOPs of modules in ViT-B

MHSA

MLP

Liner Attention

CBSA *  (ours)

CBSA (ours)

4.2 Evaluations on real-world visual tasks

We pretrain CBT models on the ImageNet-1k dataset, and finetune them on several downstream datasets. The top-1 accuracy on validation sets is reported in Table 3. In particular, our CBT-Small achieves comparable top-1 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, our CBT models remain competitive with CRATE while computing significantly fewer pairwise similarities. These results confirm that, from MSSA to 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 high-resolution images (e.g., 512 × 512 ) in Table 7 (Appendix C), showing that our methods consistently achieves superior training and inference efficiency.

To investigate the potential of applying CBSA to large-scale pretraining, we preliminarily finetune ViT models pretrained on ImageNet-21k 13 by adapting their attention blocks into the CBSA style. For comparison, we also adapt them into linear attention. Experimental results are shown in Fig. 9. Although CBSA deviates more from MHSA than linear attention and is consequently harder to adapt from pretrained ViTs, CBSA achieves comparable performance to 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 stacking CBSA layers without the feed-forward modules on 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. 8. Clearly, our CBT decoder consistently surpasses both white-box (DEPICT) and black-box (Segmenter) counterparts that rely on softmax attention. In particular, the best-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.

5 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 or memory can 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 low-dimensional 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-Broadcast Self-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 as well as new attention mechanisms in an inherently interpretable way.

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

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.

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 15979-15988, 2022.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human 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.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andrew J Ballard, Andrew Cowie, Bernardino Romera-Paredes, 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 prediction with AlphaFold. Nature , 596(7873):583-589, 2021.

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In IEEE/CVF International Conference on Computer Vision , pages 9630-9640, 2021.

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.

Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer. arXiv preprint arXiv:2103.11886 , 2021.

Yihe Dong, Jean-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.

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.13835 , 2024.

Quoc-Vinh Lai-Dang. A survey of vision transformers in autonomous driving: Current trends and future directions. arXiv preprint arXiv:2403.07542 , 2024.

Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

Ze 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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11999-12009, 2022.

Lemeng Wu, Xingchao Liu, and Qiang Liu. Centroid transformers: Learning to abstract with attention. arXiv preprint arXiv:2102.08606 , 2021.

Peng Wang, Yifu Lu, Yaodong Yu, Druv Pai, Qing Qu, and Yi Ma. Attention-only transformers via unrolled subspace denoising. In International Conference on Machine Learning , 2025.

Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, 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.

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.

Ruifeng Ren and Yong Liu. In-context learning with transformer is really equivalent to a contrastive learning pattern. arXiv preprint , 2023.

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.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , volume 119, pages 5156-5165, 2020.

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 Neural Information Processing Systems , 2020.

Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D. Haeffele, and Yi Ma. White-box transformers via sparse rate reduction. In Advances in Neural Information Processing Systems , 2023.

Ziyang Wu, Tianjiao Ding, Yifu Lu, Druv Pai, Jingyuan Zhang, Weida Wang, Yaodong Yu, Yi Ma, and Benjamin David Haeffele. Token statistics transformer: Linear-time attention via variational rate reduction. In International Conference on Learning Representations , 2025.

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.

Chong You, Chi Li, Daniel P Robinson, and Rene Vidal. Self-representation based unsupervised exemplar selection in a union of subspaces. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(5): 2698-2711, 2020.

Tianjiao Ding, Shengbang Tong, Kwan Ho Ryan Chan, Xili Dai, Yi Ma, and Benjamin D. Haeffele. Unsupervised manifold linearizing and clustering. In IEEE/CVF International Conference on Computer Vision , pages 5427-5438, 2023.

Yi Ma, Harm Derksen, Wei Hong, and John Wright. Segmentation of multivariate mixed data via lossy data coding and compression. IEEE transactions on pattern analysis and machine intelligence , 29(9): 1546-1562, 2007.

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-103, 2022.

Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In AAAI Conference on Artificial Intelligence , pages 14138-14148, 2021.

Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305.16300 , 2023.

Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In International Conference on Machine Learning , pages 399-406, 2010.

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.

Qishuai Wen and Chun-Guang Li. Rethinking decoders for transformer-based semantic segmentation: A compression perspective. In Advances in Neural Information Processing Systems , 2024.

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 European conference on computer vision , pages 124-140, 2024.

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.

Jusen Du, Jiaxi Hu, Tao Zhang, Weigao Sun, and Yu Cheng. Native hybrid attention for efficient sequence modeling. arXiv preprint arXiv:2510.07019 , 2025.

Han Guo, Songlin Yang, Tarushii Goel, Eric P Xing, Tri Dao, and Yoon Kim. Log-linear attention. arXiv preprint arXiv:2506.04761 , 2025.

Ehsan Elhamifar and René Vidal. Sparse subspace clustering. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition , pages 2790-2797, 2009.

Jie Hu, Li Shen, and Gang Sun. Squeeze-and-excitation networks. In CVPR , pages 7132-7141. Computer Vision Foundation / IEEE Computer Society, 2018.

Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In So Kweon. CBAM: convolutional block attention module. In ECCV (7) , volume 11211 of Lecture Notes in Computer Science , pages 3-19. Springer, 2018.

Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10809-10819, 2022.

Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, 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 arXiv:2209.10652 , 2022.

Emmanuel J Candes and Terence Tao. Decoding by linear programming. IEEE transactions on information theory , 51(12):4203-4215, 2005.

René Vidal, Yi Ma, and S. Shankar Sastry. Generalized Principal Component Analysis . Springer, 2016. ISBN 978-0-387-87810-2.

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.08894 , 2021.

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.

Feng Wang, Yaodong Yu, Guoyizhe Wei, Wei Shao, Yuyin Zhou, Alan Yuille, and Cihang Xie. Scaling laws in patchification: An image is worth 50,176 tokens and more. arXiv preprint arXiv:2502.03738 , 2025.

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 5998-6008, 2017.

Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In IEEE/CVF International Conference on Computer Vision , pages 7242-7252, 2021.

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 on Parsimony and Learning , volume 234, pages 72-93, 2024.

Ross Wightman. Pytorch image models. https://github.com/rwightman/pytorch-image-models , 2019.

Bolei 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 , 127(3):302-321, 2019.

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 arXiv:2507.19595 , 2025.

Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning , volume 80, pages 4052-4061, 2018.

Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning , volume 119, pages 9438-9447, 2020.

Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics , 9: 53-68, 2021.

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.

Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star-transformer. In Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 1315-1325, 2019.

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. In Advances in Neural Information Processing Systems , 2020.

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.

Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In International Conference on Learning Representations , 2024.

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-5948, 2023.

Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoc V. Le. Symbolic discovery of optimization algorithms. In Advances in Neural Information Processing Systems , 2023.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019.

Mark Everingham and John Winn. The pascal visual object classes challenge 2012 (voc2012) development kit. Pattern Analysis, Statistical Modelling and Computational Learning, Tech. Rep , 8(5):2-5, 2011.

Appendix

A Representative initialization

Besides pooling the input tokens, another intuitive choice for representative initialization is to set them as learnable (i.e., trainable) parameters, similar to the object queries in [61], the class embeddings in [50], and the registers in [62]. However, in this case we 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 = d H = 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-based initialization is employed, the rank of the attention matrix always attains its maximum, i.e., m . Nonetheless, as pointed out in [63], 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-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 ⊤ k Q = L k Σ k , A k = R k into (7):

Note that an arbitrary set of orthogonal representatives can be expressed as the principal directions under an orthogonal rotation and scaling, i.e.,

where P k ∈ O ( p ) denotes an orthogonal rotation within the k -th subspace, and Λ k is a diagonal matrix of size p . Generalizing the way to form the representatives in (4) from U ⊤ k Q = U ⊤ k ZA k to U ⊤ k Q = B k U ⊤ k ZA k , where B k ∈ R d × d , we have:

Similarly, by substituting U ⊤ k Q = P k L k Σ k Λ k , A k = R k Λ k and B k = P k into (22), we obtain

C More experimental results

Training setup. We train the CBT models in Table 3 and all models in Table 4 150 epochs with the Lion optimizer [64]. The learning rate is 2 . 0 × 10 -4 , the weight decay coefficient is 0 . 05 , and the batch size is 256 . We also incorporate a warm-up strategy over the first 20 epochs. For data augmentation, we adopt a 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 AdamW optimizer [65], a learning rate of 5 × 10 -5 , weight decay of 0 . 01 , and batch size 64 . The settings above are largely inherited from [23].

Quantitative segmentation properties. Following prior works [6, 51], we evaluate the zero-shot segmentation performance of ImageNet-1K pretrained models on the PASCAL VOC12 validation set [66]. Specifically, we assess the best-performing attention maps of the [CLS] token, but in a simplified setting that considers only three segmentation targets, instead of fine-grained 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 in Table 6.

Reduced coding rate across all layers and heads

L9

L1

L1

L2

L2

L3

L3

L5

L5

L6

L6

L7

L7

L8

L8

L9

L9

L10

L10

L11

L11

L12

L12

2

0

2

4

6

8

10

Reduction w.r.t. input tokens

w.r.t input tokens

w.r.t. representatives

40

20

0

20

40

60

80

100

120

Reduction w.r.t. representatives

4

2

0

2

4

6

8

10

Reduction w.r.t. input tokens

w.r.t input tokens

w.r.t representatives

60

40

20

0

20

40

60

80

100

Reduction w.r.t. representatives

Reduced coding rate across all layers and heads

L1

L1

L1

L2

L2

L3

L3

L3

L5

L5

L6

L6

L6

L7

L7

L7

L8

L8

L9

L9

L10

L10

L10

L11

L11

L12

L12

L12

L15

L15

L15

L16

L16

L16

L17

L17

L18

L18

L18

L19

L19

L19

L20

L20

L22

L22

L22

L24

L24

L24

D PyTorch implementation

E Limitations

Our work aims to propose a unified framework for both interpretable and efficient attention. Currently, we have proposed a unified optimization objective to derive such an attention mechanism, CBSA, which is capable to accommodate previous interpretable attention mechanisms derived from MCR 2 . Nonetheless, more fundamental and mathematical connections between our CBSA and existing efficient attention mechanism remain unexplored.

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 1 is justified in Section 3 and Section 4.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the current limitations of our work in Appendix E.

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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The assumptions can be found in Section 2 and Section 3 and complete proofs are in our Appendix.

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

All assumptions should be clearly stated or referenced in 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 a short proof sketch to 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: Our designs are clearly formulated in Section 3, and illustrated in Fig. 1. 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.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.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We present our experimental detail, such the training setup, in Appendix C. Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or 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. 6. 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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).

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

It 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 a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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, or 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 make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research conforms with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 AI.

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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: 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, or scraped datasets)?

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 a 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 a 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 or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.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 package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was 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. Crowdsourcing 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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to 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 on 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is 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 as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nSTRENGTHS\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\n\nQUESTIONS:\nnone\n\nLIMITATIONS:\n•\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", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThere isn't a limitation section in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAttention 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.\n\nSTRENGTHS AND WEAKNESSES:\nPaper 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.\n\nQUESTIONS:\n- 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)?\n\nLIMITATIONS:\n- See weaknesses\n- CBSA seems more like an alternative to TSSA than an improved method with regards to performance and computational efficiency.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}