Transformers Learn Faster with Semantic Focus
Authors: Parikshit Ram
Kenneth L. Clarkson
Tim Klinger
Shashanka Ubaru
Alexander G. Gray
Venue: NeurIPS 2025 poster ()
Keywords: transformers
sparse attention
learning convergence
analysis
Primary area: deep_learning
TL;DR: We demonstrate scenarios where sparse attention based transformer models learn and generalize faster, and theoretically characterize conditions under which this occurs.

Abstract: Various forms of sparse attention have been explored to mitigate the quadratic computational and memory cost of the attention mechanism in transformers. We study sparse transformers not through a lens of efficiency but rather in terms of learnability and generalization.  Empirically studying a range of attention mechanisms, we find that input-dependent sparse attention models appear to converge faster and generalize better than standard attention models, while input-agnostic sparse attention models show no such benefits -- a phenomenon that is robust across architectural and optimization hyperparameter choices.  This can be interpreted as demonstrating that concentrating a model's "semantic focus" with respect to the tokens currently being considered (in the form of input-dependent sparse attention) accelerates learning. We develop a theoretical characterization of the conditions that explain this behavior. We establish a connection between the stability of the standard softmax and the loss function's Lipschitz properties, then show how sparsity affects the stability of the softmax and the subsequent convergence and generalization guarantees resulting from the attention mechanism. This allows us to theoretically establish that input-agnostic sparse attention does not provide any benefits. We also characterize conditions when semantic focus (input-dependent sparse attention) can provide improved guarantees, and we validate that these conditions are in fact met in our empirical evaluations.

Transformers Learn Faster with Semantic Focus
Parikshit Ram', Kenneth Clarkson', Tim Klinger', Shashanka Ubaru', Alexander Gray?.
IBB Research, "Centaur Al Institute, Purruue University
parikshit ramiibm.....mmmmmmmmm com, klclarks@us iom,.om, iklinger@us ibm. com, shashanka. ubaru@ibm com, alexander, tra.deeeggeettttttttttttttttttt. org
Abstract
Various forms of sparse attention have been explored mitigate the quadratic computational and memory cost of the attention mechanism in transformers, We study sparse transformers not through lens of efficiency but rather terms of learnability and generalization. Empirically studying range of attention mecha- nisms, we find that depttdeeenndett sparse attention models appear to converge faster and generalize better than standard attention models, while input-agnostic sparse attention models show no such benefits phenomenon that robust across architectural and optimization hyperparameter choices. This can be interpreted demonstrating that concentrating a model's s "semantic focus" with respect to the tokens currently being considered (in the form of input-dependent sparse attention) accelerates learning. We develop a theoretical characterization of the conditions that explain this behavior. We establish connection between the stability of the standard softmax and the loss function's 's Lipschitz properties, then show how sparsity affects the stability of the softmax and the subsequent convergence and generalization guarantees resulting from the attention mechanism. This allows us theoretically establish that input-agnostic sparse attention does not provide any benefits. We also characterize conditions when semantic focus (input-dependent sparse attention) provide improved guarantees, and validate that these conditions are in fact met in our empirical evaluations.
Introduction
Transformers [1] are expressive set encoders, which when paired with positional encodings, can serve sequence encoders. The attention mechanism in transformer block allows model the long and short term dependencies in sequence in an imput-dependen manner instead of relying on handcrafted dependency modeling as recurrent (uni-directional and bi-directional) and convolutional models. The single hidden layer multi-layered perceptron (or MLP) in the transformer block introduces non- linearities enabling further expressivity. Transformers have been extremely successful in modeling natural language, and are the core blocks of various large language models or LLMs. They have also been successful vision, tabular data, and time series among various other applications.
The expressivity of attention-based transformers [2] comes with a computational overhead where the attention mechanism requires time and memory quadratic in the sequence length. To address this, various efficient transformers have been developed [3], utilizing various techniques such as fixed sparse attention patterns, low rank approximations of the attention matrix, and input-dependent sparse attention patterns. In this work, we focus on sparse attention mechanisms, both input-dependent and input-agnostic, Existing literature has studied sparse attention (inference), which in turn can speed up each training step [4], However, sparse attention has always been viewed as an approximation of the gold standard full attention.
way to speed up the forward pass
One can view sparse attention as a form of sensory gating, and this is considered an essential component of biological cognitive systems, allowing rapid learning [5, 6]. and the absence of it is often considered a marker for schizophrenia [7]- The gating is often achieved via inhibitory signals. Related observations made by Bengio [8] suggest some motivations. He makes a connection between a form of input-dependent sparse attention and the global workspace theory of consciousness in cognitive science, well as the properties of natural language sentences and
Contributions.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

