Long-Context Generalization with Sparse Attention
Authors: Pavlo Vasylenko
Hugo Pitorro
Andre Martins
Marcos Vinicius Treviso
Venue: ICLR 2026 Poster (2026)
Keywords: long-context
sparse attention
length generalisation
Primary area: generative models

Abstract: Transformer-based architectures traditionally employ softmax to compute attention weights, which produces dense distributions over all tokens in a sequence. 
    While effective in many settings, this density has been shown to be detrimental for tasks that demand precise focus on fixed-size patterns: as sequence length increases, non-informative tokens accumulate  attention probability mass, leading to dispersion and representational collapse.
    We show in this paper that dynamically sparse attention mechanisms using $\alpha$-entmax can avoid these issues, due to their ability to assign exact zeros to irrelevant tokens. Furthermore, we introduce Adaptive-Scalable Entmax (ASEntmax), which endows $\alpha$-entmax with a learnable temperature parameter, allowing the attention distribution to interpolate between sparse (pattern-focused) and dense (softmax-like) regimes.
    Our empirical evaluation on synthetic tasks and language modeling demonstrates that ASEntmax substantially outperforms softmax, scalable softmax, and fixed-temperature $\alpha$-entmax baselines, achieving up to 1000$\times$ length extrapolation on synthetic benchmarks and superior long-context generalization on language modeling while preserving short-context performance, including better perplexity trends and higher retrieval accuracies at 8$\times$ training length.

Published as a conference paper at ICLR 2026
LONG-CONTEXT GENERALIZATION WITH SPARSE ATTENTION
Pavlo Vasylenko 12, Hugo Pitorro', André F. T. Martins MMrross Treviso Instituto Superior Téenico, Universidade de Lisboa, "Instituto de Telecomunicações. PrransPerfect, "ELLIS Unit Lisbon
ABSTRACT
Transformer-batt architectures traditionally employ softmax to compute attention weights, which produces dense distributions over all tokens in sequence. While effective in many settings, this density has been shown be detrimental for tasks that demand precise focus on fixed-size patterns: sequence length increases, non- informative tokens accumulate attention probability mass, leading dispersion and representational collapse. We show in this paper that dynamically sparse attention mechanisms using d-entmax can avoid these issues, due to their ability to assign exact zeros to irrelevant tokens. Furthermore, Entmax (ASEntmax), which endows a-entmax with learnable temperature parameter, allowing the attention distribution interpolate between sparse (pattern-focused and dense (softmax-like) regimes. Our empirical evaluation on synthetic tasks and language modeling demonstrates that ASEntmax substantially outperforms softmax, scalable softmax, and fixed-temperature a-entmax baselines, achieving up 1000 length extrapolation on synthetic benchmarks and superior oong-conextt generalization on language modeling while preserving short-context performance, including better perplexity trends and higher retrieval accuracies 8x training length. Source code: https: / thub. com/ henp/deep-oo spwn/asentant asentmaxx
introduce Adaptive-Scalatee Scalable
INTRODUCTION
The transformer architecture (Vaswani et al., 2017) has become the foundation of modern large language models (LLMs), establishing new benchmarks across diverse domains. However, as re- searchers push these models toward increasingly longer contexts-from thousands to millions of -seeerl fundamental limitations emerge that can be traced to the softmax transformation used in attention. Three critical limitations stand out: representational collapse occurs due to soft- max's 's inability to maintain distinct attention patterns as sequence length grows, erasing meaningful distinctions between tokens (Barbero et al., 2024); over-squashing is exacerbated by softmax's dense probability distribution, leading to exponential dilution of gradients (Alon & Yahav, 2021; Barbero al., 2024); and attention dispersion arises from softmax's fundamental property that forces prob- ability mass to be distributed across all tokens, with attention weights necessarily approaching an uniform distribution context grows (Veličković et al., 2025; Nakanishi, 2025).
Previous approaches to address these challenges include positional encoding innovations such as ALiBi (Press al., 2022) and RoPE (Su et al., 2024), which help to mitigate position bias issues. Recent works directly target the root cause- the softmax function itself. Nakanishi (2025) proposes Scalable-Softmar to scale logits based on context length, while Veličković et al. (2025) identify fundamental limitations of softmax for sharp out-o--distribution generalization and propose to learn adaptive temperatures to control the sharpness of softmax. While effective, these solutions often require careful tuning address only subset of the challenges.
In this paper, we address the root cause of these problems by replacing softmax with c-entmax (Peters et al., 2019), differentiable sparse transformation that induces probability distributions where irrelevant tokens receive exactly zero attention. While a-entmax has been used successfully in transformers (Correia et al., 2019; Gonçalves et al., 2025), its length generalization properties, the

---

Published as a conference paper at ICLR 2026
Exant March Accuracy
100%
100%
80%
Diarrbuton
Softmax ROPE
Sattmax NAPE
Aczunny
80%
In-Detriuttio
Adipt. Temp.
Top-K (K-2)
60%
5Smax NAPE
60%
5.Maa
Emmax -11
20%
ASEntmax: NAPE
20%
ASEnimes
32K 65K
256 Sequence Length (n)
4K
Sequence Length (n)
Figure Long-context generalization on Multi-query Multi-token Associative Recall (left) and Max Retrieval (right). SSMax represents the Scalable Softmax approach by Nakanishi (2025). and Adaptive Temperature (Adapt. Temp.) represents the approach by Veličković et al. (2025). While all methods benefit from using NAPE (NoPE ALiBi), adaptive-scaling version of a-entmax exhibits the best extrapolation results, effectively handling extremely long sequences.
best of our knowledge, have never been studied. We show theoretically and empirically that o-entmax consistently helps to address challenges in long context modeling. Our key contributions include:"
Non-dispersion: We establish that o-entmax attention distributions maintain consistent focus regardless of sequence length, with entropy bounded by Olog 8) rather than approaching maximum entropy Olog n) with softmax, where < is the number of tokens with nonzero probability.
Representational preservation: a-entmax attention with sparse support can avoid represen- tational collapse and reduces the number of gradient paths from O( to O(s²), alleviating over-squashing by strengthening the gradient flow for long-range dependencies.
Adaptive-scalable a-entmax: We introduce ASEntmax, which adaptively adjust sparsity based sequence length to maintain optimal token selection even in extremely long contexts.
Empirical results: We demonstrate that ASEntmax achieves superior performance across synthetic and real-world tasks. For example, as shown Figure I, ASEntmax achieves 95.3% accuracy on associative recall at 65K tokens after training on just 64 tokens a 1000 length extrapolation.
2 BACKGROUND
2.1 TRANSFORMERS
In this work, we study (causal) transformers with sparse attention distributions created by replacing softmax with a-entmax. We present the precise mathematical formulation below, following closely the notation from (Barbero et al., 2024). Concretely, given a sequence of token embeddings X B where is the sequence length and d is the hidden dimension. transformers compute query, key, and value projections Q = XWq. K = X WW.. and V = XWy We denote with qi k; v, E RRd the d-dimensional query, key, and value vectors of the i-th token. For each query position the representation at layer for the i-th token is computed as:
Σ Piy norm (6) iSi
(o)--1)
+ +u ((-1)
0.6 = FFN(") (norm, (E) uu:
(u(²))
(6) + +u
u
(1)
where Pij are attention weights, (FN is the feed-forward network, norm(.) represent LayerNorm
(8)
0 (L) The attention
modules (Xiong al., 2020). The output is computed as y. norm3
weights P.j (2(
(
computed by applying transformation : A, to the attention
logits 2/ (q. (1) ,.. (())) //V, where 4n {p R p > 0. 1} represents the
probability simplex. Standard transformers employ the softmax function as म. this work, we study transformers by casting the o-entmax transformation.
2

---

Published as a conference paper at ICLR 2026
Softmax 1.0)
Entmax d-1.5
Entmax a-1.66
Entmax a=
Probabiity
0.75
0.25
0.00
2.5
0.0
0.0 5.0 0.0
100
Figure Visualization of a=-entmax(=/0) for different values of a. Each panel shows how probability mass is distributed among five elements of z = [2.0, 1.8, 1.6, 1.4, 1.2] the temperature parameter decreases (8 increases). The vertical lines show the temperature that leads to zero probability.
2.2 G-ENTMAX
a-entmax (Peters et al., 2019) is differentiable transformation that generalizes softmax by allowing for sparse probability distributions. For an input vector RR" and a > a-entmax is defined as:
a-cntmax(z): = [(a 1)2, - -())1=
(2)
where [+] == max(0,,) -) and T R + yields a threshold that ensures the resulting distribution sums A key property a-entmax is that tokens with scores below the threshold receive exactly zero probability, creating sparse attention patterns. When + 1 this reduces to the standard softmax function. The sparsity level increases with a, with a corresponding to the sparsemax function (Martins & Astudillo, 2016). Figure 2 illustrates a-entmax for different values of a. We provide more information on a-entmax ŞA. While a-entmax is suitable choice for sparse attention, its theoretical and empirical impact on long inputs is still unclear. In the next section, we demonstrate how it fundamentally changes the way attention behaves for long contexts.
3 THEORETICAL PROPERTIES OF Q-ENTMAX FOR LONG CONTEXTS
We analyze the theoretical properties of a-entmax that make it especially suitable for long-context modeling, focusing how it addresses the fundamental limitations of softmax.
3.1 NON- VANISHING ATTENTION PROBABILITIES
A critical limitation of softmax in transformers is that attention weights inevitably decrease as the sequence length increases. Our first result demonstrates how a-entmax avoids this issue.
Lemma 1 (Non-Vane Vanishing Attention Property). Consider scalars 01, an-1,C E R. Let x (0t, ..y an c] T R" and xx [ai...., an-1, b, R=+11 with all entries bounded. The following properties hold:
For all a 2 1, we have a-entmax(x)), 2 a-entmax(x" (n=11 In the softmax case (a 1), Barbero al. (2024, Lemma B. have shown that the inequality is always strict: softmax(x),, softmax( x )n+1.
For all a > 1, there is some bmax E R such that, for any b S bmaxo we have a-entmax(x)), a-entmax(a"" m+11
Furthermore, for 1, the difference a-entmax(x)), d--entmxxxx"))mm can take any value in [o, a-entmaxx(x))) by appropriate choice of
The proof can be found in 1. This result demonstrates a fundamental difference between softmax and o-entmax. Unlike softmax, where adding new token always reduces the attention probability of existing tokens strictly, q-entmax allows distinct behavior: the attention probability can remain unchanged. This occurs because the a-entmax's thresholding effect allows tokens with logits below certain threshold to receive exactly zero attention, letting the model focus only on the relevant tokens.
Our code will be available upon acceptance.

---

Published as a conference paper at ICLR 2026
Having established that a-entmax prevents the vanishing of individual attention weights, formalize the broader concept of attention dispersion to better understand how attention distributions whole behave as the sequence length increases.
3.2 ATTENTION DISPERSION AND CONCENTRATION
Recent work by Nakanishi (2025) and Veličković et al. (2025) has highlighted attention dispersion as fundamental limitation of softmax for long context generalization. Building upon these insights, we provide a formal definition to characterize attention dispersion and show how a-entmax naturally exhibits concentration properties that address these limitations.
Definition 1 (Attention Dispersion). Let f RR" + An denote transformation (such as softmax) mapping logits to the probability simplex An == R p,
1. I exhibits complete dispersion for any bounded sequence of logits (Zn)neN, the normalized entropy approaches I as the sequence length increases:
H(f(zin)) 1.
lim
(3)
log n
2. f exhibits concentration resilience there are bounded sequences of logits where the normalized entropy remains bounded away from I:
H(f(zin))
lim n n0000
< 1.
(4)
log n
These definitions allow us to examine how softmax and a-entmax behave as sequence length grows:
Proposition (Dispersion Properties of Attention Mechanisms). Comparing softmax and c-enimax (a > 1) attention mechanisms:
1. o-enimax can retain probability, while softmax always leaks: For any a > 1 and any logits R", there are logits z² E RR- with N > n such that:
a-entmax(2); = a-entmax(x*") Vi S
(5)
This is impossible for a (softmaxx, for which we always have softmax(z); softmax (2");
2. Softmax exhibits complete dispersion: For any fixed temperature 0 and any bounded sequence of logits (2n) (Zn)ne::
[[ssotmax(21.m/90
lim 1 1-0
(6)
log nl
3. a-entmax exhibit strong concentration resilience: When the support size grows sublinearly as |S] = Onn with 3 < 1, a-entmax maintains bounded normalized entropy:
((a-entmax((1sm))
lim R-+00
<8<1
(7)
logn
The full proof can be consulted in SC. The key takeaway of this result that the entropy of attention distributions reveals how concentrated or dispersed they are across tokens. While softmax distributions approach maximum entropy O n) as the sequence length increases (indicating complete dispersion), a-entmax distributions maintain bounded entropy O(log s) where is the support size. This allows models with a-entmax to maintain focused, low-entropy attention patterns even when processing extremely long sequences, long as the support size is smaller than the full sequence length. This non-dispen non-dispersion property means transformers with a-entmax can scale to very long contexts without the attention becoming dispersed, maintaining their ability to focus on relevant information regardless of how much additional context is present. However, attention dispersion is not the only obstacle effective long-sequence modeling.
4

---

Published as a conference paper at ICLR 2026
3.3 REPRESENTATIONAL PRESERVATION AND OVER- SQUASHING ALLEVIATION
Two other critical challenges in long-context transformers are representational collapse and over- squashing, both exacerbated by the diffuseness of softmax attention (Barbero al., 2024). More concretely, collapse means ||0: (2) / 111 +0 0 as n + oc, and over-squashing means gradients from distant inputs vanish across (n²²) paths. Here, we show that o-entmax can mitigate both:
Proposition 2 (Representational Preservation and Reduced Gradient Paths), Let > 1. Consider depth-L. I, transformer with residual connections and attention weights given by a-enimax. Suppose each attention distribution has support size at most 8 with & n. Then:
1. (Preserved representations) There exist input families b0)) EBax and v+ g and
constant e > 0 such that | [[( [a 2 for all n; i.e. o-entmax can maintain distinct token representations as + 00. In contrast, Barbero et al. (2024) shows that for softmax attention, || v (L) wt 0 as n oc.
(Alleviated over-squashing) The number of effective gradient paths scales as Oss (rather than Onn)) under softmax), alleviating over-squashing by concentrating gradient flow on fewer active paths.
Full statements and proofs are in App. $D.2-$D.3.
Empirical evidence. We corroborate Proposition 2 with controlled experiments. For representa- tional collapse, following the Barbero et al. (2024)-style one-token extension probe, we implement
(1.0, 1.5, 1.75, 2.0}. and report the L, gap between vi and u (L))
counterexample with With softmax attention, the gap rapidly decays toward with length, while a-entmax preserves a non-vanishing margin up to 128K tokens, with larger a yielding stronger preservation (App. $D.2). For over-squashing, we analyzed gradient flow through an 8-layer network on copying task, where the model must copy information across long distances. We find that a-entmax attention sustains substantially larger norms across depths and lengths, consistent with O(x4) path growth, whereas softmax degrades sharply with sequence length (App. $D.3).
ADAPTIVE-SCALABLE Q-ENTMAX (ASENTMAX)
In the previous section saw that a-entmax, for any choice of > 1, can avoid some of the pitfalls of softmax thanks to its ability to assign zero weight to many tokens, ignoring irrelevant information. But what if ignores too many tokens? Can it handle situations where many tokens are relevant and should be attended? We show in this section that indeed the model might not be able to cope with this for fixed a and a fixed temperature, and we propose practical solution.
4.1 CONTROLLING SPARSITY IN LONG CONTEXTS VIA ASENTMAX
As sequence length grows, the spread of attention logits increases- IID Gaussian logits, the expected range satisfies E[A] 2o 20 2 log n (Kamath, 2015). With fixed temperature, this makes attention overly peaky at long n. We address this with Adaptive-Scalabb a-entmax (ASEntmax). which rescales logits per head as function of context length and content:
ASEntmax(z) = a-entmaxx(6 (l)gg)))))))
(8)
where Y 8 E IR are head-specific scalars (inverse temperature). Concretely, for each head, we obtain vectors B and y whose entries contain these coefficients for each query:
B softplus X w E R", Y tanh Xw,) E (-s,s)",
(9)
where W3, W. E d are learnable, head-specific projection vectors. This characterization allows the model to learn slowly rising (y > 0) or dampening (y < 0) temperature schedule without interfering in the positional encodings (which would happen with negative values of (). Specifically, for IID Gaussian logits (00, r), when 8 = 0 and 7 -0.5 the scaling counteracts the growth of logit ranges (A,)) as context increases:
B(lggn) -0.5 A., 33log n) -0.5 2a / 2og 20.
(10)
5

---

Published as a conference paper at ICLR 2026
2, Head 11
Layer Head
Layer 6. Head 10
Head
Layer 7. Head
Ai (010-123)
015
5- Blogrt PP
Zogn (INE 20
0.0
5- /1bps//PPPE0000
plae youd 80
5 +=30ggrr geodsts
0.00
8.6
6
9.65
0.75
0.79
2.0
6.00
2+ dlog: (200 3]
0.15
0.00
140
2+ #bgeith0,00
202
V00
200 100
A00 182 200
1000
Figure 3: Learned positions per head. Besides a simple linear fit baseline (3n), we also show the fit given by 8 log and 6+ (logn)), which are used by SSMax and ASEntmax, respectively. These plots further reinforce the idea that having a parameter is beneficial for length extrapolation.
which remains constant as increases, preventing excessive sparsification. Furthermore, with this parameterization, ASEntmax can recover standard o-entmax when 3 = 0, hence allowing smooth transition between scaled and unscaled regimes. By making wg and wy -ped-speciiic and learnable, ASEntmax can adapt to the optimal scaling behavior for each head, balancing the natural concentration benefits of a-entmax with precise control over how sparsity patterns evolve with sequence length. Finally, we note that simply scaling the query-key products is appealing from a practical perspective since it allows the direct use of fast optimized kernels for a-entmax, such as AdaSplash (Gonçalves et al., 2025), without any modifications.
Empirical Analysis. To empirically validate the importance of the parameter 7 in our proposed scaling formulation, conducted experiments language modeling task using 1 120M 200-mmmree transformer trained on 5B tokens from the Fin dataset (Penedo et al., 2024). Following the methodology from SSMax (Nakanishi, 2025), we implemented learnable scaling parameters for the attention logits, but with key difference: while Nakanishi (2025) uses a global scaling parameter, we learn separate scaling parameters for each attention head, motivated by Correia et al. (2019)'s s finding that attention heads develop distinct sparsity patterns. Figure 3 presents the learned scaling behaviors for representative attention heads, along with fitted curves from different scaling models. First, we note that different heads learn significuntly distinct patterns, highlighting the need of head-specific scaling. Second, the results demonstrate that a simple log-scaling model 6 + ,3 log n provides poor fits for many heads. In contrast, the inclusion of y power provides consistently better fits across different attention heads. The complete distribution of fitted 3 and 7 values across all heads is provided in $F, and additional training details can be found in $F.1.
4.2 INTERACTION WITH POSITIONAL ENCODING
Sparse transformations do not merely reweight attention, they effectively change which edges exist in the attention graph. This makes the choice of positional encoding especially important when replacing softmax with a-entmax. Recall that a head computes logits of the form 2,ح = q, k2 + bij, where bij is the relative positional contribution (possibly zero). Under o-entmax, token j is attended by query iff it survives the query-dependent thresholding (a - 1) zi > T(Zi). Thus, positional encodings acting on Zij through bij directly affect the set of active tokens.
With NoPE (Kazemnejad et al., 2023) we have bij = 0, so (apart from causal masking) attention scores are purely -based This can be beneficial for semantic retrieval, but may be less robust under large length shifts without an explicit positional prior. With ALiBi (Press et al., 2022) we add head-specific linear distance bias bi mj 2) with my > O, inducing smooth recency preference under softmax. Under a-entmax, thresholding yields exact zeros: any position whose biased score falls below the entmax threshold is pruned. Since the ALiBi penalty decreases with distance, for finite content logits there is an effective head- and query- deeenddeeendnnt cutoff beyond which sufficiently distant tokens receive exactly zero mass, yielding a localized support. With RoPE (Su et al., 2024), relative position enters through rotations in query-key interactions, producing oscillatory -eaunce-eeeedeett score contributions; combined with d-entmax, this can lead to content-depent depended sparse supports whose active distances vary with the representations and RoPE frequencies, and the resulting sparsity can be more fragmented. We visualize this interaction in Fig. 4, and provide additional theoretical statements and empirical evaluations in $E.

