Title: Transformers Learn Faster with Semantic Focus
Venue: NeurIPS 2025
Accepted: yes
Keywords: transformers

sparse attention

learning convergence

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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).\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors already mention their limitation in large scale experiments. I think another limitation is the sparse attention methods used being too old.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThis 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.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nSee above", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n*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).\n\nQUESTIONS:\nlease refer to the Weakness part in the previous section.\n\nLIMITATIONS:\nlease refer to the Weakness part in the previous section.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nThe reviewer would love to see the concerns in the above addressed and revise the rating.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}