Transformers Learn Faster with Semantic Focus

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.

---

[..] full attention is already input-dependent sparse [..] only a matter of being complete zero or a small value.

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

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

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.

• [1] Yuan, Jingyang, et al. "Native sparse attention: Hardware-aligned and natively trainable sparse attention."
• [2] Lu, Enzhe, et al. "Moba: Mixture of block attention for long-context llms."
• [3] Gao, Yizhao, et al. "Seerattention: Learning intrinsic sparse attention in your llms."
• [4] Jiang, Huiqiang, et al. "Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention."
• [5] Tang, Jiaming, et al. "Quest: Query-aware sparsity for efficient long-context llm inference."

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

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

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

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

The paper does not provide results on large models

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

[..] 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?

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

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?

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

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

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

---

References

• Velickovic, P., et al. "softmax is not enough (for sharp out-of-distribution)". arXiv preprint arXiv:2410.01104 (2024).
• Nakanishi, K. M. "Scalable-softmax is superior for attention". arXiv preprint arXiv:2501.19399 (2025).
• Barbero, F., et al. "Transformers need glasses! Information over-squashing in language tasks." NeurIPS (2024).
• Gupta, A., et al. "Memory-efficient Transformers via Top-$k$ Attention." SustainNLP workshop (2021).
• Zeng, Q., et al. "Zeta: Leveraging z-order curves for efficient top-k attention." ICLR (2025).
• Kitaev, N., et al. "Reformer: The efficient transformer." ICLR (2020).
• Yun C., et al. "$O(n)$ connections are expressive enough: Universal approximability of sparse transformers." NeurIPS (2020).
• Strobl L., et al. "What formal languages can transformers express? A survey." TACL (2024).
• Correia, G. M., et al. "Adaptively sparse transformers." EMNLP-IJCNLP (2019).

We thank the reviewer for their acknowledgement and their time. We respond to them in the following:

Thanks for the explanation. I understand the paper is trying to explain why sparse attention is good and converging fast. But at this stage, I can not fully believe that the preassumption that sparse attention is good.

We would like to clarify that our paper is not "trying to explain why sparse attention is good and converging fast" but rather trying to understand conditions under which sparse attention can train and generalize faster, and conditions under which they will not. Our results clearly show that not all sparse attention are the same.

• Input-agnostic sparse attention mechanisms do not necessarily train/generalize faster than full attention. Our empirical results demonstrate this, and our theoretical results explain why input-agnostic sparse might not show any improved performance.
• Input-dependent heavy-hitter attention schemes can train/generalize faster than full attention, as clearly demonstrated by our empirical evaluations across many variations of the learning setup. Our theory explains why this might be happening:
  • The transformer model performance (training convergence or generalization) can depend on the semantic dispersion (Definition 2) -- larger dispersions lead to worse performance.
  • Heavy-hitter input-dependent sparse attention schemes can significantly reduce the semantic dispersion, thereby improving performance.
  • Input-agnostic sparse attention schemes do not reduce the semantic dispersion, and thus would not generally give us faster convergence or generalization.

Even with the latest NSA, we have seen a certain amount of gap in challeging downstream NIAH test (e.g. muilti-key, or multi-value). That's something can not show in the training loss. I have been fooled once using sparse attention. The loss and ppl is so good but downstream task is dispointing.

We thank the reviewer for sharing their experience with "latest NSA" (we are assuming that NSA refers to "native sparse attention" that focuses on input-dependent block-sparse attention).

It would be great if the reviewer can share a published reference regarding this comment: "Even with the latest NSA, we have seen a certain amount of gap in challeging downstream NIAH test (e.g. muilti-key, or multi-value). That's something can not show in the training loss. I have been fooled once using sparse attention. The loss and ppl is so good but downstream task is dispointing."

It is important to understand the actual sparse attention mechanism to make any claims about whether we will really see any improvements. As we have already mentioned in our rebuttal, for native sparse attention mechanism focusing on block-sparse attention, there is again no guarantee that the semantic dispersion within a block will be low, and thus, there might not be any performance improvement over full attention. So our empirical findings and theoretical analysis actually aligns with the reviewer's experience.

So converging fast and lower PPL can be miss-leading and is not a convincing evidence from my side.

We have presented both training loss and held-out accuracy for all the problems and methods (see for example in Figures 1 and 2). We are not just claiming faster training convergence, but also highlight how the training convergence also translates to faster generalization (on unseen data).

Thus, I would raise my rating if the author can provide more results and evaluation on tasks like muilti-key retrievel or any other reasonable real tasks.