---

Published as a conference paper at ICLR 2026
NoPE
ALIBI
RoPE
Altenton Weight
0.15
Altention We Weight
Athentoon Weight
0.4
Zaro weights
0.10
Dead
0.2
02
0.00
16
25
30
10 16
26
80
10
20
26
Distance from Query
Distance from Query
Distance from Query
Figure 4: Example of attention weight profiles for different positional encodings with a-entmax with a = 1.5. NoPE induces content-driven sparsity. ALiBi induces attention windows with a clear cutoff.
RoPE promotes Requency-depennde patterns with potential periodic dead zones.
Robust locality + sparse long-range retrieval. With that mind, we propose NAPE, where half of the heads induce recency bias via ALiBi, while the other half use NoPE and thus are more content-driven. This division is particularly effective under length extrapolation: ALiBi heads maintain consistent local receptive fields, while NoPE heads allows the model to retrieve distant key-value evidence when needed. Finally, ASEntmax complements NAPE by adapting logit scaling context grows, preventing the sparsity induced by the o-entmax transformation from becoming overly aggressive or overly diffuse at extreme lengths.
5 EXPERIMENTS
5.1 SYNTHETICC TASKS
A number of works have turned to synthetic tasks probing ground for transformers' length- generalization capabilities (Anil et al., 2022; Dziri et al., 2023; Zhou et al., 2024). Such tasks, like copying sequence and sorting numbers, allow precise control over training and test lengths, revealing whether a model has truly learned an algorithm that scales or merely memorized patterns within limited length. Vanilla transformers struggle this setting: they often achieve perfect accuracy on sequences up to the training length, yet fail catastrophically on even slightly longer sequences (Press al., 2022). To quantitatively evaluate our proposed improvements, we embrace this paradigm of synthetic tasks for long-sequencs testing. Concretely, evaluate our models diverse set of synthetic tasks designed to test different aspects of long-context modeling, covering both aasttion-agnottic reasoning (where token positions are not critical) and position-sensitive operations (where relative or absolute positions matter):
Petrieval-focusee tasks: These include Max Retrieval (Barbero et al., 2024), which requires identi- fying maximum values in sequences, and Multi- Multi-token Associative Recall (MQMTAR)- a variant of that proposed by Arora et al. (2024), but with multi-token keys and values values-which involves matching queries to their corresponding key-value pairs. Both tasks test the model's ability to maintain focus on relevant tokens regardless of their positions in long contexts.
Memory-dependent tasks: We evaluate models on Copy (reproducing input sequences). assesses how well the model preserves token representations and accesses specific positional information throughout the network. On this line, we also evaluate on 2Back, described in $G.
Ordering tasks: This category contains tasks such Sort (arranging tokens in ascending order) and Reverse (outputting tokens in reverse order). These evaluate compositional generalization and positional reasoning, becoming increasingly challenging as sequence length grows.
Experimental Setup. We use small decoder-only transformers and keep the number of layers as low as possible so that the results reflect the attention method's ' capabilities rather than model scale. The only exception is Reverse, where we raise until a plain softmax baseline reaches at least 1.5x the in-distribution performance. For positional information we default to NAPE (NoPE+ALiBi) where half the heads have no positional encoding (NoPE) and the other half use ALiBi with linear slopes.? 2 We treat NAPE as a practical default, not a contribution, and report RoPE and standalone
Hardd-ALiB (Jelassi et al., 2024) allows zero slopes some heads, which is equivalent to NoPE; our NAPE default mirrors this practical configuration.
7

---

Published as a conference paper at ICLR 2026
Table 1: Exact match accuracy (%) on representative tasks. For each task, we report in in-distribution sequence length n in the first column (n 64 for all tasks), followed by OOD results at increasing sequence lengths. L indicates the number of layers. Best results are bold.
MQMTAR (I.
Reverse (I.
Method
ID
2x
4x
16x
64
256
1024x
ID
1.5x
2x
4x
8x
Softmax
100.0
100.0
100.0
99.5
97.8
80.2
3.0
100.0
36.0
0.0
0.0
0.0
Top-K (K=32)
100.0
99.9
6.2
0.0
0.0
0.0
0.0
100.0
100.0
98.7
57.0
0.0
SSMax
99.9
100.0
99.9
99.6
98.3
26.7
100.0
54.6
0.0
0.0
Entmax
100.0
100.0
100.0
99.2
92.7
66.8
9.3
100.0
99.0
86.0
28.5
0.2
ASEntmax
100.0
100.0
100.0
99.7
99.6
99.0
95.3
100.0
100.0
99.8
96.4
56.7
Copy (L = 2)
Sort (L
2)
ID
2x
4x
8x
16x
32x
64x
ID
2x
4x
8x
Softmax
100.0
100.0
99.9
99.9
99.4
96.1
85.5
100.0
0.0
Top-K (K=32)
100.0
99.7
96.8
26.7
0.0
0.0
0.0
100.0
92.5
0.0
0.0
SSMax
100.0
100.0
100.0
99.9
99.6
99.3
95.8
100.0
0.0
0.0
0.0
Entmax
100.0
99.0
86.0
28.5
0.2
0.0
0.0
100.0
99.3
57.8
0.0
ASEntmax
100.0
100.0
99.9
99.7
99.4
96.3
100.0
100.0
79.7
0.0
ALiBi in App. $H for completeness. We employ 1.5 for both Entmax and ASEntmax models, and = for SSMax and ASEntmax. Further hyperparameters and ablations appear in App. SG.
Discussion. The results, shown in Table 1, reveal a critical factor for length generalization: attention sparsity. Specifically, ASEntmax dramatically outperforms others at extreme cengths-maintainn 96.4% accuracy at 256x test length on MQMTAR (vs. 80.2% for softmax) and 96.4% at 4x on Reverse (vs. 0% for softmax). Moreover, the consistent superiority of ASEntmax over basic a- entmax confirms the benefits of adaptive scaling, particularly at extreme lengths where fixed-a may become too sparse or too diffuse. SSMax performs well on the Copy task, even outperforming other methods, but struggles indicates that while scaling logits helps maintain peak attention magnitude, the explicit sparsity of a-entmax provides additional benefits by completely removing irrelevant connections. These findings are further supported by results on the Max Retricval task (Figure right), where sparse attention mechanisms demonstrate superior length extrapolation compared to dense approaches, with ASEntmax maintaining over 60% accuracy even at 4096-length sequences- dramatic improvement over standard Softmax and Adaptive Temperature (Veličković et al., 2025). Finally, Copy and Reverse show moderate generalization (up to 64x and 8x respectively), while Sort fails beyond 4x length for all methods. This pattern suggests that tasks requiring precise global ordering (Sort, Reverse) inherently more challenging for length generalization than tasks dependent on local or independent token properties. We provide per-task results, including results for other positional encoding methods such as RoPF, ALiBi, and NoPE, in SH.
more complex tasks like MQMTAR and Reverse at extreme lengths. This
5.2 LANGUAGE MODELING
To validate our approach real-world tasks, train 420M-paramees decoder-only models follow- ing the LLaMA 3 architecture (details in App. G.3) on the high-quality DCLM-Edu dataset (Allal al., 2025), for 7B tokens with context length of 2048. As in the synthetic experiments, we report results in the main paper using NAPE (NoPE+ALiBi). Results with RoPE are reported in App. H.9. We evaluate short-context performance on Lambada, HellaSwag, PIQA, ARC-C, Winogrande, and OpenBookQA, and assess long-context generalization via perplexity on ArXiv/PubMed subsets of The Pile (Gao et al., 2020) as well as RULER's Needle-in-a-Haystann tasks (Hsieh al., 2024).
Discussion. In short-context evaluation (Table 2), models perform comparably. We note that per- formance ARC-C, Winogrande, and OpenBookQA is near-random. Among the other benchmarks,
Across all attention methods tested-includde softmax-Ne NAPE consistently performs best. We study its interaction with o-entmax and positional encodings in $E, with further empirical insights in $E.4 and $H.9.
8

---

Published as a conference paper at ICLR 2026
Table 2: Downstream-task results on short-context datasets. Best results are in bold.
Method
Lambada (PPL)
Lambada
Hellaswag
PIQA
Arc-C
WinoGrande OpenbookQA
Softmax
52.4
30.9
33.1
65.1
25.6
49.5
28.2
SSMax
48.9
31.6
32.9
65.1
25.0
51.5
30.4
Entmax
47.9
32.1
32.8
24.6
50.9
29.0
ASEntmax
41.6
34.3
33.4
63.8
26.0
50.0
28.6
Table 3: Long-context perplexity on ArXiv and PubMed from The Pile. Best results are in bold.
ArXiv (ID)
ArXiv (OOD)
PubMed (ID)
PubMed (OOD)
Model
1K
2K
4K
8K
16K
1K
2K
4K
8K
16K
Softmax
21.63
17.10
13.87
12.46
12.71
19.54
18.05
15.59
15.31
18.23
SSMax
21.57
17.02
13.74
12.29
12.31
19.35
17.84
15.14
13.75
14.72
Entmax
21.36
16.85
13.36
11.04
10.07
19.10
17.64
14.79
12.86
13.02
ASEntmax
21.30
16.86
13.31
10.89
10.01
19.04
17.69
14.76
12.61
12.90
ASEntmax achieves the highest scores on Lambada (both in terms of perplexity and accuracy) and HellaSwag, while Softmax and SSMax perform best on PIQA. For long-contex: modeling (Table 3), ASEntmax outperforms all other models. On ArXiv, it shows strong extrapolation, maintaining decreasing perplexity trend even when extended to 8 the pre-trained sequence length. On PubMed, although all models struggle with 8x extrapolation, ASEntmax still leads by margin of about 1 perplexity point vs SSMax. In retrieval tasks from RULER (Table 4), nalmax-basee models perform better overall, with both Entmax and ASEntmax surpassing their softmax counterparts at the extrapo- lation 4x and 8 x. ASEntmax in particular shows strong extrapolation in passkey retrieval on the simple haystack task (S-NIAH-I), maintaining near-perfect performance even at 8x context length. On the harder S-NIAH-2 variant, despite performance dropping substantially across all models, ASEntmax remains the best performer at 2 and 4x length generalization.
RELATED WORKS
Attention Dispersion. Recent work has identified attention dispersion as fundamental limitation in softmax-based transformers (Dong et al., 2021; Zhai et al., 2023; Veličković er al., 2025). For example, Veličković et al. (2025) demonstrate that softmax attention inevitably disperses focus as sequence length increases, while Nakanishi (2025) propose SSMax to scale attention logits based on sequence length. Our approach employs a-entmax (Peters et al., 2019). which naturally produces sparse distributions by assigning exactly zero probability to irrelevant tokens. We provide theoretical guarantees that o-entmax maintains bounded normalized entropy as sequence length increases a property softmax fundamentally lacks. Our ASEntmax further improves a-entmax with learnable, context-dependent scaling, leading to consistent gains over SSMax across diverse tasks.
Representational Collapse and Over-Squashing, Studies analyzing attention patterns in neural networks have noted that increasing depth and context length can induce representational degenera- tion (Dong et al., 2021; Noci et al., 2022; Arroyo et al., 2025). In particular, Barbero et al. (2024) prove that with softmax attention, token representations become indistinguishable as sequence length increases and gradient paths grow as Onn)), causing exponential signal dilution. Our analysis shows, theoretically and empirically, that a-entmax can address both limitations by maintaining distinct token representations and by reducing gradient paths ((s") for increasing sequence lengths n.
Positional Encodings. The design of positional encodings plays central role in enabling trans- formers to generalize to long contexts. ALiBi (Press et al., 2022) introduced linear attention biases with fixed slopes that encourage recency, and has since inspired several extensions that either mitigate softmax-reland issues (e.g., Hard-ALiBi (Jelassi et al., 2024)) or induce various learned decay patterns (KERPLE (Chi et al., 2022), FIRE (Li et al., 2024)). More recently, Stick-Breakin Attention (Tan et al., 2025) can be interpreted as a dynamic variant of ALiBi, in which the slopes are input- dependent and calculated adaptively from spans of tokens. Our work incorporates NAPE (NoPE

---

Published as a conference paper at ICLR 2026
Table Retrieval performance on RULER benchmark. All models use NAPE positional encoding and were trained on 2048-token contexts. Best results are in bold.
S-NIAH-I (ID)
S-NIAH-I (OOD)
S-NIAH-2 (ID)
S-NIAH-2 (OOD)
Model
1K
2K
4K
8K
16K
IK
2K
4K
Softmax
100.0
99.4
94.2
11.4
100.0
100.0
4.8
SSMax
100.0
99.8
99.2
92.0
75.2
99.4
99.2
64.4
14.8
Entmax
99.8
99.8
89.0
21.6
1.2
99.6
99.4
64.8
7.2
ASEntmax
99.6
100.0
100.0
99.8
97.4
99.4
99.4
83.2
25.4
ALiBi). In this formulation, ALiBi heads combined with o-entmax create hard attention windows similar Hard-ALiBi (Jelassi al., 2024). NoPE, in turn, can learn positional bias (see App. H.I), and when combined with ASEntmax, it induces learnable, input-dependent recency bias.
Attention Scaling, Recent work has shown that scaling attention logits is key for maintaining sharp attention in long contexts. Methods like YaRN (Peng al., 2024) and the entropy-awart approach of Zhang et al. (2024) use dynamic logit scaling-ofii often with modified RoPE- to stabilize attention during training on extended contexts. Scalable-Softmax (Nakanishi, 2025) and SWAN (Puvvada et al., 2025) apply a log n-based logit scaling to control dispersion without requiring post-training adaptation. InfoScale (Li et al., 2025a) derive scaling rules from the principle of entropy invariance, while Scale-invariant Attention (Anson et al., 2025) introduces rosition-depender transformations to balance attention across both local and global contexts. Across these methods, adaptive scaling consistently improves extrapolation to longer sequences. Our ASEntmax builds on this line of work by introducing context-dependen, learnable scaling within the a-entmax framework, enabling sparse, focused attention context length increases.
Sparse Attention. Previous sparse attention approaches include structured patterns like Longformer (Beltagy et al., 2020) and BigBird (Zaheer et al., 2020), as well as adaptive methods like a-entmax (Peters al., 2019), top-k related methods (Gupta al., 2021; Zeng et al., 2025), and chunk- approaches (Mohtashami & Jaggi, 2023; Hu et al., 2025). A complementary line of work targets efficient long-context inference via structured or dynamic sparsification, e.g., attention sinks (Xiao et al., 2024), dynamic sparse prefill accelerators (Jiang et al., 2024; Lai et al., 2025; Li et al., 2025b), and block/permutation sparse schemes (Chen et al., 2025; Xu al., 2025: Zhang et al., 2025). In contrast to many inference-time sparsification strategies, we study more general modification the analyze how it affects length generalization through
choice of the attention transformation-ane dispersion, representational collapse, and over-squashing. Finally, we note that ASEntmax employs AdaSplash (Gonçalves et al., 2025) for computing a-entmax attention efficiently and thus inherits its runtime properties. Hence, these inference-orienter sparsification methods are largely orthogonal to ASEntmax and can be combined for speeding up inference while retaining theoretical benefits.
7 CONCLUSIONS
In this paper, present principled approach long-context modeling by replacing softmax with a-entmax in transformer attention. Our theoretical analysis demonstrates how this simple change addresses three fundamental limitations: it avoids attention dispersion through naturally sparse distributions, prevents representational collapse by maintaining distinct token representations, and alleviates over-squashing by reducing gradient paths from O(n²) to Oss)), where & n is the number of tokens with nonzero probability. We further introduce Adaptive- Sdaptive-ccalable Scalable a-entmax (ASEntmax), which adaptively adjusts sparsity based on sequence length for each attention head and query input. Our empirical results confirm these theoretical predictions across both synthetic and real-world tasks. On synthetic benchmarks, ASEntmax achieves 95.3% accuracy on associative recall at 1000 the training length, substantially outperforming softmax and existing alternatives. On language modeling with 420M -arammeeeeter models, ASEntmax maintains decreasing perplexity at 8x the training context length and achieves 97.4% retrieval accuracy 16K tokens after training on only 2K tokens. These findings suggest that addressing the fundamental mathematical limitations of transformer attention mechanisms provides direct path to robust long-context generalization.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
To facilitate the reproducibility of our results, we make our code publicly available at https:// github. com/ deee-epin /atn//seetaa... For efficient softmax attention, we have used FlashAttention- 2 (Dao, 2024), and for a-entmax we relied on AdaSplash (Gonçalves et al., 2025). All theoretical results presented in $3 accompanied by complete proofs in Appendix $C and &D. For synthetic task experiments, we provide detailed task descriptions, data generation procedures, model architectures, and hyperparameters in Appendix I-.1- Our language modeling experiments publicly available datasets: DCLM-Edu (Allal et al., 2025) for training and standard benchmarks (LAMBADA, HellaSwag, PIQA, ARC-C, Winogrande, OpenBookQA), The Pile (Gao al.. 2020) (ArXiv, PubMed), and RULER (Hsich et al., 2024) for evaluation, with complete model specifications and training details provided in Appendix $G.3.
ACKNOWLEDGMENTS
We thank the SARDINE Lab members for reviewing this paper and providing helpful feedback. This work was supported by the Portuguese Recovery and Resilience Plan through project C645008882- 00000055 (Center for ResponsibleAl), by the project DECOLLAGE (ERC-2022-CoG 101088763), and by FCT/MECI through national funds and when applicable co-funded EU funds under UID/50008: Instituto de Telecomunicações.
REFERENCES
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlíček, Agustín Piqueres Lajarín, Vaibhav Srivastav, Joshua Lochner, Caleb Fahlgren, Xuan-Son Nguyen, Clémentine Fourrier, Ben Burtenshaw, Hugo Larcher, Haojun Zhao, Cyril Zakka, Mathieu Morlon, Colin Raffel, Leandro von Werra, and Thomas Wolf. Smollm2: When smol goes big data-centric training of a small language model, 2025. URL https:tp///. //tpxvv/.vvvv//////// abgg/bs///a// 22502 /owp.//522.0000 2273377777777....
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations, 2021. URL attps://openree
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 35:38546-38566, 2022.
Ben Anson, Xi Wang, and Laurence Aitchison. arXiv:2505. 17083, 2025.
Scale-invariant attention.
arXiv preprint
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re. Zoology: Measuring and improving recall in efficient language models. The Twelfth International Conference on Learning Representations, 2024. URL https: weeeanreieee.ne//founiiii.e
Álvaro Arroyo, Alessio Gravina, Benjamin Gutteridge, Federico Barbero, Claudio Gallicchio, Xiaowen Dong, Michael Bronstein, and Pierre Vandergheynst. On vanishing gradients, over- smoothing, and over-squashin in gnns: Bridging recurrent and graph learning. arXiv preprint 0rXiv:2522.108// 2025.
Federico Barbero. Andrea Banino, Steven Kapturowski, Dharshan Kumaran, João Madeira Araújo, Oleksandr Vitvitskyi. Razvan Pascanu, and Petar Veličković. Transformers need glasses! informa- tion over-squashing language tasks. Advances in Neural Information Processing Systems, 37: 98111-98142, 2024.
Federico Barbero, Alex Vitvitskyi, Christos Perivolaropoulos, Razvan Pascanu, and Petar Veličkovié. Round and round we go! what makes rotary positional encodings useful? In The Thirteenth International Conference on Learning Representations, 2025. URL https: htttps///opnreviieee net/ forum??d-