symbolic Al representations used in planning and reasoning [9], "stipulating that elements of conscious thought are selected through an attention mechanism (such as the content-base attention mechanism we introduced in Bahdanau et al. [ [10]) and then broadcast to the rest of the brain, strongly influencing downstream perception and action as well as the content of the next conscious thought". As the "elements" or weight vectors being attended to are often discussed as semantic concepts, one can refer to the same phenomenon as "semantic focus" and explore its possible benefit to learning efficacy. Motivated by this, consider the following question in this paper: "Can sparse attention in transformers be beneficial in terms of learning convergence and generalization, in comparison to full attention?". To this end, we share the following findings:
($3) Focusing on benchmarks of structured languages designed to evaluate capabilities of transform- ers [4, 1 I1, and controlling for all involved hyperparameters, we make two empirical observations: - Sparse attention with input-agnose
sparsity patterns empirically struggles with expressivity (as implied by Yun et al. [2, 12]), and does not show benefits in terms of learning convergence and generalization even when equipped with enough expressivity (via global tokens [13, 14)).
Sparse attention with a specific form of input-dependent sparsity pattern that limits the attention to the top attention scores - the heavy-hitters (such as top-k attention [15, 16]) - are empirically as expressive as the standard full attention, and can converge significantly faster during training, while generalizing as well as, and at times better than, the full attention model. These improve- ments hold across various hyperparameters, both related to the architecture (such as the number of heads per transformer block, the number of transformer blocks, the MLP activation function), and the optimizer (such as the initial learning rate, and the learning rate decay).
($44 We then try to theoretically understand why this might be happening, and characterize conditions under which sparse attention can provide better learning convergence and generalization guarantees. Our analysis is based on two critical insights:
For any X-Lipschitz learning objective (with respect to the learnable parameters), the convergence rate and algorithmic stability [17] of (stochastic) gradient-descen based algorithms are dependent on Lipschitz constant A, with smaller values implying better convergence and stability guarantees; better stability implies better generalization [18]. We show that the Lipschitz constant of a transformer-base model is tied to the
of the softmax in the attention mechanism better input-stability implies better Lipschitz constant. Thus, we establish how the input-stability of softmax directly affects the learning convergence and generalization.
The sparsity pattern of the sparse attention affects the overall learning convergence and general- ization through its effect on the input-stability of the softmax. The input-stability of the (sparse) softmax is closely tied to the range or the semantic dispersion of the values (the query-key dot- dot-products) over which the softmax is applied (formally discussed in definition 2) - larger dispersion implies worse input-stability. While input-agnostic sparsity patterns do not neces- sarily improve the dispersion over the full-attention model, input-dependent sparsity that only focuses the heavy-hitters can significantly improve this dispersion, thus implying improved input-stability. This effectively translates to an improved Lipschitz constant, thus convergence and generalization guarantees. We also empirically validate that the dispersion and the estimated Lipschitz constant of input-dependent sparse attention show improvements over full attention.
Related Work. Transformers have been studied from various aspects since their conception. Various sparse attention transformers have been developed to improve their computational complexity [3, 19] with both input-agnostic sparsity [20-23, 14, 13] and input-dependent ones that focus the attention on the keys with the highest dot-product scores - the heavy hitters- while explicitly ignoring the rest [15, 24, 25, 16]. Various benchmarks empirically study the efficiency [4] and capabilities [11, 26, 27] of transformers. These capabilities are also theoretically characterized under various computation models [28, 29J. Please see more detailed literature survey around transformers in appendix C.
2 Problem Setup
In this section, we detail the problem setup, introducing the notation, and presenting the details of the transformer-based model, the training data and the learning loss.
Notation. We denote the index set [n]] A {1.....m n) for any natural number N. We use X for input sequences of token indices v E [D] in a vocabulary V of size D, and for labels or targets. We use x E IRd for a token embedding vector and X E RdxL for the sequence (matrix) of L token embeddings. For any vector v, &s to denote its i-th entry, and ||v|| to denote
2

---

its Euclidean norm. For a matrix W, we denote its (-th entry as Wi,, i-th column as W. and i-th row as W.. We use || and || to denote the spectral and 12,1 norms of W. For a tuple = (W ..... w()) of n matrices, we let ||0|| = maX [W We consider a learning problem with imput sequences X [ut....,,zz] E [D] of length exactly L with its i-th entry 4 denoting the 4:-th token in a vocabulary V. with outputs E y. For a learnable function f X y with learnable parameters 6, we explicitly write the function as fo(X) with X X.
Transformer block. For L length sequence of token embeddings X E IRd with the i-th token embedding denoted as X R., let TF:; Rdx Rdx denote transformer block with learnable parameters 0- = (W, V. P, R) with W, V E Rddd P. R. E RR Xd defined as:
TF (X) LN X R o (PX))
and X = LN(X + VX softmax(X (X (X));
MLPP (X)
Aw.v(X)
where LN : IRd + R is the token-wise (columnwise) layer normalization (or LayerNorm), and R (PX) denotes the token-wise single hidden layer MLP: : Rd - Rd. The columnwise softmax(.) of the dot-products D = X WX between the query and key matrices, combined with the value matrix VX, denotes the dot-product self-attention A g + Rd We consider single head attention here for the ease of exposition, but our analysis can be easily extended to multi-headed attention (see appendix F.4). While Vaswani et al. [1] utilized ReLU as the activation
the MLP, subsequent works [30] have used other activations such as GELU [31] and ELU [32]. Furthermore, many different variations of the transformer block has also been utilized in literature.
Masked softmax. A common modification of this transformer block is the replacement of the softmax
with sparse masked softmax which has an associated masking function R4XL {0, 1) L with M m mD) for a dot-pre poo-product matrix D. The (j, i)-th entry Ajı of the post-activation attention matrix A- softmax(D) for standard and masked attention is given as follows:
exp(Dja)
exp(Drt)= Maa
Aji Σ eD.))
Aje Σ exp(Dpx)) My's
Complete model. The model is defined as fe [D] - 5 with token and position embed- dings T R and E E Rdxt respectively, 7 transformer blocks each with parameters g(e) = (W()), v(O, P((), R()), E [r]. and readout linear layer with weights & E Yxxd us- RL,, where Y is the dimensionality of the output 5 (for example, the
ing token projection vector number of classes in output domain })) The i-th token U, E [D] in the input is initially embedded as x (0) TT,, + E, using the token and position embeddings:
(3) Here E R is the (fixed) token projection vector we can set the w = [0, 0, ..... 0,1] 1] to select the last token to make the final prediction, and w = (3/1)1, , uses the average of the tokens (along the sequence length dimension), where 1, is the all-one L dimensional vector. The O in f-(-) denotes the tuple of all the (learnable) model parameters, that is θ A (T, (() ,((T), 2). Here assuming that the position encodings are not learned, but that can also be incorporated in our study.
fe(X) = ()) X (t) TFacc) (x (x,)),,, Vt (, [7].
Training, Given a set S of n sequence-output pairs (X, v), X E (DI' y E y for training, and a per-sample loss function f: y' x ý - RR, the learning involves solving the following empirical risk minimization or ERM problem:
min ( 010)
((O) A n2
Σ (ey (e(-) defined equation (3)). (X
In the sequel, we will study, first empirically and then theoretically, (i) the convergence rate of stochastic gradient descent for this learning problem, and (ii) the generalization of the learned model.
In our experiments, we consider supervised learning with Y as a set of labels, but our analysis applies to any y where we have a scalar loss e: y' x 5 + R. where is the model output space: 3 c R13l for multi-clase classification with cross-entropy loss, and
y C IT'h for m-output regression with mean-squared loss. Witt queries QX, keys KX,, the scores (QX) (KX)= = X (Q K)X; we denote W = QK" Foor example, instead of the transformer block described in equation (1), there are versions that modify the location where the LayerNorm is applied: (Xx X + MLPp (LN(X)) and X = X Aw.v(LN(X)).
3

---

3 Empirical Observations
20
3.6
04-
In this section, we focus on em-- pirically ablating the effect of the different forms of sparse atten- tion on the ERM convergence and generalization. For this purpose, we ensure that all hyperparame- ters (architectural and optimiza- tion) are the same between the standard full attention, and the var- ious sparse attention mechanisms. We consider total of 8 tasks from LRA [4] and the NNCH benchmark [11]; we present re- sults from 3 of the tasks here (results for all the tasks can be found in appendix E); we also present preliminary results on NLP next-token prediction task in appendix E.3. Details on the tasks and the sparse attention choices, along with the hyperparameter (ar- chitectural and optimization) selection procedure and our compute resources are discussed in ap- pendix D. First, we will present the comparison between the standard full attention and various sparse attentions. Then, focusing on full attention and heavy- we will study the effect of hyperparameters (or the lack thereof) on their relative behaviors.
60
58:
50
100
350
Sparthe
(a) List Ops I ReLU (b) Even Pairs
ReLU
(c) Miss Dup ReLU
Figure I: Learning convergence and generalization curves for full attention and various sparse attention based models. Each column corresponds
task. The legend is the same across all datasets - B5) denotes banded attention with a band size of 5: BAND(5):1 denotes the same with a single global token. BL)C(5) denotes block local attention with a block size of 5; BLOC(5):1 denotes the same with single global token. TOPK(5) is top-k attention with k = 5. Top row: Training cross-entropy loss trajectories lower better. Bottom row: Generalization performance on held-out set as training progresses higher is better. Further results for 8 tasks with different mask sizes and global tokens is presented in appendix E.
style input-dependent sparse attention,
We present our first set of experimental results in figure I, comparing the overall learning convergence and generalization of full attention based models to those using sparse attention. Here we use ReLU for the MLP in the transformer block as in the original configuration [1]. These results are aggregated over 10 repetitions, and we present the median performance and its inter-quartile range. We present results for a single sparsity level (mask size) here; please appendix E. for more variations.
Observation I. Input-dependent heavy-hitter sparse attention speeds up learning convergence while input-agnostic sparse attention do not show any improvement over full attention.
The results in figure (top row) show that the input-agnostic sparse attention often converges slower than full attention. They also often struggle with expressivity in the absense of the global tokens, as seen for both block local and banded attention with Even Pairs and Missing Duplicates. This is expected as per the expressivity results of Yun et al. [12]. The inclusion of the global token addresses this issue. In contrast, the training loss of top-k attention converges significantly faster than full attention. Top-k attention shows improvements (in terms of achieving 95% training accuracy) over full attention ranging between 1.37x (121 epochs vs 167 epochs) with ListOps to 8.83 (6 epochs vs 53 epochs) with Even Pairs (see table in appendix F. 1). In all cases, top-k attention is able to be as expressive as full attention without the need for any global tokens. This consistent faster training of top-k attention in terms of the number of optimization steps needed to converge is not something discussed in existing literature to the best of our knowledge.
Observation II. Input-dependent heavy-hitter sparse attention generalizes faster during learning.
The results in figure 1 (bottom row) show that the input-dependent sparse attention achieves similar (Even Pairs and Missing Duplicates) or better (ListOps) holdout accuracy when compared to full attention. Furthermore, it attains this generalization level much earlier during the training process. Note that nepu-deeendentt heivy- top-k achieves better empirical generalization performance both in terms of the highest holdout accuracy during the training trajectory, and the final holdout accuracy. The latter highlights that the faster ERM convergence of input-dependent sparse attention does not lead to overfitting. In fact, with the ListOps task, the final holdout accuracy with full attention drops from around 35.1 = 0.6% to 28.9 + 1.4%, while the drop with top-k attention is only from 36.3 + 0.3% to 31.3 + 0.9% (see table 3 in appendix E for complete results). In general, the top-k attention based transformers also have comparitively lower variations in their performance as evidenced by the fairly tight inter-quartile ranges of the trajectories of the loss and accuracy.

---

Next we study the effect of the dif- ferent hyperparameter choices on
25 10000
the relative performances. First
FULL
we change the activation in the
MLP of the transformer block to evaluate whether the differences in empirical performance are due
0.-
60.
to the attention or the MLP in the
block. Then, ber of blocks and the number of
vary the num-
50
100
300
300
80
160
35
102
116
230
250
2oada
Cooda
heads in the model. Finally, vary various optimizer hyperpa- rameters such the learning rate
(a) List Ops GELU
(b) Even Pairs GELU
(c) Miss Dup GELU
and its scheduling as well as the optimizer itself.
0.0
Observation III. The gain from the input-dependent heavy-hitter sparse attention in terms of convergence and generalization is unaffected by the choice of the activa- tion function a in the MLP of a transformer block.
50
100
350
300
100
207 330
Sparhh
6oodh
(d) List Ops Mish
(c) Even Pairs Mish
(f) Miss Dup Mish
Figure 2: Same experimental setup figure with GELU ac- tivation (top rows) and Mish activation (bottom rows). See
We present the performances of the different attention mecha- nisms with the GELU activation [31] in figure 2 (top 2 rows) and with the Mish activation [33] in figure 2 (bottom rows), while keeping all other hyperparameters exactly the same as in figure to ablate the effect of the change in the MLP activation. Comparing these results to figure 7, see that there is not a lot of qualitative difference in the performances both in terms of learning convergence and generalization, indicating that the difference in performance is due to the difference in the attention mechanism of the transformer block.
additional results in appendix E.2.
Observation IV. The improvement of input-dependent heavy-hitter sparse attention over full attention is not affected by the number of transformer blocks, and increases with the number of heads in each transformer block.
We study the effect of varying the number of transformer blocks (figure 3a) and the number of attention heads per transformer block (figure 3b). We have again fixed all other hyperparameters as in figure ablate the effect of the considered architectural changes. Here we focus on learning convergence of full and top-k attention (with k 5) for ListOps. The results indicate: (i) Top-k attention continues to converge faster than full attention across all number of blocks tried (r {6, 10, 15, 22)).. The relative performance is not affected by the number of blocks. (ii) Top-k attention continues to and
converge faster than their full attention variants as the number of heads increase from The convergence of the full attention model slows down while the convergence of top-k attention stays almost the same, thus increasing the relative improvement with the number of heads.
Observation V. The improvement of the input-dependen heavy-hitter sparse attention holds across varying optimizer hyperparameters, especially for hyperparameters that have the most promising convergence.
We present the effect of varying SGD parameters for full and top-k attention ListOps. In figure 3c, we fix the decay rate to 0.99 (as in figure 1) and vary the initial learning rate. Smaller initial learning rates (0.66 and 1) have the best convergence for both full and top-k attention, with top-k converging faster. For larger initial learning rates (1.5 and 2.25), convergence slows down for both, and the difference between full and top-k attention is less pronounced, though top-k appears to be slightly better, especially initially. In figure 3d, we fix the initial learning rate to 1.0, and vary the decay rate. For slower decay rates (0.9999 and 0.999), the overall convergence for both methods slow down though top-k continues to converge faster. For faster decay rate of 0.9, top-k initially appears to outperform full attention with a big margin. However, both methods prematurely stall as the learning rate becomes too small.
5

---

NBLOCKS:6
NBI 1BLOCK5:10
NBLOCKS:15
OCKOC2
FULL TOPS55))
50
100
150
200
5.0
300
150
200
0
30
120
150
200
50
100
150
200
Epochs
Epartss
Epocta
Epocts
(a) Varying number of transformer blocks (with I attention head in each block).
NHEADS2
NHEADS:8
TOP
Traking
3.50
202
130
150
200
3.00
200
Epochs
(b) Varying number of attention heads per block (with total 10 blocks).
LR:0
INIT
LR:2.25
100
150
100
(c) Varying initial learning rate (with decay rate of 0.99).
LR DECAY RATE:0.9
DECAY RATE:0.99
DECAY RATE:0.999
DECAY RATE-0.9999
Traning
PULL
10PO51
Eportss
Epecte
Eports
(d) Varying decay rate (with an initial learning rate of 1.0).
INIT 1R:1e-05
INIT LR0000001
INIT UR:0.001
INIT LR:0.01
2.0
2.0
Triring LU3
1.8
1.0
1.6
0.5
1.4
0.0
0.0
100
200
100
200
100
200
100
200
Spachs
Epochs
Esocts
Epochs
(c) Varying initial learning rates for Adam [34].
Figure 3: Comparison of full and top-k attention training loss trajectories for varying hyperparameters, both architectural (number of transformer blocks and attention heads), and optimization (initial learning rate, learning rate decay, and optimizer) with the List Ops task.
Observation VI. The improvement of the deput-dependent heavy-hitter sparse attention over full-attention also holds for the Adam optimizer with varying learning rates, especially for hyperparameters that have the most promising convergence.
While all our previous results use SGD, in figure 3e we evaluate whether the relative performances translate to the more widely used Adam optimizer [34] ListOps. We evaluate various learning rates, and see that the learning rate that provided convergence for SGD (initial learning rate around 0.1-1.0) lead divergence with Adam (see 4th column in figure 3e). With smaller learning rates, we see that the improved convergence of the input-dependent heavy-hitter sparse attention persists.
4 Theoretical Understanding
The empirical observations in the previous section demonstrate that input-agnostic sparse attention can struggle with expressivity, and does not show any consistent benefit over full attention. In contrast, input-dependent heavy-hitter top-k attention shows significant advantages. In this section, we want to theoretically understand why this might be happening. We begin by considering the factors that affect the convergence and generalization of SGD based training.

---

Standard SGD analysis show that, for a 01 -iipschitz and d-y-mmoohh finite- sum (non-convex) ob- jective, with learning rates T7 at the i-th step, learning converges to a e-stationary point in T steps where c (agaa) (o nt (E n)). Different choices of 7R; i E [T] (such as n/i or n/ for some constant > 0) provide different convergence rates (such as O(1/ log(T)) or Q(1/VT)). As we control for the learning rate and its scheduling in our experiments, and ensure that all models start learning from the same initial parameters, the main distinction between the different forms of attention could be the Lipschitz constant aj and the smoothness constant 02- Note that, with non-smooth activation function like ReLU, we are effectively performing stochastic sub-gradient descent, where the guarantees are much weaker but still depend on the Lipschitz constant.
Generalization error is the difference between empirical risk (on the training samples) and the true risk (over the population). A low training error combined with low generalization error implies strong performance on unseen data. Utilizing the seminal work [17] on algorithmic stability, Hardt et al. [18, Theorem 2.2] show that learning with e-stable randomized algorithm guarantees expected generalization error of at most e. Furthermore, for 0r -Ipsschcii and 022Smoothh finite-sum nonconvex objective, the T step SGD algorithm with per-step learning rate 1Ι S n/i is s-uniformly stable with E O (r9az)//+=zr)) + 1/a2n) + [18, Theorem 3.1 As before, the distinguishing factors between our models pertinent to generalization are the Lipschitz and smoothness constants.
Based this intuition, will focus on the Lipschitz constant. First, we will try to characterize how the behavior of the softmax- - specifically the input stability of softmax- - in the attention mechanism of the transformer block affects the Lipschitz continuity of the overall learning objective.
Definition 1. A masked softmax is f-input-statle stable if Vz. 2 Rd, ||softmax(z) softmax < E E|z - 1|11. 1. The attention A with parameters W, V is stable with respect to its input and parameters V X E R W, W, V, V E Rdddd with constants (x( (C) depending on E:
(A). Aw.v(X))1 12,, S X X 12,11
(5)
(Aw.v(X) - (X).v(X)l2.11 S (E(()) W - WII
(6)
(Aw.v(X) - ( S AvIIV-
(7)
We will precisely characterize the values of the constants in the above definition (C, Ax(5). Aw (E).
Av) for the different (masked) softmax and corresponding (masked) self-attention in the sequel. However, we define them here to highlight how the stability of softmax affects the stability of the self-attention A, and how this affects the Lipschitz continuity of the learning objective in equation (4)
with respect to the model parameters θ = (T, 0(2), O(T), +b). For completeness, we first need to establish the stability properties of the MLP component of a TF block (see proof in appendix F.1):
Lemma 1. Assuming that the MLP activation a is A Lipschitz with (0) = O, and the MLP
parameters have norms bounded by B > O, that is ||P|| < B and |R] S B, the token-wise MLP and LN operations are stable with respect to their input and model parameters as follows Vx, x' E
Rd ||x||, |||| < E, P, P E TRduuxd xd R., R. E Rddaex xd. with nxx B² >o, 1pp 1n X BE:
||MLPp P. R (x) MLP P (x)1 S 7x x X|I,
[MLPp.R(x) MLP P.R (x)]l S Sp |MLPP.R(x) MLPP,R(X)) (x)] S MR R R|I, [|LN(x) LN(x)]I S x x|].
Pll
(8)
The Lipschitz continuity of the LayerNorm has been previously established in Kim et al. [35]. Given this, we can establish the following results for a transformer block (see proof in appendix F.2):
Theorem 1. Given definition I and lemma I, a transformer block TF with learnable parameters 0- (W, V.P. is (o(E)-stable with respect its learnable parameters A with
Xo(E) ÇLN (GIN(1 qx 7x) (AW(E)) + Xv) + [[pr + nR)),
and TF is Xx(E)-stable with respect to its input X with xx(() Ciy + nx) + Ax(E)), where we explicitly note the dependence on the stability & of the (masked) softmax operation. Thus, for any parameter tuples 0, 0 and inpul X, X, we have
(X) (22()) S Xe(E)) Øl, [ TFo(X) S Ax(E))XX XI|. (10)
Thus, we establish the following for our model with 7 transformer blocks (proof in appendix F.3):
7

---

6, Rd 2T
6, Ro 2T
2g
R4
RaRgRs
RaReRs
Ri F10
Gr
+
+1=
(a) Standard attention dispersion 6g
(b) Banded attention dispersion 6,
2
6n o 21
2g
2g 210
210 98
+
+1 +r
(c) Causal banded attention dispersion 8,
(d) HH attention dispersion on & separation As
Figure 4: Semantic dispersion ö (definition 2) and heavy-hitter (HH) attention semantic separation A (definition 3): For a sequence of length L = 10, we demonstrate the concepts for query token X.s. Let 2j X WX. denote the jth query-key dot- dot-product. (a) Figure 4a shows that in full attention, the Ejs can range between - - and +r=² (theorem 1), giving us a semantic dispersion 6, 2 2
In general, we cannot expect a tighter bound on 6,. (b) Figure 4b shows the same example for an input-agnostic banded masked attention, where the query token X. only attends to succeeding key tokens X. X.7,X,,,, ). while the rest are masked (L). Here, the dispersion 6, = ôg R 21 no better than full-attention. (c) Figure 4c shows the example with imput-agnostic causal banded attention mask where token X.6 only attends to the preceding key tokens X.5, X.4, X.,,,., X.3. Here, this masked attention has a small dispersion 6, 2 better than that of full-attention 6, 2 However, there is usually no way to ensure that a condition where 8, < 6, will exist. (d) Figure 4d shows the example with input-dependent HH attention, where only the high values unmasked, and there is a significant semantic separation An between the masked and unmasked dot-products. Here, we can expect significantly smaller semantic dispersion dn < 2 implying on < d..
Theorem 2. Assuming that the per-sample loss function & in (4) a-Lipschitz and || S 1 with (11))L, under the conditions of definition I and theorem 1, the learning objective L in (4) is Lippsiitt with respect the learnable parameters θ (T, (((),
0(T), D). where
xc(5) = ( (= Ax(E) (1 (1+ ((xx(E)-1)
Ao(E)
(11)
and [c()) - C(O)) Ac(1)11--1 6|| for any set of model parameters θ G.
This characterizes how the Lipschitz constant of the learning loss, and thus the convergence rate of the SGD based ERM, is tied to the input-stability constant & of the (masked) softmax. Thus, based on theorem 1, the larger the values of Aw(E), (x) and λv in definition 1, the larger the Lipschitz constant of the training loss. We will characterize these quantities in the sequel. To understand the effect of sparsity on the stability of softmax, we begin with the stability of the standard full softmax and the standard full attention in the following (see lemma 4 in appendix G.1);
Lemma (adapted from Li et al. [36] Lemma B.1). For any z, Z E R- with max; (1 (z--z)) 2z) S 8.
and |ax - E,) I 8, for a positive constant 8 O, we have the following:
||softmax(z]l S e /L, ||softmax(z) - softmax(2)ll S (e*/L)ll 2/|1.
(12)
A critical factor in the softmax stability is this quantity 5 that is the upper bound on the difference between the largest and smallest values over which the softmax is applied. In the context of dot- product self-attention, it corresponds to the difference between the largest and smallest query-key droduuusscts for any query. We term this as semantic dispersion, and define it precisely as follows:
Definition 2. For a (spurse) atlention transformer block with L length input sequences X E L and a mask M E {0. 1 (imput dependent or inpul agnostic), we define the per-query semantic dispersion
O such that, for any query token X i [L] the maximum difference between the largest and smallest unmasked query-key dot-products is bounded from above by d. That is, for any input sequence of token representations X E Rddx mask M E {0, 1] LXL and attention parameters W E Rdxd, for all query tokens i E [L], we have
scalar
max J. M,,
(X.
WX,, X. WX
(13)
We discuss this definition with examples in figure 4. We now establish the stability of standard softmax and self-attention A by characterizing &, Ax (E), (w in definition I in terms of the semantic dispersion as follows (see theorem 7, appendix G.I for details):
8

---

Theorem 3 (partially adapted from [36] Lemma B.2). Assuming that the per-token Euclidean norms are bounded as ||* S Evi E [L], and the parameter norms are bounded at [| - I and [IVI| S Y, and the per-query semantic dispersion is bounded by 5g >> 0. Then the standard softmax is Es-stable with & = e L, and the standard attention is stable as in definition I with (i) xx E. LL + = e 1(2r=² + +1,, (ii) (W(Es) = X. = c² TLE (iii) Xv = LE
The semantic dispersion 8, plays significant role in (x()s) and (5.) with larger values implies higher per-transformer-bloe stability constants (E( and Xx(Ex) in theorem 1, As discussed in figure 4a, we cannot expect this dispersion 6, to be significantly smaller than 21
Next we study the stability of input-agnostic regular k- k-sparse attention transformers, where each query attends to exactly k keys, and each key is attended to by exactly k queries. This form includes banded attention [20], block-local attention [21] and strided attention [22, 23]; random attention [14] satisfies this only in expectation. We establish its properties as follows (see theorem 8 in appendix G.2):
Theorem 4. Consider self-attention with a k-regular inpul-agnostic mask M. Assume that the per-token Euclidean norms are bounded as [|x:i| S Evi E [L]. the parameter norms are bounded at |W|| S T, | S Yr and the per-query semantic dispersion is bounded by 6,- Then the masked softmax is -stable with &,--. eer jk, and the attention is stable as in definition I with (i) Xx (Ex)
( + 1) = é T(2r=* + 1), (ii) Xw(E)) = E, LAK= = TL (iii) xv = LE.
This shows that this input-agnostic sparse attention provides guarantees very similar to full attention except for the ς term, implying significant improvement in stability only if the per-query semantic dispersion 6, is sufficiently small relative to the full attention dispersion dg; see one such situation in figure 4c. In general, the dispersion 6, would be small only if the per-query dot-products somehow align with the sparsity patterns - with temporal locality based patterns (banded, block-local), the dot-products for nearby keys (in terms of sequence position) would require have small dispersion; with strided patterns, the dot-products for keys matching the stride regularity should span a small range. These conditions are too restrictive, and thus, in general 6, I 6₂ 2 (see figure 4b).
In input-dependent heavy-hitter sparse attention, for any query token i E [L], we mask all but the highest values of X WX, and there is a significant gap between the unmasked dot-produce X_ WX. for the unmasked keys j with Mfi- 1, and the masked dot-product X WX.3' for the masked keys j', Myy; = O. Unlike the regular k-sparse attention, here each query attends to k keys, but each key can be attended to by anything between 0 and L queries, making the analysis of input-agnostic regular sparse attention (theorem 4) inapplicable. To study these heavy-hitter sparse attentions, we formalize a notion of semantic separation between the masked and unmasked keys:
Definition 3. For a sparse attention transformer block with L length input sequences X E R and an dependenn heavy-hitter mask M E {0, 1]LxI, we define the per-query semantic separation as a scalar A > O such that, for any query token X E [L] the minimum difference between the pair of masked and unmasked query-key dot-products is bounded from below by A. That is, for all query tokens i E [L]. with unmasked key j and masked key j', we have
<
min Vf M, (X. WX.j - X. W..)))
(14)
The notion of semantic separation is visualized in figure 4d. We present the stability of the heavy-hitter attention in the following (see theorem 9 in appendix G.3):
Theorem 5. Consider the self-attention with a k-heavy-hitter input-dependent masking function m, applied columnwise to the dot-product matrix to get a mask matrix M E {0, 1LL Assuming the following: (i) For any query-key pairs X, E Bdx and parameter W E Raxd, the k-heavy-hitter mask M = =X WX) (applied columnwise) has a minimum per-query semantic separation of An. (ū) A maximum of Ak., B querry tokens attend to a single key token, (iã) The per-token Euclidean norms are bounded as ||X,;|| < Evi E [T.], and the parameter norms are bounded at [[]] S T, [V]| < Y, and (iv) The per-query semantic dispersion is bounded by ôn. Then the masked softmax En- Stable with En //)) (1 + 1/ar),, and the sparse attention is stable as in definition with
)
B
(x = Tk 21 (B + 1) + 1+
= (3 (1 + 1) 1//n)/
(15)
AW(EN) = 26, TLAE3 26 T1 + 1/2r), =E
9

---

With the heavy-hitter atten- tion, we would expect the per-query dispersion dn to
STAMSADD
1/0
STA STAMDNAD
STANDAAD
0.75
673-
0.8
be significantly smaller than 5g especially for small k. In appendix G.4, we explic- idly characterize the condi- tions under which the sta- bility constants for input- dependent sparse attention (theorem 5) show improve- ments over full attention (theorem 3). We see that moderate reduction in the dispersion (ō,, 6,) al- low for significant improve- ments in Åw even for small separation An. while im- provements in Ax are more moderate.
0.29
E25
0.98
1.00
TOrK
192
TDPK
0.79
679
0.38
659-
0.5
C25-
To see how these stabil- ity constants affect the loss landscapes, also visual- ize them in figure 5 (top and middle rows) utiliz- ing the techniques proposed in Li al. [37] (see ap- pendix G.5). We see that the contours on the loss sur- faces of full attention model are somewhat asymmetric see for example, around the center in figure 5b, fig- ure 5c, and moderately in figure 5a. In contrast, the loss surfaces of the heavy- hitter top-k attention model quite symmetric, especially around the origin, which corresponds to the final learned model.
Ordrange
(a) List Ops
(b) Even Pairs
(c) Miss Dup
Figure 5: Top and middle rows: Loss surfaces of the models with full attention (top row) and top-k attention (middle row) for the tasks considered figure with the corresponding hyperparameters utilizing the filter-normalized version of the loss landscape visualization. The (0,0) grid point corresponds to the final trained model - the optimum. Bottom row: Distribution of the estimated Lipschitz constants com- puted in the random directions used to generate the loss landscapes. We report the distributions on the vertical axis in terms of the 50-th (dot- ted), 75-th (dash-dotted), 95-th (dashed) and 99-th (solid) percentiles (lower is better). On the horizontal axis, we denote the distance of the parameters from the optimum on the grid, and visualize how the distributions vary with the distance.
We also utilize the loss surface to approximately estimate the Lipschitz constant across the loss landscape (see details in appendix G.5). We plot the distribution of these estimates in the bottom row of figure for varying distance from the optimum. We see that near the optimum (the final trained model), the distributions of these estimates are close for both the models. However, as we move farther away from the trained model, the distributions change significantly, and top-k attention provides a smaller Lipschitz constant estimate compared to full attention all percentiles of the distribution. This indicates that, empirically, the loss for top-k attention has a more favorable Lipschitz continuity compared to full attention, which in turn implies both faster convergence and better generalization guarantees. Thus, our stability- basedd theoretical investigation in this section appears to align with empirical observations in section 3.
5 Conclusion
In this paper, we study the potential advantages and drawbacks of sparse attention over standard atten- tion beyond the currently studied computational perspective. Our empirical findings, characterized by our theory, show that (i) input-agnostic sparse attention can in general only provide computational ben- efits, but (ii) inpur-dependent heavy-hitter sparse attention can provide significant improvements over full attention in terms of learning convergence and generalization. We hope that this motivates further use of heavy-hitter sparse attention at scale with transformer based models. We discuss the limitations of our work in appendix A.
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. URL pttpps://rroceeddngg neurips...... cwtppsccccpppp//p//eee /017/ ftto:/ffee2255777ee10003311ee
pdf. 3,
[2] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of Conference on Learning Representations, 2020. URL attps://openrevieee wtt.rrvvviienee/ffrrrrr. net/forum?
functions? In International
id-ByxRMONtvr. 1, 44
[3] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6), 2022. ISSN 0360-0300. URL
3530811. 25,
[4] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. International Conference on Learning Representations, 2021. URL
httpp:://penreviee..eet wtw..net///oommmmdd=//lllll grC2k. 1, 2, 4, 27, 28, 29
[5] LA Jones, PJ Hills, KM Dick, SP Jones, and P Bright. Cognitive Mechanisms Associated with Auditory Sensory Gating. Brain and cogmition,, 102:33-45, 2016. URL httph:://///tt sci sciencedirect. com/ sttp:/ccccnccccceeeeeeee ctt.c///ciieeeeaaatacceee//////s /ttpa//cleee/////772222211000000
1
[6] Bernd Fritzsch. The Senses: Comprehensive Reference. Academic Press, 2020.
[7] LL Judd, McAdams, B Budnick, and DL Braff. Sensory gating deficits in schizophrenia: new results. The American journal of psychiatry, 149(4):488-49)) 1992. URL https://pubmed. ncbi. nlm. nih. www..gov/155eee gov/1554034/.
[8] Yoshua Bengio. The consciousness prior, 2019. URL https://arrivv.org//b///// 08568.
[9] 3rd Neuro-Symbolic AI Summer School, 2024. URL https://narroyyablllc.itttt io/nsss2024. 2
[10] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2016. URL https:
org/ abs//409 0473. 2
[11] Gregoire Deletang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, and Pedro Ortega. Neural networks and the chomsky hierarchy. In The Eleventh International Conference on Learning Representations, 2023. URL 29
2,4, 28,
[12] Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O(n) connections are expressive enough: Universal approximability of sparse transformers. Advances Neural Information Processing Systems, 33:13783- 13794, 2020. URL https:// neurips ltapparapeeeesssppllllpaaaaaeeee fil.files/peeee fiwes/paper////// haappa//0000 /2020/ 111/8887588005088002228888838333
.df. 2, 4, 28
[13] Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Ete: Encoding long and structured imputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 268-284, 2020. URL httts:://cclatthie 2020. emnlp-main. 19/. 2,27 27
org/
[14] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020. URL htwph::tpp///tee httops///oooeeeedngg neurips. cctpaper_fffaaeeeeeeeeeeee saaeer_filssppppeeee Paper.pdt 2,

---

[15] Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Nemorr-eefficieee transformers via top-k attention. In Proceedings of the Second Workshop on Simple and Efficient Natural Language Processing, pages 39-52. Association for Computational Linguis- ties, 2021. doi: 10. 10.18653/v1/2022 1ssstain...stlppp--.. 1.5.. URL htt.httts//accaaanollogg.. www.e//aannooloooooo2/00000 sustainlp-1.5 1.... 2, 25, 27, 29
[16] QIUHAO Zeng, Jerry Huang, Peng Lu, Gezheng Xu, Boxing Chen, Charles Ling, and Boyu Wang. ZETA: Leveraging z-order curves for efficient top-k attention. In The Thirteenth International Conference on Learning Representations, 2025. URL https://
net//oorum?2cc-joVeeeeEE 2
[17] Olivier Bousquet and André Elisseeff. Algorithmic stability and generalization performance. Advances in Neural Information Processing Systems, 13, 2000. URL
https://proceedingss neurips. wwtpccpppppeeeeeeeee saat2aaleecattt88770084fa.f.dc Paaeeeerfaaaappr..ree http/pppr_fileeee_ fileffilel//lleelf htt.:/fflaaaaappeaeaaeee http////r...lleeppprr/////////// httpa//pppeeee/////eeeee
2.7
[18] Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of
stochastic gradient descent. In International conference on machine learning, pages 1225-1234. PMLR, 2016. URL huttps://arxiii http://xxivv.rrrgg
.or. wwt/ootg///oooooooooooo.ppp...tt httf//159..022. 01240. 2,
[19] Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transform- ers. Al Open, 2022. URL hattps://wet http::::/wwwsscceeneeiliiee sciencedirectt. wot./sccence//rtticleee.. $2666651022000146 2, 27
[20] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference on machine learning, pages 4055-4065 PMLR, 2018. URL https://proceeecingg mlr htt///80//aaa/ ht0/////aa//....1/ 1armar18a.pde 2, 9, 27,29
[21] Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 2555-2565, 2020. URL
org/
2020. 2020. www.fifiiiinsnnneeee emm 232/ 9, 29
[22] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/100, 2020. URL https://///aat http://axxvvoooo/aaa/a org/ 05150. 9, 27
[23] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019. URL abs/1904. 10509. 2, 9, 27
org/
[24] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computa- tional Linguistics, 9:53-68, 2021. URL https: mit. addu/aac//articl//bo. 20055//3330000333/.... 00353/ 2, 25, 27,29
10. 20.1162/tacc_a___e
[25] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URL https: //tp:///opppeeieeeeeeeeeee
2, 25, 27, 29
[26] Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-to teeequeete recurrent networks. In International Conference Machine Learning, pages 2873-2882, PMLR, 2018. URL
mlr wtt.:ggw.w...r.essmememmm press/v80/
lake18a. html. 2, 28
[27] Najoung Kim and Tal Linzen. COGS: A compositional generalization challenge based on semantic interpretation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 9087-9105, 2020. URL httpp://aaaee http:///aaaantoogyoyyoooo org/2020. emnlp-main. 731, 731.pdf. 2, 28
[28] Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What formal lan- guages can transformers express? a survey. Transactions of the Association for Computational Linguistics, 12:543-561, 2024. URL https::t... pettp:://diee //aps:://ddrtttte/ mit. aait.edu/tacl http:/talccarrrreee www.e//aa//arrieeeeeeaaee. /arttcle///t//........ 1 1162/tacl 0000000060000000 120983. 2, 28
12

---

[29] Parikshit Ram, Tim Klinger, and Alexander G Gray, What makes Models Compositional? A Theoretical View. In Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence, pages 4824 4832, 2024. URL. https: //ww://wwwwwww/ www.wwjca................. ijjaii org/ pttp/pooceedings/222222 0533. pdf.
[30] Jacob Devlin, Ming-Woi 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 Linguis- tics: Human Language Technologies, Volume (Long and Short Papers), pages 4171- 4186. Association for Computational Linguistics, 2019. doi: 10. 10.188633/1//1911160 URL
https://aclanthoolgg. www.llogy.rrg//SS1
[3]] Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, abs/1600.00815, 2016. URL http://arxiv, org/absl 1606. 08415. 3, 51
[32] Djórk-Aréé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (ELUs). arXiv preprint 1rXiv:1511.07289, 2015. URL https://arxiii wttp:///////....orrrrr/ 07289.
[33] Diganta Misra. Mish: A self regularized non-monotonis neural activation function. CoRR, abs/1908.0860 2019. URL http: org/ /a//abb/////// 08681. 5, 32
[34] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, In The Third International Conference Learning Representations, 2015. URL http::ttt://:iii/e org/ prrg/b///// 1412. 6980.
[35] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning, pages 5562-5571. PMLR, 2021. URL https://arrt http///rxvvvv/v/// . .rg/ owt/p/t//////oootooo/.ggp.o. /2006. 04710. 7,
[36] Yingcong Li. Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabaio, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 19565-19594 PMLR, 23-29 Jul 2023. URL https: http://ppoccee...e//e////o mlr, wwwps/eee22222 /0000///121111 html. 8, 9, 28, 45, 46
[37] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018. URL https::/
nwurips.cc. wwtpcc/pppppeeaeeeeeeeeee www.fffiweeeelee /llls//pperr/018/E1ie a4 00113036600000000777666600--2220 .dff 10, 62, 63
[38] Ido Amos, Jonathan Berant, and Ankit Gupta. Never train from scratch: Fair comparison of long-sequence models requires data-driven priors. In The Twelfth International Conference on Learning Representations, 2024. URL https: http:::://///pevvveeeeee 25
[39] Mary Phuong and Marcus Hutter. Formal algorithms for transformers. arXiv preprint wnXiv:207.002338 2022. URL https://arxiv.orgg. orggppff/20f/ 092388 pdf. 27
[40] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn atten- tion. Proceedings of ICML, 2020. URL http:///proceedings. mlr. wtwpaiiinnn.....rrseeeeeeeeeeeee prepss/v199/aay00aa tay20a. pdf. 27
[41] Ankur Sikarwar, Arkil Patel, and Navin Goyal. When can transformers ground and compose: Insights from compositional generalization benchmarks. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods Natural Language Processing, pages 648 648-669, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.1 12605//110022222mmmmmmmo https: / acllaatoolttooyyc..c 22g22022. emnlp-main 41. 28
URL
13

---

[42] Róbert Csordás, Kazuki Irie, and Juergen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of transformers. In MranieeFa Moens, Xuanjing Huang. Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference Methods Natural Language Processing, pages 619-634, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 10.18//3/v 1/2021. emnlp-main.49, URL https: ptttcp::/aaaaaathhee http://olaaattoogg...oooo /tgg//2//1 emnlp- .ttpmmain.4333 28
Empirical
[43] Santiago Ontanon, Joshua Ainslie, Zachary Fisher, and Vaclav Cvicek. Making transform- ers solve compositional tasks. In Smaranda Muresan, Preslav Nakov, and Aline Villavi- cencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Pupers), pages 3591-3607, Dublin, Ireland, May 2022. URL
Association for Computational Linguistics. doi: 10. 1/2665//1/2022.ac1--ooo https:/
aoo-o/............oooo 251. 28
[44] Yichen Jiang and Mohit Bansal, Inducing transformer's s compositional generalization ability via auxiliary sequence prediction tasks. In Marie-Francing Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 0253-6265 Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguisties. doi: 10. 101158855122221 emnlp-main.505. URL https: http//acllaatt a/cle//ooolaaaogogooooooooo org/2021 emnlp-main. 505. 28
[45] Andrew Drozdov, Nathanael Schärli, Ekin Akyürek, Nathan Scales, Xinying Song, Xinyun Chen, Olivier Bousquet, and Denny Zhou. Compositional semantic parsing with large language models. In The Eleventh International Conference on Learning Representations, 2023. URL https:/ pttps://oppnrevies: www.feeeffeeoe ft/mm 2 28
[46] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of trans- formers to recognize formal languages. In Proceedings of the 2020 Conference on Empir- ical Methods in Natural Language Processing (EMNLP), pages 7096-7116, 2020. URL
https: bttts://aclane http:///aaaaaaloogo..a//ao org /2200.2000.22 emnlp-main. 576/. 28
[47] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 1761177 2020. URL. htttts//p///rrctt mit tttcl//aticl///e/iiii...
[48] Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. Transactions of the Association for Com- putational Linguistics, 10:800-810 8100008-00000 2022. URL
ojs/ ojs//ndexx.... /idexx.dex.... /in//.ppp/////////////
tacl/article/vier/
[49] William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant- depth threshold circuits. Transactions of the Association for Computationaal Linguistics, 10:843- 856, 2022. URL https://trane orge org/ojs/ index.ndxxi/e httpp//aaa/atttaieee///eeeeeeeee 28
[50] David Chiang and Peter Cholak. Overcoming a theoretical limitation of self-attention. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics org/
(Volume 1: Long Papers), pages 7654-7664, 2022. URL https://aclaattoot 2022.ac1- acl- long. 527/
[51] David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In Internatioonal Conference on Machine Learning, pages 5544-5566 PMLR, 2023. URL http:::/pproceeeddngg mlr. wwwpppeesssveeee w212/ chiang23a.html html. 28
[52] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, PMLR, 2023, URL https: .mw.aingmmmm.....mmmmmm....mm..m mll. www.m..w..////a.... http:////22///ooooaallaaseas/es/ pdf. 28
[53] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2023. URL. htte:p::/pprcceeiinng neurips.com cc/ ccc/papr__eee ffl_fills/// www.ls/ppppr/////222/ pppppppppp////////////// iile/ b2b3e1 22beeaadeeaacaaa77dd000000003032 Paper Paper-Conference Conference, .df.. 28
14

---

[54] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2023. URL httptttps://pproeeediin mew. cc/ cc/paper_ _files/ Confeernce pdf.
28
[55] Ruiqi Zhang. Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context, Journal of Machine Learning Research, 25(49):1-55, 2024. URL httpp://www jmlr www.papers////oooooooooo..eeseer http//oolm////////2 10-2/23-2222222 1042. pdf. 28
[56] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for atten- tion models? Advances in Neural Information Processing Systems, 33:15383-15393, 2020. URL https://proceeddngss neurips. httpppper.fllleepprrr//0/0000000 W0B57f66d4100337496666664633 Paper. .ppper.pdf.. 28
[57] Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than SGD for transformers. OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop),
2022. URL https: //tp:////p/e/eeeeeeeeeeee net/ dd
[58] Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact lo- cal neural network geometry? Advances in Neural Information Processing Systems, 36:8305-8384, 2023. URL httpps:/ neurs neurips cc/ cc/paper_ft feaapr_fileee faresrffilepp_eeeeeee fillespppppr///// ooa.r//rerrrreee-eeaeeooorrrreee Conference. pdf.
[59] Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be, In The Eleventh International Conference on Learning Representations, 2023, URL /ttp::///ppereviee. wwt.:neeet/tiiiiiieeo a665rrmogg
[60] Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization), The Twelfth International Conference on Learning Representations, 2024. URL https::/opperrviiae net/forum? id=0uI5415ry7. 28
[61] Bingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou, Taiji Suzuki, Jun Zhu, and Jianfei Chen. On the optimization and generalization of two-layer transformers with sign gradient descent. In The Thirteenth International Conference on Learning Representations, 2025. URL https:
/ //openreview net/ffrummidd-teeeee d== 28
[62] Nikita Nangia and Samuel Bowman. Listops: A diagnostic dataset for latent tree learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Student Research Workshop, pages 92-99, 2018. URL https: //aclanthology. org/ NN/NNNN88111111 29
[63] Mary Ann Marcinkiewicz. Building a large annotated corpus of English: The Penn Treebank. Using Large Corpora, 273:31, 1994. 33
[64] Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. EMNLP 2018, page 66, 2018. URL attps://aclanthology wrg./.rrgg.. 201222 pdf. 33
[65] Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguisties (Volume Long Papers), pages 1715-1725, 2016. URL org/P16-1162 pdf. 33
https://aclanthoon
[66] Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems, 31, 2018. URL http:tttts://prrceeediiaa newriinnnn.eeeu.s.eeeee...rriiee neurips.cc/pppppppe wwcpppapr_ffleeeeee /ttps//ieee///pe//////111..
5411977baa770019eba4455777222223 Papp:ee2222-2-e----e-eerrrr pdf. . 40
15

---

[67] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learn- ing via prrmeearaattiiiiaiionn In International conference on machine learning. pages 242-252. PMLR, 2019. URL https:: https://proceeddnges mlr prw.ss//99/aallenn-eeeeee zhu19a/ milennzhuu99appm
[68] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine learning, 109:467-492, 2020. URL https: / Fliin.sprring com/ http//ootent/ppf////// 10010000000090/000111111 05839- 6. pdf.
[69] Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21(13)):1-48, 2020. URL https: //::://jmmm /jmrr. csail mitt..mit....... wtt.mmtted/ppppeeeemm httpppppps//ollueeeeeese 19- 6366 19- 636. ..f..df..... 40
16

---

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: The two main contributions discussed in the introduction are detailed in section 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 generalize 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 limitations of our work in appendix A.
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 encouraged create 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 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 or images are taken in low lighting. Or a speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 grounds for rejection, 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 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 proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All the theoretical results presented in section 4 are stated in complete form with assumptions and detailed proofs in appendix F and appendix G.
17

---

Guidelines:
(c) If the contribution is either be a way to access this model for reproducing the results or a way to reproduce 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 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 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 perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
The details for the empirical evaluation are provided in appendix D. We also provide the code with necessary documentation to reproduce the experimental results in this GitHub repository.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer this question will be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided 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 novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 replicate the results, access hosted model (e.g., in the case of large language model), releasing model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions provide some reasonable avenue for reproducibility. which may depend on the nature of the contribution. For example
(a) If the contribution primarily 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.
model (e.g., large language model), then there should
the model (e.g., with an open-source dataset or instructions for how construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case closed-source models, may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
18

---

Justification:
We will provide the code with necessary documentation to reproduce the experimental results.
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/ ceallic//uddess ButdggdsCoodddSumiiiieeee "oodeSumiisionn onPolicy) 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 central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https: http//guide///oeuubeeeeee
//nips.
my) for more details.
The authors should provide instructions 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 ver- sions (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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The details for the empirical evaluation are provided appendix D. We will also provide the code with necessary documentation to reproduce the experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper a level of detail that necessary to appreciate the results and make sense them. The full details can be provided either with the code, in appendix, or supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All the experiments are repeated 10 times, and we present all results with inter-quartile range based on these 10 repetitions.
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, 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, 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 the standard deviation or the standard error of the mean.
19

---

It is OK to report 1- I-sigma enror bars, but one should state it. The authors should preferably report 2-sigma enror 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables or 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We discuss the computational resources required in appendix D.
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: http///eerrpppepeeee/eee/er wwwpeep...pppuuuaepepppp htt.ai///Ehhccccccennneeeeeeee
Answer: [Yes]
Justification: This paper studies properties of learning with transformers and the effect of sparse attention both theoretically and empirically, and 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 deviation from the Code of Ethics.
The authors should make sure preserve anonymity (e.g., if there is special consideration due 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: This paper studies properties learning with transformers and the effect of sparse attention. This has no direct societal impact.
Guidelines:
The answer NA means that there 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, alone deployments. However, if there a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point
20

---

out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is needed point out that 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 models, providing defenses addition to attacks, mecha- nisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper deseribe 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: This paper poses no risks misuse as it focuses on empirical behaviour and theoretical analyses.
Guidelines:
The answer NA means that the paper poses 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 best faith effort.
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: The creators of the original models were appropriately cited, and the libraries used for the implementations briefly discussed in appendix D.
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 URL..
The name of the license (c.g.. CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms 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, paperswi thcode. com/ 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. If this information 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?
21

---

Answer:
[Yes]
Justification:
The code used for the empirical evaluations will be provided with appropriate documentation.
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 sub- missions 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. 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 details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve research that requires IRB approvals or equivalents.
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. 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 as the institution conducting the review.
16. Declaration of LLM usage
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 is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were not used nor were a core component of this paper.
Guidelines:
22

---

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 L.M policy (https: pettps://neuripee cc htw.rferrrrenee/222....2. for what should or should not be described.
23

---

Appendix
Table of Contents
A Discussion of Limitations
25
B Table of Symbols
26
Related Work
D Details Experimental Setup
29
E Additional Empirical Results
30
E.1 Detailed Evaluation
30
E.2 Effect of MLP Activation Function
31
E.3 Natural Language Processing Evaluations
33
F
Softmax Lipschitz Continuity: Technical Details
39
E.1
Proof of Lemma
................................
39
F2
Proof of Theorem
40
F3 Proof of Theorem 2
42
F4
Multi-headed Attention
44
Role of Sparse Softmax: Technical Details
45
Standard Softmax based Attention
45
G.2
Regular Input-agnostic Sparse Softmax based Attention
49
G.3
Heavy-hitter Input-dependent Sparse Softmax based Attention
54
G.4 Comparison of Bounds between Full and Heavy-hitter Attention
61
G.5 Loss Surfaces and Estimated Lipschitz Constants
62
24

---

A Discussion of Limitations
Our results imply that there is value in pursuing input-dependent sparse attention [25, 24, 15] in real world LLMs given that they would be both computationally cheaper while having improved generalization guarantees. However, we would like list some limitations of our work:
(1)
Our empirical results are limited to benchmarks developed to study transformers under a controlled setup, and do not speak of their capabilities (in terms of improved training speed, equivalent expressivity and improved generalization) in the wild as we are unable to perform such experiments scale. 'he potential advantages of this input-dependent sparse attention at scale remains open question, though our theoretical results and accompanying preliminary experiments provide strong motivation.
(2) We study transformers in a supervised learning setup with an encoder-only architecture, where the models are trained from scratch. We do not consider the effect of pretraining, which has been shown to be quite useful with transformers [38], and we do not cover how our results would transfer sequence-to-sequenee to-sequence learning setup with an encoder-decoder architecture (though the now common decoder-only architecture can be easily analyzed in our framework).
(3) our empirical evaluations, we have considered a few representative input-dependent and input- agnostic sparse attention to validate our theoretical results. However, there are various other sparse attention mechanisms [3] that we have not considered our empirical evaluations.
(4) Our analyses establish upper bounds for the worst case performance (convergence rate general- ization error) for various forms of full and sparse attention, and we compare these upper bounds in our discussion to understand relative behavior. We do support our discussion with empirical evaluations. Furthermore, our study is focused on in-distribution generalization, and does not consider the commonly studied problem of length generalization.
(5) As with any theoretical analysis involving neural networks, we acknowledge that there might a gap between the theoretical constants (such Lipschitz constant weight utilize and the practical estimates of those constants empirically seen with these models. However, much of analysis is adaptive in nature, where improved value of such constant can be directly incorporated for improved guarantees.
upper bounds)
25

---

B Table of Symbols
Table Problem and transformer model specific symbols discussed in section 2.
Symbol
Meaning
X
Input string of token indices
f
Ground truth function
y
Output f(X)
V
Vocabulary
Number of tokens in the vocab Sequence length
L
Sequence embeddings
W
QKK K
Query-key projection matrix Value projection matrix
V
MLP first layer weights
R
MLP second layer weights
LN
Layer normalization
MI.P activation
M
Mask matrix
Initial token embeddings
Positional embeddings
w
Token projection vector
Readout layer weights
((0)
transformer block parameters
θ
Full model parameters
fe
Learned model with parameters +
k
Number of unmasked keys per query in sparse attention
Table Analysis specific symbols discussed in the theoretical analysis in section 4.
Symbol
Meaning
E
(Sparse) softmax input stability
(x
Input stability constant for self-ateention Stability constant w.r.t. to parameter W for self-attention
Aw (E)
Av
Stability constant w.r.t. parameter V for self-attention
Ao(E)
Per-transformer- block parameter stability
Ax(E)
cnrtrransformer-llock input stability Learning objective Lipschitz constant
Ac(E)
62
Per-query maximum semantic dispersion for standard softmax attention
&
Per-query maximum semantic dispersion for regular input-agnostic k-sparse attention Per-query maximum semantic dispersion for heavy-hitter k-sparse attention Per-query minimum semantic separation for heavy-hitter k-sparse attention The maximum number of queries that attend specific key The "sink ratio"
8n
A
3k
B
Per token embedding Euclidean norm upper bound Maximum spectral norm of the query-key projection matrix W Maximum spectral norm of the value projection matrix V
I'
T
26

---

(a) Full
(b) Banded
(c) Strided
(d) Block-local
(e) Banded+Gi
(1) top-k
Figure 6: Visualizations of dot-product based attention scores matrices, which along with the value matrix VX, gives the attention-based token updates A(X) (see equation (1) in section 2). The horizontal axis denotes keys and the vertical axis queries. The color intensities denote the value of the attention scores (higher intensities denote higher scores), and the white entries in the matrices corresponds to masked entries. Figure 6a depicts standard full attention score matrix; figure 6b, figure 6c and figure 6d depict various input-agnostic sparse attention score matrices. Figure 6e shows the use of global tokens (attention scores are shown in orange) in conjunction with banded attention (scores are shown in blue), with the last two tokens being the global tokens all tokens attend to and attended by these global tokens. Note that the per-query semantic dispersion (see definition 2, figure 4) of the unmasked attention scores in the input-agnostic masks would be similar in general to that of standard attention. Input-dependent masked attention such as top-k attention (shown in figure 6f) can have a much smaller semantic dispersion compared to standard attention.
C Related Work
In this section, we cover literature on efficient transformers, and the theoretical and empirical investigations on the capabilities and limitations of transformers. Finally, we will also briefly discuss the existing research on optimization with transformers.
Efficient transformers with sparse attention. The transformer architecture [1] has had tremendous impact in various fields such language modeling, vision and tabular data, and spurred new research into the development of architectural variants or X-formers [3, 39, 19]. Many of these have been developed to address the quadratic computational complexity of the attention mechanism in a transformer block with respect to the context length (the number of tokens in the context), with the goal of increasing the context length. One common technique is sparsify the attention mechanism. Usually each (query) token in the context attends to all other (key) tokens as in figure 6a, leading to the quadratic cost. Instead, we can limit the set of key tokens attended to by any particular query token. input-agnostic sparsification strategies include attending (i) within a window as in figure [20] or a block as figure 6d [21], (ii) in a strided manner as in figure 6c [22, 23], (iii) to random tokens [14], or (iv) to only a small number of global tokens and these global tokens attend to all other tokens [13, 14]; this is often used conjunction with other forms sparse attention as shown in figure 6c. Input-dependent sparsification strategies include (i) using a scoring mechanism and attending only to the highest scoring tokens as in figure 6f [40, 15], or (ii) clustering [24] or hashing [25] tokens into buckets and attending only to in buckucket tokens. Surveys such as Tay et al. sparse attention mechanisms
[3] and Lin et al. [19] cover various other forms. These input-depended focus the attention the keys corresponding to the highest dot- doo-product scores the heavy hitters while explicitly ignoring the remaining keys. Sparse attention is considered in all these cases as a way to speed up the attention mechanism in the transformer block during the forward pass without significantly deteriorating the downstream performance, with the standard full attention being the gold-standard. The Long-range Arena or LRA [4] serves one such benchmark comparing different efficient transformers to the standard transformer.
In contrast to above, we theoretically study the effect of sparse attention based transformers the learning or empirical risk minimization (ERM) convergence of the whole model (containing multiple transformer blocks), and the in-distribution generalization of the model obtained via ERM. We attempt to characterize conditions under which sparse attention might show improvements over full attention.
Empirical evaluations of transformer capabilities. While benchmarks such as the LRA [4] focus on the efficiency and in-distribution generalization, transformers have also been thoroughly evaluated on benchmarks studying specific forms of out-of-distributton generalization such as compositional generalization and length generalization. Compositional generalization benchmarks
27

---

such as COGS [27] and SCAN [26] consider sequence-to-sen to-seqeence translation problems, and they have been used to highlight the inability of transformers to systematically generalize [41]. However, subsequent work such as Csordás et al. [42]. Ontanon et al. [43] have demonstrated ways in which transformers can systematically generalize. The Neural Networks and Chomsky Hierarchy or NNCH benchmark [11] considers language transduction tasks from different formal language classes such as regular, deterministic context-fiee and context-sensitive languages. This benchmark studies the ability of various models (including transformers) to length generalize that is, generalize to longer input sequences when being trained in length limited manner. There has also been a lot of research on improving the performance of transformer based models on these out-of-distribution generalization benchmarks leveraging auxiliary tasks [44] and ohain-of-though prompting [45].
In our work, we focus on the theoretical analysis of the ERM convergence and the in-distribution generalization of models based on multiple transformer blocks, and empirically validate our theoretical insights utilizing these above benchmarks. We consider one multiclass classification task from the LRA benchmark [4] and subset of the tasks from the NNCH benchmark [11] that can be posed as supervised classification problems.
Theoretical treatment of transformer capabilities. Given the widespread success of transform- ers, there have been various theoretical studies on the capabilities and limitations of transformers. One line of research focuses on the ability of transformers to express (and thus recognize) formal languages [28]. Some of these works study transformers with hard attention [46-49], while others consider the more commonly used softmax attention [50, 51]. Another line of research has focused on understanding the capabilities of transformers as algorithms [36], demonstrating how transformers can, under specific parameter settings, perform in-context gradient descent for linear regression [52] or in-context clustering [53], and how easily can such parameters can be found [36, 54, 55]. Yun et al. [12] focus on universal approximation of sparse attention transformer for scquence-to-sequeme problems, and establish conditions on the sparsity pattern that ensure desired expressivity given enough number of transformer layers.
Viewing hard-attention as a form of input- seput-dependent sparse attention, these existing expressivity results [28] are complementary to our focus on learning convergence and in-distribution generaliza- tion for models using multiple sparse attention based transformer blocks existing hard-attentin expressivity results discuss whether sparse attention transformers are expressive enough for the task at hand. Our study here focuses on how quickly and sample efficiently can such transformers learn the task, and how the attention sparsity pattern plays role.
Optimization with transformers. There has been a lot of work on understanding the optimization of transformers in terms of the benefit of adaptive methods such as Adam over non-adaptive SGD [56- 60]. However, the focus there is to understand why optimizers such as Adam converge significantly faster than SGD with transformer models; no such consistent difference has been established for previous architectures such as convolutional or residual. Li et al. [61] recently present an analysis of the training dynamics with SignGD for single transformer block model for specific noisy binary classification problem, working in the "feature learning framework", and empirically demonstrating that the dynamics of SignGD and Adam are quite similar, thus making SignGD a useful proxy for analyzing Adam.
Our study complementary to this line of work where we study the effect of sparsity in attention to non-adaptive SGD convergence and generalization. We also consider a more general sequence learning problem with multiple transformer blocks.
28

---

D Details on Experimental Setup
The code and results for the paper are available in this GitHub repository.
Tasks. We consider the List Operations or ListOps task [62] from the LRA benchmark [4] with sequence lengths between 500 and 600 both for training and testing because we are evaluating in-distribution learning and generalization. This is a 10-class classification problem. We select this task over the other tasks in the LRA benchmark because (i) this is task where transformers have better than random performance (around 30-40% compared to a random 10% performance), but there is still significant room for improvement, and (ii) can control the length of the input sequences and still have a meaningful problem, which is not as straightforward with the other document or image processing tasks LRA. From the NNCH benchmark [11], we consider tasks that can be solved as a binary classification problem- - Parity, Even Pairs, and Missing Duplicates, and 4 tasks that can be solved as a multi-class classification problem - Cycle Navigation, Stack Manipulation, Modular Arithmetic with Brackets and Solve Equation. Parity, Even Pairs and Cycle Navigation are regular languages. Stack Manipulation, Modular Arithmetic and Solve Equation are deterministic context-free languages, while Missing Duplicates is context-sensitive language. For the NNCH tasks, we consider input sequences of length 40 both for training and testing; Deletang et al. [11] train on the same length but test on longer to evaluate out-of-distribution length generalization. For all the tasks, we utilize training / holdout sets of sizes 5000 / 2000.
Sparse attention. While there are various sparse attention mechanisms (as we discussed in ap- pendix C), we will consider representative subset for our empirical evaluations. For input-agnostic sparse attention, we choose banded attention (figure 6b [20]) and block-local attention (figure 6d [21]), with varying band and block sizes respectively. For input-dependent heavy-hitter sparse attention, we choose top-k attention (figure 6f [15]). The main motivation for selecting top-k over LSH based [25] or clustering based [24] input-dependent sparse attention is that we can then easily ensure that the input-dependent sparse attention attends to exactly the same number of tokens as in the input-agnostic ones - that is, the number of nonzeros in each column of the attention score matrix is exactly the same across all sparse attention patterns we consider. We also consider versions of these input-agnostis sparse attention with varying number of global tokens (figure 6e). Note that, as we have highlighted before, the number of learnable parameters is exactly the same between the model using stan- dard full attention and the one using sparse attention. A minor difference is with global tokens where we also learn their initial global token embeddings. For this reason, we use exactly the same hyperpaaraneters for the full and sparse attention versions of the same model to ablate the effect of the sparse attention.
Compute resources and experimental setup. All our empirical evaluations are performed on a Intel i7 Core CPU (16 threads, 64GB memory), and Nvidia V100 GPU (8GB memory). Each experiment was executed with 10 random seeds and all results are aggregated across these 10 trials. Each trial took around 55 hours ListOps: 21.5, Parity: 10, Missing Duplicates: 2.5, Even Pairs: 1, Stack Manipulation: 2, Modular Arithmetic: 6, Solve Equation: 6, Cycle Navigation: 7.5 - for a total of 550 hours for each of the 3 activation functions considered. Ablation of additional hyperparameters took another 160 hours. The implementation is in Pytorch 2.2 with CUDA 12.4. We implement our own attention block to handle different forms of sparse attention.
Hyperparameters. For the NNCH tasks, we considered the transformer architecture used in Deletang et al. [11] with (i) T = 5 transformer blocks, (ii) embedding dimension d = 64 and (iii) the MLP hidden layer dMLP- 64, but with a single head (instead of 8) and a dropout of 0.01. The final classification layer uses the average of all the token representations after the final transformer block. For the ListOps task, we utilize the same architecture but use T = 10 transformer blocks for the initial experiment. We also consider varying number of heads and blocks in experiments studying the effect of hyperparameters. For all problems, we use the SGD optimizer and the StepLR learning rate scheduler with a decay rate of 0.99 for ListOps and 0.9995 for NNCH tasks and decay period of 1 epoch. For the NNCH tasks, we use an initial learning rate of 0.1, while we use 1.0 for ListOps. The number of epochs is selected to ensure that standard full attention transformer is able to consistently achieve 100% training accuracy (and thus, the ERM has converged). Thus, we use 100 epochs for Even Pairs, 200 epochs for ListOps and Stack Manipulation, 250 epochs for Missing Duplicates, 600 epochs for Modular Arithmetic and Solve Equation, 750 epochs for Cycle Navigation, and 1000 for Parity.
29

---

3.4
3.4
0.2
0.0
0.0
100
150
200
40
40
105
100
290
1.0
100
202
460
510
930
1000
Spodks
(a) ListOps
(b) Parity?
(c) Even Pairs
(d) Missing Duplicates
Figure 7: Learning convergence and generalization curves for full attention and various sparse attention based models. Each column corresponds to a task; we present 4 tasks here and 4 more in figure 8. The legend is the same across all datasets BAND((5) denotes banded attention (figure 6b) with a band size of 5; BAND(5):1 denotes the same with single global token (figure 6e). BLOC(S) denotes block local attention (figure 6d) with a block size of 5; BLOC(5):: denotes the same with single global token. TOPK(5) is top-k attention with k: = 5. Top row: Training cross-entie tross-entropy loss trajectories lower better. Bottom row: Generalization performance on held-out set as training progresses - higher is better. Further results with different mask sizes and different number of global tokens presented in figure (training cross-entropy), figure 10 (training accuracy), table (generalization) and table 4 (convergence). For the Parity task, all forms of attention have poor generalization, with held-out accuracy as low as random guessing (50% for binary classification).
E
Additional Empirical Results
E.1
Detailed Evaluation
In this subsection, we present a detailed view of the results presented in figure 7 and figure 8, where we evaluate different mask sizes (number of nonzeros in each column of the attention matrix) and the number of global tokens included with the input-agnostic sparse attention patterns. We present the trajectories of the training cross-entropy loss figure and figure 11, and the trajectories of the training accuracies in figure 10 and figure 12. In table 3, we present the best accuracy on the held-out set for each of the sparse attention patterns and contrast with that of the full attention model. Table presents the number of epochs (aggregated over the 10 repetitions of each experiments) required by each attention pattern to (i) achieve at least 95% training accuracy for the first time (if at all), and (ii) achieve the best held-out accuracy.
The results in figure and figure 11 (along with figure 10 and figure 12) show that the input- sparse attention has a slower ERM convergence than standard full attention, being unable to reach even 95% training accuracy with the ListOps and Even Pairs tasks. With the input-agnostic sparse attention, having the global tokens helps convergence in almost all cases, being critical for convergence in the NNCH binary classification tasks (Parity, Even Pairs and Missing Duplicates), especially with the block local attention. In contrast, the ERM convergence of the top-k attention is significantly improved over the standard full attention in all tasks, with improvements (in terms of achieving 95% training accuracy) over standard attention ranging between 1.37 x (121 epochs vs 167 epochs) with ListOps to 9.5x x (6 epochs vs 53 epochs) with Even Pairs (see table 4 for further results on this).
The results in table show that, in almost all cases, the input-dependent sparse attention has similar (Even Pairs and Missing Duplicates) or better (ListOps) holdout accuracy than the standard full attention. This is true both in terms of the highest holdout accuracy during the training trajectory, and the final holdout accuracy. The latter highlights that the faster ERM convergence of input- dependeen sparse attention does not lead to overfitting. In fact, with the ListOps task, the final holdout accuracy with standard attention drops from around 35.1 1 + 0.6% to 28.9 + 1.4%, while the drop with top-k attention is only from 36.3 + 0.3% to 31.3 + 0.9%. In general, the top-k: attention based transformers also have comparitively similar or lower variations in their performance. This set of results align with our theoretical result that the improved stability of the input-dependent sparse attention translates
30

---

LANE
2.0
14
0.
100
150
700
tpechs
12
200
620
220
602
2.00
4000
Spada
Epecra
Spects
Epecna
(a) Stack Manipulation
(b) Modular Arithmetic!
(c) Solve Equation!
(d) Cycle Navigation' F
Figure 8: Same figure 7 with 4 more NNCH tasks. Further results with different mask sizes and different number of global tokens is presented in figure 11 (training (ross-entropy) and figure 12 (training accuracy). For the Modular Arithmetic, Solve Equation and Cycle Navigation tasks, all forms of attention have poor generalization, with a held-out accuracy as low as random guessing (20% for each of these 5-class classification tasks).
Table 3: Generalization performance (higher is better) for standard full attention (highlighted in green) and sparse attention. We report the sea aggregated over the 10 trials (same as figure and figure 10). The first set of columns show the best holdout accuracy obtained across the training trajectory, while the second set show the holdout aceuracy at the end of training. We highlight methods that have not reached 95% training accuracy at the end of training in blue; among the remaining methods, the highest mean each column is shown in bold. See figure for the naming of the attention mechanisms.
Attention
Best holdout MissingLops
Final holdout MissingDups
Listpps
EvenPairs
ListUps
EvenPairs
2000
0.00
1:40
99.98.amm
aco
BandedlS)
34.47.40.35
58 38.55 1.34
S2.57 cas
28 1.70
54.044:13
30.344 1.42
Randed(9)
10.55
96 06.51.1100 1.00
52.0 12.06_ 1.07
1
95.35
50.42.g
Randed[5)411
34,75:20:33
2:334
81.86 22.56
30.50.00.55
99.62-20.35
$1.40._;226
Binded9441
35.54-4-0.53
99.81 19,81400 +-013
64.47+
31.06
99 99,40:4-.34
63.07 20.59
35.25
00.92 20.12
75
50,99
99.80 99.80
75 38.4882444
HCC
35.29
99.0
66.20 22.13
31.50::141
Riklocal(S)
34.85-40.42
$7.87 1.36
98 c.98
2.100---.12
usi
Riklocal())
34.75
12
51.79.g 6.94
28.59 1.31
2:.50
50.28
35.20:+0.55
28 3:22
85.62.+ 21
29.73 4:
98.47
15.10.2:22
34.63-40433
99 +-0.44
71 71.134, 23.56
30.57 1.07
98 08.58
70.33
35.55-2006
99 96
66.55
29.97 1:35
99 99.86
65
35.53
66.34 22.05
1.35
19.12
65.45 33.63
Topk(5)
16.02.40.59
100,00 20000
0.00
31.06-1075
90.94.mmm
n.00
Topk(9)
36.25 $0.29
100.00 Lo00
100.00 -: 0.00
guss
99.95 2..s
100.00 E. 0.00
to matching or better generalization error. This does hold with the Parity, Modular Arithmetic, Solve Equation and Cycle Navigation tasks. However, note that these are tasks for which al forms of attention have very close to random performance (which is 50% for a balanced binary classification problem and 20% for a 5-class classification problem), and thus none of the attention mechanisms are generalizing well. The inability the input-agnostic sparse attention to obtain high training accuracy within the training budget translates to low holdout error, especially with the Even Pairs task.
E.2 Effect of MLP Activation Function
Here, we present a detailed view of the results presented in figure 2, where we evaluate different mask sizes (number of nonzeros each column the attention matrix) and the number of global tokens included with the input-agnostic sparse attention patterns. We present the trajectories of the training cross-entropy loss with the GELU activation [31] figure and figure for all tasks, and their corresponding training accuracy trajectories in figure 14 and figure 18. Similar results for 4/8 tasks
31

---

30
(a) ListOps (Deterministic context-free)
BANDSOG
BLC.OCAL
TOPK
100
12
750
100
100
1200
(b) Parity (Regular)
BONING
PEPACEN
HEDON
TOM
180
35
100
75
(c) Even Pairs (Regular)
0.4
(d) Missing duplicates (Context-sennstivve)
Figure 9: Training cross-entropy (vertical axis, lower is better) vs number of epochs (horizontal axis) across different tasks and sparse attention forms aggregated across 10 repetitions. Each plot contains the training curve for the standard transformers (in black). Sparse attention: Banded (column 1), banded with global token (column 2), banded with 3 global tokens (column 3), block-local (column 4), block-local with global token (column 5), block-local with 3 global tokens (column 6), top-k attention (column 7).
namely ListOps, Parity, Even Pairs and Missing Duplicates with the Mish activation [33] in the MLP block are presented in figure 15 (training cross-entropy) and figure 16 (training accuracy).
In all cases, the qualitative results do seem the change much from the previous results with the ReLU activation in the MLP block presented in figure 9 and figure 11 (and corresponding figure 10 that (i) input-agnostic sparse attention models continue to train and generalize comparitively to the full attention model, and (ii) input-dependent heavy-hitter sparse attention models continue to converge faster (and generalize similarly or better) than the full attention model.
and figure 12). The overall trend continues
The input-agnostic sparse attention models continue converge comparably to full attention with ListOps and Missing Duplicates while falling behind in Even Pairs. One marked difference here is that, with ListOps, the full attention model initially converges slower than the other sparse attention models (compare figure 7a with figure 2a and figure 2d). This is more marked with the Mish activation. However, finally the full attention model convergence catches up the input-agnostic sparse attention models. This initial slowdown in the convergence is also reflected in the initial lower generalization accuracy. In contrast, the input-dependent heavy-hitter top-k attention continues consistently converge faster than full attention in terms of the training loss for both these MLP activations, with very little differences from the results with ReLU activation. This form of sparse attention also achieves better generalization performance earlier in the training process. This indicates that the difference performance is probably due to the differences in the attention mechanism and not an artifact of the MLP block configuration.
32

---

(a) ListOps (Deterministic context-(ree)
GANGED
BURLCCA
HULDCALLGG
10M
BOOT
140
250
742
12
1200
(b) Parity (Regular)
KOVEND
W.C.DOPA
TOPE
BOOK
18
75
180
15
124
90
(c) Even Pairs (Regular)
(d) Missing duplicates (Context-sensitive)
Figure 10: Training accuracy (vertical axis higher is better) number of epochs (horizontal axis) across different tasks and sparse attention aggregated across 10 repetitions (median (line) and inter-quartile range (shaded region)). Each plot contains the training curve for the standard full attention transformer (in black). Sparse attention are as follows with each with = 5 and k = 9 nonzeros in each row of the attention score matrix column 1: banded, column 2: banded with global token, column 3: banded with global tokens, column 4: block-local, column 5: block-local with global token, column block-local with global tokens, column top-k attention.
E.3 Natural Language Processing Evaluations
We consider a preliminary experiment with the Penn Tree Bank [63] natural language dataset where we use the context of tokens predict the next token. The text is tokenized using the SentencePiece tokenizer [64] with BPE (byte-pair encoding) [65] and a vocabulary size of 4096. We consider a transformer with embedding size of 32 and MLP hidden dimensionality of 128, varying the number of transformer blocks with a single attention head per block. We train the model for 50 epochs with SGD, We consider full attention and top-k attention with & 5 and report the token misclassification cross-entropy on the training set at increasing number of epochs in figure 19.
As the results indicate, the top-k attention mechanism continues to converge faster that full attention even in this NLP task for varying number of transformer blocks. We also consider a more challenging version of Penn Tree Bank with larger vocabulary of 10000 tokens. Here we also consider larger transformer model with an embedding dimension of 128 and ML.P hidden dimensionality 512, and vary the number of transformer blocks with 4 attention heads in each block. The training loss curves are presented in figure 20, and demonstrate that the input-dependent sparse attention continues to converge faster than the full attention transformer model.
33

---

Table 4: Additional generalivation/conveeggg
results: We note the number of iterations required during the training (i) to reach 95% training accuracy, with denoting that we do not reach that training accuracy, and (ii) to reach the highest holdout accuracy. For training that reach 95% % training accuracy, the smallest in each column is highlighted in bold.
Attention
Iterations 95% training accuracy
Iterations to best holdout accuracy
ListOps
MussingDups
EvenPairs
ListOps
MissingDups
EvenPairs
Standard
167218
96 53
6211
174
64
Bundi
63
38
66
Band(9)
114430
68
236
45
Band
114+13
$94-2
236
83
Band(9)+G1
72
74::21
231
514557
S)+3
137
62
227
81
Bund(9)+G3
80
69 69
224
70
Biklocal(S)
71+11
54
67-+26
Biklocal(9)
57
3
$7.226
53++
238
85 1:
BIk
154
54
235
B[k
134
69
230+19
86-+11
B1k
146
60+17
234
62:t36
Topk(5)
131+
48
6+3
40
226
18
Topk(9)
122+ s
1343
41+7
193:;99
19
0.
(a) Modular Arithmetic (Deterministic contexl-free)
BANGGD
BARDED663
0J100L-G)
TOPK
1.5
(b) Solve Equation (Deterministic context-free)
BWVAAD
SANDING
BACKET BARADA60 4G2
KULLOCA
TOPE
BADIN
BOOT
BODT
1.0
300
106
600
300
60
600
200
655
800
205
40
60
100
40
225
180
(c) Cycle Navigation (Regular)
Figure 11: Training cross-entropy (vertical axis, lower better) vs number epochs (horizontal axis) across different tasks and sparse attention forms aggregated across 10 repetitions. Each plot contains the training curve for the standard transformers (in black). Sparse attention: Banded (column 1), banded with global token (column 2), banded with 3 global tokens (column 3), block-local (column 4), block with global token (column 5), block-local with global tokens (column 6), top-k attention (column 7).
34

---

DANICC
DA1023+-L
SANDED
TOPK
(a) Modular Arithmetic (Deterministic context-free)
0450 G1
C3
BU/ACOAL
BURCCALL
TOPK
0.3
6.9
(b) Solve Equation (Deterministic context-free)
BLAUCOL
BOOT
BACKS
600
200
0.0 3.2 0.0
000 200
100
680
(c) Cycle Navigation (Regular)
Figure 12: Training accuracy (vertical axis higher is better) number of epochs (horizontal axis) across different tasks and sparse attention aggregated across 10 repetitions (median (line) and inter-quartile range (shaded region)). Each plot contains the training curve for the standard full attention transformer (in black). Sparse attention are as follows with each with k = 5 and k: 9 nonzeros in each row of the attention score matrix column 1: banded, column banded with global token, column 3: banded with global tokens, column 4: block-locas column 5: block-local with global token, column block-local with global tokens, column top-k attention.
(a) ListOps (Deterministic context-free)
SANDED
BLC BLC.NOL
B.SECAA ES1
RUL 101
TIPE
26
14
100
76
15
150
100
12
96
300
100
24
1200
(b) Parity (Regular)
BAHCCA
TOPC
180
15
100
78
75
30
(c) Even Pairs (Regular)
02
(d) Missing duplicates (Context-sensitive)
Figure 13: Same as figure 9 with GELU activation in the MLP component of the transformer block.
35

---

SANDCCHOL
DANDED+
BUACOLL
BJUCOA C
TOPK
ROOT
BOO
BACKL
130
100
210
10
30
180
18
16
(a) ListOps (Deterministic context-[ree)
(b) Parity (Regular)
(c) Even Pairs (Regular)
BWNDED +01
SANDCD403
BLCLOCAL
TOPE
13
350
150
202
350
352
132
153
10
1.0
(d) Missing duplicates (Context-sensitivee)
Figure 14: Same as figure 10 with GELU activation in the MLP component of the transformer block.
SANDED
7UC.CCL-03
-011
BUCIN
05
300
300
110
180
300
(a) ListOps (Deterministic context-free) frree)
(b) Parity (Regular)
BAACCAL
REGOCAL
(c) Even Pairs (Regular)
KEPLICA +G
TORK
150
180
200
300
150
120
250
200
15
200
120
(d) Missing duplicates (Context-sensitive)
Figure 15: Same as figure 9 with Mish activation in the MLP component of the transformer block.
36

---

(a) ListOps (Deterministic context- free)
DAMEED
DAIEED-C
BAIDED+
SUNOCAA
EUGOCALICS
TOM
76
78
(b) Parity (Regular)
BAVICAD
BLALCCAL
TOPK
RACT
120
72
150
15
(c) Even Pairs (Regular)
BANOOD
DANGCD-
BLCUCCAL
BUALCOLL
TOPE
200
(d) Missing duplicates (Context-sensitive)
Figure 16: Same as figure 10 with Mish activation in the MLP component of the transformer block.
GENOLD
JANCED 4G2
BUNLOGN
BL6.DCAL-
BLALCCAL-G
TORK
CAN
650
280
40
200
350
aa) Modular Arithmetic (Deterministic context-free)
1.5
00
(b) Solve Equation (Deterministic context-free)
1.5
(c) Cycle Navigation (Regular)
Figure 17: Same as figure 11 with GELU activation in the MLP component of the transformer block.
37

---

ALCEOCAL
80
(a) Modular Arithmetic (Deterministic context-free)
(b) Solve Equation (Deterministic context-free)
155
80
30
225
(c) Cycle Navigation (Regular)
Figure 18: Same as figure 12 with GELU activation in the MLP component of the transformer block.
NBLOCKS:8
NBLOCKS:16
NBLOCKS:32
NBLOCKS:64
FULL
Training Loss
3.5
3.0
3.0
3.0
2.5
2.5
2.5
Epochs
Epochs
Epochs
Epochs
Figure 19: Training loss convergence for full attention and top-k attention with a small transformer for 50 epochs with Penn Tree Bank on vocabulary of size 4096.
:
NBLOCKS:10
NBLOCKS:15
FULL TOPK(S)
Training Loss
2
0
50
100
0
50
100
50
100
Epochs
Epochs
Epochs
Figure 20: Training loss convergence for full attention and top-k attention with larger transformer for 120 epochs with Penn Tree Bank on a vocabulary of size 10000.
38

---

Softmmx to Lipschitz Continuity:
Technical Details
F.1
Proof of Lemma
Lemma 3. Consider the following assumptions:
(M1) The MLP activation is X Lipschitz with (0 = 0.
(M2) The MLP parameters have norms bounded by > O, that is |P] s and |||| S B.
(M3) The input to the MLP is bounded by E 0, that is ||x|| SE
Then the token-wise MLP and LN operations are Lipschitz with respect to their input and model parameters as follows Vx, x' E RR |x|], || - P.P E R xd R. R E R xd,
||MP ( MLP MLPPP.nXX)) S 1/x x x||,
(16)
[MLPP.R((M
MLP P.LP (x)]l S 7/r| P P|l,
(17)
||LP P.P.P.RX (x) MLPP.R (x)ll $ 7n R -R|l,
(18)
[NN()) (xx - <LN||x x
(19)
where yx = B²>o. np 7R = BE.
Proof. First, the Lipschitz property of the LayerNorm (and the corresponding value of (LN) has been previously established in Kim et al. [35, Appendix NJ. With LayerNorm LN IRd - R defined as
x ((iee 2.)
LN(x)
aa+bl b,,
(20)
2
V à 24 - 2(Eie(a) mii)
where a and b are the scale and shift hyperparameter. Then LayerNorm is Lipschitz with CLN =
||alloo (dd -21/a a in equation (19).
For equation (16), have the following:
MLPP (x) MLPp R (x)]
R o (Px) I R (Pxx) (Px) (21)
BA P x)|| BX P X x S B B²>, |
x||
(22)
where we use the assumption (MI) that [o(z) oz))ll - A |z - E/1. and assumption (M2) that |P||| s B and ||]|| S B.
For equation (17), we have the following:
[|MLPp.(()) MLP -
T a (Px)]l < [ (Px)
(23)
BA, (P
p)x| s BA Pll.
(24)
since ||x|| < E for all tokens as per the assumption (M3).
For equation (18), have the following:
[|MLPp.(x) MLPp ELPP.F. (x)ll
R T a [Px) - R T (Px)]l
(25)
R RPx) < R (Px))
(26)
R R| (Px) x(0)]
(27)
R || Px B. - R|
(28)
since |||| < for tokens as per assumption (M3) and (0) = 0 as per assumption (MI).
Note that the σ (0) holds for standard activations such as ReLU(x)) max(x, lr, 0) and GELU(z) (x) where
- [0, 1] is the cumulative density function of the standard Gaussian distribution.
With activations such ReLU, the MLP(x) R (Px) often positive homogeneous such that, for any 0, we have R o (Px) aR a(a 1Px), leading symmetries in the paramter space,
39

---

and making analysis of optimization algorithms challenging [66-69], and some results focus on the convergence under specific conditions. However, most convergence rates depend on the Lipschitz- ness of the ReLU network, and the Lipschitz constant is not affected by this positive homogeneity as long as we assume that the matrices R. have bounded norms (which we do). As an example, note the following for the Lipschitz-ness with respect to P:
(MP - MLP R _ max ||LPP a P MLP, a a- P (x)l
(29)
max
||a a(a
Px) -oR' a( Px)]]
-RR"o(a
(30)
I max
||aR|llo( PPx - -o(a [Px)]l
(31)
max R
Ao (P - [PP))
(32)
max B P PI
(33)
P Pll,
(34)
where we see that the effect of the is cancelled out and we get the result in lemma 3. Similar result can be shown for Lipschitz-ness with respect to R.
F.2 Proof of Theorem
Theorem 6. Given definition I and lemma 1, transformer block TF with learnable parameters 0 (W, V,P, R) is \o(E)-Lipschitz with respect its learnable parameters with
Xo(E) ÇiN (SLN(1 + 7x) 7x)(((())) + Av) + I(nr + nn)).
(35)
and TF is LEpsssiis with respect to its input X with
xx(E) <iN(1 + nx)(L + Ax(E)),
(36)
where explicitly note the dependence of the Lipschitz constant with respect learnable parameters )o(E), and input )x(()) to the Lipschilz constant of the (masked) softmax operation.
Proof. Let 0- (W, V P. R) and 0 (W, V, P. R). Then, have the following:
| XX - TFa (1,1 = [F W.W.v.R...KK (X) - TF TFw.V.r V.P. R
(1),221
(37)
S |T.... X) TF TF..V.P.R (X)112,
(Ti)
+ |
- TF (X)112,1
(T2)
+ [F R X TF (1)12,1
(Ta)
+ | IF X TF (X) 12,11
(T.)
First, processing equation (Ti), let us denote with X LN(X Aw.v(X)), then
(Ti) IIFFw.vP,n( TFW.V. (x))l2
(38)
||LN(XX + MLPP. LN(X + MLPP.R R. (8))1l2.
(39)
Σ || + MLPP (x.)) LN(X, + MLP P.P.R((X,)) (X.))
(40)
xE(1.
ÇI ( MLP ALPP.R xI)I (using equation (19))
(41)
KE[r.]
SNNTR R- - R| = LLINTAI R||
(using equation (18)).
(42)
EE[A]
Handling equation (T2) in similar fashion, we have
40

---

(72) IIFF.VVP,R(X) (X) T. (x)1221
(43)
|ILN(X + MLP (LP).R(X)) LN(X+
+ MLPP, (1))
(44)
Σ |ILN(X M.RP.R (X.)) LN(X, + MLP P. r (X.)))I
(45)
SE[L]
SLN GLN|
P.R (X.)) (using equation (19)
(46)
SE[L]
[LNTIP||P PI = LLLNTP Pll
(using equation (17)).
(47)
SE[4]
For equation (Ty), let us denote with X' LN(X Aw.v(X)). Then we have
(Ta) = =|ITFw.V...RR TF (FW.
(48)
=|LLN(X+ MLP MLPP.R LN(X* MLP M.PP.R (X' 12,,1
(49)
Σ IILN(X, MLP P.R (X. - LN(X + MLP P R (X))
(50)
VELL]
SLN - X.) + (MLPP.((..) MLP R (X.))I
(using equation (19))
(51)
VE[L]
E (1 Ix) X.3I (using equation (16))
(52)
EILL]
E RE[L]
(LN(1 + x + Aw.v(X).)) - LN(X,, + Aw, (X),(())))
(53)
Σ (1 vx SuN/AA//))) Aw.v(X).)) (using equation (19))
(54)
iE(M]
(LN(1 + v - (X)la
(55)
s CiN(1 + TII VII (using equation (7) in definition 1)-
(56)
For equation (T4), let us denote with X" = LN(X + Aw, (X)). Then we can follow the same
procedure for equation (I') and get the following:
(Ta) = I|TF TFW.V...())).
(57)
-ILNNXX' MMP L + MLP P.P.P.(.. (X"))
(58)
Σ ||LN MLP (X.)) LN(X" MLP (X))
(59)
vE(L]
Σ X X3) + (MLPp. (X.) - MLP n (X") (using equation (19))
(60)
(E(A]
Σ
GIN(1 xx)]l XX1I (using equation (16))
(61)
[E[M]]
(L + +x))) + Σ
- LN(X,, + A (X)+()))
(62)
RE[L]
Σ
(L mX Aw, (X)all
(using equation (19))
(63)
rE(L]
<LN(1
+ NX) (X) -
(64)
41

---

I [LN(1 + nx CW ( W- W|
(using equation (6) in definition 1).
(65)
Putting these all together, we have
(TF((X) (xF((X)la))) 1l,,1
(66)
LL (7R R R|| + MP PPl)
+ CiN + nx) (AvlIV
( W WIl
(67)
S SIN (næll0 - 0| mp|10 - 01) + CiN + nx) (Avll0 - 01 + Aw - Ø1)
(68)
= CiN (SLN(1 + [xx)(A (E) + Av) + L(np + n))) 100 - 0||,
(69)
where used the definition that, for matrix tuples 8, 6, ||0 011 max W -V VP
Pll. || R.}. This gives us the desired result in equation (35).
For inputs X, X, let X LN(X + Aw. v (X)) and X' = LN(X + Aw. (X)) Then we have the following:
(TF (Fx(X)122, IILN(X + MLPP.R(X) LN(X + MLP M.PP.R (X)))
(70)
=E
| + MLPP,R LN(X. + MLP (X.))I
(71)
[E[L]
Σ
[ - X (MLPP.R(X.) MLP (x.)))
(72)
(using equation (19))
SiN(1 X X
(using equation (16)
(73)
{LN(1 + +x) + Aw.v LN(X, + (X).(X).))
L
(74)
Ci 7x) (X,, - X.,) + -ww.v(X)s-w.v(P)))
[
(using equation (19))
s (i + xx) - -X) + (Aw,v Aw,v(
(76)
[1 (I)X X1 |Aw.v()) Aw.x()22.))
(77)
S (x)+)+xx) (1 + Axx -X- |
(78)
(using equation (5) definition 1),
which gives us the desired result in equation (36).
F3
Proof of Theorem
Corollary 1. Consider the following assumptions:
(LI) The sample wise loss l equation (4) a-Lisschitz.
(L2) The final readout layer weights are norm-bounded as |||b|| < 1 and the per-token output of
each transformer block is norm bounded as [|x(0 < for all E [L] and E [r].
(L3) The sequence aggregator is = (1/L)11.
Under the above assumptions and the conditions of definition and theorem 1, the learning objective L in equation (4) is Xe(E)-Lipschitz with respect to the learnable parameters θ (T, g(1) "..)............... ((7), ②), where
Ac(E) = a (E+Ax(0)) (1+ L(xx(E)
-(xx(6-11))
Xo(E)
(79)
42

---

$ and θ Proof. Let us first denote the model parameter tuples as 6 = (T, (T, O(T), (T, "....... 0(T), 2). Let X = [T tE...... T oz + =E, and X (0)) = [T us E..... T uz + +E, denote the initial token embeddings for the same input X ( , V v parameters θ and ④ respectively. Note that we are not learning the position encoding E in our setup. For any / = I,...,,TT T, let X TFate (X((-1)) and X = TF (X ((-1),, both defined recursively.
[D] with model
Then, using the loss function in equation (4), we have the following:
(C( - [(¹) =
1
Σ
(Mw,
((,.
(80)
(X.))eS
1
( - fe(x))
(81)
n (X.y)ES
where we utilized the assumption that é is a-Lipschitz. Focusing on the [f((X) fe(X)] term in equation (81), we see the following:
(X) fa(X)
x (w)) ((((W))
(82)
+
1
( x(3)
+( ))
1
Σ x
(83)
L
L
(1
x!'1
||
+
+| 1 क
(84)
L
x x 12211 + 111
(85)
where we utilized the assumption that | and || EVi E [4]. Considering the [|
(()) in the right-hand-side of equation (85), and noting the recursive definition of (
TF ac (X (t-)),, we have the following:
XX
X 12,1
(86)
Too (TF (TFoc) (x())))) TF
(TF aco (x))
(87)
2,1
TF (+.- (TFaa) (X (0))) - TF (TF (X (x)))
(Pi)
TFam (-.. (TF. (x(+=1)))) - (F (TF [X
(P2)
2,1
1=1
+ Tooc (x(1))) - TFg X (11)
(Px)
Utilizing the Ax Liisssiiss of each transformer block with respect to the input (as per theo- rem equation (36)), and applying it recursively through the transformer blocks, can bound equation (P1) as:
(P)) Xx
x 2211 -x((x) Σ
TT T..
(88)
x Σ |T T -X
(89)
For equation (P2), we will again utilize the X Lipschizzess of each transformer block with respect the input recursively to get the following:
(Pa)
TF ( (.T-.TFo gcr (X(+-)))) TF (. (..- (X
(90)
2.1
43

---

(4) T- || occ) TFc (x TFace (X (X
(91)
2.1
V
(2))) 11
(92)
where we utilize the (e()))Lipschitzneee of the transformer block with respect to the parameters in the last inequality.
We can use se(E)-Lipschirzness of each transformer block with respect to the parameters (as per theorem 1, equation (9)) to bound equation (Py) with λo ( - 1|-
Substituting this, equation (89) and equation (92) in equation (85), we have
(X) (x()
(93)
1 $ T
xx Xx - -001) (0(3)1106)-000))
(x LIIT
- Ao(E)
+ Φ111
(94)
Ax (x +11+ (x(()) Σ {x 611) +x(())10-
$
+310 - Al
(95)
(L+ -E+
Ao(E)
))
|10 110- ō11.
L(xx(E) -
(96)
Finally, substituting the above in equation (81) gives us:
! Σ (=+Ax(S) Xx(E)² (1 (1+ L(Ax(E) -(xx(())11)) (23)--1)) )) |10-@| n (X,Y)E
20(E)
(C(e) -[(]]
(97)
Sa E ++)x(()) (1 L(Ax(E) -1) L(x -1)) 100 10- 611
Xo(E)
(98)
This gives us equation (79) in the statement of the corollary.
F.4 Multi-headed Attention
As per Yun et al. [2, Section 2, equation 1J, we can write multi-headed (self) attention h heads in our notation as:
MHA (wo), vo.,H.) H(O) kejn
(X)
Σ H Amaa vi (X).
(99)
KE(A)
where (+) E Idd are the head-aggregator matrices. Here we are assuming that each head is of size d, same as the dmodet- This is for ease of exposition, as we can introduce new variable for head size
and get the same guarantees.
Now, for each of the heads
[b]. let assume the following (as in definition 1): Awwr yio (X) Awn, v X 112,1 & Ax (c) X - X|12,11
(100)
Awio yo (X) Awt,, a (X)1l2, S Aw ( W -WIl-
(101)
Then, we can show the following for multi-headed attention, assuming |['ll| S A for all i e [A]:
(X) MHAA(wa) vo, vO, (X)
(102)
H (Awc), v (X) Awo, vo (X))
(103)
2.1
Awa, (X) (x)))o()))222. S A -1|2.11
(104)
44

---

Thus, the stability of multi-headed attention with respect to its input is preserved as with a single head, but with additional constant factors.
Furthermore, in terms of Lipschitz-ness with respect to its parameters, such as W, we can see that
(waa V6, NEejA (X) (Wia) v HO)ee (X)112,1
(105)
h
H (
(Aw (X) Awo, vo (X))
(106)
h
All (X) - Awo, (X,
(107)
S AAw (5)
IWO
W.
(108)
[E
where we utilize equation (101) for the (()) parameters for each of the heads. This shows us that we establish results for multi-headed attention analogous those we study for single head attention. The driving factors continue to be Ax (g) and Aw which are tied to the properties of the masked softmax functions. However, both the terms get multiplicatively magnified with increasing number of heads, and thus any improvement in the stability of the masked softmax function will get more pronounced the number of heads increase. This intuition supported by our results figure 3b.
G Role of Sparse Softmax: Technical Details
G.1 Standard Softmax based Attention
Lemma (adapted from Li et al. [36] Lemma B.I). For any z, z R- with
max (JE(Z]
2q 5, and
max 2; 2j S 8, A.E[L]]
(109)
for a positive constant 8 O, we have the following:
é
cd T, - ø||1
|/sftmax(2)]l>
||softmax(z) - softmax(Z)]1 S
(110)
T,'
Proof. For any E R_L,, without loss of generality, let the first entry 21 be the largest, and the second entry be the smallest. By equation (109), S 6. With softmax(z). the first entry $1 will the largest. Thus:
exp(21)
[oofmma(()) loc = S1 exp(21) El=n exp(2:)
(111)
(21
exp(21 22 <
exp(o)
(112)
exp(z1) + El exp(22)
exp(2) - 22) + (L-
L
Now we can write softmax(z) - - softmaxx (z) as an aggregation infinitesimal steps along the gradient of the softmax in the direction z - 2:
softmax(z) softmax(Z)
sesotmaa + + dz)
(113)
|softmax(z) softmax(Z)]1 S
V softmaxx(z + E(Z - =z)ddElli
(114)
V V-sotmax ee(z z))llide.
(115)
Considering the ||
/osotmaxx(2 + e(z 2))l term, and denoting z(e) z + e(2 2) and s(e) =
softmax(z(3)) (z()), we have
45

---

< ( = [[diag(s(E)) - s(e)s(E) -2)l11
(116)
Σ (s(E); - ((E))) - 21)
Σ
s(t))(())((2 - 22))
(117)
je
[s())((2) -2)1.
(118)
since all the (()),,(E)) [0, 11. Noting that s(s); z |softmax(s())]l>> S exp(8)/L, we have
exp(o)
V sof)max(z(e))
S
(cxp(8)/L(2+ -z)l
1 1-
(119)
L
Thus
|sftmaa(() softmax
V |Vssomaax(2 E(2 =2)]llde
(120)
exp(o)
1z E||ide exp(o)
1z 2/11.
(121)
thus giving us equation (110).
Theorem 7. Consider the self-attention operation A R - B with input X of L token representations and parameters W, V E Rdxdd Consider the following assumptions:
(SI) The pok-token Euclidean norms are bounded as |X S Evi E [L], and the parameter norms are bounded at [[]] S and || S SY
(S2) The per-query semantic dispersion (definition 2) is bounded by 8,, that is:
Vi [L], max (.''(LL)
JX,, X.. WX.) < 8..
(122)
Then the standard softmax E. E--Lipschitz with E.= -L,/L, and the standard attention is Lipschitz
with respect to its input and parameters as following for any input pair X, X,X X R with || [X+]] S
Ivi E [L], and parameter pairs W, W. V. V E Rdd with || < T, [| < and [[V] S
T, V. < T:
(Aw.v(X) Aw.v X 2,1 (.TT((TT++1)) X 2,1,
(123)
(Aw.v(X) Aw.v 12,,1 ( W],
(124)
[Aw.v(X) Aw,v(X) S LE VII
(125)
Proof. Now, given the upper bound the per-query semantic dispersion 6, in equation (122), can apply Lemma 4 with ō = 6., giving us a E,-Lipschitz softmax with &. = exp(6.))I.
Next, we can show equation (123) utilizing lemma 4 and adapting Li et al. [36, Lemma B.2).
|Aw,v(X) Aw.v (X)12,1 VX X WX) VXfoftmax X wX 1|2,1 (126)
VVXoftmax X WX) VXsoftmax X WX WX)
(A1)
|VXsoftmax X WX VXsoftmax(X
(A2)
VXsoftmaxx WX) VXsoftmax(X WX) ( (Ag)
We will handle each of the equation (A1), equation (A2), and equation (Ag) individually. We
will use az to denote the j-th entry of softmax(X WX.)). and aji, āzf denote the j-th entry
46

---

of softmax(X
X T ( and softmax(X X :X.;) respectively. Note that, by lemma 4 and equa-
tion (122), all Ozzs azz. ā.i S E, exp(6,)/L. L.
L VV(X X)softmax(X T WX.))l
-A = IV(X X )softmaxx X (x) Σ
(127)
(x - x,))a.1
(128)
1=1
1=1
L
1. L Σ (x x,y)a S [I
<
ΣΣ
| X,1llad
(129)
i=1
j=l
i=1
TE.
[ X.y1l TE Σ
X X1|2,11 TE.LIIX X112,1,
(130)
where we utilize the fact that [V| < Y.
(A2) VV ssottmax X WX) softmax X WX)] 12.11
(131)
|||stmmaax(XX T wX.) - softmax(X T WXx)ll
(132)
X (a - azrill
(133)
j=1
L
L L
|X,1|las- - S TE
Σ
aa
and
(134)
i=1 ==
YE
|softmax X XWW.TWTL (X..) softmax(X X WXX))l1
(135)
1
S TEC
W(X x,)11= TEC.
ΣΣ
T ,
( Xx,II
(136)
1=1
< TEG
X - -x11= Σ
|X,, - X,11
Σ
(137)
f=1 /=1
$ YEG, Σ
X X Σ X1,1)
(138)
S T PL Σ - = TE - 112,,
(139)
utilizing equation (110) and the assumption that ||W]| < I' and |X, < E for all i e [4].
(Aa) = || VV [softmax(X T WX) - softmax(X X wX)] ||.,
(140)
V X//softmax(X" wX..) - softmax( X wX.))]]
(141)
X x (a, 3,))l
(142)
47

---

E 7
L
X ||as - Fzr S T=
agj - azd
(143)
==1
j=1
TE
|softmax X WX:))) softmax X WX
(144)
$ T
X WX, X wX = TEC.
(X:g X,1) T WX
(145)
i=1
i=1 j=1
< TE Σ
Σ
X - x = TEc. Σ
|
E1x X
(146)
i= ==
i=1
Y Σ
[ 1,,,,1
(147)
S TEC// X||2,1
Σ
X TEL X
(148)
Combining the individual bounds on equation (A1), equation (A2), and equation (Aa), we have the following bound as per equation (123):
Aw. (X) Aw. (2) S S + +1]|X X|12,11
(149)
For equation (124), we note the following:
||w.v Aw, (2),,,,1
VXsoftmax(X X WX) VXsoftmax(X W
(150)
(X(oftmaxx X wX,)) softmax X WX.)))
(151)
X(softmax X WX.)) softmax X WX.))
(152)
Denoting 0jf. as the j-th entry of softmax(X WX.)) and ājs as the j-th entry of softmax X WX.), and using the assumption that || < T, we have
(Aw,v(X) Aw.v(X)l221 -
r
(a3 - 0,) Y
[1a X,1)
1=1
1=1
(153)
T
ajj - X
(154)
i=1 j=1
Y
Σ
softmaxxX X WX..) softmax(X WX.))li,
(155)
where we use the assumption that X
S E. Now, utilizing lemma have
L.
[Aw,v(X)) Aw.v(X)l2,, S TE
Σ
E X WX.; X Xww.ll
(156)
i=1
48

---

L L
T
X
WX.; - WXX
(157)
i=1 j=1
1,
TE Σ
Σ
| w
- W
x,
(158)
i=1
<TE's ΣΣ
| w -X,WII
(159)
i= j=1
S rE'S. ΣΣ
X ||l - w/l &.
W]I, (160)
where we utilize XX,] < E twice, thus giving us equation (124).
For equation (125), we note that
(A) (w.v(X)]l2, |VVSsofmax((X wX) VXsoftmax(X WX)
(161)
II( V)Xsoftmax(X WXX))
(162)
< |
Xsoftmax(X X WX.))l
(163)
Noting the fact that Xsoftmax(X WX.)) is a convex sum of the columns of X, its maximum Euclidean norm is bounded by maximum Euclidean norm of the individual columns, max, X which itself by bounded from above by E. This simplifies the right-hand-side above LE V giving us equation (125).
Remark 1. For the Lipschitz constants in definition 1, Ax(E.) = ELLT + 1) T(21) +), Aw(E,) TL = e and Av = LE with & = exp(6.)/L and 8, defined in equation (122). Under the assumption (S2) of theorem 7, 8, S 21
G.2 Regular Input-agnostic Sparse Softmax based Attention
Lemma 5. Given mask b E (0, 1 with k nonzeros, define the i-th entry of the masked softmax softmaxp : IR_ + St, for an input z E IRd as:
exp(z;)b;
softmax,(2); Σ exp(zz)bj
(164)
Now, for any 2, 2 E IR_ with
max
2: -
5. and
max A =,
2, - 5
(165)
for a constant 8 O, we have the following:
|softmaab (z)ll (2))
k
(166)
[|sotmaax (z) softmax, (2)tmaxb(2)]11 < e
h b (z 2) s 6 k =1 E||1.
where 3 denotes the elementwise multiplication of two vectors.
Proof. For any 2, E R. and fixed mask b E {0, 13 with k nonzeros, let 2(b), x[b] E R denote the k-dimensiionl vectors corresponding to the unmasked entries of z, z. Then, utilizing lemma 4 for softmax operation over a k-length vector with equation (165), we have the following:
49

---

(z)]ll = |> 1 exp(o)
(167)
k
Furthermore,
|softmaxb (z)
softmax, (2)ll ||softmax [b]]
softmax(21b)) bb)]l
(168)
exp(6)
[z[b] xbb] ||
(169)
k
exp(6)
lb @ (z 2)
(170)
1:
cxp(6)
z z
(171)
k
where the last inequality from the fact that { distance between masked vectors is smaller than the ( distance over the full vectors.
Theorem 8. Consider the relf-attention operation A Rdxx + RoxL with input X of L token representations and parameters W, V E Rdxd utilizing k-regular input sparse agnostic masking
function IR I. +
where m(D MVD RLXL Consider the following
assumptions:
(RI) The per-token Euclidean norms are bounded as |X & EVi E [L], and the parameter norms are bounded || [W] S
and [ S Y.
(R2) The per-query semantic dispersion (definition 2) is bounded by 8r, that is:
Vi E [L], J'' (L;
max
(X. WX, - -XJ. 6,.
(172)
Then the masked softmax is ---iishchitz with &r < and the masked attention is Lipschitz with
respect to its imput and parameters as following for any inpul pair X, X E RRx with [X S
Ivi E [L], and parameter pairs W, W.W, TRdxd with || < T, W < and [[] s
Y, V T.
(Aw)vXX) Aw.v (X) 2.,1 s T. + X
X 2.1
(173)
(A.v(X)) Aw.v (X) / W/|,
(174)
(Aw.v(X) (x),((X)l2.111 S L3 -VI.
(175)
Proof. Now, given the upper bound the per-query semantic dispersion 8, in equation (172), can apply Lemma 5 with d = 8,, giving us a E,-Lipschitz softmax with & exp(6,)/k.
Note that, given k-rrgulaar input agnostic masking function and the corresponding mask matrix M. know that, for any column M,, E [4 Mji k. and for any row M. E E Mij k the mask matrix has k nonzeros in each and each column. We de- note the masked softmax with a mask matrix M of a dot- dot-proouct matrix D e R. as softmax (D), defined as the columnwise masked softmax, which itself is denoted softmaxM, (D..) and defined in equation (164).
For equation (173), we proceed as follows:
Aw.v(X) Aw. 122.1
VXsoftmax (X WX) VX
X (X
WX ||,,1
(176)
S VXsoftmaxM (X WX) VX X WX)
[2,,
(B1)
+ VXSoftmaxx X WX) VXsoftmaxM WX)12a
(B₂)
+ VXSoftmaxx X
WX) VXSoftmaxm (X WX) ||2,1-
(Bg)
50

---

We will handle each of the equation (B1), equation (B²), and equation (By) individually. We will use Jj to denote the j-th entry of masked softmaxM. (X W..)), and ajis ā denote the j-th entry of softmaxM, (X T ;)))) and softmaxM, (X WX.;) respectively. Note that, by lemma and equation (172), all aje, Ojt, aji, àji I & /xp(6,)/k.
(B,) = IV(X X)softmaxM(XX wx)|12,1
Σ IV(X X)softmaxM, (X :XX;)]
(177)
L
I
Σ
(X.j- - )ayaa S V
(X. - Xzjlajill
(178)
j=1,M,;=1
i=1 j=1,M,
1
|X X |1/ S TE,
X X,ا
(179)
i=
M=1
i= j=1,My;=1
L
L
T&
I(M = 1) - Xx, 11 TG Σ
I( = 1)|X, - x,1
(180)
jml
L
= -,
>
|X - -x,;112
I(M = 1))
(181)
j=l
1m1
I.
TG
>
|X
- ||,,,1 |
(182)
where we utilize the fact that [ Y, and the row sum of the mask matrix is exactly equal to k.
(B2) ||VX softmaxx (X T WX) - softmaxM (X T WX) 11211
(183)
VX V||softmaMM. (X WX.)) softmaxM, (X :XI
(184)
Ky(an IVI Σ Σ
(185)
j= 1,=11
M,M
b
$ T=
las- -asl
assl
(186)
jim M,,
L TE
|softmaXM. (XT TWX,) softmaxM, (X TWX,)ll
(187)
L
$ TES Σ
M O X W(X, - X,)112
(188)
= T3G, Σ
Σ
W(X : - xx,)I
(189)
1= j=I,M
< TEG Σ
W X, - X.1l
(190)
i=1 I,
L
= TEE Σ
X X
x.J.W Wil
(191)
M=
51

---

< TES,
Σ
X
X IIlill
Σ
|x
(192)
1=I,M,1=11
S rE r
Σ
|x X
rE 8112,1,
(193)
utilizing equation (166), the assumption that || S r, X S E for all i E [4], and that the column sum of M is k.
(Ba) VV softmaxM (X WX) softmax (X T [X) 12,1
(194)
X|softtaaM.. (X WX.) softmaxM. (X wXlIll
(195)
Σ
III
Σ
-
(196)
i= M,,=M
$ Y=
رaاs -
(197)
j= 1,M
TE
|softmaMM, (X WX.)) softmaxM, (X TWX,)ll
(198)
S TEC,
M 0 X WX, X :X,))l
(199)
i=1
= TEG, Σ
Σ
(X.j -x,))xX. wx.1
(200)
1=1 j= M,,=M
4
S <TE Σ
Σ
-x--X., WX
(201)
1=1 j= M,,=1
L
= TES, Σ
|XxIll
Σ
| - X,,1
(202)
j=1,
L
S YEE Σ
|
Σ
[ - X
(203)
1=1
j=1.M,;:11
<TE I
| Xz,
(204)
M=
S <TEEST 1=1
ΣΣ
I(Ms = =1][|, - X,11
(205)
||x-j -X,1l F. 1IIM==1)
-
Σ
(206)
|X - =E 111,,11 -= Σ
(207)
Combining the individual bounds on equation (B1), equation (B2), and equation (Bz), we have the following bound as per equation (173):
52

---

[Aw.v(X)-Awwv((//) AX)l(X)l2,, & 1)] X11,,11
(208)
For equation (174), we the following:
(X) v (X)12,,
VXsoftmaxM (X WX) VXsoftmaxx WX)122,
(209)
VX(sotmmax,, (X (X.)) softmax (X WXx)))
(210)
$ X (softmaxx (X WX.)) softmaxM, (X WX.
(211)
Denoting Oji as the j-th entry of masked softmaxM X WX.)) and āz. the j-th entry of the masked softmaxM. (X :X.)), and using the assumption that || < T, we have
(X.)
Aw,v
ST
(a.i -
Sr
Σ
Σ
I(ay; az. X
(212)
1=1
j=1,M
1=1 j=
r
agi نهو |
(213)
1=1 j=1,
< YE
|softmaxM,
(X WX.)) softmaxm, (X WX.)))
(214)
where we use the assumption that |X, & 1. Now, utilizing lemma 5, we have
(Aw.v(X) - Aw.v(X)]l2) S TE Σ
M, O (X wX, - X Wxx))l
(215)
i=1
YEG,
WX - Wx.
(216)
j=1,M;;=1
S YEG,
| x
(217)
M=1,M;:11
STE'S
X W X!WI
(218)
J=
< <
| Wl
(219)
M=
$ < LK W W|I
(220)
where we utilize |X,,|| < 3 twice, thus giving us equation (174).
For equation (175), we note that
AAw,v(X) - Aw.v(X)l2.1 = VVXooftmaMM (X WX) VXsoftmaxM (X WX)
(221)
53

---

Xsofmmax (X WX,))l
(222)
XsoftmaxM, (X WX.))l
(223)
Noting the fact that XsoftmaxM... (X .X.) ) is a (sparse) convex sum of the columns of X, its maximum Euclidean norm is bounded by maximum Euclidean norm of the individual columns, max; Xxjll. which itself by bounded from above by E. This simplifies the right-hand-side above to LE Vil. giving equation (175).
Remark 2. For the Lipschitz constants in definition 1. (x( = E + 1) ( + 1), (W(Er) kLKE = TLL and Xy LE with & exp (xp(6,)/k and 6, defined in equation (172). Under the assumption (R2) of theorem 8, 6, 2
Remark Note that, with + L,, which corresponds to standard- -ooftaa based atlention, 8, + 8s. and the results in theorem 8 reduce to the results in theorem 7.
G.3 Heavy-hitter Input-dependent Sparse Softmax based Attention
Lemma 6. Given k-heavy-hitte masking function m . + {0,1) such that, for any z E R4, with the corresponding mask b m(z), the number of nonzeros in b is exactly k, and
2: - 2
(224)
JE( 1,6
where A denotes the smallest gap between the heavy-hitter unmasked values in and remaining masked values. Furthermore, for any z, E R with corresponding input dependent masks = m(z) and b- m(z) respectively,
max LjE (1
2j S
and
-z 6,
(225)
for a constant 8 > O. Denoting the combined masked vector as c = bvb with : = I[(b; = 1vb; = 1), we have the following:
e
softmax
soffaaa ||softmax softmax, (2)]l1 S (1+ 1/++1/4 1/2)
llc (z 2)11 (226)
k.
k
where 3 denotes the elementwise multiplication of two vectors.
Proof. For any Z, E R with input dependent masks b, E {0, with nonzeros, let [b. 2bb] E R denote the k-dimensional vectors corresponding to the unmasked entries of z, Then, utilizing lemma 4 for a softmax operation over k-length vector with equation (165), we have the following:
|softmax = ||sotmax(z((b)ll>> S exp(o)
(227)
h
Furthermore,
softmaxx (z) softmax, (z)ll
< ||softmax - softmax, [z)]l + |softmax, (z) softmax; (z)]l
(228)
exp(o)
b (z 2)|1 + |sotmmax (2) softmax, (z)]l
(229)
where we utilized lemma with the fixed mask b. Now the second term is the masked softmax with two different masks b and b on the same input 2. Then the maximum change between the two masked softmax occurs when the entries that go from being masked to being unmasked (or vice versa) the i e [L] such that b, @ b, - have the highest values. That is,
54

---

[|softmax((z) softmax,(2)) 5
|softmax(()) softmax;(2))
(230)
Σ
||sotmmaxb (z) oo
(231)
iE(L] du 86,
< b b| ||oftmax (2)
(232)
Let k ||b 1o-ibii b//1 denote the change in the mask when the input to the mask changes from
Without loss of generality, assume that z is such that b, = 1, Vi E [k] that is the first k entries of z are the heavy-hitters. Similarly, assume that for z, the corresponding mask b overlaps with the last k" entries of b and b; Ivi E [k - k" + 1, 2k k"]. Given that k' = ||b bll1 2(k k"), we can show that k" (k- k k2)
Note that, by equation (224), we know that there exists thresholds t, f E IR such that
2 2 t E [k] 2 م l - (k+
z: f 2 22 [k k" 1, k"]
2 م A E [1, k - k")
12 k2 + 1,
Now we will just consider the first (2k k") entries of and We see that
(1 - + A). i 11 k* k"] (-t+ A f) [k + 2k k"
(z+ 2,)
(233)
Thus the & norm between such two z and ž is lower bounded as
2
Cc (z 2)]11
Σ
22
E/|
(234)
K"
2k-k"
(z - 2:)1
((zi - E:)l +
Σ
1( - 2:)
(235)
#M
#=
2k-k" k"
1(4- 14-=1)1 +
1(2 - -;:]
(236)
k=++1
2k-k"
(t- 7) + + + Σ
1[t ( 1) A
(237)
k=k+1
(k k²) (1(t 8) A + (t 8) Al)
(238)
Denoting (t f) as E. consider the term (|: + A
+ - A and note that A > We can see
that, if |s| < A, the term is equal to A + | + A
2A.
= 2A. If | > A, then term is equal to
|s] + A + | A 2 2A.
Thus we have
| (z - 21 2 (k-
Al)
(239)
2 2(k - A'A = |bb - b
(240)
giving us ||b 11b-00 bllI < (1/A)]cc
(z 2) |. Utilizing this in combination with equation (229) and
equation (232), we have
exp(
(z 2)|1 + (1/2) exp(o)
||softmaxp (z) - softmax;(*)]11 <
cc (z 2)[1
k
(241)
55

---

exp(d) k (1+ @ (z 2)
(242)
(z 2)]1 & ||e (z 2)[l1 b contained with c. This gives us the desired result in equation (226).
since |lb
Theorem 9. Consider the self-attention operation Rdx - RăxL with input X of L token
representations and parameters W, V _ddd utilizing a k-heavy-hitter input-eepenent masking function m RL + 10,1] applied columnwise to the dot-product matrix to get mask matrix M E {0, 1] Consider the following assumptions:
(HI) For any query-key pairs X, X E BdxL, the k-heavy-hitter mask M = m WX) (applied columnwise) has a minimum per-query semantic separation (definition 3) of Dn. > that
Vi E [L],
min M. (X -X, WX.) An.
(243)
(H2) A maximum of Bk, query tokens attend to a single key token, that is, ||M S Bk for any i [L].
(H3) The per-token Euclidean norms are bounded as ||X Evi E [L], and the parameter norms are bounded at [W] I and |[V Y.
(H4) The per-query semantic dispersion (definition 2) is bounded by on, that is:
Vi E [L],
max My [XX, X.J WX.) S 8n.
WX,
(244)
Then the masked softmax is &n Si-iischhitz with &n = (<>h/k) (1 + 1/aw), and the masked attention Lipschitz with respect to its input and parameters as following for any input pair X, X d with || < Ivi E [4]. and parameter pairs W, W,V,V E Iaxd with ||
I || S
and [ S T, | T:
B 1+ 1/2. )
AAw (X) Aw. (8)122 S Th (B++)
|X 1X-X11
112,111
(245)
(x)) - AX)v(X)]l2. S T W -W],
(246)
(Aw.v - AX)v(X)]la S V
VII -VV--VI
(247)
Proof. Now, given the upper bound on the per-query semantic dispersion 6n in equation (244), and the per-query semantic separation AS in equation (243), we can apply Lemma 6 with = 8n and An. giving us -I.lichccitzz softmax with &a = exp(5,)(1 + 1/a.)/k.
Note that, given a h--ter input-dependent masking function m and the corresponding mask matrix M, we know that, for any column M., i E [L], プ Mys k. However, unlike the k-regular input-agnostic mask, for any row M., i E [4], E/ Mij + k. Here, we will utilize assumption H2 which states that, for any row M.:. Σj Mj I .k.
For equation (245), we note that the mask matrix is npput-dependent, and thus will denote as mask matrices M, M, M for the following dot not-product matrices (X WX). (X wX) and (X TTWW WX) respectively. That is, M = m(X WW), M = m(X T WX), M = m WX), where the masking function is applied columnwise to the dot- dot-product matrices. Given this, we proceed as follows:
Aw.v(X) - Aw.v(X)ll2
VXSoftmaxM (X WX) V X WX ||,,1
(248)
S VXsoftmaxM (X WX) VX X WX)
|2,1
(C1)
+ VXSoftmaxx WX) VXsoftmax M
(X)12
(C²)
+ VXSoftmax X WX) VXsoftmaxx XX WX) ||2,1-
(C²)
56

---

We will handle each of the equation (C1), equation (C2), and equation (Ca) individually. We will use 4j to denote the j-th entry of masked softmaxM. (X W.,)., and ajis ā to denote the j-th entry of softmax (X T (X;)) and softmax (X WX;;) respectively. Note that, by lemma and equation (244) in assumption H4, all ajz aje, ā. S exp(6x)/k = 6a/(1 + 1/an).
(G1) = |V(X X)softmax (X Wx)]12,1 =
IV(X X)softmaM (X T WX.))l (249)
(X, X,yيررر S IV
(X.: - X,)aazil
(250)
j=1.M,==1
j=1,M,x=1
< M Σ
[X - X,11llaز S r
X - Xz1
(251)
h :=1 1=
En
= =r
T(M5: = 1)|1x, - X,,
(252)
1+ 1/A,
4 Σ
L
En
T(Ms: = 1)|XX, - Xx,1
(253)
1+ 1/AI
i=
L
L
r
En
|X.- X,1
Σ T(M,; = =1)
(254)
1+ 1/Ai
j=1
En 1/A.
Tk 1+ 1/2.,
S r
X 8,y118k=
|X- X 1 2,11
(255)
where we utilize the fact that || S Y, and the row sum of the mask matrix is upper bounded by Bk from assumption H2.
We handle equation (C2) in the following manner:
(C2) IV [softmaxM (X WX) softmax (X WX)] Ila,1
(256)
VX[ (X : - softmax (X WXx)]l
(257)
IV
Σ
Xz(asi - -1:))l
(258)
Σ
Σ
|x,s|laje - aa,i
(259)
I, 1, Myamlvvv Iv f
< TE
lay- aael
(260)
TE
|sotmaaM (X WX.) - softmax (x WXx))l1
(261)
< YEG
[[MM, O - -x,il12
(262)
=TEG
Σ
Σ
X R.X,X.-XX
(263)
j= 11,M,-IVMM,.
57

---

S TECA
Σ
| X.
(264)
i=1 1, IvM
= TEch Σ
| - X,il
Σ
(265)
1=1
1= = IvM ml
T Σ
|X X|I
Σ
Xx,||
(266)
i=1
L < T
Σ
|X | X.1|
-215) = TT X|
(267)
im1
utilizing equation (226), the assumption H3 that ||| T. |X,|| s E for all i [4]. and that the column sum M is k thus Σ/= I(M = vv M,. = 1) I 2k.
We handle equation (C3) the following manner:
(C2)
IV [softmax (X WX) softmax (X wx)] Ila,1
(268)
| [||softmaxx
(x WXX:) - - (X WXx))ll
(269)
L
IV
Σ
X - -zx:)l
(270)
jml M My
Σ
| azpl
(271)
JiI,M,IIVMM,
< <E
(272)
j=1 M
-E
|softmax M..
WX.) softmaxM. (X WX
(273)
< TEC
[I,, O WX - X (XX:)l
(274)
TEE,
Σ
Σ
(X,-X.)) x,)) Twx..
(275)
1
< TEW Σ
[ 1--x.yl1 x.I.I
(276)
= TEC Σ
WXX:ll
[X X
(277)
J=1,M,;=IVMM
TEG
Σ
X
Σ
X -x-1
X.11
(278)
58

---

Σ
|x:;; X5o1
(279)
1=1, 3=I,,;;=IvA N1y=1
< < Σ
Σ
IMM1 = V Mas = 1)11X,3 - X1,5
(280)
==
|| - X,yl Σ 1(M1 V Mai Mjj= =1)
=TE?SAL'
Σ
(281)
5=1
E Σ
||, X,5ll
Σ I(M3i 1) Σ 1(Ma = 11
(282)
j=1
<TERGA1' Σ
| - -X,y (8k + 3k)
(283)
sr
Σ = X - X |
(284)
Combining the individual bounds on equation (C1), equation (C2), and equation (C²), we have the following bound as per equation (245):
3 AAw.v - Aw, v (X) S TEAh (B (1+1)+ 1)+ 1+ 1/A,
|X - X 12.111
(285)
First, let us denote the input-dependent mask matrices as M and M for the dot-product matrices (X WX) and (X WX) results. Thus M m(X WX) and M mxX" WXX. Utilizing this
for the left-hand-size of equation (246), we note the following:
(X) Aw,v (2) VXsoftmaxM X (X) - VX (X Wx)
(286)
VX(softmaxM. (X wX.)) softmaxM, (X WX.))
(287)
Σ X (softmax.. (X WX.)) softmaxM, (XTXX WX.)))l
(288)
Denoting 0zz as the j-th entry of masked softmaxM (X WX.;) and āzz the j-th entry of the masked softmax (X WX.), and using the assumption that ||V|| < Y, we have
AAw.v(X) Aw. (X)la.
Σ
(ajz
X.jll
(289)
1=
,=
r
Σ
[aas X
(290)
1=1
laj- - a
(291)
1=1 j=1 IVM
59

---

S TE>
softmaxM, (X T WX.)) softmaxm.. (X wX
(292)
where we use the assumption that |X S E. Now, utilizing lemma 6, we have
[Aw.v(X) - (x)l(X) S TE Σ
En (M,,
V
O - X xWx))l
(293)
= TESh Σ
Σ
WX w
(294)
1= 1
Σ
< YEGR
[ - -.J.WilII..l
(295)
M, M
S <
X W WWXXWWi
X/WI
(296)
< S
X W W|
(297)
1, M,,
S S W W
|
(298)
j=1,M,,=IvMM I,,, =1
S &
W
(2K = 2 TLE W
W|
(299)
where we utilize |X,|| < E twice, thus giving us equation (246).
Denote the input-dependent mask matrix as M = m(X T WX) for the dot-product matrix (X T WX). we can express equation (247) following:
|Aw.v(X) - Aw.v(X)ll.1 = VXsoftmaxM WX) VXsoftmaxM (X [x)
(300)
(VV V) Xsoftmaxx (X Wx,))l
(301)
< V XsoftmaxM, (X Wx.))l
(302)
Noting the fact that XsoftmaxM, XX| WX.) is a (sparse) convex sum of the columns of X, its maximum Euclidean norm is bounded by maximum Euclidean norm of the individual columns, max, X,iI,, which itself by bounded from above by E. This simplifies the right-hand-se hand-sie above to L Vil, giving equation (247).
Remark 4. For the Lipschitz constants in definition
Ax(Ex) = Tk 21= 3) + (1+ 1/-.)
(303)
8
= ((sp(6))((111//))
21 +
(1+ 1/A.)
60

---

Table 5: Bounds for ç. (E), (C), Ay from definition 1 for different forms of attention. Note that xy LE for all forms of attention, and thus elided from this table.
Attention
(x)
Full (thcorem 3)
dx
T' +)
eLTLS
& regular (theorem 4)
fr
<'"T(2rE" +1)
k-heavy-bitter (thoorem 5)
+ (1 1ow)
e 2 Du (o,))
TLE
and Aw (En) = 2 TLKE = 2exp(6,) TLE + 1/an) and Xy = LE with &n = exp(ox)(1 1/a,)/ /k and on defined in equation (244). Under the assumptions (HI) and (H3) of theorem 5n S 2 A..
G.4 Comparison of Bounds between Full and Heavy-hitter Attention
To compare the stability constants for all different forms of attention, we have put them together in table To characterize the conditions when the stability constants for the heavy-hitter sparse attention provides improved guarantees over full attention, we have the following result:
Corollary 2. Consider the definitions and conditions of theorem 3 and theorem 5. Further assume that (i) the maximum per-query semantic dispersion for standard attention is og S 2T while that of heavy-hitter attention is ón. c,dg, and (ii) the heavy-hitter minimum per-query dot-product separation is Dr = c2d..or some positive constants c1, C2. Then (w(En) < Aw (W(C)) when
1+ 1 <1.
C1 + 5g
log2
(304)
c25.
and (E(En) < (Ex) when
21 3) (1
1 c252
ei + 5,
log
+
log (2T=² +1 <
(305)
8,
Proof. We arrive at equation (304) by comparing (5.). exp( TLE in remark with 1/a.). We arrive at equation (305) by comparing Ax(E.)
Aw (En) = 2 exp(,,) (LE +
exp (.) (21' + 1) in remark 1 with (Ex) defined in equation (303) in remark 4.
Based on this result, we want the constant ei (corresponding to the semantic dispersion) to be small and the constant 62 (corresponding to the semantic separation) to be large. However, this condition also depends on the full attention dispersion 8,. We present this relationship for (W vs Aw (Cx) in figure 21. For small values of 8g, C2 needs to be quite large and e₁ needs to be quite small for (E.) once 5, is large enough, the condition in equation (304) holds for almost all values of c1, C2. This indicates that it is relatively easy to satisfy the condition for heavy-hitter sparse attention to have better stability constant (E with respect the learnable attention parameter W than the Aw (E.) for full attention. However, the conditions for Xx (Ex) < Ax(E.) in equation (305) are bit more restrictive it depends 8 which corresponds to the number of query tokens that might attend to the same key- the attention sink ratio. This relationship visualized in figure 22. While for small values of 3 (column 1-3) and large enough ò,, almost all values of C1, C2 satisfy equation (305). However, as the value of / increases, the conditions are only satisfied for large values of 8, and small enough c1. We present the distribution of the semantic dispersions, semantic separations and sink ratios different datasets computed over the whole training set with the trained model in table 6. Overall, it shows that the full attention dispersion 8, is usually significantly larger than the heavy-hitter attention dispersion dn. We present different percentiles of the values seen over all queries in all training points across all transformer blocks in the model. Based on these values, we also plug them into the conditions equation (304) and equation (305) in corollary 2 and report the left-hand-see -ide values in the table. We
1.0
0.8
0.6
0.4
Aw(EE), However,
0.2
0.25
0.50 0.75 1.00
C1
Figure 21: Relation- ship of c1, C22 5, in equa- tion (304).
For any value of 8,, the region above the line denotes values of C₁, C2 for which Xw < Aw (S.).
61

---

6
2=2
== B=
= 32
64
0.1-
62
0.23
0.15
1.00
0.2
2.00
2.75
0.30
1.0
0.30
1.52
0.20
C
6
Figure 22: Relationship of C1, C2, 6,, in equation (305). For any value of 6g and 8. the region above the line denotes values of C1, C2 for which Ax (5n) < (Ex) In these figures, we assume 6₃ = 2 so that we just need to vary 8, without considering different values of and E.
see that, in almost all cases, the hidd--d--- values are lower than 1, implying that the heavy-hitter attention has better guarantees, which aligns with the empirical results we see in figure 7. This especially true when we only consider the values using the 95-th percentile values for semantic dispersions, sink ratios, and the 5-th percentile values for the semantic separations, which is the most relevant quantity as we have been studying the worst-case stability constants. There one case where the values are not less than I, counter to what we see in the empirical evaluations of figure 7. Note that we are evaluating these conditions at the optimum (the final trained model) instead of over the whole parameter space. For that reason, it is important to look at the whole loss surface which we do in the sequel.
Table 6: Empirical distribution of the semantic dispersions 8g, 8n semantic separations An and for different datasets. For each metric, we report the 75-th, 90-th, 95-th and 100-th (maximum) percentile (except for A, for which we report the 25-th, 10-th, 5-th, 0-th (minimum) percentile as it is lower bound). The left-hand-siee (LHS) of equation (304) and equation (305) are computed using the corresponding percentile values. Note that for this set of results, k in heavy-hitter sparse attention.
Dataset
Full attn dispersion 6,
dispersion 5%.
separation
Sink ratio
ListOps
[8.61, 18.5, 29.4, 87.8]
[3,51, 6.74, 9.67, 2]
[0.016, 0.005, 0.002, Ic-9]
[0.2, 0.6, 3.0, 119.6]
Purity
[8.30, 10.1, 11.2, 19.4]
3.13 3.78, 9.16] [2.31, 3.13
[0.062, 0.022, 0.011, [e-6]
[1.6, 2.6, 6.6]
EvenPairs
[2.03, 4,73, 9.44, 14.6]
[1.03, 2.84, 5.50, 8.25]
[0.009, 0.003, 0.002, 3e-8]
[1.2, 5.2, 8.0
MissDup
9.25,
23.91
12.36, 4.25.4.8. 10.5]
[0.018. 0.006, 0.003, Io-7]
[1.4. 8.0]
Dataset
IHS (304)
LHS (305)
ListOps
[0.97, 0.69, 0.56, 0.57]
[0.90, 0.67, 0.59. 0.61]
Parity
[0.70, 0.76, 0.80, 1.22]
[0.72, 0.81, 0.86, 1.29]
EvenPParrs
[3.17. 1.98, 1,31, 1.80]
13.02, 2.10, 144. 1990
MissDup
[1.53, 1.09, 0.67, 1.41]
[1.53, 1.15, 0.71, 1.20]
G.5 Loss Surfaces and Estimated Lipschitz Constants
Beyond understanding the relative behavior of the aforementioned stability (and thus Lipschitz) bounds, we also empirically visualize the training loss landscapes for 4 of the tasks in figure 23. We use the version of transformer block with the ReLU activated MLP (for loss surfaces of transformer blocks with GELU activated MLPs see appendix G.5 in figure 25). We utilize the techniques proposed Li al. [37]. Given the training model parameters O we pick two random directions 01 and 02. and then plot the training loss L(O + rú, + yd2) at the grid point (I, yI, 2, [], The grid points are computed as a granularity of E = 0.005 in both axis, that is, E, 1}. We utilize the computed loss at each grid point to generate contour plots (a heatmap visualization of the loss surface is provided in appendix G.5 in figure 24). Note that the grid point (0. 0) corresponds to the loss Z()) of the trained model. The contours on the loss surfaces full attention model are somewhat asymmetric see for example, around the center in figure 23c, figure 23d, and moderately in figure 23a. In contrast, the loss surfaces of the heavy-hitter top-k attention model are quite symmetric, especially around the center.
a,y E (-1,-1 =, 2
'Note that the random directions are "filterwise normalized". which means that each matrix of parameters is normalized independently.
62

---

-033
-0.0
-015
-0.
1.2
-.0
-05
63
1.00
-1.
-AP
6.5
-25
46
-3.0
-55
T5P
T0PC
12
1.0
(a) ListOps
(b) Parity
(c) Even Pairs
(d) Missing duplicates
Figure 23: Loss surfaces of the models with full attention (top row) and top-k attention (bottow row) for each of the 4 tasks considered in figure with the corresponding hyperparameters utilizing the notea-normalized version of the loss landscape visualization techniques proposed in Li al. [37]. Note that the (0,0) grid point corresponds to the final trained model.
SIANANI
STANDARD
53
59
75-
15
35-
50
X04
3x
10
22
25
12
120-
55
18
159
175-
379
195
119
50
240
124
20-
189
129
126
55
244
154
159
D5
25-
175
155
220
(a) ListOps
(b) Parity
(c) Even Pairs
(d) Missing duplicates
Figure 24: Loss surfaces as in figure 23 but in the form of heatmaps instead of contour plots.
Beyond visualizing the loss surface in 2-dimensions, also utilize the loss surface to approximately estimate the Lipschitz constant of the model in the selected random directions E d2. Given the loss values C(0 + zd1 + yd2) grid points (I, y). we compute the following ratios at neighboring horizontal and vertical grid points as an estimate of the Lipschitz constant Xc in theorem 2:
L(O + x01 + y02) L(O + (a + +)01 + y02)
and
<110211
(306)
[C( + x0, + y02) L(O x0 + (v + e)02)
<|10211
plot the distribution these estimates figure for the loss surfaces in figure 23 of of the tasks for varying grid ranges E (0, 1] with x, E [-r,r]. We plot the 50-th, 75-th, 95-th and 99-th percentile values of these estimates of the full attention model and the heavy-hitter top-k attention model. We see that near the trained model (small values of the grid range r), the distributions of these
63

---

-033
-0.0
-011
-0.
1.2
-.3
-05
1.0
-A0
66
110
-25
-.5
-.0
35
12
-1.0
(a) ListOps
(b) Parity
(c) Even Pairs
(d) Missing duplicates
Figure 25: Loss surfaces in figure 23 of the models with full attention (top row) and top-k attention (bottow row) for each of the 4 tasks considered in figure 7 and table 4. Note that both forms of
attention now utilize the MLP with GELU activation for all tasks (as opposed to ReLU activation in figure 23).
STANDARD
sc
75
15
X04
3x
100
22-
25
12
125-
106
159
472
17
11.
50
240
124
20.
103
z2
129
126
35
140
199
D0
25
175
155
120
2200
(a) ListOps
(b) Parity
(c) Even Pairs
(d) Missing duplicates
Figure 26: Loss surfaces as in figure 25 but the form of heatmaps instead of contour plots.
estimates are close for both the models. However, as we move farther away from the trained model (large values of in the horizontal axis), the distributions change significantly, and the top-k attention model provides smaller Lipschitz constant estimate compared the full attention model across all percentiles. This indicates that, at least empirically, the loss for the top-k attention model has a much more favorable Lipschitz constant compared to that of the full attention model, which in turn implies both faster convergence and better generalization guarantees. Thus, our stability-based theoretical investigation in this section appears align with our empirical observations section
64

---

05
65
64
(a) ListOps
(b) Parity
(c) Even Pairs
(d) Missing duplicates
Figure 27: Distribution of the estimated Lipschitz constants computed in the random directions utilized to visualize the loss landscape in figure 23 for full attention and top-k attention each of the tasks considered in figure 7 with the corresponding hyperparameters. We report the distributions of the estimated Lipschitz constants in the vertical axis in terms of the 50-th (dotted), 75-th (dash-dotted), 95-th (dashed) and 99-th (solid) percentiles (lower better). On the horizontal axis, we denote the radius of the ball around the parameters of the final learned model, and visualize how the distributions vary as the ball radius increased.
65


Reviews:
[{"review_id": "HlfhznYgOq", "reviewer": "Reviewer_spof", "summary": "The paper compares input-dependent and input-agnostic sparse attention in terms of learnability, convergence, and generalization, both empirically and theoretically. Firstly, through extensive experiments, the authors show that input-dependent sparsity enhances convergence and generalization, whereas the input-agnostic counterpart does not. Then, the paper theoretically argues that the top-k mask reduces the semantic dispersion, which improves the soft-max’s input-stability and lowers the Lipschitz constant of the training loss, leading to better algorithmic stability and faster convergence.", "questions": "The reviewer would love to see the concerns in the above addressed and revise the rating.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Clarity\n\n1. Overall, the presentation is too dense which makes it hard to follow, and I strongly recommend the authors improve readability. This is explained in detail as follows.  \n   1. The figures 1,2,3,5 are very small, and their labels and legends are impossible to read. Further, each figure has many subplots, which at least to me is overwhelming.   \n- I am privileged by being able to zoom in with an electronic device, but a reader could have been reading it from a hard copy and not been able to read the figures.    \n- Consider suppressing some information to make space for the figures. For instance, in the caption of Figure 1, “The legend is the same across all datasets” seems removable without major damage. As another example, each row in figure 3 has four plots. Do we really need all four plots right here to make the point? It seems that two per row would be enough in the main text, and more figures could be left in the appendix.  \n  2. A large number of notations are defined throughout the paper, and it is sometimes difficult to remember all of them. Further, some notations have conflicting definitions.   \n- Let us take $\\\\beta$ for example. In page 6, $\\\\beta$ is used as the smoothness constant, but **the same symbol $\\\\beta$ is used on page 9 (line 352\\) for a different meaning**. Notably, for the latter usage, the symbol was not given a name **until 50 pages later** in the paper where $\\\\beta$ is now called sink ratio.   \n- It would be great if there is a centralized table summarizing the symbol and its high-level meaning, and a link to where it is best contextualized. This would not only help readers keep track of the notations, but it also prevents overloaded definitions.     \n2. Minor typo: line 165 “absense” \n\n## Significance\n\n1. Theorem 5 relies on the assumption that $\\\\beta$ (the maximum number of queries attending to a single key) is bounded. While Table 4 gives empirical percentiles of $\\\\beta$, the theoretical guarantees depend on its worst-case value. Could the authors comment on the maximum observed values of $\\\\beta$ across the experiments?  \n2. How does the per-step computational cost of calculating the top-k mask compare to the baselines? The reason for this question is that, while top-k mask has better performance given the same number of steps, it may have a higher time-per-step. Therefore, a fairer comparison is to fix a time budget (instead of number of steps) for all the baselines and see their performance.", "was_revised": true, "final_justification": "The detailed clarification by the authors has addressed my concerns. In the comment, I have encouraged the authors to integrate their explanations in a revised version of the paper. I am leaning towards accepting the paper.", "rebuttal": "We truly appreciate the reviewer's time and effort in thoroughly reviewing our manuscript and pointing out notational and typographical errors. We have tried to address all their comments in the following. We are thankful for their support, and we are happy to answer any further questions that can motivate the reviewer to raise their ratings.\n\n--------- \n\n>   _Overall, the presentation is too dense which makes it hard to follow, and I strongly recommend the authors improve readability._\n\nWe apologize for the dense presentation, and we thank the reviewer for their helpful suggestion. We were trying to present an in-depth empirical analysis to the extent possible to highlight the robustness of the relative behaviour across experimental settings. We will update our manuscript to incorporate the reviewer's suggestions. \n\n> _A large number of notations are defined throughout the paper, and it is sometimes difficult to remember all of them. Further, some notations have conflicting definitions._\n\nWe thank the reviewer for thoroughly reading our manuscript and identifying this notation overload. We really appreciate it! We will correct the discussion in the beginning of section 4 on page 6 to consider $\\alpha_1$-Lipschitz and $\\alpha_2$-smooth objectives (instead of the current $\\alpha$-Lipschitz and $\\beta$-smooth objective) to disambiguate the use of $\\beta$ later.\n\nFurthermore, we will add a longer version of the following table to clarify notation:\n\n  | Symbol | Meaning |\n  | ------ | ------- |\n  | $X$    | Input string of token indices |\n  | $f$    | Ground truth function |\n  | $y$    | Output $f(X)$ |\n  | $\\mathcal{V}$ | Vocabulary |\n  | $D$ | Number of tokens in the vocab |\n  | $L$ | Sequence length |\n  | $\\mathbf{X}$ | Sequence embeddings |\n  | $\\mathbf{W} = \\mathbf{Q}^\\top \\mathbf{K} $ | Query-key projection matrix |\n  | $\\mathbf{V}$ | Value projection matrix |\n  | $\\mathbf{P}$ | MLP first layer weights |\n  | $\\mathbf{R}$ | MLP second layer weights |\n  | $\\textsf{LN}$ | Layer normalization |\n  | $\\sigma$ | MLP activation |\n  | $\\mathbf{M}$ | Mask matrix |\n  | $\\mathbf{T}$ | Initial token embeddings |\n  | $\\mathbf{E}$ | Positional embeddings |\n  | $\\boldsymbol{\\omega}$ | Token projection vector |\n  | $\\boldsymbol{\\phi}$ | Readout layer weights |\n  | $\\theta^{(t)}$ | $t$-th transformer block parameters |\n  | $\\Theta$ | Full model parameters |\n  | $f_\\Theta$ | Learned model with parameters $\\Theta$ |\n  | $k$ | Number of unmasked keys per query in sparse attention |\n  | $\\xi$ | (Sparse) softmax stability |\n  | $\\lambda_X(\\xi)$ | Input stability constant for self-attention  |\n  | $\\lambda_W(\\xi)$ | Parameter $\\mathbf{W}$ stability constant for self-attention |\n  | $\\lambda_V$ | Parameter $\\mathbf{V}$ stability constant for self-attention |\n  | $\\lambda_\\theta(\\xi)$ | Per-transformer-block parameter stability |\n  | $\\lambda_{\\mathbf{X}}(\\xi)$ | Per-transformer-block input stability |\n  | $\\lambda_{\\mathcal{L}}(\\xi)$ | Learning objective Lipschitz constant |\n  | $\\delta_s$ | Per-query maximum semantic dispersion for standard softmax attention | \n  | $\\delta_r$ | Per-query maximum semantic dispersion for regular input-agnostic $k$-sparse attention | \n  | $\\delta_h$ | Per-query maximum semantic dispersion for heavy-hitter $k$-sparse attention |\n  | $\\Delta_h$ | Per-query minimum semantic separation for heavy-hitter $k$-sparse attention |\n  | $\\beta k$ | The maximum number of queries that attend to a specific key |\n  | $\\beta$ | The \"sink ratio\" |\n  | $\\Xi$ | Per token embedding Euclidean norm upper bound |\n  | $\\Gamma$ | Maximum spectral norm of the query-key projection matrix $\\mathbf{W}$ |\n  | $\\Upsilon$ | Maximum spectral norm of the value projection matrix $\\mathbf{V}$ |\n\n\n> _Could the authors comment on the maximum observed values of $\\beta$ across the experiments?_\n\nIn the worst case, $\\beta$ can be as large as $L / k$, where $L$ is the length of the input sequence, and $k$ is the number of unmasked keys per query token. This can happen especially in the later transformer blocks where the model starts aggregating all relevant information in certain tokens for putting together the final output, and all queries attend to these \"aggregating keys\". In all our experiments, $L \\in [40, 600]$ and the smallest $k$ is 5, thus, the worst case $\\beta \\in [8, 120]$.\n\nHowever, note that $\\beta$ shows up only in the input-stability constant $\\lambda_X(\\xi_h)$, and not in the parameter-stability constants $\\lambda_W(\\xi_h)$ and $\\lambda_V$ in Theorem 5 for the heavy-hitter attention. From our detailed comparison between full attention and heavy-hitter attention in Appendix F.4, $\\beta$ is within a logarithm term; see equation (277) in Corollary 2 where the relative performance of full and heavy-hitter sparse attention depends on $\\beta$ approximately as $O(\\log \\beta)$. Thus $\\beta$ has a smaller effect than the full attention dispersion $\\delta_s$.\n\nAt the risk of introducing further notation, we believe that the dependence on the worst case $\\beta$ can be improved to depend on the distribution of the per-key sink-ratios (as well as the other parameters) in Theorem 5. For this reason, we present the different percentiles of the actual dispersions, separations and sink ratios.\n\nWe also want to apologize for a mistake in Table 4, where we actually report the distribution of $\\beta k$, the number of queries attending to a specific key, and not the sink ratio $\\beta$ distribution, thus the actual values of $\\beta$ are even lower as they would be divided by $k = 5$. Below, we present an updated version of Table 4 (Top) where we report both $\\beta k$ and $\\beta$ for different percentiles (including the worst case $\\beta$ value):\n\n| Dataset |  $\\delta_s$ [75%, 90%, 95%, max]  | $\\delta_h$ [75%, 90%, 95%, max]  |  $\\Delta_h$ [25%, 10%, 5%, min] |  $\\beta k$ [75%, 90%, 95%, max] |   $\\beta$ [75%, 90%, 95%, max] |\n| ----  | ----- | ---- | ----- | ---- | ---- | \n| ListOps | [8.61, 18.5, 29.4, 87.8] | [3.51, 6.74, 9.67, 28.2] | [0.016, 0.005, 0.002, 1e-9] | [1, 3, 15, 598] |  [0.2, 0.6, 3.0, 119.6] |\n| Parity | [8.30, 10.1, 11.2, 19.4] | [2.31, 3.13, 3.78, 9.16] | [0.062, 0.022, 0.011, 1e-6] | [8, 13, 16, 33] | [1.6, 2.6, 3.2, 6.6 ] |\n| EvenPairs | [2.03, 4.73, 9.44, 14.6] | [1.03, 2.84, 5.50, 8.25] | [0.009, 0.003, 0.002, 3e-8] | [6, 17, 26, 40] | [1.2, 3.4, 5.2, 8.0] |\n| MissingDup | [4.63, 9.25, 17.1, 23.9] | [2.36, 4.25, 4.88, 10.5] | [0.018, 0.006, 0.003, 1e-7] | [7, 15, 21, 40] | [1.4, 3.0, 4.2, 8.0 ] |\n\nThe actual worst-case $\\beta$ values are close to $L/k$. However, the worst-case $\\delta_s$ values are also much higher than their 95-th percentile values, thus making the overall bound still favourable for heavy-hitter attention. In the following, we report the updated left-hand-sides of equation (276) and (277) as in Table 4 (bottom), including the worst case values for all the involved parameters (note that the LHS (277) values are slightly lower than the ones reported in the paper as we corrected the use of $\\beta$ instead of $\\beta k$ though the effect is minor as $\\beta$ is in the logarithm):\n\n| Dataset |  LHS (276)  |  LHS (277) |\n| ----  | ----- | ---- |\n| ListOps | [0.97, 0.69, 0.56, 0.57] | [0.90, 0.67, 0.59, 0.61] |\n| Parity  | [0.70, 0.76, 0.80, 1.22] | [0.72, 0.81, 0.86, 1.29] |\n| Even Pairs | [3.17, 1.98, 1.31, 1.8] | [3.02, 2.10, 1.42, 1.90] |\n| Missing Dup. | [1.53, 1.09, 0.67, 1.41] | [1.53, 1.15, 0.71, 1.20] |\n\nAs can be seen, even with ListOps, where $\\beta \\approx 120$, the bounds are quite favourable for heavy-hitter attention because of the outsized effect for the larger standard full-attention semantic dispersion $\\delta_s \\approx 8$ compared to $\\delta_h \\approx 28$.\n\n> _How does the per-step computational cost of calculating the top-k mask compare to the baselines?_\n\nThank you for this great question. Naively computing the top-$k$ mask would require $O(L \\log k)$ per-query (as opposed to the $O(L)$ cost per query for standard attention). Gupta et al. (2021) focus on the memory efficiency of top-$k$ attention without improving the runtime of the attention. However, various advancements have sped up the computational cost of top-$k$ attention both in terms of runtime and memory overhead such as Zeng et al. (2025). This has shown significant speedup over pytorch native attention implementation, as well as the optimized FlashAttention, with the speedup increasing with context length. Thus, top-$k$ attention can now be significantly faster than standard full attention in terms of actual runtime. Furthermore, other heavy-hitter attention such as Reformer (Kitaev et al., 2020) can also speed up the actual attention runtime with locality sensitive hashing.\n\nOur focus has been on a complimentary angle where we study how sparsity affects optimization convergence, and try to provide theoretical intuition as to when and how sparse attention can converge faster than standard full attention. In situations where top-$k$ (or any other $k$-heavy-hitter attention) would converge faster than standard full attention, the overall runtime would improve cumulatively with faster sparse attention and fewer steps to convergence.\n\n------ \n### References\n- Gupta, Ankit, et al. \"Memory-efficient Transformers via Top-$ k $ Attention.\" SustainNLP workshop (2021).\n- Zeng, Qiuhao, et al. \"Zeta: Leveraging z-order curves for efficient top-k attention.\" ICLR (2025).\n- Kitaev, Nikita, et al. \"Reformer: The efficient transformer.\" ICLR (2020)."}, {"review_id": "6YrALkmEsT", "reviewer": "Reviewer_PEBb", "summary": "This paper shows an interesting case that transformers with sparse attention sometimes learn faster than full attention. The authors try to explain this by showing that sparse attention has smaller semantic dispersion, resulting in smaller Lipschitz constants, thus learns faster.", "questions": "lease refer to the Weakness part in the previous section.", "limitations": "lease refer to the Weakness part in the previous section.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "*Strengths*\n1. This paper provides extensive experiments to show that sparse attention sometimes learns faster than full attention.\n2. This paper provides detailed theoretical analysis for their experimental findings through the lens of Lipschitz constants and smoothness.\n\n*Weaknesses*\n\nI have some concerns for this paper.\n\n1. The paper aims to show that sparse attention is more stable than full attention, i.e., the constant in Definition 1 is smaller for sparse attention. Their analysis shows this is true when semantic separation is large and semantic dispersion is small. This only happens when the model is well trained and the loss is small, i.e., the model is close to the optimal point. For an attention model with random initialization, the semantic separation should be 0 as the attention score is equal for each position. In this case, at initialization, the loss landscape for sparse attention (top-k attention) is less smooth than full attention.\n\n2. Although in section 4 lines 243-259, the authors try to connet the convergence rate with Lipschitz constants, I still feel it is somehow far-fatched, a stable optimization algorithm doesn't mean this is a fast algorithm, especially when the following analysis consider the case to loss is close to the optimal point.\n\n3. Note that in Figures 1, 2, 3, there is a clear loss plateau at initialization, and sparse attention can break this plateau faster compared with full attention. However, if we ignore this plateau, we can see that the convergence rate for sparse attention and full attention is close. In other words, as I analyzed in (1), I think the reason why sparse attention learns faster than full attention is that, at initialization, the loss landscape for sparse attention is sharper, thus the model can soon escape the plateau at initialization, so we can observe sparse attention learns faster.\n\n4. In Figure 4 in [1], there is no clear difference that sparse attention learns faster. Could the authors provide some explanation?\n\n[1] Gupta, Ankit, et al. \"Memory-efficient Transformers via Top-$ k $ Attention.\" arXiv preprint arXiv:2106.06899 (2021).", "was_revised": true, "final_justification": "I maintain my position that sparse attention learns faster primarily because the loss landscape for top-k attention is sharper, allowing it to learn more quickly at initialization. The main experiments in this paper, Figures 1 and 2 clearly demonstrate that the primary loss difference emerges at initialization. The authors provide no experimental evidence that contradicts this intuition that top-k attention with random initialization, compared to full attention, is less stable. I still question aspects of the paper's main claim.", "rebuttal": "We truly appreciate the reviewer's time and effort in reviewing our manuscript and sharing their insightful intuitions. We have tried to address all their comments in the following. We are thankful for their time, and we are happy to answer any further questions that can motivate the reviewer to raise their ratings.\n\n--------- \n\n> _**Weakness (1).** This only happens when the model is well trained and the loss is small, i.e., the model is close to the optimal point. For an attention model with random initialization, the semantic separation should be 0 as the attention score is equal for each position. In this case, at initialization, the loss landscape for sparse attention (top-k attention) is less smooth than full attention._\n\nWe thank the reviewer for sharing their thoughtful intuition. The main theoretical advantage of heavy-hitter attention over standard full-attention would occur when the full-attention semantic dispersion $\\delta_s$ is higher than the heavy-hitter semantic dispersion $\\delta_h$, while the heavy-hitter semantic separation $\\Delta_h$ is not too small. We discuss the precise conditions when these hold in Appendix F.4. However, the semantic separation $\\Delta_h$ can be quite small, not only at initialization but also at the optimum (corresponding to the trained model) as we list in Table 4 in Appendix F.4. Even when the semantic separation $\\Delta_h$ is quite small (very close to zero), the stability constants $\\lambda_W(\\xi_h), \\lambda_X(\\xi_h)$ of heavy-hitter sparse attention can be better than the constants $\\lambda_W(\\xi_s), \\lambda_X(\\xi_s)$ for full attention because of the exponential dependence on the semantic dispersion.\n\nThe stability constants $\\lambda_W(\\xi), \\lambda_X(\\xi)$ of the attention operation, full or sparse, and thus the Lipschitz constant $\\lambda_{\\mathcal{L}}(\\xi)$ of learning objective $\\mathcal{L}(\\Theta)$ in equation (4) (page 3) depends on both the dispersion and separation. The reviewer posits that the stability constants (and thus the Lipschitz constant of the learning objective) would be higher for heavy-hitter sparse attention than for standard full-attention far away from the optimal points (the trained model), and thus at initialization.\n\nWe present the loss landscapes for full-attention (Figure 5 top row) and top-$k$ attention (Figure 5 middle row) with the center of the landscapes at (0,0) denoting the optimum. The loss landscapes of top-$k$ attention do not appear to be sharper than full-attention far away from the optimum (the center at (0,0)). We also present the estimated Lipschitz constants as we move farther away from the optimum (Figure 5 bottom row). The results indicate that, at the optimum --- the trained model, the estimated Lipschitz constants for both sparse and full attention are quite similar. However, as we move away from the optimum model, the estimated Lipschitz constants increase, with the Lipschitz constants for full-attention growing faster than those of the top-$k$ attention, implying that the loss landscapes are equally sharp near the optimal point, but the full-attention loss landscape gets sharper than top-$k$ attention as we move away from the optimum, counter to the reviewer's intuition. \n\nWe do acknowledge a caveat that the Lipschitz constants are estimated. However, both our theory and empirical evaluations (and the actual observed convergence results presented in Section 3) imply that heavy-hitter sparse attention has more favourable loss landscapes for optimization.\n\n> _**Weakness (2).** Although in section 4 lines 243-259, the authors try to connet the convergence rate with Lipschitz constants, I still feel it is somehow far-fatched, a stable optimization algorithm doesn't mean this is a fast algorithm, especially when the following analysis consider the case to loss is close to the optimal point._\n\nWe apologize for the confusion here. We do not claim or imply that stable optimization algorithms are faster algorithms. Rather, we claim that both the SGD convergence rate bounds for a learning problem and the algorithmic stability bounds of a learning algorithm depend on the Lipschitz constant of the learning objective --- these are established results. Thus, an improved Lipschitz constant provides improved guarantees for both SGD convergence and algorithmic stability guarantees.\n\nFurthermore, we do not make any assumption that the \"loss is close to the optimal point\", nor do the established results require this assumption. We also discuss the behaviours of the different attention mechanisms close and far from the optimal points in our response to weakness (1) above.\n\n> _**Weakness (3).** Note that in Figures 1, 2, 3, there is a clear loss plateau at initialization, and sparse attention can break this plateau faster compared with full attention. However, if we ignore this plateau, we can see that the convergence rate for sparse attention and full attention is close. In other words, as I analyzed in (1), I think the reason why sparse attention learns faster than full attention is that, at initialization, the loss landscape for sparse attention is sharper, thus the model can soon escape the plateau at initialization, so we can observe sparse attention learns faster._\n\nWe thank the reviewer for their thoughtful insight. We have already discussed the behaviours of the different attention mechanisms close and far from the optimal points in our response to weakness (1) above. It is an interesting hypothesis that the initialization of full attention is on a \"loss plateau\". First, we would like to note that all models, using full or sparse attention, have the same number of parameters, and start learning from the same initialization. Second, we repeat each training 10 times from different initializations (and seeds for sampling training batches), and we present results (the learning curves) aggregated over these 10 trials.\n\nWe note that while the training loss curve for full attention appears to \"plateau\" initially (for example in Figures 1 and 2), the full-attention training loss is actually changing (both increasing and decreasing because of an unfavourable loss landscape) as the training progresses. The effect of this change can be also seen from the variation (relatively small but present) in the heldout accuracy curves, indicating that the full-attention model is changing. This is usually different from a loss plateau where the loss gradient is very small, and thus the model changes are very small.\n\nFurthermore, for evaluations where the full-attention model does not have an initial apparent plateau, such as in Figure 3(e) with the Adam optimizer, top-$k$ sparse attention continues to converge faster (except in the case where both models diverge because of too large of an initial learning rate). \n\nFinally, we consider a small experiment with the PennTreeBank natural language dataset where we use the context of tokens to predict the next token. In this case, full attention does not have that apparent \"plateau at initialization\" but heavy-hitter sparse transformer continues to optimizer faster. For a 8-block transformer, trained for 50 epochs we see the following token misclassification training cross-entropy:\n|Method / Epoch|5|10|15|20|25|30|35|40|45|50|\n|-|-|-|-|-|-|-|-|-|-|-|\n|Full|3.29|3.08|2.97|2.89|2.83|2.79|2.74|2.72|2.69|2.67|\n|top-$k$ (% gain)|3.21 (2.4)|2.98 (3.2)|2.87 (3.4)|2.79 (3.5)|2.74 (3.2)|2.70 (3.2)|2.65 (3.3)|2.62 (3.7)|2.59 (3.7)|2.56 (4.1)|\n\nWith a 16-block transformer model for 50 SGD epochs:\n|Method / Epoch|5|10|15|20|25|30|35|40|45|50|\n|-|-|-|-|-|-|-|-|-|-|-|\n| Full | 3.35 |3.13 |3.00 |2.93 |2.86 |2.81 |2.77 |2.73 |2.69 |2.67 |\n| top-$k$ (% gain) | 3.21 (4.2) |2.98 (4.8) |2.87 (4.3) |2.79 (4.8) |2.73 (4.5) |2.67 (5.0) |2.63 (5.1) |2.59 (5.1) |2.56 (4.8) |2.52 (5.6) |\n\n>  _**Weakness (4).** In Figure 4 in [1], there is no clear difference that sparse attention learns faster. Could the authors provide some explanation?_\n>\n> _[1] Gupta, Ankit, et al. \"Memory-efficient Transformers via Top-$k$ Attention.\" arXiv preprint arXiv:2106.06899 (2021)._\n\nRegarding Figure 4(a) in Gupta et al. (2021), we really appreciate the reviewer making an appropriate connection with the ListOps dataset. There are of course experimental differences in terms of learning rates, schedulers, optimizers, training data, plotting granularities and such. But beyond those, we believe there are a few important points:\n\n- First, these plots show the eval set and test set accuracies, and not the training loss. So we can compare this Figure 4(a) in Gupta et al. (2021) to Figure 1(a) bottom row in our manuscript. Upon this comparison, we can see that in both cases, the held-out accuracies of top-$k$ attention with ListOps improves initially faster than full-attention. So these results are somewhat aligned (barring the plotting granularity).\n\n- Second, our results consider a much smaller value of $k=5$ in Figure 1(a) of our manuscript compared to Figure 4(a) in Gupta et al. (2021) where $k=128$. In our experiments, we see that a small $k=5$ is sufficient in terms of expressivity, achieving 100\\% training accuracies for different problems and hyperparameters. The smallest $k$ considered in Gupta et al. (2021) is actually $k=64$. Smaller $k$  implies higher reduction in the semantic dispersion $\\delta_h$ compared to the full-attention semantic dispersion $\\delta_s$, and thus, better learning convergence guarantees.\n\n- Finally, Gupta et al. (2021) consider a transformer block which makes use of the top-$k$ operation not only in the attention, but also in the MLP (see **Feed-forward as attention** paragraph in Section 2.1 in Gupta et al. (2021)), thereby giving us different transformer blocks. Thus, it is hard to directly compare these two different transformer models.\n\n------------\n### References\n-  Gupta, Ankit, et al. \"Memory-efficient Transformers via Top-$k$ Attention.\" arXiv preprint arXiv:2106.06899 (2021)."}, {"review_id": "0IpCg0Q3lT", "reviewer": "Reviewer_XrxX", "summary": "This paper investigates the role of sparse attention in Transformers from both theoretical and empirical perspective. The authors empirically show that input-dependent sparse attention, such as Top-k masking, leads to faster training convergence and better generalization compared to both full attention and input-agnostic sparse variants (e.g., block, banded). Theoretical analysis links this behavior to improved softmax stability and a tighter Lipschitz constant for the training objective. The work is rigorous, well-motivated, and presents both theoretical insight and empirical validation for an underexplored benefit of sparse attention mechanisms.", "questions": "See above", "limitations": "See above", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This paper presents strong and consistent empirical evidence that input-dependent sparse attention improves convergence and generalization. Theoretical analysis is solid, linking attention sparsity to improved softmax stability and tighter Lipschitz bounds. Results are robust across multiple tasks, architectures, and hyperparameters.\n\nHowever, I have several concerns before this paper can be accepted.\n\n1. Lack of literature search. The papers should discuss more about the current advanced sparse attention techniques (e.g., Performer, Longformer, Reformer) as discussed in [1].  Is it possible for the authors to compare these techniques in synthetic data experiments? I am curious about the performance of these techniques. \n\n2. Small-scale synthetic tasks. I appreciate the detailed experiments. A minor suggestion is that the authors could consider some real data experiments.\n\n3. How the choice of $k$ affects performance, generalization, or the theoretical bounds? \n\n[1] Tay Y, Dehghani M, Bahri D, et al. Efficient transformers: A survey[J]. ACM Computing Surveys, 2022, 55(6): 1-28.", "was_revised": true, "final_justification": "I will keep my scores.", "rebuttal": "We truly appreciate the reviewer's time and effort in thoroughly reviewing our manuscript and raising insightful questions. We have tried to address all their comments in the following. We are thankful for their support, and we are happy to answer any further questions that can motivate the reviewer to raise their ratings.\n\n--------- \n\n> _Lack of literature search. The papers should discuss more about the current advanced sparse attention techniques (e.g., Performer, Longformer, Reformer) as discussed in [1]._\n> \n> _[1] Tay Y, Dehghani M, Bahri D, et al. Efficient transformers: A survey[J]. ACM Computing Surveys, 2022, 55(6): 1-28._\n\nWe thank the reviewer for the useful pointers, and we apologize for the succinct literature survey in our main paper. We do provide a detailed literature survey in Appendix B where we discuss Reformer (Kitaev et al., 2020) and Longformer (Beltagy et al., 2020) alongside other sparse attention transformers. We thank the reviewer for bringing up Performer. We chose to not include Performer in our discussion as we focus on sparse attention transformers. Instead of sparsifying the attention matrix, techniques such as Performers (Choromanski et al., 2020) consider a low-rank approximation of the post-softmax attention matrix using (orthogonal positive) random features (Rahimi and Recht, 2007). Our current discussion and analysis will not be applicable to low-rank approximating schemes such as Performer; we leave such a study for future work.\n\n> _Is it possible for the authors to compare these techniques in synthetic data experiments? I am curious about the performance of these techniques._\n\nWhile we did not explicitly consider the input-agnostic strided sparsity pattern from Longformer in our experiments, we do make use of the global tokens from Longformer alongside the input-agnostic banded and block local sparse attention models. These global tokens are critical for the expressivity of the input-agnostic sparse attention transformer models as discussed by the theoretical results of Yun et al. (2020), and also validated in our empirical experiments. However, the inclusion of these global tokens in input-agnostic sparse attention does not speed up the training convergence of input-agnostic sparse attention when compared to full-attention --- see BAND:5(1) and BLOC:5(1) in Figures 1 and 2 in the main paper, and further empirical analyses in Figures 9-18 in Appendix D for varying values of $k$ and varying number of global tokens for input-agnostic sparse attention transformers. We also briefly discuss the implications of Theorem 4 (for regular input-agnostic sparse attention) for strided sparse attention as in Longformer in Lines 332-333.\n\nIn contrast, input-dependent heavy-hitter sparse attention such as top-$k$ attention retains necessary expressivity, and does not need the global tokens. Furthermore, top-$k$ attention transformers converge significantly faster than full-attention transformers, highlighting the need for input-dependence for faster convergence. This is corroborated in our theoretical analyses, where the theoretical bounds for input-agnostic sparse attention in Theorem 4 do not necessarily show any improvement over full-attention results in Theorem 3.\n\nWhile Reformer would have been a good candidate for empirical evaluation of input-dependent heavy-hitter sparse attention transformers, we limited our experiments to top-$k$ attention for reasons discussed in Appendix C, lines 1129-1133:\n\n\"The main motivation for selecting top-$k$ over LSH based or clustering based input-dependent sparse attention is that we can then easily ensure that the input-dependent sparse attention attends to exactly the same number of tokens as in the input-agnostic ones -- that is, the number of nonzeros in each column of the attention score matrix is exactly the same across all sparse attention patterns we consider.\"\n\nWe believe that is necessary to ablate the effect of input-dependent vs input-agnostic sparse attention. Otherwise, with methods such as Reformer, a query can often attend to more that $k$ keys depending on the number of key tokens in the same hash bucket as the query token. In such a scenario, it is not clear whether the improved expressivity and convergence of input-dependent sparse attention over input-agnostic ones is due to the input-dependence or due to the ability to attend to more that $k$ keys.\n\n> _A minor suggestion is that the authors could consider some real data experiments._\n\nWe consider a small experiment with the PennTreeBank natural language dataset where we use the context of tokens to predict the next token. The text is tokenized using the SentencePiece tokenizer with a vocabulary size of 4096. We consider a transformer with 8 single-headed blocks, embedding size of 32 and MLP hidden dimensionality of 128. We train the model for 50 epochs with SGD (we plan to run it for longer but are reporting intermediate results here). We consider full attention and top-$k$ attention with $k=5$ and report the token misclassification cross-entropy on the training set at increasing number of epochs:\n    \n| Method / Epoch | 5 | 10 | 15 | 20 | 25 | 30 | 35 | 40 | 45 | 50 |\n| -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- |\n| Full | 3.29 |3.08 |2.97 |2.89 |2.83 |2.79 |2.74 |2.72 |2.69 |2.67 |\n| top-$k$ (% gain) | 3.21 (2.4%) |2.98 (3.2%) |2.87 (3.4%) |2.79 (3.5%) |2.74 (3.2%) |2.70 (3.2%) |2.65 (3.3%) |2.62 (3.7%) |2.59 (3.7%) |2.56 (4.1%) |\n\nWe also considered a larger model 16 single-headed transformer blocks for the same experiment for 50 epochs with SGD:\n\n| Method / Epoch | 5 | 10 | 15 | 20 | 25 | 30 | 35 | 40 | 45 | 50 |\n| -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- |\n| Full | 3.35 |3.13 |3.00 |2.93 |2.86 |2.81 |2.77 |2.73 |2.69 |2.67 |\n| top-$k$ (% gain) | 3.21 (4.2%) |2.98 (4.8%) |2.87 (4.3%) |2.79 (4.8%) |2.73 (4.5%) |2.67 (5.0%) |2.63 (5.1%) |2.59 (5.1%) |2.56 (4.8%) |2.52 (5.6%) |\n\n\n\nWe see that even with natural language problems, heavy-hitter sparse attention is able to optimize the loss faster (in terms of the number of optimization steps) than full attention. We do note that these are preliminary results as we have not been able to verify these results across multiple different settings due to the time constraints during this rebuttal process.\n\n> _How the choice of $k$ affects performance, generalization, or the theoretical bounds?_\n\nThe choice of $k$ in sparse attention affects the overall behaviour in multiple ways:\n\n- First, unless the number of transformer blocks are significantly high, larger values of $k$ imply better expressivity for sparse attention transformers (Yun et al., 2020). This is something we see with input-agnostic sparse attention transformers in our additional experiments in Appendix D. For example, in Figure 9, we see how the improved expressivity of larger $k$ with input-agnostic banded sparse attention and block-local sparse attention allows the learning to progress instead of stalling prematurely. In terms of learning convergence with regular input-agnostic sparse attention, we actually discuss how our bounds in Theorem 4 approaches the results for full-attention in Theorem 3 as $k$ grows to $L$, the complete sequence length (which would be full-attention); see Remark 3 on page 52 of our manuscript.\n\n- However, when expressivity is not an issue, as we see in our empirical experiments with heavy-hitter top-$k$ sparse attention, increasing $k$ *can* have an implicit effect of increasing the heavy-hitter semantic dispersion $\\delta_h$ (see for example, Definition 2 and Figure 4). Larger semantic dispersion implies larger bounds for the Lipschitz constant of the learning objective, implying slower SGD convergence. This behaviour is again manifested in our additional experiments in Appendix D. For example, in Figure 9, we can see that top-$k$ attention with $k=5$ sometimes converges slightly faster than top-$k$ attention with $k=9$. Larger Lipschitz constants also imply larger algorithmic stability bounds, and thus worse generalization bounds.\n\n-------------\n\n### References\n- Kitaev, Nikita, et al. \"Reformer: The efficient transformer.\" ICLR (2020).\n- Beltagy, Iz, et al. \"Longformer: The long-document transformer.\" arXiv preprint arXiv:2004.05150 (2020).\n- Choromanski, Krzysztof Marcin, et al. \"Rethinking Attention with Performers.\" ICLR (2020).\n- Rahimi, Ali, and Benjamin Recht. \"Random features for large-scale kernel machines.\" NeurIPS (2007).\n- Yun, Chulhee, et al. \"$O(n)$ connections are expressive enough: Universal approximability of sparse transformers.\" NeurIPS (2020)."}, {"review_id": "hkvTZCAYWp", "reviewer": "Reviewer_yPUz", "summary": "This paper investigates whether sparse attention in transformers improves learnability and generalization, beyond just computational efficiency. Empirically, input-dependent sparse attention  accelerates convergence and enhances generalization compared to full attention, while input-agnostic sparse attention offers no benefits. Theoretically, the authors link this to \"semantic focus\": input-dependent sparsity improves the stability of the softmax operation by reducing semantic dispersion.", "questions": "1. From my experience, when training with sparse attention, the loss and some benchmark can show similar or better results than full attention. However, the sparse attention will fail in the challenging **downstream tasks** such as Needle in a Haystack. Could the authors provide any explanation using the proposed theory on this issue? \n\n2. Since the paper discusse top-k or k-heavy-hitter attention. Could you provide any theorical or empirical guidence on how to set the k in practice? Or I can ask in a different way, how can we deterimine the sparsity before seeing the input?", "limitations": "The authors already mention their limitation in large scale experiments. I think another limitation is the sparse attention methods used being too old.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The authors aim to provide a theoretical explaination of why sparse attention improves learnability and generalization, which is different from most of the research that focusing on efficiency.\n\nWeakness:\n- I think full attention is already **input-dependent sparse** as most of the softmax score is very low. It's only a matter of being complete zero or a small value. As a result, I guess what interest me more is whether sparse attention can be both efficient while being better than full attention? To be hardware efficient, there are many new input dependent sparse attention methods such as NSA[1], MoBA[2],  SeerAttention[3], MInference[4], Quest[5], and etc that provide block-wise sparsity. The current sparse attention methods used in the paper are not state-of-the-art. \n- The paper does not provide results on large models. It's questionable whether their conclusion still holds at scale. \n\n\n\n\n[1] Yuan, Jingyang, et al. \"Native sparse attention: Hardware-aligned and natively trainable sparse attention.\" arXiv preprint arXiv:2502.11089 (2025).\n\n[2] Lu, Enzhe, et al. \"Moba: Mixture of block attention for long-context llms.\" arXiv preprint arXiv:2502.13189 (2025).\n\n[3] Gao, Yizhao, et al. \"Seerattention: Learning intrinsic sparse attention in your llms.\" arXiv preprint arXiv:2410.13276 (2024).\n\n[4] Jiang, Huiqiang, et al. \"Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention.\" arXiv preprint arXiv:2407.02490 (2024).\n\n[5] Tang, Jiaming, et al. \"Quest: Query-aware sparsity for efficient long-context llm inference.\" arXiv preprint arXiv:2406.10774 (2024).", "was_revised": true, "final_justification": "I get better understand standing of the experiments and insights from the paper. However, I agree with other reviwers' comment that the paper is too heavy to read. It would be better if the author can improve the writing and make it well organized.", "rebuttal": "We truly appreciate the reviewer's time and effort in thoroughly reviewing our manuscript. We have tried to address all their comments in the following. We are happy to answer any further questions that can motivate the reviewer to raise their ratings.\n\n---\n\n> _[..] full attention is already input-dependent sparse [..] only a matter of being complete zero or a small value._\n\nAs softmax is intended to be a regularized argmax (with the Shannon entropy regularization), softmax attention is already meant to be **almost** input-dependent sparse, with the difference being that the low values are not exactly zero. However, as our empirical evaluations indicate, this can lead to significant differences in behaviour between full softmax-attention and heavy-hitter top-$k$ sparse attention, and the theoretical analysis explains why this might be occurring, highlighting the effect of the semantic dispersion in full softmax-attention.\n\nSoftmax-attention will always put a strictly non-zero probability mass on each of the keys. This property of softmax has already been shown in literature as the source of **attention dispersion**, where the softmax output tends to approach a uniform distribution instead of looking like the one-hot distribution of an argmax (Velickovic et al., 2025; Nakanishi, 2025). Such behaviour of softmax can lead to **representational collapse** in transformers where all tokens end up with the same representation (Barbero et al., 2024). Thus, it seems that full softmax-attention can behave quite differently (and undesirably) in practice from its originally intended smoothed-argmax form.\n\n> _To be hardware efficient, there are many new input dependent sparse attention methods such as NSA[1], MoBA[2], SeerAttention[3], MInference[4], Quest[5], and etc that provide block-wise sparsity. The current sparse attention methods used in the paper are not state-of-the-art._\n>\n> - _[1] Yuan, Jingyang, et al. \"Native sparse attention: Hardware-aligned and natively trainable sparse attention.\"_\n> - _[2] Lu, Enzhe, et al. \"Moba: Mixture of block attention for long-context llms.\"_\n> - _[3] Gao, Yizhao, et al. \"Seerattention: Learning intrinsic sparse attention in your llms.\"_\n> - _[4] Jiang, Huiqiang, et al. \"Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention.\"_\n> - _[5] Tang, Jiaming, et al. \"Quest: Query-aware sparsity for efficient long-context llm inference.\"_\n\nWe thank the reviewer for these state-of-the-art input-dependent hardware-efficient blockwise-sparse attention schemes. We completely agree that top-$k$ attention is not necessarily the state-of-the-art input-dependent sparse attention scheme. We would like to clarify that our theoretical analysis is focused on \"heavy-hitter\" sparse attention, where the keys with the highest attention scores are the only ones that are not masked. There are many other ways to make the sparse attention input-dependent without necessarily being a heavy-hitter style sparse attention.\n\nWe consider top-$k$ attention in our experiments as it is one of the most basic forms of input-dependent sparse attention mechanisms. We believe that it is important to understand some phenomena at the simplest level first to see how consistent it is, and if there is some theory that can explain it. The above schemes [2-5] are designed for inference, and thus their effect on training convergence is unexplored; they are applied to pretrained models trained with full attention. Note that all schemes [1-5] utilize a \"top-k\" over block scores (each with their own scoring rule). Given that we are selecting blocks (in an input-dependent manner), there is no guarantee that the semantic dispersion within a block will be low. We hope to extend our theory to better understand input-dependent block-sparse schemes in future work.\n\nGupta et al. (2021) originally consider top-$k$ attention, and focus on the memory efficiency of top-$k$ attention without improving the runtime of the attention. However, various advancements have sped up the computational cost of top-$k$ attention both in terms of runtime and memory overhead such as Zeng et al. (2025). This has shown significant speedup over pytorch native attention implementation, as well as the optimized FlashAttention, with the speedup increasing with context length. So top-$k$ attention can now be significantly faster than standard full attention in terms of actual runtime. Furthermore, other heavy-hitter attention such as Reformer (Kitaev et al., 2020) can also speed up the actual attention runtime with locality sensitive hashing.\n\nOur focus has been on a complimentary angle where we study how sparsity affects optimization convergence, and try to provide theoretical intuition as to when and how can sparse attention converge faster than standard full attention. In situations where top-$k$ (or any other $k$-heavy-hitter attention) would converge faster than standard full attention, the overall runtime would improve cumulatively with faster sparse attention and fewer number of steps to convergence. If the sparse attention is also hardware-efficient, then that benefit will be in addition to the faster convergence --- each training step will be faster, and we will need fewer training steps.\n\n> _The paper does not provide results on large models_\n\nThe reviewer is accurate in their assessment here, and this is something we have already listed in our limitations. Our small scale experiments and theoretical analyses provide a strong motivation to attempt these at scale. However, we believe that it is important to understand some phenomena at a small scale to see how consistent it is, and if there is some theory that can explain it. In our thorough small-scale experiments, we see that the behaviours are quite consistent --- heavy-hitter sparse attention does consistently converge faster than full-attention. Furthermore, while we do not consider large models, we do study the effect of increasing the size of the models in terms of the number of transformer blocks, or the number of heads per block; see Figure 3 and Observation IV.\n\n> _[..] when training with sparse attention, [..] show similar or better results than full attention. However, the sparse attention will fail in the challenging downstream tasks such as Needle in a Haystack. Could the authors provide any explanation using the proposed theory on this issue?_\n\nIt is true that the input-agnostic sparse attention can fail on associative recall tasks such as Needle in a Haystack. This is expected as the input-agnostic sparsity pattern can easily miss the relevant associations. This is also related to the expressivity results of Yun et al. (2020) for sparse attention transformers. On the other hand, heavy-hitter input-dependent sparse attention can do quite well with associative recall tasks --- as long as the dot-product attention scores can appropriately model the associations, a heavy-hitter sparse attention can appropriately perform associative recall. As one example, Zeng et al. (2020) show that top-$k$ attention can match the associative recall performance of standard attention (see Figure 2(a) in Zeng et al. (2020)), and small values of $k$ is sufficient (see Figure 2(d) in Zeng et al. (2020)). However, we cannot make any general claims here since \"downstream performance\" does not only depend on the use of full or sparse attention, but also on what data/tasks the model is trained on.\n\n> _Since the paper discusse top-k or k-heavy-hitter attention. Could you provide any theorical or empirical guidence on how to set the k in practice? Or I can ask in a different way, how can we deterimine the sparsity before seeing the input?_\n\nAs with many architectural hyperparameters, we think that $k$ is a problem dependent hyperparameter. We are usually not able to select the number of transformer blocks, the number of attention heads, the MLP hidden-layer dimensions, and other architectural hyperparameters without \"seeing the input\", and we believe this to also be the case with $k$. Ideally, for a given architecture, we would like to make $k$ as small as possible while still retaining the necessary level of expressivity (that is, be able to achieve perfect training accuracy).\n\nAt a more high level, softmax is a smoothed version of argmax for the purposes of differentiability, and thus attention is actually crafted to look for the argmax, implying that $k \\sim O(1)$ should suffice. This is the motivation behind studying the theoretical properties of the argmax or hardmax attention transformers, where $k=1$ (Strobl et al., 2024).\n\nAs an alternate mechanism, there are heavy-hitter sparse attention transformers that can adapt the sparsity level to the input and the problem (Correia et al., 2019), but that scheme also requires the setting of various other problem-dependent hyperparameters. These are an interesting class of sparse attention transformers, and we hope to extend our theoretical framework to such adaptive sparse attention mechanisms in future work.\n\n-----------\n### References\n- Velickovic, P., et al. \"softmax is not enough (for sharp out-of-distribution)\". arXiv preprint arXiv:2410.01104 (2024). \n- Nakanishi, K. M. \"Scalable-softmax is superior for attention\". arXiv preprint arXiv:2501.19399 (2025).\n- Barbero, F., et al. \"Transformers need glasses! Information over-squashing in language tasks.\" NeurIPS (2024).\n- Gupta, A., et al. \"Memory-efficient Transformers via Top-$ k $ Attention.\" SustainNLP workshop (2021).\n- Zeng, Q., et al. \"Zeta: Leveraging z-order curves for efficient top-k attention.\" ICLR (2025).\n- Kitaev, N., et al. \"Reformer: The efficient transformer.\" ICLR (2020).\n- Yun C.,  et al. \"$O(n)$ connections are expressive enough: Universal approximability of sparse transformers.\" NeurIPS (2020).\n- Strobl L., et al. \"What formal languages can transformers express? A survey.\" TACL (2024).\n- Correia, G. M., et al. \"Adaptively sparse transformers.\" EMNLP-IJCNLP (2019)."}]

Author rebuttal:

Meta-review: This paper makes a strong case that input-dependent sparse attention (like top-k) helps transformers learn faster and generalize better than both standard full attention and simpler input-agnostic sparse methods. The authors back this up with solid, consistent experiments on smaller-scale tasks and a theoretical argument that this works by focusing the model's "semantic attention," which improves optimization stability. While the theory's bounds are somewhat loose and the experiments aren't large-scale, the core finding is robust and insightful. The reviewers' concerns about the loss landscape and scalability were adequately addressed in the rebuttal. Despite its limitations, the paper offers a valuable new perspective on sparsity, focusing on learning speed rather than just computational efficiency, which justifies acceptance.
Decision: Accept (poster)