Low semantic dispersion, and thus improved performance, is usually guaranteed for heavy-hitter attention such as top-$k$ attention. As we discussed in our rebuttal, Zeng et al. (2025) have already demonstrated clearly that top-$k$ sparse attention performs well for needle in a haystack associative recall problems. Furthermore, the tasks we have considered such as ListOps (list operations) is a problem requiring (hierarchical) multi-key/value attention as it requires the retrieval of all the keys/values corresponding to operands within the enclosing brackets for the query corresponding to the operator to appropriately solve the problem. This multi-key/value retrieval in the long sequence needs to be done hierarchically through the whole input. And given the nature of the problem, the attention span -- the distance between the query token index and the required key token indices -- can be quite large (as already discussed in Tay et al., 2022). The input-dependent nature of top-$k$ allows it to model this desired behaviour.

Thus, we have already shown the necessary empirical evidence, and provided theoretical support for this empirical findings. Furthermore, existing work has already demonstrated the multi-key retrieval capabilities of top-$k$ sparse attention. Our contribution is not the top-$k$ attention, but understanding when and how heavy-hitter sparse attention outperforms standard full attention (and when sparse attention does not outperform full attention).

Finally, to reiterate, the goal of our study is to understand conditions under which sparse attention can be beneficial in terms of learning convergence and generalization. Our results, supported by our theory, show that not all sparse attention schemes would show benefits all the time, and there are precise conditions --- sparse attention mechanisms that ensure low semantic dispersion --- that lead to improved performance over standard full attention. Input-agnostic sparse attention mechanisms do not ensure such a condition (and this is also possibly why "native sparse attention" with input-dependent block-sparsity does not ensure improved performance). In contrast, input-dependent heavy-hitter attention can ensure low semantic dispersion, and our empirical evaluations demonstrate improved performance over standard full attention.

---

References

• Zeng, Q. et al. "Zeta: Leveraging z-order curves for efficient top-k attention." ICLR (2025).
• Tay Y, et al. "Long range arena: A benchmark for efficient transformers". ICLR (2021).
• Nangia, N. and Bowman, S. Listops: A diagnostic dataset for latent tree learning. NAACL: Student Research Workshop (2018).

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.

---

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

[1] Tay Y, Dehghani M, Bahri D, et al. Efficient transformers: A survey[J]. ACM Computing Surveys, 2022, 55(6): 1-28.

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

Is it possible for the authors to compare these techniques in synthetic data experiments? I am curious about the performance of these techniques.

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

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

While 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:

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

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

A minor suggestion is that the authors could consider some real data experiments.

We 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:

We also considered a larger model 16 single-headed transformer blocks for the same experiment for 50 epochs with SGD:

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

How the choice of $k$ affects performance, generalization, or the theoretical bounds?

The choice of $k$ in sparse attention affects the overall behaviour in multiple ways:

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

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

---

References

• Kitaev, Nikita, et al. "Reformer: The efficient transformer." ICLR (2020).
• Beltagy, Iz, et al. "Longformer: The long-document transformer." arXiv preprint arXiv:2004.05150 (2020).
• Choromanski, Krzysztof Marcin, et al. "Rethinking Attention with Performers." ICLR (2020).
• Rahimi, Ali, and Benjamin Recht. "Random features for large-scale kernel machines." NeurIPS (2007).
• Yun, Chulhee, et al. "$O(n)$ connections are expressive enough: Universal approximability of sparse transformers." NeurIPS (2020).

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.

---

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.

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

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

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

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

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.

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

Furthermore, 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.

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.

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

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

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

Finally, 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:

With a 16-block transformer model for 50 SGD epochs:

Weakness (4). In Figure 4 in [1], there is no clear difference that sparse attention learns faster. Could the authors provide some explanation?

[1] Gupta, Ankit, et al. "Memory-efficient Transformers via Top-$k$ Attention." arXiv preprint arXiv:2106.06899 (2021).

Regarding 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:

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

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

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

---

References

• Gupta, Ankit, et al. "Memory-efficient Transformers via Top-$k$ Attention." arXiv preprint arXiv:2106.06899 (2021).

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.

---

Overall, the presentation is too dense which makes it hard to follow, and I strongly recommend the authors improve readability.

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

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.

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

Furthermore, we will add a longer version of the following table to clarify notation:

Could the authors comment on the maximum observed values of $\beta$ across the experiments?

In 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]$.

However, 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$.

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

We 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):

The 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):

As 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$.

How does the per-step computational cost of calculating the top-k mask compare to the baselines?

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

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

---

References

• Gupta, Ankit, et al. "Memory-efficient Transformers via Top-$k$ Attention." SustainNLP workshop (2021).
• Zeng, Qiuhao, et al. "Zeta: Leveraging z-order curves for efficient top-k attention." ICLR (2025).
• Kitaev, Nikita, et al. "Reformer: The efficient transformer." ICLR (2020).