---

Published as a conference paper at ICLR 2026
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05150, 2020.
Yonatan Bisk, Rowan Zellers, Ronan Le bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05),7432-7449, Apr. 2020. doi: ht09//aai.v3440...oo org/ index.. php/ TaAT/articleevvee//22222
URL
https://ojs .tw.aaaaaii.aa............
Mathieu Blondel, André FT Martins, and Vlad Niculae. Learning with fenchel-young losses. Journal of Machine Learning Research, 21(35):1-69 2020.
Yaofo Chen, Zeng You, Shuhai Zhang, Haokun Li, Yirui Li, Yaowei Wang, and Mingkui Tan. Core context aware transformers for long context language modeling. In Forry-second International Con- ference on Machine Learning, 2025. URL /ttttts://opperrveee
net/ f formm? MAHP
Ta-Ch Pa-Chung Chi, Ting-Han Fan, Peter J Ramadge, and Alexander Rudnicky. Kerple: Ker- nelized relative positional embedding for length extrapolation. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 8386-8399, Curran Associates, Inc., 2022. URL https: attppe///proceedingss neurips cc/paper fumper_ffless filef/lls//leelle wawee//ppppr/22///1111 9T44.881161146441163336666883333
Paper- Conference, pdf.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoming challenge. 2018. URL https:: / //txx/axvvvv/////////ee//aaaaavl org/ abs 11803.0007.... 05457.
Gonçalo M. Correia, Vlad Niculae, and André F. T. Martins. Adaptively sparse transformers. In Kentaro Inui, Jing Jiang, Vincent Ng. and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 2174-2184, Hong Kong, China, November 2019. Association for Computational Linguisties. doi: 10.18653/v1/D19-122 URL https://aclantholl.g/ org/D19-11 -0t////1112//....
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL https://
openreview www.eeeewwttnreeeeeee ddmmm22222nd222.2
Yihe Dong, Jean nan--Batistte Cordonnier, and Andreas Loukas, Attention not all you need: Pure attention loses rank doubly exponentially with depth. In International conference machine learning, pp. 2793-2803, PMLR, 2021.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers on compositionality. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
URL htttps://openrevieee httt://orrmmmmm////cckkrr
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2104 000027.7 2020.
Nuno Gonçalves, Marcos V Treviso, and Andre Martins. Adasplash: Adaptive sparse flash at- tention. In Forty-second International Conference Machine Learning, 2025. URL https: //openreview wwt.:eettfoooaa Pi
Aaron Grattafiori, Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle
12

---

Published as a conference paper at ICLR 2026
Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olamo, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee. Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafilld Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Contala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang. Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang. Pengwei Li, Petar Vasic, Peter Weng. Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Pumit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Railcanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang. Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang. Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou. Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang. Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang. Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaci, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang. Due Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery. Elconora Presani, Emily Hahn. Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers. Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpemn Grant Herman, Grigory Sizov, Guangyi. Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Aina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang. Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li. Kiran Jagadeesh, Kun Huang. Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manay Avalani, Manish
13

---

Published as a conference paper at ICLR 2026
Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vattkkvv, Mikayel Samvelyan, Mike Clark. Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj. Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani. Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng. Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang. Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng. Sungmin Cho, Sunny Virk, Suraj Subramanian. Sy Choudhury. Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar. Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang. Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang. Zhiwei Zhao, and Zhiyu Ma. The llama herd of models, 2024. URL https: / /arxiiv.vvv org/ 220s 21783.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory--fficient transformers via top-k attention. In Nafise Sadat Moosavi, Iryna Gurevych, Angela Fan, Thomas Wolf, Yufang Hou, Ana Marasović, and Sujith Ravi (eds.), Proceedings of the Second Work- shop on Simple and Efficient Natural Language Processing, pp. 39-52, Virtual, November 2021. Association for Computational Linguistics. doi: 10.18653/w1/201...so attps://aclantholee
URL
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL https:://///ttt //ppr//ee/e/.e...e//////m nwt/forumm id= kIoBbc76Sy.
Xiang Hu, Zhihao Teng, Jun Zhao, Wei Wu, and Kewei Tu. Efficient length--enerrllizblle attention via causal retrieval for long-context language modeling. In Forty-second International Conference on Machine Learning. 2025. URL https://openreete
net
Samy Jelassi, Stéphane d' d'Accoli, Carles Domingo-Enricn Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 15400, 2023.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying, arXiv preprint arXiv:2402.01032. 2024.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo. Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin- Yew Lin, et al. Minference I.0: Accelerating pre- pre-filling for long-context Ilms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-52515, 2024.
Gautam Kamath. Bounds on the expectation of the maximum of samples from gaussian. http: www../aauaakaaattett....... com/ wwwm/ooooottteeeeee htt///rttiins///// weittingggguusiiii gawssiannnnnn. aaaesiiinnaaaaaaaaaeeee pdf, 2015. Accessed: 2025-05-13,
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36:24892-2444 24928, 2023.
14

---

Published as a conference paper at ICLR 2026
Xunhao Lai, Jianqiao Lu, Yao Luo, Yiyuan Ma, and Xun Zhou. Flexprefill: A context-aware sparse attention mechanism for efficient song-sequence inference. In The Thirteenth International Conference on Learning Representations, 2025. URL http:///pppeeeeeeeeeeeettttteree Of OfjllberrT
Kewei Li, Yanwen Kong, Yiping Xu, Jianlin Su, Lan Huang. Ruochi Zhang, and Fengfeng Zhou. Information entropy invariance: Enhancing length extrapolation in attention mechanisms. arXiv preprint arXiv:2501. 250100870, 2025a.
Shanda Li. Chong You, Guru Guruganesh. Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang. Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. In The Twelfth International Conference on Learning Representations, 2024. URL
net teee//oorummidrrR0ogg-
Yucheng Li, Huiqiang Jiang. Chengruidong Zhang. Qianhui Wu, Xufang Luo, Surin Ahn, Amir H. Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang. and Lili Qiu. MMInference: Accelerat- ing pre-filling for long-context visual language models via modality-aware permutation sparse attention. In Forty-second International Conference on Machine Learning, 2025b. URL httpps://operrevie net frrmm/rrrmmmmmmm m==.epf.ATT....
Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Purandare, Gokul Nadathur, and Stratos Idreos. Torchtitan: One-stop pytorch native solution for production ready LLM pretraining. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //openreview net/ wormmriiudeiiee 1dY..
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang, Exposing attention glitches with flip-flop language modeling. In Thirty Conference on Neural Information Processing Systems, 2023. URL https: /estps://operreee /app//peerrveeeeeeee www.eeeweeiwwweeeeeeeeeeeeeeeeee nnt/ id
Andre Martins and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. Maria Florina Balcan and Kilian Q. Weinberger (eds.). International Conference on Machine Learning (ICML), volume 48 of Proceedings of Machine Learning Research, pp. 1614-1623, New York, New York, USA, 20-22 Jun 2016. PMLR. URL http: //proceedings.. mlr. press//48/martinss h.mml.
André FT Martins, Marcos Treviso, António Farinhas, Pedro MQ Aguiar, Mário AT Figueiredo, Mathieu Blondel, and Vlad Niculae. Sparse continuous distributions and fenchel-young losses. Journal of Machine Learning Research, 23(257):1-74, 2022.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and lun'ichi Tsujii (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2381-2391, Brussels, Belgium, October-November 2018. Association for Computational Linguistics. doi: 10.18653/v1/D11. 1-111000.1111. // //aclanthology...
URL https:
Amirkeivan Mohtashami and Martin Jaggi. Random-access infinite context length for transformers. In Shirth- Conference on Neural Information Processing Systems, 2023. URL https: // //ppnrevieww
Ken M Nakanishi. Scalable-softmax is superior for attention. arXiv preprint 19X1v:2501.19399 2025.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27
2022.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume
1: Long Papers), pp. 1525-1555 1534, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 1.. N0.118663/vll 1/P16 144, URL ht /tttppp://alaaaaai http//allaatttth /tt////aaaattooyyy.a.ea wrtg/16--1111 1144/.
15

---

Published as a conference paper at ICLR 2026
Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL /ttp:////////eieeeeeeeeee net/ffrummic id= n6SCkn2QaG.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, 2024. URL https: http:///pppeeeeeeeeeeeeeee
Ben Peters, Vlad Niculae, and André F. T. Martins. Sparse sequence-to-sequeree models. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1504-155 1-19, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/ v1/P19-1146. URL https:tp/// : / wwwweb...bb....... orggantholloy//P9-1111
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables imput length extrapolation. In International Conference on Learning Representations, 2022. URL
httprrrewweeeeeeteeeeeeeeeeeeeee wete/forum1id-R8s@PPGGG. https://openreviex
Krishna C Puvvada, Faisal Ladhak, Santiago Akle Serano, Cheng-Ping Hsieh, Shantanu Acharya, Somshubra Majumdar, Fei Jia, Samuel Kriman, Simeng Sun, Dima Rekesh, and Boris Ginsburg. SWAN: An efficient and scalable approach for long-context language modeling. In Christos Christodoulopoulos, Tanmoy Chakraborty, Carolyn Rose, and Violet Peng (eds.). Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pp. 2424-2438, Suzhou, China, November 2025. Association for Computational Linguistics. ISBN 979-8-891 332-6. doi: 10.18653/v1/2023 12116665/112022.emmpppp main. 123.. URL https: attps://aclannthol.gg/ wrw.e/aaatttooooooo/////0000 org/2025. emnlp-me mapp.main... 123/.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge scale. Proceedings of the AAAI Conference Artificial Intelligence, 34(055):872-884400 Apr. 2020, doi: 10. 1.0999/ai..33400.......... w3w....6433333300 63399, URL https://ojs. aaai beai.org/indeed intex.nddx...e php/ AWWAIAAAaaiii.aaa htteeiiccceveeeee999/////
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Shawn Tan, Songlin Yang, Aaron Courville, Rameswar Panda, and Yikang Shen. Scaling stick- breaking attention: An efficient implementation and in-depth study. In The Thirteenth International Conference on Leurning Representations, 2025. URL https: ttttps://openee http///pppeeeeieeee/eeee/ee feot//frrmmddd r8J3DSD5KF.
Constantino Tsallis. Possible generalization of boltzmann-gibbs statistics. Journal of Statistical Physies, 1988.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural infor- mation processing systems, 30, 2017. URL https:://aapeeeeee /ttpp/aappassae/e/// httpps.cc/appppp ptcc//ppp///////////// ///0/hhh/hh//// $ffe223s4reee9111dddScccccooose -at.raaccccceeeaa-aateeeeeetttt Abstract. html..
Petar Veličković, Christos Perivolaropoulos, Federico Barbero, and Razvan Pascanu. Softmax is not enough (for sharp size gemeralisation). In Forty-second International Conference on Machine Learning, 2025. URL httpp::
Xinyi Wu, Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie. On the emergence of position bias in transformers. In second International Conference on Machine Learning, 2025. URL https://openrrevex..u http://pppenrreeeeeeeeeeoe/errne wwr.ee/forrm?idie (ee
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024, URL https: : http://ppppeeeeeeeeeeeeee wtt..eeoooooommdddd/T//
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International conference machine learning, pp. 10524-10533. PMLR, 2020.
16

---

Published as a conference paper at ICLR 2026
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Mar- tin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang. Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.). Proceedings of the 2024 Confer- ence of the North American Chapter of the Association for Computational Linguistics; Human Language Technologies (Volume I: Long Papers), pp. 4643-4663, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10.18653/v1/2022 wow.:224..aallooo...222... URL
org.oog///2 maacl-long 260/
Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo. and Song Han. XAttention: Block sparse attention with antidiagonal scoring. In Forty-second International Conference on Machine Learning, 2025. URL /ttps////ppeeeeereeeeeeeeee atte/foorumiid-KGBBttt..
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Yang, al. Big bird: Transformers for longer sequences, Advances in neural information processing systems, 33: 13:33123-77 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Mårquez (eds.). Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v 1/116555/11/19-1472 URL
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: https:///ppereview.ro wttp:/iiwwweeeetommmmmiiia j9VVzueEbG.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference Machine Learning, pp. 40770-40803, PMLR, 2023.
Shuhai Zhang, Zeng You, Yaofo Chen, Zhiquan Wen, Qianyue Wang, Zhijie Qiu, Yuanqing Li, and Mingkui Tan. Curse of high dimensionality issue in transformer for long context modeling. In Forty-second International Conference on Machine Learning, 2025. URL /ttpps//opperreieer net./forrum..d-CmmmpppGDGG
Yikai Junlong
Junlong Li, and Pengfei Liu. Extending lims' context window with 100 samples. arXiv preprint arXiv:2401.07004 2024.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M. Susskind. Samy Bengio, and Prectum Nakkiran. What algorithms can transformers learn? a study in length generalization. In The Twelfth International Conference Learning Representations, 2024. URL nwt nortffoum??? Ass TuHnmHIX.
17

---

Published as a conference paper at ICLR 2026
Softmax
Top-2 softmax
1.5-entmax
Sparsemax
0.8
0.8
0.6
0.6
0.5
0.6
0.4
0.4
0.4
0.0
0.0
0.0
0.0
2 0
-
22
Figure 5: Visualization of a-entmax for different values of a. We also include top-k softmax with
k 2 for completeness. Each panel shows how po varies for the input [0, 21, 22].
A Q-ENTMAX TRANSFORMATION
The a-entmax transformation of score vector RR is defined as follows (Peters et al., 2019):
a-entmax(2) arg max p (p), An E PEλn
p > 0, 0.1 (11)
where (p) is the Tsallis(a) entropy (Tsallis, 1988), defined
(p) = 1) E, pf),
a +
(12)
Ej log
a a=
Solving the optimization problem above corresponds to finding a threshold so that p* sums Given 7, we can easily evaluate o-entmax follows:
a-entmax [(a -1]=
(13)
where [++ the ReLU function. Figure illustrates how q-entmax behaves for different choices of a. From the equation above, it is clear that coordinates with (a 1)2; S r become exactly zero. In other words, d-entmax yields dynamic sparsity, where the pattern zeros depends on the input, with controlling the propensity of sparsity: larger values leads to increasingly sparser outputs.
Relation softmax and sparsemax. The a-entmax family continuously interpolates between softmax and sparsemax. When a - 1+, the Tsallis entropy H. converges to Shannon entropy and a-entmax converges to the usual softmax mapping. At a 2, we have Hz(p) x ||p|| (Gini entropy), and the resulting optimizer coincides with sparsemax (Martins & Astudillo, 2016).
g-logarithm and g-exponential view. A useful way to interpret o-entmax through g-deformed exponentials (Tsaliis, 1988). For q # 1, define the qeexponential
cxp,(x) = | (1 q)
(14)
with exp; (z) exp(x) obtained by continuity. q := 2 a so that 1-a=-------- and T-g T-T
1
Then the closed form of o-entmax can be written as a normalized q- g-exponential of a shifted score:
a-entmax(2); exp2-a(z: cc c) for scalar chosen so that Σ a-entmax(2): 1. (15)
This makes the connection with softmax explicit softmax corresponds the q - 1 limit (no truncation), while for a > 1 (i.e., < 1) the ReLU operation yields exact zeros (sparsity).
a-entmax computation. Order-based algorithms have only been proposed for (Martins & Astudillo, 2016) and a = 1.5 (Peters et al., 2019). In contrast, root-finding algorithms apply for all values of a. Letting max(=) and following Peters al. (2019), have
0 < 1-a
(16)
Blondel et al. (2020) propose a bisection or binary search approach to finding r* Similar in spirit, Gonçalves et al. (2025) introduces GPU-oriented solver of a-entmax that uses hybrid Halley- bisection method, which combines the fast local convergence of higher-order root-finding with the convergence guarantees of bisection.
18

---

Published as a conference paper at ICLR 2026
B MODEL DEFINITION AND NOTATION
In this work, we study (causal) transformers with sparse attention distributions created by replacing softmax with a-entmax. We present the precise mathematical formulation of our model below, following closely the notation from (Barbero et al., 2024).
Let Q XWe K XWR, V XWr Ruxd be the query, key, and value projections of the imput embeddings respectively, where n is sequence length and d the hidden size. We denote with qi kss 2; I the d-dimensional query, key, and value vectors of the i-th token. For a single attention head, transformers compute the representation of the i-th token through the following layer-wise transformations:
(0) PG norm
(8)
D0) (1)--1) + 0 0/!-1)
(17)
(F)
norm (2)
( ) n,
(18)
y: norma
(19)
where & is the later index, Pij represents the attention weights, FFN()) Rd - Rd represents
(
the feed-forward network, and norm," (2) norm? (t) and norm; normalization functions. The final representation y; is computed after applying L transformer layers. For next-token prediction tasks, the model output typically depends solely on yn the final representation of the last token. The attention weights Pij are computed by applying a transformation æ R An as follows:
(2) Pij (2(10)
(20)
where 2, E R is the vector of logits for token at layer f, with elements (0) = (q. K ///
The function ı maps these logits to a probability distribution over the n tokens, with An denoting the probability simplex. In standard transformers, i is the softmax function:
exp(e;)
softmax(z))
(21)
Kui exp(2*)
In our approach, replace softmax with a-entmax (SA). We group the attention weights into attention matrix at the E-th layer, defined element-wise as [p(0] == Piş This is a row-stochastic lower triangular matrix that can also be interpreted as probabilistic directed graph. Finally, when incorporating positional information, we modify the attention logits computation according to the chosen positional encoding strategy:
NoPE: (0) = (g. k
(0) \/Va
(0) ALiBi: 4j (g: k) 1/Vd + m (j i). where E R is slope hyperparameter.
(0)
RoPE: 2ng (q. 1 R - $1,00 where R dxd is a block-diagonal rotation matrix.
(8)
C NON-DISPERSION OF Q-ENTMAX
critical problem in long-context modeling is the dispersion of attention, where relevant signals get diluted across increasingly long sequences. For clarity, we begin by examining how a-entmax behaves with two-level logits, and then proceed to define dispersion more rigorously and how a-entmax naturally counteracts this issue.
Lemma 2 (Threshold Behavior for Two-level Logits). Consider logits z E I2"h where k tokens have value M and (n - k)) tokens have value m with m.
'Following Barbero et al. (2024), we omit the linear projections used to compute the vectors from the output of previous layers for clarity; however, this does not impact our derivations and conclusions.
19

---

Published as a conference paper at ICLR 2026
(a-1; a-1 only the tokens with value M receive non-zero
1. For a when A == M attention. The threshold converges to [(z) (a 1 k k-(n-1), and each high-value token receives attention & while others receive zero attention. As a consequence, o-entmma maintains constant attention weight of (+)) on high-value tokens regardless of the total sequence length n.
2. In contrast, softmax (with fixed lemperature 0 0) necessarily disperses with attention weights of ( as n increases. For softmax to maintain concentration of at least e E (0, 1) on the k high-value tokens, the required logit difference must grow logarithmically with nt:
A
11 k
In
(22)
k
1-c)
Proof. We prove the two parts below.
Part (i): Let be the support set. A token is in S if and only if 24 a-1
[(2)
where 7(2) satisfies:
Σ 11a ((a)))=-(3))==--- =1.
(23)
For our two-level distribution, this becomes:
Σ - 1Ma
- 11--1mm-((a) 1a- 1)m
= 1.
(24)
1:2
For only tokens with value M to receive non-zero attention, we need:
(a 1)M r(2)> 0 and (a 1)m 7(2) s 0.
(25)
Rearranging: (a (a-11) ()m S T(2) < (a M. In this regime:
k ((a - 1)M - (2))
(26)
Solving for r(2):
((a (JM - -r(=) 1
(27)
(a 1)M T(2) = (a-1))
(28)
T(Z) = (a - 1)M - k -k (a-1)
(29)
For this threshold to satisfy r(2) 2 (a 1) 1)m, need:
(a 1)M k (a-1) 1) 2 (a 1)m
(30)
(a
M
>
(31)
a 1
k (a-1)
A
(32)
Thus, when A (e-1) an attention of Therefore, when A a-1
a-1 only the tokens with value M receive non-zero attention, each receiving (a(-)) the attention weights for a-entmax are:
=
a-entmax(2)) = {
(33)
if m.
These weights remain ( for high-value tokens regardless of n., demonstrating that a-entmax maintain constant attention on important tokens even as sequence length grows, as long as is fixed oc.
20

---

Published as a conference paper at ICLR 2026
Part (ii): For softmax with temperature 8 > 0, the attention weight for tokens with logit M is:
exp(M//)
softmax(z/0); exp((M/0) + (n k)) exp(m/0)
(34)
For softmax to maintain concentration of at least on the h high-value tokens combined:
kexp(M/0)
k kexp(M/0)+ (n k) exp(m/0) 2c.
(35)
Through algebraic manipulation:
k exxp(M/0) 2 c exkexp(M//)) (n k) exp(m/0)
(36)
(1 c)k (xp(e1/p))) 2 c(n - k) exp(m/0)
(37)
kexp(M/0)
c
2 -
(38)
(n b) exp(mm/0)
b exp(A/0) 2 - c
(39)
n n k
exp(A/0) 2 k - c
c
(40)
Taking the natural logarithm:
A
k
(41)
k
This shows that as grows, the required for maintaining concentration with softmax grows logarithmically with In contrast, for a-entmax, assuming we have a k that is fixed as n grows, the condition A (-(--1) is independent of n, enabling constant focus regardless of sequence length.
We now prove Proposition which concerns the concept of dispersion presented in Definition 1.
Proof. We address each claim in turn. For bounded sequences (Zn)nEN, we assume m, M R with m S M such that m S 2; I M for every E N.
Part (i) a-entmax can retain probability, while softmax always leaks: For 1, consider logits z E IR" and an extended sequence z" E R. with > n, where all additional elements have values below the threshold, 2 S 7(x)/(a- 1) for i n. By the vanisiiihing attention property of o-entmax (Lemma I), these additional elements receive exactly zero probability, resulting in:
a-entmax = a-entmax(2") ) Vi S.n. n.
(42)
This demonstrates that a-entmax can produce identical distributions despite arbitrarily different sequence lengths, maintaining the same concentration regardless of whether we have a distinct number of tokens. In contrast, for softmax (a = 1), Barbero al. (2024) proved that adding any element to the sequence strictly decreases the probability assigned to existing elements, making such invariance impossible.
Part (ii) Complete dispersion of softmax: For softmax with constant temperature 0 > 0, the attention weights for bounded logits can bounded as:
exp(m/0)
exp(M/0)
S softmax(z1.m/8) S Σ exp(2x
(43)
Σj= exp(2,/0)
Since Σj=1 exp(z,/0) exp(m//0 and Σj= exp(z,/00 /0) - n exp M M/0), have:
exp(m//) < softmax(z1m100) < exp(M//)
(44)
exp(M/O)
exp(m/0)
21

---

Published as a conference paper at ICLR 2026
This simplifies to:
1
A
1 exp (
exp
softmax
(45)
n
n
where A M is bounded.
These bounds show that as n - oo, all softmax weights are ((1/n). For the entropy:
H (ssottmx((1:m//0000
softmax(21.//88) /0), log softmax (21m// logn. n.
(46)
Thus, lim
H (softmax(=/00) log
1, showing complete dispersion.
Part (iii) Strong concentration resilience a-entmax: First, we focus on two-level case from Lemma 2, where k tokens have logit value M and (n k) tokens have value When
M m la 17 only the tokens with value receive non-zero attention:
2 = M
a-entmax(Z1:m):
R
(47)
if m. 0
The Shannon entropy of this distribution is:
1
1
(aaeentma(ZIIm))
log K
log k.
(48)
h
The normalized entropy is:
H (o-entmax (2)m))
log k log
(49)
log
For fixed ls as n + oc, this ratio approaches O, confirming concentration resilience.
For cases where the support grows sublinearly as k [8] O(n)) for some 1, the Shannon entropy is bounded by:
H ((0-entmax(21...) log O(log n² ((3logn).
(50)
The normalized entropy is therefore:
H (a-eentmx(21_a)) < 3 <
lim
(51)
log
This confirms that the normalized entropy remains strictly bounded away from 1, even with growing support, as long as the growth is sublinear.
This proposition shows that the entropy of attention distributions reveals how concentrated or dispersed they are across tokens. While softmax distributions with bounded logits must approach maximum entropy O(log n.) as sequence length increases (indicating complete dispersion), o-entmax distributions can maintain bounded entropy Ollog k) where l: the support size. This allows models with c-entmax maintain focused, low-entropy attention patterns even when processing extremely long sequences.
Moreover, this proposition demonstrates that a-entmax attention distributions have remarkable property: they do not necessarily disperse as sequence length increases. They can maintain identical attention patterns regardless of context length. This non-dispersion property means transformers with a-entmax can scale to very long contexts without the attention becoming diluted, maintaining their ability to focus relevant information regardless of how much additional context is present.
22

---

Published as a conference paper at ICLR 2026
D REPRESENTATIONAL COLLAPSE AND OVER-SQUASHING
D.1 PROOF OF LEMMA 1
Adding element to the sequence of logits only redistribute probability mass, ) 2 a-entmax(2"" [=)+++11++++)+) must always hold, with equality iff t-entmaxxx*)), = 0.
Since softmax (a cannot return zeros, we must have strict inequality for a =
For a > 1, we need to find the value bmax such that, for any b S bmax. a-entmaxx holds. From the definition of o-entmax equation 2, token receives non-zero probability iff (a 1)2, > T(z). where T(z) ) is the threshold ensuring the sum of probabilities equals 1, Therefore, for the token in the extended sequence x² receive zero probability (thus not affecting other probabilities), we need:
0
(a (a-1)b 1)b S T(x*).
(52)
We know that r(x") 2 7(x) in general for a-entmax, as shown by Peters et al. (2019); Martins et al. (2022, Lemma 3; Proposition 4). Therefore, sufficient condition is:
(a 1)6 T(x).
(53)
Solving for b, we get:
Σ T(x) a 1
(54)
Thus, we can define bmax r(x) For any b S bmax, the token position n in x* receives zero a-1
attention, meaning doesn't 't affect the normalization. Therefore, r(x") Tar), which means that the condition equation 54 is both necessary and sufficient, and:
a=entmaa(x));; [(a= 1)c [(=:)] -(0==))-7(0)]
o-eetmaaxxx") (x))) h+1- (55)
By choosing different values of b such that b < bmax. we can control the change in threshold r(x") and consequently the difference a-entmax(2)), a-entmax (x*) nn+1 can be large as a-entmax((x).-
D.2 PROOF OF PROPOSITION 2 FOR REPRESENTATIONAL PRESERVATION
We prove the first part of Proposition 2 by exhibiting the counterexample below, following the synthetic construction introduced by Barbero et al. (2024).
Proposition 3 (Counterexample Representational Collapse with a-entmax). Let ()x-))xd be sequence of embedding vectors, and define:
v (0) [V. val IRd (0(0) Ua,
Va Va] (x++))xd
(56)
where the final token v. < Rd is repeated.
For appropriate choice of embeddings and a there exists constant e independent of n such that:
(e vntill
-(L)) c>>
(57)
after transformer layers with a-entmax attention, demonstrating that representational collapse is not inevitable.
In contrast, Barbero et al. (2024) proved that for softmax attention, |101
(L)
v,((1]
||I
+ do for any such construction.
Proof.
We prove this through explicit construction.
Since Vi:n-1 1 and both sequences end with v., the attention logits computed by the final
(0)
(0)
tokens are:
x1) [u U1 e Un U Val.
(58)
2
(1) [u_ Ul.... V. Un- U Va. v. Ual.
(59)
23

---

Published as a conference paper at ICLR 2026
Consider the specific embedding choice where:
Da chosen such that vg 8aa = @ for some & E R.
0t for i ,...., n are chosen such that v, v, = for some value Φ.
This construction yields the following attention logits:
x = [b, b, ... 6, 6].
(60)
2n1 [b,b, ... b.
(1)
(61)
Specific counterexample. Consider d = 1, a = 2.0, o = 0.5, and = We can construct inputs as U2 v0.5 and v; = for 1, ,,.., n 1. The attention logits are:
2 (1) [0. 0......0 0,0.5).
(62)
(1) 00.0.... ..... 0,0.5,0.5)
(63)
x+1
For 2.0 (sparsemax), the attention distributions are: sparsemax (=(")) [Pb, Pb. pb Pnl (dense)
(64)
+(1) sparsemax (2 n+1 = [0,0, 0, 0, 2 (sparse)
1
(65)
where << ps < Pm < 1. As + oc, we have Pn. p" where:
p' (aa 1) - - [(2 (0.5 0) (0.5)* 0.5.
(66)
The average representation is ú - Σi 0 0, and therefore:
(1) lim [|0 2 = (1-p") val = 0.5) - vo.5 = 0.5 x v0.5 R 0.354. (67)
n-000
This demonstrates that the L1 difference remains bounded at approximately c 0.354, independent of sequence length n. This specific example already establishes the existence of a counterexample to representational collapse with o-entmax. We now extend this result to prove Proposition 3 for general constructions of v and V..
With values of & such that b + o-entmax produces the following distributions:"
General Construction.
1),
a-entmax( (x) [p()
P2
())
P
(1)
(2)
(68)
a-entmax +0)) [0,0, ..... 0,
1
(69)
2'2
2
(1)
In particular, since the first n - positions share the same representation, we have pl
P2
1 =1 / 1) pb5' 0, with 0 < pn (1)
1. This leads to representations:
( = Pb 01 + .... + Pbb 1 + ph va
(70)
(1) Ue + Ua Va.
(71)
Let, ū: n Σm=1 vs denote the average of the first block of vectors. Taking the Ln-norm of the representations difference:
1-1
| (1)
(1 ( + p: (. vall,
(72)
= (1 - p (1) _ valli
(73)
We need to show that the above expression does not tend to 0 as n + co. To that end, we need (i)
lim Pn (1) p" < 1, and (ii) lim, | |
0.
'The upper bound ensures dense output for (1 following Lemma with r (a 1) The lower bounds ensure sparse output for (11)
2, following Lemma 2 with & 2.
24

---

Published as a conference paper at ICLR 2026
First condition. We need to choose parameters that as oo, the original sequences remains dense and the extended sequence is in the sparse regime. From our analysis with Lemma and 2, this requires:
2 (a-1)
S - a --1
(74)
a a-1
From the a-entmax definition, we have:
(p)))a-11 (a- 1)b T, (v (1 = (a - 1) T,
(75)
(76)
where p3"h is the probability for each b token and ph is the probability for the & token. Subtracting the first equation from the second:
( (pb - a = (a () - b).
(77)
From the normalization constraint (n 1)ph (2) + Pn 1:
(1) 1 1
Pb
(78)
Substituting:
(1) a-1
(p
1
Pn
-aa-1)(0 b).
(79)
11
We know from the dense regime that pn (1) Ph satisfies:
- p², with < <p" < 1. Thus, as n oc, the probability
- p¹
a-1
(1) - p" where p" solves (p")o-1 lim
= (a - (a) ()(6-b).
(80)
n-000
11 -
Since
1-p"
+ 0 as nt + oo, we get:
(p²)a-1 = (a 1))( b)
(81)
Therefore:
p" [[a 1)(6 6)] --1
(82)
The inequality holds because @ - so (a - () 1. The key insight is that as n -> oc,
the small probabilities on the tokens become negligible, but their sum (n - 1)pp - p= remains finite, each individual Pb (1) 0.
Second condition. construction:
Choose the input sequence such that ö n-T ER-1 * Var and use the
V.. = Voei where e is the first standard basis vector.
J/ e1 + 62 for i = 1,..., 1.
Note that this construction satisfies the logit constraints:
0 U. = Vo.
+0.1 b
(83)
T U/ (vo)²
(84)
25

---

Published as a conference paper at ICLR 2026
Softmax
Enimax g=1
Entmax g=1
Entmexx 6=2
Norm
2
Sequence Length
Sequence Length
Figure 6: L1 norm of representation difference between original sequence and extended sequence after 6 transformer layers, with constant prefix (left) and random prefix (right). With softmax (a = 1.0), representation difference rapidly approaches zero, demonstrating inevitable collapse. a-entmax (a 1.0) maintains bounded differences even at extreme sequence lengths.
The average representation is:
n-1
8;
C1
C2.
(85)
71
v@
Since v. = Voei, we have:
0 valli
v eeel +
Voei
b b-
+1>0 0
(86)
The bound is strictly positive because b * by the logit difference requirement, and because of the constant +1 term from the é₂ component. Therefore, lim | valli A 0. For the case = 1, we can use the simpler construction = vØ and e = V where the non-collapse condition b F Vo, which follows from * Φb In contrast, as shown by Barbero et al. (2024), the resulting representations become increasingly similar as n + with softmax (a 1.0). regardless of the input content, leading to representational collapse.
becomes verifying that V
Empirical Verification of Representational Preservation. To empirically validate our theoretical analysis, we conducted the following experiment: we implemented the counterexample construction using identity projection matrices for queries/keys/valess and tested two scenarios with = 1:
1. Constant prefixes: = 1 and Φ = 1.2.
Random prefixes: b (1 1) and Ø 1.2.
Using a 6-layered transformer with residual connections, we experiment with increasing sequence lengths {128, 256. 16384) and {1.0, 1.5, 1.75, 2.0}, and compute |/ ( || Figure shows how representational differences evolve across increasing sequence lengths. As estab- lished in the previous counterexample, while softmax attention inevitably leads representational col- lapse in long contexts, o-entmax can maintain distinct representations even as sequence length grows.
D.3 PROOF OF PROPOSITION 2 FOR EVEERSSQASHING ALLEVIATION
The following proposition demonstrates how a-entmax helps alleviate the problem of over-squashing. the exponential dilution of gradient signals through deep networks. For clarity, we follow the same set of assumptions as Barbero et al. inde)=nineeeeneence of attention coefficients from the values and approximation of layer normalization with a constant factor.
Proposition 4 (Over-squashing Alleviation with a-entmax). Consider an I-layer transformer-like model where the attention distribution for each head computed by a-enimax with 1. For token n the final layer, let ( E Bd be its hidden representation and
yn normy(u!!
(87)
26

---

Published as a conference paper at ICLR 2026
be its final normalized outpul. The sensitivity of yn to the initial embedding (0) of token experiences less over-squashing with o-enimax than with sofimax attention. Specifically, if the support size of the o-entmax attention distributions is (5
< n for tokens j across all layers then the number of gradient paths from token to token n reduced from (n to O(*)), and consequently helping to alleviate over-squashing by providing stronger gradient signals.
Proof. We begin by expanding By. through the chain rule. Since yn norm(uh" (1 and on is 80) the output of I transformer layers, we have:
dyn 1 Qu
(88)
dv
B3 ov;") (0)
where 2g
1 accounts for the normalization assumption. Expanding the gradient through all layers:
do
Oon (L) deEL-1
(L-1)
doy
(89)
Qv: (0)
(1-1) du Di
(L-2)
do, (0)
koki-
Due to causal masking, the only non-zero terms occur when i S hi s h2 S ... S 1z-1 S n. For each pair of adjacent layers, we have:
Ou (e((+))
On
+1
(90)
dv (9) 3/00
2o (t)
k
where u (8) Exer Pi.k
(0
(8)
(6)
and Pj.k are the attention probabilities computed using a-eettmx. + 0j
Thus, for k S j. we have:
du
(2) 6g
Pia )
(91)
0u 10)
where 65,x is the Kronecker delta and reflects the contribution from the residual connection, which happens when & j. For simplicity, let p/. layers, we obtain:
8j.k. Taking the norm and combining all
Öyn io (0)
E
((-11) (1-22 Pn,ke_, Phk-,,,k-,,-- P
(92)
1121k22
AL
12
where C - 1 II
Ge x
+
is a constant independent of the sequence length.
The crucial distinction between softmax and a-entmax lies in the attention probabilities (f]) a-entmax with a > 1, many tokens receive exactly zero attention. Specifically. we define the ( support set for the j-th token as 5/6) {k Pj.k 0 and h < j}, then p), 0 for all k & s/9. (e) Consequently, Pj.k when k & S and j + k (i.e., when there is no contribution from either attention or the residual connection). This means we can rewrite our bound as:
pj, For
Byn
Σ
(L-1) (L-2) Pa.ke-PPi Phi-
Pk..
(93)
0o (0)
K (k)) kI-1ETh--tkku-.))
where we precisely characterize the gradient flow paths via the sets Ti and (ekk-11), which identify tokens that receive non-zero gradient contributions:
T = {1: E {i, + l.....,.bh nJ & 5:0) or i}
(94)
Te(ke-1) = {k E (k-11, K + 1, 1....,,. n k E s(e-1) or = =k-1 for (95)
These sets have the following meaning:
27

---

(f) Published as a conference paper at ICLR 2026
Avg. Gradient Paths
Avg. Gradient
6
Entmax a=1.3
100
4
Entmax a=1.75
Entmax a=2.0
2
2
3
5
6
7
0
2
5
6
7
Layer
Layer
Figure 7: Left: Layer-wise gradient norms an 8-layer transformer with sequence length n = 256. a-entmax with 1.0 maintain substantially stronger gradient signals, especially in earlier layers, compared to softmax. This demonstrates how o-entmax alleviates over-squashing by enabling more effective gradient flow through the network, with gradient norms up 6x higher than softmax. Right: Visualization of the average number of non-zero gradient paths in 8-layer transformer layer, showing how d-entmax creates fewer paths, which helps to alleviate over-squashing compared to softmax, which always selects all possible paths (within machine precision).
Ti represents the tokens in layer & = 1 that receive non-zero gradients from token i in layer t 0, either through attention (A: E (( or via the residual connection (k i).
T represents the tokens in layer that can receive non-zero gradients from token kg-1 in layer 1, either through attention (k E s(f=1) ) or via the residual connection (k: kz-
The causal constraint (A 2 for Ti and is 2 he-1 for Te) is explicitly incorporated in these definitions to account for the causal attention mask. Importantly, the cardinality of these sets directly corresponds to the potential gradient paths through the network. While softmax attention would yield [T((k
n 1 paths from each token, a-entmax's sparsity ensures (T) = | +
Hence, if the support size ( = < n and assuming that i tokens are always in the support
due to the residual connections, this reduces the number of terms in the sum from (nn)) Oss"). drastically reducing the total number of gradient paths. Furthermore, as a direct consequence of Lemma 1, since a-entmax may concentrate probability mass on fewer tokens, the non-zero Pj.k values can be larger than with softmax. In such cases, the gradients along the remaining paths will be stronger, helping to further alleviate over-squashing by concentrating gradient flow on important tokens.
Empirical Verification of Over-squashing Alleviation. To empirically validate our theoretical prediction that a-entmax reduces gradient paths from O(n²) ) to O(s"), we conducted controlled experiment using a delayed copying task. In this task, the model is presented with sequence consisting of prefix of random tokens, followed by separator token, after which it must reproduce the prefix different attention mechanisms on sequences of length 256, where the model must copy information from the beginning of the sequence to predict tokens after the separator.
natural long- range dependency. We trained 8-layer transformers with
Figure 7 shows the average gradient norm per layer when backpropagating from the loss to each layer input. Softmax exhibits consistently low gradient norms across all layers, indicating severe gradient dilution. In contrast, a-entmax variants maintain substantially stronger gradient signals, particularly in the earlier layers of the network. This confirms that gradient information propagates more effectively through the network with a-entmax, preserving signal strength even when flowing through multiple layers.
The right part of Figure quantifies the average number of non-zero gradient paths per layer. With softmax, nearly all possible connections remain active within numerical precision, creating O(n²) paths per layer. This compounds across layers, resulting in O(n 5) total paths. d-entmax dramatically reduces the number of active paths, with stronger sparsity (higher values) creating even sharper reductions. Notably, in the first layer less than tokens are kept active on average. This empirically confirms theoretical claim that o-entmax prunes the computational graph to (s*²) paths.
28

---

Proposition 5 (No Positional Encoding with a-entmax). Let & be the causal mask graph and p(() E === represent the causal attention matrix at layer & computed row-wise using a-entmax Published as a conference paper at ICLR 2026
E INTERACTION WITH POSITIONAL ENCODING
Here, we aaddy-theoreticllly and empirically- how d-entmax with a > 1 (sparse) interacts with different positional encoding methods. We follow the same model definition and notation in $B,
E. No POSITIONAL ENCODING (NoPE)
We adopt the same set of assumptions as Wu et al. (2025). Namely,
A1 There exists E IR such that max
l2.
W (t) K [1a] C.
KEN
A2 The sequence {
is bounded.
The first assumption tells that key and query weight matrices are bounded. The second assump- tion ensures that the node representations' trajectories across t layers stays within fixed interval [-C²,
(f) (a subscript) with (1,2] or softmax (soft subseript). In particular, let Pij denote the (i.j) probability entry in p()), Further, let p(()
.())..(())) represent the product of attention matrices through layer f, which we call the cumulative attention matrix, which captures how information from tokens in the input layer flows to tokens in layer & through the composition of attention operations.
For softmax, Wu et al. (2025) have shown that
P Psoft.11 80200
lim
=
for all <iS n.
Under assumptions A1. any indices < S
n, with a-entmax we have:
1. Edge deletion: For every layer f, an edge (j, E G is present in the dynamic graph (! if and
only if (a 1) (8) T(z) ). where 2/g (g. (e) k)) and T(Z) (0)) is the entmax threshold.
Otherwise, P Pa.
(t) 0 and the edge is removed for that layer.
2. Modified attention patterns: Unlike softmax, a-entmax with >> 1 creates sparse attention patterns by completely removing some connections. For tokens that survive the threshold, the behavior of their atlention weights depends on:
How far the token's logit sits above the threshold
The relative differences between logits
For tokens that remain connected through the dynamic attention graph go at all layers, the cumulative attention still exhibits a decay pattern:
-00 Pa.s CC(1
(96)
where diz depends on the connectivity puttern of the dynamic atlention graph. This decuy rate differs from softmax due to edge pruning and the redistribution of probability mass.
3. Disrupted limit behavior: Unlike softmax, a-entmax does necessarily converge the first token:
(a) If for every layer f, there exists at least one directed path from token to token i in the dynamic graph (!)) then:
lim Pa.1
(97)
t-000
29

---

Published as a conference paper at ICLR 2026
(b) If some layer En, directed paths from token I token i deleted in g(fo)
then:
0 lim Pa.l < 1,
(98)
E000
with the exact limit determined by the structure of the strongly connected components formed in the dynamic graph.
Proof. (i) Edge deletion: For o-entmax, coefficient is non-zero if and only if its pre-activation exceeds the layer-specific threshold. The stated condition follows directly from the definition of a-entmax:
Pa.ij [(a - 1) (z
(99)
where [2]+ = max(0, z).
(ii) Modified attention patterns: For a-entmax, the attention weights are determined by thresholding:
p = [ 1)2, - -
(100)
T Let x-1 for simplicity. Consider two tokens with logits 24 2 2y both in the support. The ratio of their probabilities is:
P3
2j
(101)
D
For softmax, the ratio is:
soft
(2; 2z)
(102)
no
The comparison between these ratios depends how far 2j sits above r'. Let A 2j and b = 2j T'. Through algebraic manipulation, we can show:
solt
A
b
(103)
p
p
c(a-I)A-
This means the relative behavior of attention weights in o-entmax compared softmax depends on the specific configuration of logits and thresholds, rather than following simple universal relationship. Since tokens that remain connected must distribute probability mass among fewer options (due to pruning), there exists some 0 < 8ig < 1 such that:
(0) Pa.ij S C(1
(104)
The specific value of dey depends on the connectivity pattern of the dynamic attention graph and may differ significantly from the softmax case due to edge pruning and probability redistribution.
(iii) Disrupted limit behavior.
(a) Case where paths to token the first token persist:
Suppose that for every layer f, there exists at least one directed path from token 1 to token in the dynamic graph g( the unique "center node" defined by Wu et al. (2025). For any token j > the geometric decay established in part (ii) applies:
Pa.i < C(1
5ty)
E op.
(105)
Since the row sums of (()) must equal 1 (as it is product of row-stochastie matrices), and all entries (()) Pa.ij with j > approach 0, we have:
lim Pa.1
-(0) lim 8
P Pa.ij
(
1.
(106)
8-000
30

---

Published as a conference paper at ICLR 2026
(b) Case where paths to token are cut: The key difference from softmax arises when edge deletion creates a configuration where token 1 cannot reach token i. Let &y be the first layer where paths from token to token are removed in g(fo) Let c(f6) C {1, ..., n be the set of tokens in the same strongly connected component as token i in g!an), By our assumption, 1 & cf At layer lo. the
attention probability is distributed only among tokens in c/fo).
Σ
(6a) Pa.i = and Pa,11 0.
(107)
jEC
For all layers l > lo, the multiplication by ensures (())
and therefore:
lim Pa.l 1.
(108)
The exact limit depends on the structure of the strongly connected components formed in the dynamic graph through subsequent layers. fact, the limit can be exactly zero whenever all paths from token 1 to token are removed from G, since layer to-
This proposition demonstrates a fundamental difference between softmax and (-entmax transformers: while softmax inevitably leads to concentration of attention the first token, o-entmax can potentially disrupt this position bias through its ability to dynamically prune edges in the attention graph. This provides a theoretical foundation for using sparse attention mechanisms mitigate position bias in transformer architectures.
E.2 ALIBI
We start by recalling the definition of ALiBi from (Press et al., 2022).
Definition (ALiBi Positional Encoding). Let H be the number of attention heads. A general form of ALiBi bias for head h E {1,2, .., H.} can be defined as:
(
mm (j- - i) if
00
(109)
otherwise,
where my E R., is the slope parameter for head h, defined as mn = 2 h.
Now, we consider how interacts with a-entmax.
Proposition 6 (ALiBi with a-entmax). Consider a-entmax attention with the ALiBi bias from the
above definition. Assume the raw attention logits 2ig E [h) 2(h) for all i, Let
(h)
(h)
((h)
2max
2min
a-1
(110)
mi
Then, any token with (i - j)
d
receives zero attention from token i head h.
Proof. For a token j to receive non-zero attention from token i with a-entmax, we require:
(a- - 1)(2)) + +b()) >T
(111)
where the threshold ensuring normalization. Since slopes are positive (m, > 0), have:
(a (a-1)(z))))
1)(2 - (imm) >T.
(112)
In the extreme case where only the closest token receives attention (single-support case), we can solve for exactly:
1 (a ()(?())
7)
+ (a ()(h)) -) = =
- a
(113)
31

---

Published as a conference paper at ICLR 2026
Since the logits drop by at least my per position due the ALiBi bias, have 2max Amin mi which gives us:
(A)
(b)
(h) min
1
mh
(114)
For a token with distance (i - j), even with the maximum logit zmax:
(a - (2 (i - j)mm) I (a 1)(2ml mn) - 1.
(115)
Solving for (i -)::
(h)
min + m + n-1
(1 1)
(116)
m
Hence,
(h)
(h)
(h)
2max
"min
a +1
(117)
max
mh
This proposition establishes that with a-entmax and ALiBi positional bias, there exists a head- dependent hard cutoff distance d beyond which tokens receive exactly zero attention. This creates an adaptive but bounded attention window that depends on both content relevance (2 = and the sparsity parameter a, naturally limiting the effective context without requiring explicit truncation. This property allows the model to focus computational resources on a relevant window of tokens, which can be particularly valuable for efficiently processing long documents.
E.3 ROPE
To analyze the interaction between RoPE (Su et al., 2024) and a-entmax, we first establish our notation. Following Barbero al. (2025), we consider queries qi E IR and keys k2 E Rd, where and are token positions in the sequence. RoPE decomposes these vectors into d/2 two-dimensional
chunks, denoted as qi (k) R² and k/))) E 2 for l E {l...., d/2). Each chunk rotates a different frequency 2k- where 0 (typically 10,000) is the base wavelength parameter.
RoPE applies position-dependeee dependent rotations through matrices p(gx)' transform the original queries and keys. The resultant raw attention logit between query position and key position is: In standard transformer with RoPE.
4/2
Zij
Σ
(k) (a )(gx))
)
(118)
where p(gx) is the 2D rotation matrix with frequency 9s.
Query-Key Interaction with RoPE. Let Øişk the angle between the original (unrotated) vectors
q: and k (k) That is:
(k)
(q, (k)
(k)
k
cos((15k)
(119)
11
||2
112
For single 2D chunk, since rotation preserves vector magnitudes, the contribution to the raw score is:
(k) (q p(..))))))))))))))))))) (o
[((b) || (x) 12llp(g)))" lla cos( + guli (()--)))
(120)
||q:
x ll
k (1 1|2 cos(Óujk gu i))-
(121)
(1) and k: (k) As shown in Proposition 3.1 of Barbero et al. (2025), RoPE allows for maximal attention at any arbitrary distance. However, RoPE combined with a-entmax creates a hard boundary on attention distance due to the thresholding effect, which we analyze next.
where Φitk is the original angle between q:
32

---

Published as a conference paper at ICLR 2026
Approximation for Small Angles. First, note that the angle-sum sum expansion for cosine as follows:
cos( + ( - i)) = cos( cos(gn i)) sin(Oizk) (i i)).
(122)
Further, note that for small angles ØijR we can use a second-order Taylor expansion for go
i):
cos(g - - 1)) = (
(123)
2
Finally, applying the dot-produee
d/2
4/2
L2 -)2
Zij
qi
(k)
lla [kg 1l2 cos(
Σ
11
(k)
1l2 [k] 112 cos(
sin terms
2
(124)
d/2 Σ k=1
2 max
Ci - -))
(125)
Here, we simplified the last step by focusing on the quadratic decay from the cosine term while omitting the sine terms sin sin( (j - i). For semantically aligned tokens where Digh R 0, the sine term's is contribution is minimal since lim sin(x:) O.
Proposition 7 (Maximum Attention Distance for RoPE (Small-Angle Regime)). Let Jmax max; gle be the maximum frequency in RoPE. Within the small-angle domain where
ji-
2gmax
so that all rolational angles 0 go (ū (i-j) -) satisfy 101 S 3. and assuming Zmax
r(zi)
there exists
critical distance dmax beyond which tokens receive exactly attention:
r(z.)
dmax
max
a-I
(126)
N
E k= CA
Proof. For a token to receive non-zero attention under a-entmax, we must have 2ig > r(a:) Sub- stituting the decay pattern from Equation 124, which is valid within the small-angle domain where cosine can be approximated using Taylor expansion cos Ro 1 2
ga
4/2
Σ cigi j) > T(Zi)
(127)
Zmax
1m
Rearranging for (i - 2)
Zmax [(2,)
(i j)²
a-1
(128)
I d/2
Ch.
Taking the floor of the square root gives us d
Note that this analysis applies to the first attention window. Due to the periodicity of rotation operations, distances beyond x for any frequency component kt the attention pattern may exhibit additional windows of non-zero attention, which we address next.
'cos(x) 1 22/2 higher order terms.
33

---

Published as a conference paper at ICLR 2026
Proposition 8 Frrequency-Speciffe Cutoff for RoPE). For each frequency component k. let By.
(a-1))lla (1aal Since at least one token must receive non-zero attention for o-entmax to yield a valid probability distribution, 3x S must hold for at least one component. Assuming 3k E [-1,1] fcovering all possible cosine values), there exists sequence of distances {dk, In at which its contribution to attention crosses the threshold.
T(Z:)
The first such distance is:
T(Zi)
d
arccos
(129)
gk
(a - 1)
1|a [k)
Due to the periodicity of cosine, subsequent threshold crossings occur at approximately:
2 + dx.o
dR
TI TEN N.
(130)
9k
Furthermore, dk,0 is non-increasing in a and inversely proportional gk-
Proof. For a single frequency component k, the contribution to the raw score from Equation 120 is:
(a
(k) (l || (k: 12 11211 cos(ge - -i) + disa)).
(131)
For this to exceed the attention threshold under optimal alignment (Øiga O, which maximizes the contribution):
1)llg; 1211 (k) l2 cos(an 1)) > T(Zi)
(132)
T(2,)
cos(gk - i)) (a 1) (k)
(133)
ll2
Taking the arccos of both sides and dividing by gives the first threshold crossing distance da.0:
1
T(2;)
jl
arecos
(134)
gk
(a- 1)llq Ik
Due to the 2x 2x-periodicity of cosine, subsequent threshold crossings occur at distances 2anidox for integers > 0. Moreover, 7(2,)/(a 1) grows increases, making the arccos term smaller and consequently decreasing dx.0- The inverse proportionality to is evident directly from the formula.
These theoretical analyses of RoPE with q-entmax reveal two interesting takeaways. First, differ- ent frequency components in RoPE naturally create attention windows of different widths. High- frequency components (large (jk) produce very narrow windows focused on local context, while low-frequency components (small gı)) enable attention over longer distances. Second, the sparsity pattern induced by the combination of RoPE and a-entmax is not uniform but varies across frequency components, creating more complex attention structure than simple distance-based decay methods like ALiBi.
E.4 COMPARISON BETWEEN POSITIONAL ENCODING METHODS WITH Q-ENTMAX
The positional encoding scheme used in transformer has significant implications for how attention behaves over long contexts. Our theoretical approach from previous subections reveals that o-entmax interacts with positional encodings in ways that fundamentally alter attention behavior compared to softmax. We summarize our theoretical findings next, along with an empirical analysis within a controlled experimental setting.
34

---

Published as a conference paper at ICLR 2026
Head
Head
NaPE
KOb ded Oar sortod
Aieerted Cantrol
Guary Posilan
Query
Query Postoo
Quary Poneton Head
Heed!
125
120
Kaunder Centrdls
Centraa
60 Postton
Figure 8: Per-head attention centroids across different positional encoding methods, Each panel represents attention head's behavior. The dashed line shows the identity function (attending to self). NoPE heads consistently exhibit an early-token bias, ALiBi heads maintain proximity-basee attention, and RoPE heads display more varied patterns with irregular fluctuations.
Theoretical Analysis. With NoPE (Kazemnejad et al., 2023), softmax transformers (without MLP layers) develop an implicit bias towards the first tokens as depth increases, as shown by Wu et al. (2025). c-entmax disrupts this behavior through its ability create disconnected attention graphs. By assigning exactly zero attention to some connections, it may remove the implicit bias encouraging attention to concentrate on early tokens.
When combined with ALiBi (Press et al., 2022), a-entmax transforms the smooth linear de-
cay into a hard attention window. For tokens separated by distance d > (h) where ((h)
(Zmax 2min + + +1 attention weights become exactly zero. This creates an adaptive but bounded attention window that depends on the input (Zmax =min) and the sparsity parameter a.
With RoPE (Su et al., 2024), a-entmax induces frequency-dependent sparsity. Each frequency component & has critical distance de beyond which its contribution falls below the attention threshold. This creates multi-scale attention pattern where nearby tokens interact through all frequency components, while distant tokens interact only through low-frequency components.
These interactions between positional encodings and a-entmax have important practical implications, and further motivate the introduction of our hybrid approach, NAPE (NoPE ALiBi). By creating natural, content-adaptin attention windows, NAPE combine the benefits of sparse, focused attention with awareness of token positions, allowing models effectively balance local and global information processing. This provides a principled alternative to manually designed sparse attention patterns like sliding windows dilated attention, with the advantage of adapting to content relevance rather than using fixed patterns.
Empirical Analysis Setup. We simulated sequence of length 128 with attention heads using a-entmax (a 1.5). For each query position we generated a vector of raw attention scores (logits) for all positions j < i according to:
2ij
base
proxx -11
noise Pij
(135)
Juse U(Zmin- Zmaxx represents the base content-based affinity, 2/ where 2ng
peexx = 0.2 Zmax
2min)(1 -1 0.5 introduces a proximity bias, and = ~ NO, o²) adds Gaussian noise. The parameters were set to Zuin -500, Zmax 5.0, and *ij mixture of content-basee based attention (random component), a mild inherent bias toward nearby tokens (proximity component), and natural variation (noise component). We then applied different positional encoding methods to modify these base logits. Finally, we calculated the attention distribution using a-entmax and determined the attention centroid for each query position as centroid; = Σi=1 Pij. where Pij o-entmax(z;)),
0.5. This formulation models realistic
35

---

Published as a conference paper at ICLR 2026
Values
Vallues
1.50
1.26
0.26
Head
Head
Figure 9:
Heatmaps for 3 and y per head for Scalable a-entmax.
Empirical Results. The read-specific analysis in Figure 8 reveals distinct behaviors across posi- tional encoding methods when combined with a-entmax. While seems NoPE exhibit weak bias towards earlier positions, it also shows a modest variability, indicating more disperse attention. ALiBi clearly creates a consistent recency bias, with centroids following slightly below the identity line, maintaining low variability that indicates focused attention. RoPE demonstrates centroid patterns similar NoPE but with lower entropy (higher variability). suggesting focused attention in more distant positions. These observations may explain why NAPE-the the hybrid NoPE+ALiBi well in practice, since ALiBi heads provide consistent positional structure focused on recent context, while NoPE heads can contribute complementary via early-token and semantic focus, creating more balanced attention mechanism than either approach alone. In fact, as we show in HH.1, models equipped with NoPE are flexible enough and can acquire relative positional encoding, thus also supporting the original hypothesis Kazemnejad et al. (2023). Therefore, NAPE has the ability to encourage short-span focus with ALiBi alongside learning more longer-span focus that guided via semantic information with NoPE
F ADAPTIVE-SCALABLE a-ENTMAX (ASENTMAX)
F.1 LEARNING SCALERS FOR LANGUAGE MODELING
To verify our scaling approach, we follow the setup from (Nakanishi, 2025) and trained a language model with learnable scales p; for each i-th position. However, we do independently for each head in the model. Specifically, we train a 12-layer transformer with approximately 120 million parameters. We set hidden size to 768, attention heads to 12, MLP intermediate size to 2048, learning rate to 6 x 10-4, weight decay to 0.001, batch size to 1M tokens, and sequence length to 1024. Each head thus contains 1024 learnable parameters (1 per position). Finally, we train on the Wib dataset for total of billion tokens.
In Figure 3 we show how the 5n Bn (log n) 1 scaling performs well for different attention heads. In contrast, removing 7 as done by Nakanishi (2025) leads to severe degradation fit for several heads. Specifically, the linear fit has an overall I22 0.12, the log fit has R² -14 (severe underfitting), and our log with y exponent has R²2 = 0.17. The full set of Br and Th learned by our approach are shown in Figure 9.
Effect of Negative 7he Experiments on the Copy task, shown in Table 10, suggest that, without scaling, c-entmax can hurt performance, leading a noticeable drop in accuracy in the OOD scenario. Introducing an adaptive temperature, however, substantially mitigates this effect. We hypothesize that Copy requires less sparse attention patterss, which can be accomplished by applying a negative power to the logarithm function. We confirm this hypothesis in Figure 13, which shows that ASEntmax learss negative values of 7n, in all heads, resulting in more spread-out attention distributions.
36

---

Published as a conference paper at ICLR 2026
G EXPERIMENTAL DETAILS
G.1 SYNTHETIC DATA
Following the data diversity assumptions of Zhou et al. (2024), we generate a large number of samples-between between 10 million and 50 million, depending task complexity. See Table for training/test details on for each task along with model hyperparameters. The 2Back and Local Count tasks are token-clasiification tasks, while the remaining tasks are generative. In Figure 10, we show examples of the tasks we introduce in this work.
2Back. In this classification task, the model must predict the class of the token that appeared two positions earlier. Via this task, we examine the ability of models equipped with NoPE to learn relative positional bias and assess their behaviour in out-of-distribution scenarios (see H.1).
Local Count. The Local Count task a classification task in which the model must predict the number of times a word has occurred far. We restrict the vocabulary size to 16, allowing multiple clusters of the same word to appear within a sequence multiple times. This increases the task's 's difficulty, the model must distinguish between different clusters of identical words. We sample the number of repetitions for each cluster uniformly from (111, 48) to test whether models equipped with NoPE can learn longer focus span than observed in the 2Back task.
Flip-Flop (Liu et al., 2023). Flip-Flop uses sequences made up of alternating instructions "write", "ignore", and "read" (w, i.r). Each instruction is paired with a bit (0 or 1). The model must memorize the bit if it follows w, and recall the last stored bit when it encounters an For example, in the string "i 0 w i 1 1 w 1 i 0 i 1 r", the correct output is 1. We conducted experiments on two datasets with different "write" instruction probabilities: 10% sparse, 80% dense.
Max Retrieval. We follow Veličković et al. (2025) in constructing the Max Retrieval dataset and the model architecture.
Multi- Queti-Query Multi-Token Associative Recall (MQMTAR). MQMTAR is a retrieval task in which models must produce sequence of multi-token values corresponding to the queries provided at the end of the input sequence (see Figure 10). MQMTAR employs three special tokens: (1) O for empty space; (2) 1 as the key-value delimiter; and (3) 3 as the query delimiter, which is also used to separate values in the target sequence. We set the lengths of both keys and values to be 2 tokens, resulting in 5 tokens per key-value pair in the input. The number of queries is 4, and the density of key-value pairs is 80% of the total number of tokens. Finally, the size of the alphabet is 256, but we limited the size of the key/value vocabulary 10K. This results in total of 100K key-value pairs.
Sort, Copy and Reverse. These well-known tasks for testing models' length generalization (Kazemnejad et al., 2023). We use a small vocabulary size of 32 to generate more sequences with repeated tokens, since models must handle such repetitions increasingly as sequence length grows.
G.2 MODELS FOR SYNTHETIC TASKS
All synthetic tasks are trained with a decoder-like transformer. We evaluate models in extreme settings by using as few layers as possible as our aim is to test the attention mechanism coupled with our positional-encoding strategy, rather than the scaling capabilities of transformer. However, for the Reverse washh proved particularly challenging for softmax- based models increment the layer count until the softmax baseline generalizes to at least 1.5x the in-distribution length.
For experiments with RoPE, we use the Hugging Face implementation from LLaMA 3 (Grattafiori et al., 2024), which includes RoPE scaling. Since our sequences are relatively short, the base frequency is set to its default value of 10,000. To improve length extrapolation in RoPE-based models, we apply a scaling factor of 16, which we found to be optimal for Flip-Flop under 4x extrapolation (see Table 11); factors of uses a slope of m = t.- where h is the head index. We employ 8 attention heads for all tasks except Copy and MQMTAR, where 16 heads yields performance boost across all models.
32 degrade performance. For NAPE, each ALiBi head
37

---

Published as a conference paper at ICLR 2026
Multi-query Multi-token Associative Recall k1
vi
K2
v2
Input 0461 1 0
Target: 268278
2Back
Input 0 0
2
3
Target: 0
2
3
Local Count
Input
8 2
Target: 1
2 3
Figure 10: Examples of the introduced tasks. MQMTAR: Each digit is token; the alphabet size 10, and the number of queries 2. 2Back: special token added at the beginning of the sequence
the model has something predict the first two positions; the vocabulary size is 10. Local Count: The maximum number of repetitions 4, and the vocabulary size is 8.
In case of ASEntmax, 8y and 7 are computed per token i via linear projections followed by activations softplus and tanh respectively, allowing to scale attention adaptively based on content. For our experiments with a-entmax. unless mentioned otherwise. Furthermore, we use the Gemma2 implementation from Hugging Face, but disabling dlding-window attention in all layers. For experiments with a-entmax, replaced FlashAttention with AdaSplash (Gonçalves et al., 2025).
1.5 as the default value for the a-entmax models,
For optimization, we use the AdamW with default betas and cosine learning-rate scheduler with warm-up, setting 10K warm-up steps. Given the large training corpus, we do not employ dropout or weight decay. In addition, we use bfloat16 in all experiments. All models are relatively small (2-10M parameters) and fit on a single GPU.
We observe that even when models are 100% accurate in-distribution. they still require significantly more training; in some cases, the training loss reached as low 10-8 Therefore, the best checkpoint selected based on performance at 8x the in-distribution sequence length. In some cases such as the Sort task and models with RoPE, where generalization up to 8x was not possible, as an intermediate metric and 2x or 4x in-distribution sequence length. We perform evaluation with 1 K samples per sequence length. 2Back and Local Count are evaluated using accuracy, as they are classification tasks, whereas for the remaining tasks, we use exact match accuracy accuracy-assigning 1 only if the entire predicted sequence matches the reference and 0 otherwise. We report results for the single hest-performing model selected from experiments conducted with multiple random seeds (3) and various learning rates.
use BLEU
For some tasks, we also report results for SEntmax-wii which learns 3 and y directly, without linear projections nonlinearities and also for ASSMax-hiih applies our adaptive-scaling strategy to Softmax. Finally, for some tasks we also experiment with Stick-Breaking Attention (SB, Tan et al. 2025), which corresponds different positional encoding strategy.
G.3 LANGUAGE MODELING
We train 420M-parameter decoder-only models following the LLaMA-3 architecture, using the following hyperparameters: (1) model dimension: 768; (2) number layers: 24; (3) number of heads: 12. For RoPE and FFN layers, we the default LLaMA-3 settings. The training is conducted using the torchtitan library (Liang et al., 2025). We train for a total of 13.5K steps, with 1.35K steps allocated for warm-up. The learning rate is set to 3 10 1 and the cooldown phase corresponds to 10% of the total steps with a minimum learning rate of 3 10 where we apply cosine decay scheduling. We use the AdamW optimizer and train with bfloat16 precision. The global batch
38

---

Published as a conference paper at ICLR 2026
Table 5: Task details and hyperparameters.
Task
Samples
Length
Batch
Vocab.
Heads
Layers
Hid. dim.
Int. dim.
2Back
10M
32-64
128
256
512
Local Count
10M
64-128
128
16
3
128
512
Flip-Flop
10M
32-64
128
256
512
Copy
20M
32-64
128
32
16
2
256
1024
Reverse
30M
32-64
128
32
6
256
512
MQMTAR
50M
32-64
128
256
16
4
512
1024
Sort
40M
32-64
128
32
2
256
1024
size 524K tokens. For tokenization, we employ the LLaMA-3.2-ID tokenizer. For NAPE experiments, we choose original exponential ALiBi slopes.
We the following benchmarks for short context evaluation: Lambada (Paperno et al., 2016), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020), ARC-C (Clark et al., 2018), Winogrande (Sakaguchi et al., 2020), and OpenBookQA (Mihaylov et al., 2018). For length generalization, use 500 documents from each dataset (ArXiv and PubMed), with all documents at least 16K tokens long, and 500 samples (default number of samples in RULER) for each length for S-NIAH-1 and S-NIAH-2. We evaluate models with RoPE in two settings: vanilla RoPE and RoPE with Adjusted Base Frequency (ABF) (Xiong et al., 2024), where the base frequency is set to 500,000, i.e., 50x used in pre-training. Results for language modeling are in Appendix H.9.
H DETAILED RESULTS
In this section, we provide a more detailed analysis of each task.
H.1 2BACK
Table 6: Accuracy (%) on 2Back.
ID
Oft-o--Distribution
Model
64
128
256
512
1024
2048
4096
RoPE
Softmax
100.0
100.0
100.0
99.3
81.4
63.4
41.1
SSMax
100.0
100.0
100.0
99.8
98.5
90.4
69.0
Entmax
100.0
98.2
94.8
83.7
65.5
45.7
31.3
ASEntmax
100.0
100.0
100.0
95.0
61.2
36.2
22.1
NoPE
Softmax
99.9
83.2
51.4
30.1
18.5
12.7
9.7
SSMax
100.0
80.5
47.2
27.2
16.9
11.8
9.2
Entmax
100.0
77.3
50.5
33.2
20.7
13.7
10.2
ASEntmax
100.0
90.3
55.1
36.5
24.3
16.2
11.5
NAPE
Softmax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
SSMAx
100.0
100.0
100.0
100.0
100.0
100.0
100.0
Entmax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
ASEntmax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
Following the hypothesis of Kazemnejad al. (2023) that NoPE can learn relative positional bias, we conducted experiments on the simple 2Back task, in which the model must predict the class of the token two positions earlier. Table 6 shows that models equipped with NoPE achieve perfect in-distribution performance. Moreover, the attention maps in Figure (left) confirm that NoPE indeed acquires relative positional encodings, thus supporting the hypothesis. However, the
39

---

Published as a conference paper at ICLR 2026
OOD attention maps (right part) reveal that, sequence length increases, the recency bias diffuses unevenly across positions. Such behavior detrimental tasks requiring attention to a fixed-size local context (e.g., associative recall, previous instructions in code, n-grams in text). By contrast, ALiBi constrains attention to a local window irrespective of content. Moreover, we observe that ASEntmax partially mitigates diffusion in the attention maps (bottom right), which is also reflected in the accuracy gains shown in Table 11. In our design of positional encoding, NoPE + ALiBi (NAPE), half of the attention heads employ a faster-decaying variant of ALiBi to enforee short-span focus, while the remaining heads use NoPE which can 1) learn focus that spans longer and depends on position 2) guide attention semantically.
H.2 LOCAL COUNT
As observed (Table 7), models using ALiBi perfectly solve the task, which is unsurprising given that ALiBi induces recency bias. Furthermore, the results for ALiBi and NAPE suggest that models rely exclusively on ALiBi heads in case of NAPE. With NoPE, however, the model is challenged because identical tokens develop mechanism to locate the current cluster. Figure 12 indicates that by the third layer, the NoPE model exhibits relative positional bias. Combined with the bias observed in 2Back, this indicates that NoPE models can acquire various content-based recency biases that differ from those induced by ALiBi or RoPE. Finally, for Local Count, we observe no improvement NoPE models when using attention scaling.
distinguishable the very first layer. Therefore, the model must
Table 7: Accuracy (%) on Local Count.
ID
Model
128
256
512
1024
2048
4096
RoPE
Softmax
100.0
99.4
91.6
55.2
31.1
17.3
SSMAx
100.0
100.0
81.3
42.6
23.0
13.4
Entmax
100.0
99.9
89.3
47.1
24.9
14.1
ASEntmax
100.0
100.0
79.1
41.4
22.4
13.1
NoPE
Softmax
99.1
71.7
36.5
18.3
9.2
4.6
SSMax
99.1
71.4
36.8
18.6
9.3
4.7
Entmax
99.8
80.8
45.6
25.0
13.8
7.7
ASEntmax
99.6
78.1
42.6
22.5
11.9
6.4
ALiBi
Softmax
100.0
100.0
100.0
100.0
100.0
100.0
Entmax
100.0
100.0
100.0
100.0
100.0
100.0
NAPE
Softmax
100.0
100.0
100.0
100.0
100.0
100.0
SSMax
100.0
100.0
99.9
99.9
99.8
99.8
Entmax
100.0
100.0
100.0
100.0
100.0
100.0
ASEntmax
100.0
100.0
100.0
100.0
100.0
100.0
H.3 MAX RETRIEVAL
Solving this task requires extremely concentrated attention distribution. Such distribution can be achieved by either lowering the temperature for softmax or increasing the entmax parameter a. As Table 8 shows, increasing yields substantial performance gains. However, if a becomes too large, the distribution may collapse to one-hot vector, causing entmax to lose all gradient signal and hindering learning (c.g. 16). Instead. this issue can be alleviated by scaling entmax based on the sequence length. With this approach, ASEntmax with a = 1.5, learned 3. and elevated achieves substantially improved performance the task. We also note that Top-K K with K=2 works roughly as
40

---

Published as a conference paper at ICLR 2026
Softmax
Entmax 1.5
ASEntmax 1.5
Figure 11: Comparison of attention maps for the 2Back task. Left: In-distribution, sequence length 64. Right: Out-of-distribution, sequence length 512 (for visualization clarity, we applied max pooling with a window size 4 and stride 4). The maps are shown for the second layer for all models. We can observe that diagonal patterns are less distorted with a-entmax. Moreover, ASEntmax mitigate dispersal of the diagonal pattern up to 3x the in disttributtion sequence length and make it less distorted up to 8x.
41

---

Published as a conference paper at ICLR 2026
Figure 12: Attention maps of Entmax model Local Count. Left: Layer 1. Right: Layer 3; We observe a local pattern: attention weights fade as relative distance increases. Input sequence: (1..1 x10, 2. x4, 3. x10, 2. 2 x4, 4. 4 x10, 2. x4) x4
well as SSMax. However, lowering K to I is not possible since this would essentially reduce Top-K to an argmax operation, non-differentiable function.
Table 8: Accuracy (%) on Max Retrieval
Out-of-Distribiin
Model
16
32
64
128
256
512
1024
2048
4096
Softmax Veličković et al. (2025)
98.6
97.1
94.3
89.7
81.3
70.1
53.8
35.7
22.6
Adapt. temp. Veličković et al. (2025)
98.6
97.1
94.5
89.9
82.1
72.5
57.7
39.4
24.9
Softmax 0 dd
99.2
98.5
96.7
93.2
86.7
73.5
54.4
36.4
24.1
Softmax = 0.1
99.5
99.0
97.8
95.1
89.6
77.9
60.2
41.2
28.5
Softmax 0 0.0004
99.2
98.4
97.0
94.2
89.4
81.8
71.4
58.4
43.4
SSMax
99.4
98.9
97.8
95.9
92.3
85.0
74.7
59.9
44.7
Top-K, K 2
99.4
98.6
97.6
95.4
91.3
85.0
75.3
62.4
48.3
Top-K, K 4
98.8
97.7
95.0
89.5
79.1
64.6
48.9
38.9
32.4
Entmax 1.5
99.4
98.8
97.4
94.7
89.9
80.1
65.1
50.0
36.8
Entmax a 2
99.5 99.1
98.0
96.0
92.1
84.5
72.0
58.4
44.6
Entmax 4
99.5
98.9
97.7
95.9
92.1
84.8
75.2
61.4
46.9
Entmax a 16
99.6
99.4
98.7
97.5
95.2
91.0
82.8
70.3
53.4
Entmax a 32
99.4
98.7
97.5
95.5
91.5
83.8
72.6
57.5
41.7
Entmax a = 64
99.1
98.4
96.8
93.9
88.7
78.6
64.6
45.5
28.1
ASEntmax, 1.5, Stearns 7 1
99.5
99.0
98.1
96.3
93.1
86.1
76.2
61.9
44.5
ASEntmax, a 1.5, Btaarns 7 = 2
99.6
99.2
98.4
96.9
94.4
89.0
81.4
69.5
55.1
ASEntmax, a 1.5, Iearrs 7 3
99.6
99.4
99.0
98.0
96.0
92.4
85.9
76.1
62.7
ASEntmax, a 1.5, Bearns = 4
99.3
98.7
97.6
95.4
91.3
84.6
73.6
59.7
45.9
H.4 MQMTAR
We observe the same pattern across all tasks: despite theoretical extrapolation to 16xvia RoPE scaling, RoPE models poorly generalize beyond 4x. Moreover, although models with ALiBi can extrapolate up to 64x, ALiBi's limited span inevitably leads to performance degradation on very long sequences. However, NAPE provides substantial boost in all models. As in the copy task,
42

---

Published as a conference paper at ICLR 2026
Entmax alone underperforms Softmax, but adaptive scaling in ASEntmax makes the model superior, extending the generalization to an impressive 1024x. We also conducted experiments with ASSMax to demonstrate that, despite adaptive scaling benefits and improved performance in comparison with the model without scaling, softmax dispersion still causes a significant performance drop on very long sequences.
Table 9; Exact match accuracy (%) on Multi-query Multi-token Associative Recall.
ID
Out-of-Distribution
Model
64
128
236
512
1024
2048
4096
8192
16K
32K
65K
RoPE
Softmax
100.0
3.1
0
0
SSMAx
99.8
6.2
0
0
0
0
0
0
Entmax
99.8
49.4
0
ASEntmax
100.0
66.9
0.8
0
0
NoPE
Softmax
100.0
30.6
0
0
0
Entmax
100.0
26.1
0
SSMax
100.0
39.1
0
0
0
0
0
ASEntmax
100.0
58.3
0
ALiBi
Softmax
100.0
100.0
99.5
99.0
98.0
93.9
77.8
38.8
6.8
0
Entmax
99.5
97.5
90.6
75.4
44.7
11.7
1.1
0
NAPE
Softmax
100.0
100.0
100.0
99.4
99.5
98.2
97.8
90.2
80.2
34.2
3.0
SSMax
100.0
100.0
99.9
99.9
99.8
99.6
98.3
97.4
90.6
74.0
26.7
ASSMax
100.0
100.0
100.0
99.7
99.8
99.7
98.7
94.2
85.2
72.8
21.7
Top-K. K 16
100.0
100.0
99.9
41.8
5.8
0.3
0
0
0
Top-K, K, K
32
100.0
100.0
99.9
45.9
6.2
0.8
0
0
Entmax
100.0
100.0
100.0
99.9
99.2
97.8
92.7
86.2
66.8
35.8
9.3
ASEntmax
100.0
100.0
100.0
100.0
100.0
99.7
99.6
99.2
99.0
97.8
95.3
SB Attention
100.0
100.0
100.0
99.7
99.8
99.2
94.6
68.6
12.4
0.4
0
H.5 COPY
As shown by Jelassi al. (2023), transformers generalize to the Copy task especially with appropriate positional encodings. Table 10 shows that the Softmax transformer generalizes up to 64x the ID length. Notably, SSMax outperforms all other models which might suggest that scaling is crucial for this task.
As we can also see, without scaling, Q-entmax can hurt performance, leading to a noticeable drop in accuracy in the OOD scenario. Introducing adaptive temperature, however, substantially mitigates this effect: ASEntmax matches Softmax performance outperforming Entmax times of the sequence length. We hypothesize that Copy requires less sparse attention patterns, which can be accomplished by applying negative power to the logarithm function. We confirm this hypothesis in Figure 13, which shows that ASEntmax learns negative values of y in all heads, resulting in more spread-out attention distributions.
H.6 FLIP-FLOP
We first conducted an ablation study to evaluate model performance with various RoPE scaling factors (Table 11). Although the random baseline accuracy for Flip-Flop is 50%, our generative training setup with vocabulary of 7 tokens (4 main and 3 special) can yield accuracies below 50%. Therefore. we treat accuracies at or below 50% as poor and select scaling factor scaling factor defines the expansion multiple which the model must generalize. Throughout all
as optimal. The RoPE
43

---

Published as a conference paper at ICLR 2026
Head
Head
Head
Head
Head
Head
Head
Hoad 8
00%
2014
1.99
1.90
00%-
49%
20%
Figure 13: Distributions of y per head and layer for ASEntmax trained on Copy,
Table 10: Exact match accuracy (%) on Copy task.
ID
Model
64
128
256
512
1024
2048
4096
RoPE
Softmax
100.0
2.8
0
0
0
SSMax
100.0
0
0
0
0
ASSMax
99.9
19.9
0
0
0
Entmax
100.0
34.3
0
ASEntmax
100.0
5.3
0
0
0
0
NoPE
Softmax
56.3
0
0
SSMax
56.1
0
0
0
0
0
00
Entmax
34.6
0
0
0
0
ASEntmax
45.8
0
0
0
ALiBi
Softmax
100.0
99.8
99.8
98.8
98.3
93.9
26.8
Entmax
100.0
100.0
96.6
14.6
0.1
NAPE
Softmax
100.0
100.0
99.9
99.9
99.4
96.1
85.5
SSMax
100.0
100.0
100.0
99.9
99.6
99.3
95.8
ASSMax
99.9
99.8
99.7
99.3
97.5
91.1
72.8
Top-K, K 16
100.0
99.9
86.3
0.6
0.0
0.0
0.0
Top-K, K 32
100.0
99.7
96.8
26.7
0.0
0.0
0.0
Entmax
100.0
99.0
86.0
28.5
0.2
0
0
SEntmax
100.0
100.0
99.9
99.0
96.2
69.7
6.5
ASEntmax
100.0
100.0
99.9
99.7
99.4
96.3
86.6
SB Attention
100.0
100.0
100.0
99.6
98.8
48.6
0.0
experiments, however, we observe that RoPE models poorly generalize sequence lengths in-distribution length.
While Flip-Flop Flop considered a challenging task for testing length extrapolation (Liu al., 2023), we found that ALiBi and NAPE strategy almost perfectly solves both the sparse and dense variants. Surprisingly, RoPE models generalize better with the sparse variants.
H.7 REVERSE
From Table 14, which to our knowledge, represents the largest extrapolation reported. Moreover, RoPE models fail sequence length of 96. Although NAPE improves Softmax and SSMax, it does not enable
that ASEntmax with NAPE achieved impressive 8x length generalization
44

---

Published as a conference paper at ICLR 2026
Table 11: Exact match accuracy (%) for ablation of LLaMA RoPE scaling on Flip-Flop (sparse)
ID
cut-of-Distribution
Model
Factor
64
128
256
512
1024
2048
4096
Softmax
100.0
79.9
54.4
51.5
48.8
50.8
50.8
Softmax
100.0
99.6
33.8
11.2
3.3
0.5
0.0
Softmax
100.0
100.0
72.6
0.2
0
0
Softmax
16
100.0
99.9
97.3
36.7
0
Softmax
32
100.0
99.2
71.6
51.3
51.1
49.3
49.2
Table 12: Accuracy (%) on Flip-Flop (sparse).
ID
Out-of-Distribution
Model
64
128
256
512
1024
2048
4096
8192
16K
32K
RoPE
Softmax
100.0
99.9
97.2
36.6
0.0
0.0
0.0
SSMax
100.0
99.8
91.8
77.8
52.2
22.0
39.2
Entmax
100.0
99.9
89.0
64.0
50.6
50.6
55.1
ASEntmax
100.0
99.8
98.9
51.4
51.4
50.2
49.2
ALiBi
Softmax
100.0
99.9
99.8
99.9
99.9
100.0
99.7
99.7
99.9
99.7
Entmax
100.0
99.9
99.8
99.8
99.8
99.9
99.7
99.7
99.7
99.7
NAPE
Softmax
100.0
99.8
99.6
99.3
99.7
99.6
99.6
99.6
99.3
99.4
SSMax
100.0
99.8
99.9
99.8
99.8
99.7
99.8
100.0
99.9
99.6
Entmax
100.0
99.9
99.8
99.9
99.9
100.0
99.7
99.7
99.8
99.7
ASEntmax
100.0
99.8
99.6
99.3
99.5
99.3
99.5
99.7
99.6
99.5
Table 13: Accuracy (%) on Flip-Flop (dense).
ID
Out-of-Distriuttion
Model
64
128
256
512
1024
2048
4096
RoPE
Softmax
100.0
70.2
62.2
53.2
49.2
50.3
53.1
SSMax
100.0
69.1
60.4
53.2
48.5
51.0
53.1
Entmax
100.0
80.4
73.6
60.3
49.3
51.2
53.1
ASEntmax
100.0
100.0
100.0
49.6
48.9
51.1
53.1
NAPE
Softmax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
SSMAx
100.0
100.0
100.0
100.0
100.0
100.0
100.0
Entmax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
ASEntmax
100.0
100.0
100.0
100.0
100.0
100.0
100.0
generalization beyond 1.5x the in-distribution length; however, applying adaptive scaling Softmax (ASSMax) enables performance at 2 the in-distribution length.
H.8
SORT
Table 15 demonstrates superiority of a-entmax on Sort, with two-layer models generalizing almost perfectly 2x under both NoPE and NAPE configurations. Furthermore, Softmax models with NAPE experience a performance decline relative to their NoPE counterparts, and adaptive scaling degrades performance for both Softmax and a-entmax (we also report results for NoPE + SEntmax
45

---

Published as a conference paper at ICLR 2026
Table 14: Exact match accuracy (%) on Reverse.
ID
Model
64
96
128
256
512
RoPE
Softmax
100.0
0
0
0
0
SSMax
100.0
0
0
0
0
ASSMax
100.0
0
0
Entmax
100.0
0
0
ASEntmax
100.0
NoPE
Softmax
100.0
0
0
0
000
SSMax
100.0
0
0
Entmax
100.0
77.1
0
ASEtmax
100.0
74.4
0
ALiBi
Softmax
100.0
00
Entmax
100.0
96.1
78.5
NAPE
Softmax
100.0
36.0
0
0
0
SSMax
100.0
54.6
0
Top-K. K 16
100.0
99.4
83.3
0.0
0.0
Top-K, K
100.0
100.0
98.7
57.0
0
Top-K, K I
100.0
100.0
98.8
68.9
0
ASSMax
100.0
98.7
62.4
0
Entmax
100.0
99.0
86.0
28.5
0.2
SEntmax
100.0
100.0
98.1
51.4
0.0
ASEntmax
100.0
100.0
99.8
96.4
56.7
SB Attention
100.0
100.0
99.1
0.0
to be convinced). However, combining NAPE with adaptive scaling enhances a-entmax. This pattern suggest that sparsity, adaptive scaling, and NAPE can act complementarily.
H.9 LANGUAGE MODELING
We present comparison between RoPE and NAPE models standard benchmarks in Table 16, and on length generalization in Table 17. NAPE alone improves short-context results on LAMBADA, HellaSwag, and PIQA. Furthermore, length generalization capabilities only emerge when using either RoPE with ABF or NAPE. While SSMax, Entmax, and ASEntmax with RoPE + ABF achieve near- perfect generalization on the S-NIAH-1 task, NAPE combined with scalable models demonstrates more consistent performance across all sequence lengths in both S-NIAH-1 and S-NIAH-2. Perhaps more importantly, we find that RoPE with ABF leads to decline in-distribution performance. For instance, Lambada perplexity/accuracy drops, does accuracy HellaSwag, PIQA, ARC-C, and the in-distribution S-NIAH tasks. Moreover, both Softmax ABF and SSMax ABF fail dramatically on S-NIAH-2. Given that NAPE shows more consistent performance across different tasks and sequence lengths, we believe stands out as a more robust positional encoding approach than RoPE
I AI ASSISTANTS
We used Cursor during development, and ChatGPT during paper writing for correcting grammar and polishing sentences.
46

---

Published as a conference paper at ICLR 2026
Table 15: Exact match accuracy (%) on Sort.
ID
Out-of Distr--Dittiibutinn
Model
64
128
256
512
RoPE
Softmax
100.0
0
0
SSMax
100.0
0
0
0
Entmax
100.0
0
ASentmax
100.0
0
0
0
NoPE
Softmax
100.0
97.6
46.6
0
SSMax
100.0
96.3
29.8
0
Top-K, K = 16
99.7
48.0
0
0
Top-K,
32
100.0
92.5
Entmax
100.0
99.9
66.2
0
SEntmax
100.0
99.4
47.7
0
ASEnmtax
100.0
97.5
20.8
0
ALiBi
Softmax
99.9
0
Entmax
99.2
NAPE
Softmax
100.0
0
0
0
SSMax
100.0
0
0
ASSMax
100.0
99.5
9.4
0
Entmax
100.0
99.3
57.8
0
ASEntmax
100.0
100.0
79.7
0
SB Attention
100.0
36.5
0
Table 16: Downstream-tass results on short-context datasets encompassing comparison between RoPE and NAPE,
Method
Lambada (PPL)
Lambada
Hellaswag
PIQA
Are-C
WinoGrande
OpenbookQA
Softmax
59.7
30.3
32.5
64.3
26.3
50.8
28.6
SSMax
54.5
30.9
32.3
63.4
51.7
31.0
Entmax
53.0
31.1
32.9
62.9
25.8
49.9
27.8
RoPE
ASEntmax
56.5
29.9
32.2
62.4
25.5
52.3
27.8
+ ABF, Softmax
126.3
19.0
31.5
63.8
24.6
51.8
28.8
+ ABF, SSMax
96.8
21.6
31.9
63.2
24.7
51.1
30.2
+ ABF, Entmax
101.9
20.3
32,3
62.2
25.5
50.6
27.2
+ ABE, ASEntmax
99.0
20.0
29.4
62.7
24.1
52.0
29.2
Softmax
52.3
30.9
33.1
65.1
49.5
28.2
NAPE
SSMax
48.9
31.6
32,9
65.1
25.0
51.5
30.4
Entmax
47.9
32.1
32.8
63.6
24.6
50.9
29.0
ASEntmax
41.6
34.3
33.4
63.8
26.0
50.0
28.6
47

---

Published as a conference paper at ICLR 2026
Table 17: Retrieval performance on RULER benchmark
S-NTAH-1
S-NIAH-2
OOD
ID
OOD
Model
IK
2K
4K
8K
16K
IK
2K
4K
8K
Softmax
99.8
79.0
0.0
0.0
0.0
99.6
11.4
0.0
0.0
SSMax
99.0
83.0
0.0
0.0
0.0
99.6
53.6
0.0
0.0
Entmax
100.0
79.6
0.0
0.0
99.6
0.0
RoPE
ASEnntmax
99.8
87.2
0.0
0.0
0.0
99.0
83.6
0.0
0.0
ABF, Softmax
99.2
97.2
93.4
75.4
75.6
0.0
0.0
0.0
0.0
+ ABE, SSMax
98.2
98.0
97.6
97.4
98.0
4.4
+ ABF, Entmax
98.4
98.2
99.4
100.0
100.0
98.8
89.0
64.8
32.4
+ ABF, ASEntmax
100.0
99.6
99.6
99.6
94.0
98.6
83.6
30.8
7.2
Softmax
100.0
99.4
94.2
100.0
4.8
NAPE
SSMax
100.0
99.8
99.2
92.0
75.2
99.4
99.2
64.4
14.8
Entmax
99.8
99.8
89.0
21.6
1.2
99.6
99.4
64.8
7.2
ASEntmax
99.6
100.0
100.0
99.8
97.4
99.4
99.4
83.2
25.4
48


Reviews:
[{"review_id": "Gk9tEUNzDl", "reviewer": "Reviewer_xkyo", "summary": "The attention scores, even transformed by a highly distorting function (such as softmax), exhibit one unfortunate trait – they tend to have high entropy and get closer to uniform distribution, if the sequence length grows significantly. This work aims to analyze and offer a solution for  this problem by building on previous literature on a specific family of functions – Entmax – which have favorable properties, most notably sparseness. The authors propose a novel variation of this function, ASEntmax with learnable weights which allow regulating sparseness level with the increase of sequence length. \n\nThe scope of this work is essentially theoretical, as it examines long-range properties of softmax and Entmax family of transforms, successfully proving several and expanding our understanding of softmax limitations and methods to circumvent them. However, the paper also for the first time (as I’m aware of) empirically validates (AS)Entmax activations in causal language modeling on a scale of 420M parameter models.\n\nOverall, I believe this is a solid work and it should be accepted.", "questions": "- Can ASEntmax be implemented as an efficient Triton kernel, similarly to AdaSplash (https://openreview.net/forum?id=OWIPDWhUcO), without incurring numerical instabilities? \n\n- Will you open-source the model weights for all variants of Llama pre-trained on DCLM data?\n\n- Proposition 2.1 (Preserved representations) – can the opposite be proved for softmax? If not, the statement of this proposition should be altered to include softmax.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* The problem of length generalization of Transformers is well-motivated and relevant.\n\n* The paper is thorough and rigorous in math, although I didn’t check all proofs and derivations. \n\n* The conclusions derived from principled theoretical analysis, are also plausible intuitively and corroborated by evidence.\n\n* The empirical validation is convincing, and the results are promising. Specifically, Entmax and ASEntmax show superior out-of-distribution generalization capabilities as compared to Softmax in causal language modeling on two long-context datasets.", "weaknesses": "* It would be helpful to also compare different attention activation functions for language modeling perplexity as in Table 3, but with industry-standard RoPE embeddings. Appendix H.9. does not provide such comparisons, but it seems that the RoPE models are already pre-trained, and you would only need to test their perplexity on these long-context benchmarks. \n\n* The paper could be more self-contained. To fully comprehend the material, it is necessary to familiarize oneself with several previous works about Entmax. The exposition in Appendix A is scarce and for the most part repeats section 2.2. I would appreciate if definitions of Tsallis Entropy, q-expential functions, and a brief recap of sections 2.1-2.2 of Sparsemax paper (https://proceedings.mlr.press/v48/martins16.pdf) and sections 3.1-3.3 of original Entmax paper (https://aclanthology.org/P19-1146.pdf) could be provided.\n\n* The paper doesn’t establish which values of $\\alpha$ and the functional form of $\\tau(\\cdot)$ are used in the experiments.\n\n* There are no comparisons of speed of calculation between the 4 considered algorithms. An earlier work (https://proceedings.mlr.press/v48/martins16.pdf) alludes that run-time complexity for at least some of the Entmax variants is $O(L \\log L)$ instead of softmax’s $O(L)$, so it warrants a test for understanding how real-life throughput is impacted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CqhaVnaGro", "reviewer": "Reviewer_HXz2", "summary": "The paper introduces Adaptive-Scalable Entmax (ASEntmax), a sparse attention mechanism that improves long-context generalization by \"zeroing out\" attention (in an exact manner) on irrelevant tokens.\nOn multiple datasets and benchmarks, ASEntmax outperforms all baselines.", "questions": "- how is adaptive scalable entmax different from both scalable softmax and vanilla entmax ? is there a version of entmax that is only \"adaptive\" and not \"scalable\"? i understand entmax enables to turn of attention on some tokens below a certain threshold; scalable softmax amplifies attention if logits exceed by some amount compared to minimum logit; why can't vanilla entmax work by itself? why is the scaling solution proposed by SSMax required in this paper?\n\n- i didn't understand the motivation behind figure 3. why should scalers learned in figure 3 help in long-context generalization?\n\n- is QK-norm used before attention? can this fix some of the problems faced in long-context due to bounded logits?\n\n- is there any finetuning used before testing on RULER ?\n\n- why do the results vary a lot when using RoPE vs when using NAPE (e.g. in Appendix H) ?\n\nI am definitely willing to improve the score if these are addressed :)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- a very simple change to existing attention mechanism\n- works well across different datasets and tasks, showing consistent length generalization compared to baselines", "weaknesses": "- there are some understanding problems that i'm facing, refer to the questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2NLNJgRLar", "reviewer": "Reviewer_5G4R", "summary": "The paper makes attention genuinely sparse with α-entmax and then auto-tunes temperature by context length via ASEntmax so the model still locks onto key tokens in very long sequences. The theory shows softmax spreads probability mass as sequences grow while entmax keeps entropy low and prevents weights from vanishing on important positions. The experiments mix controlled synthetic tasks with language modeling and point to steadier long-context behavior without hurting short-context quality.", "questions": "1. What happens when the logit distribution departs from IID Gaussian or becomes extremely aligned, and can you show counterexamples and diagnostics. \n2. How should this combine with Top-k, chunk routing, or RAG pruning, and are there simple ordering/budgeting rules to follow. \n3. Can you provide an end-to-end compute profile and an early-exit strategy so systems can auto-tune ASEntmax intensity by sequence length.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper cleanly connects attention dispersion to representation collapse and over-compression with precise definitions and propositions.\n\n2. The method plays nicely with NAPE and existing fast-attention kernels like FlashAttention/AdaSplash, which makes it engineer-friendly. \n\nOverall,I think it is solid, but I'm not fully sure my evaluation is correct.", "weaknesses": "1. Entropy bounds rely on bounded-logit or near-Gaussian assumptions that may break on heavy-tailed or high-contrast distributions. \n2. Deployment-critical profiles for memory, throughput, and tail latency are thin, especially for large models and ultra-long inputs. \n3. Interplay with RoPE extensions, retrieval-style sparsity, or SSM hybrids is discussed but not systematically tested. \n4. Robustness beyond NAPE and specific hyperparameters is unclear, so portability to other position encodings or task families needs work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t4l3IEO5Ju", "reviewer": "Reviewer_7xk8", "summary": "This paper addresses the long-context generalization limitations of transformer architectures by replacing the traditional softmax attention with ASENTMAX, a differentiable sparse transformation that assigns exact zero attention to irrelevant tokens, achieving up to 1000× length extrapolation and superior long-context perplexity.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper provides theoretical analyses to address core limitations of softmax attention. It formally proves that alpha-entmax maintains non-vanishing attention probabilities.\n2. ASEntmax introduces head-specific, learnable scaling parameters to adjust logits dynamically. \n3. The empirical evaluation is thorough and diverse, covering synthetic tasks (retrieval-focused, memory-dependent, ordering) and real-world language modeling.", "weaknesses": "1. It is unclear whether ASEntmax can be integrated with FlashAttention-2, a widely adopted library for efficient softmax attention. Given FlashAttention’s popularity in industrial and academic settings, clarifying ASEntmax’s compatibility with FlashAttention (e.g., modifications needed for sparse logit scaling, runtime overhead comparisons) is critical for its practical adoption.\n2. The paper does not explicitly confirm whether ASEntmax inherits the theoretical properties of alpha-entmax in Section 3. For instance, does the learnable temperature parameter disrupt the thresholding effect of ω-entmax that ensures zero attention to irrelevant tokens?\n3. The theoretical motivation for ASEntmax’s scaling mechanism (Section 4.1) relies on offsetting the range growth of IID Gaussian logits. However, real-world attention logits are unlikely to follow IID Gaussian distributions—they are influenced by semantic relevance, positional biases, and layer-wise transformations. \n4. Proposition 2 claims that alpha-entmax alleviates over-squashing “via stronger gradient signals,” but the reasoning is incomplete. \n5. In the MQMTAR task, ASEntmax achieves 1024× length extrapolation, but the paper relies on standard positional encodings (e.g., NAPE, ALiBi). Existing literature shows that many positional encodings struggle with extreme extrapolation due to degraded positional awareness. The paper does not analyze whether NAPE alone can support 1024× extrapolation or if its performance is entirely dependent on ASEntmax. \n6. Table 4 shows ASEntmax achieves near-perfect performance (97.4%) on S-NIAH-1 (OOD, 16K tokens) but only 25.4% on S-NIAH-2 (OOD, 8K tokens). The paper does not explain this stark disparity. Is S-NIAH-2 is presumably a harder task？\n7. The related works should include a discussion about the broader landscape of sparse attention mechanisms designed for long-context modeling (e.g., [1-7]).\n---\n[1] Efficient streaming language models with attention sinks. ICLR 2024. \n\n[2] MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention. NeurIPS 2024. \n\n[3] Core Context Aware Transformers for Long Context Language Modeling. ICML 2025. \n\n[4] XAttention: Block Sparse Attention with Antidiagonal Scoring. ICML 2025. \n\n[5] FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference. ICLR 2025. \n\n[6] Curse of High Dimensionality Issue in Transformer for Long Context Modeling. ICML 2025. \n\n[7] MMInference: Accelerating Pre-filling for Long-Context VLMs via Modality-Aware Permutation Sparse Attention. ICML 2025", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)