α-Entmax: Sparse Generalization of Softmax
- α-entmax is a differentiable sparse transformation that converts logits into probability distributions with exact zeros for low-scoring entries.
- It generalizes softmax (α→1) and recovers sparsemax at α=2, offering a tunable, thresholded power-law transformation that enhances model focus.
- Despite computational bottlenecks in threshold determination, its sparse mechanism benefits applications in language modeling, long-context attention, and hierarchical routing.
Searching arXiv for primary and papers on α-entmax to ground the article in the relevant literature. Searching arXiv for "alpha entmax". -Entmax is a differentiable sparse transformation from logits to a probability distribution on the simplex that generalizes softmax by allowing for sparse probability distributions. In the form used across modern neural sequence models and attention mechanisms, it is written as a thresholded power transformation,
where is chosen so that the probabilities sum to one, and coordinates below threshold are mapped to exactly zero. As , it reduces to softmax; at , it yields sparsemax; and for it gives a smoothed sparse transform. The operator is especially relevant when dense softmax tails are undesirable, including large-vocabulary prediction, long-context attention, and dense open-vocabulary vision, but its threshold computation has historically been a central computational bottleneck (Tezekbayev et al., 2021).
1. Definition and variational form
-Entmax is commonly presented in two equivalent ways. The first is as a closed-form thresholded map:
with and chosen so that
0
This makes 1-entmax a probability distribution, but unlike softmax, many coordinates can be exactly zero (Tezekbayev et al., 2021).
The second is as the solution to an entropy-regularized optimization problem using Tsallis 2-entropy: 3 with
4
This formulation makes explicit that 5-entmax is the sparse analogue of the usual softmax-as-entropy-regularized argmax perspective (Zohra et al., 8 Jul 2026).
The interpolation parameter 6 controls sparsity. The literature summarized here states three standard anchor points: 7 corresponds to softmax-like behavior, 8 yields sparsemax, and values in 9 give a smoothed sparse transform (Tezekbayev et al., 2021). The practical consequence of the thresholded form is equally standard: a token receives nonzero mass iff 0, so low-scoring entries can be discarded exactly rather than merely downweighted (Vasylenko et al., 19 Jun 2025).
This exact-zero behavior is the defining distinction from softmax. Softmax is strictly positive everywhere, so every class or token gets nonzero mass. 1-Entmax instead preserves simplex normalization while enabling finite support. In applications, this means that probability mass can be concentrated on the most relevant classes, and irrelevant classes can be assigned exactly zero probability (Tezekbayev et al., 2021).
2. Information-geometric and entropy-theoretic background
The mathematical substrate underlying 2-entmax predates modern neural attention. A foundational antecedent is the theory of relative 3-entropy, denoted 4, which generalizes the usual relative entropy and yields power-law rather than exponential structure in projection problems (Kumar et al., 2014).
In that framework, 5 is defined using escort distributions. If 6 with densities 7, and 8, then
9
With 0 and
1
the generalized relative entropy is
2
where 3 is Csiszár’s 4-divergence. The paper explicitly proves that
5
so 6 is a genuine generalization of KL divergence (Kumar et al., 2014).
For projection onto a linear family,
7
the forward 8-projection 9 of 0 has density 1 satisfying the power-law relation
2
with 3 on 4. This is not exponential-family-like but a power-law family member, and the thresholded support condition 5 on a set 6 already exhibits the sparse / thresholded behavior that later becomes central in sparsemax and 7-entmax (Kumar et al., 2014).
The same theory establishes a Pythagorean property: 8 for projection of 9 onto convex 0. It also connects minimization of 1 against the uniform distribution with maximum Rényi entropy, and therefore also with the Tsallis maximum entropy principle because Tsallis entropy is a monotone transform of Rényi entropy (Kumar et al., 2014).
This suggests a clear conceptual lineage: modern 2-entmax inherits a power-law geometry in place of the exponential-family geometry associated with KL and softmax. The earlier theory does not discuss neural attention, but it shows that minimizing a generalized entropy divergence naturally yields sparse, thresholded, power-law distributions, precisely the structural phenomenon exploited by 3-entmax (Kumar et al., 2014).
3. Sparsity mechanism and computational bottlenecks
The principal attraction of 4-entmax is that it retains normalization on the simplex while introducing an input-dependent threshold. Because
5
the operator acts as a content-aware selector rather than a merely sharper dense normalizer (Zohra et al., 8 Jul 2026).
This threshold, however, is also the main source of computational overhead. Unlike softmax, where normalization is a simple log-sum-exp over all logits and can be implemented in highly optimized CUDA kernels, 6-entmax requires identifying the active support and solving for the threshold that makes the outputs sum to one. In the original formulation, computing 7 involves sorting the logits. Sorting over a large vocabulary is expensive, especially in NLP settings with tens of thousands of classes. That is why standard 8-entmax is significantly slower than softmax in both training and inference (Tezekbayev et al., 2021).
For general 9, the threshold is generally found by bisection. Special cases admit more structured solvers. The sparsemax case 0 uses
1
with the threshold found by sorting logits 2 and finding the largest active set size 3 such that
4
then
5
For 6, an exact sort-based algorithm exists because the exponent becomes an integer (Zohra et al., 8 Jul 2026).
The computational problem can also be posed as root finding. One must solve
7
where 8 is continuous and strictly decreasing, so the root is unique. This is substantially less hardware-friendly than softmax’s additive normalizer
9
which can be accumulated online in one pass (Gonçalves et al., 16 Apr 2026).
A common misconception is that sparsification benefits arise only from sharpening or temperature scaling. In dense open-vocabulary CLIP, an ablation reported that a matched-support random mask and temperature scaling both underperform entmax, showing that which entries are removed matters, not just how many. The benefit is attributed to a data-dependent cutoff rather than to mere saturation (Zohra et al., 8 Jul 2026).
4. Output-layer use, losses, and the 0-ReLU alternative
In large-vocabulary generation, 1-entmax is motivated by a limitation of softmax: by producing a dense probability distribution each token in the vocabulary has a nonzero chance of being selected at each generation step, leading to a variety of reported problems in text generation. The key benefit of 2-entmax is that probability mass is concentrated on the most relevant classes, and irrelevant classes can be assigned exactly zero probability, which is attractive for language modeling and machine translation (Tezekbayev et al., 2021).
A practical complication is that normalized entmax is slower than optimized softmax. One response is the output-layer substitute 3-ReLU, an unnormalized sparse transform defined as
4
where now 5 is a constant, not a function of 6. The associated loss is
7
with
8
The paper proves the gradient identity
9
which mirrors the familiar “prediction minus target” structure (Tezekbayev et al., 2021).
The same work emphasizes that 0-ReLU is a “shifted ReLU raised to a power”: for 1, 2-ReLU reduces to standard ReLU when 3, and for 4, it becomes a smoothed version of ReLU. The Jacobian is diagonal,
5
which enables the cancellation with the derivative of Tsallis entropy that yields the gradient identity (Tezekbayev et al., 2021).
Empirically, the method was evaluated on Transformer-based NMT across IWSLT’14 De6En, WMT’14 En7De, and WMT’13 En8Ru. The main BLEU results reported were softmax: 9, sparsemax: 0, 1.5-entmax: 1, 1.5-entmax (top-100): 2, and 1.5-ReLU: 3. The same study also reported that softmax assigns the empty translation a much higher probability too often, whereas sparsemax, entmax, and 4-ReLU drastically reduce this issue (Tezekbayev et al., 2021).
A notable empirical difference is that 5-ReLU is less sparse than 1.5-entmax, yet this lower sparsity correlates with slightly better BLEU in these experiments. The paper argues that 6-entmax’s threshold 7 grows during training, increasing sparsity, whereas 8-ReLU keeps 9 fixed (Tezekbayev et al., 2021).
5. Attention mechanisms and long-context modeling
Within transformer attention, 00-entmax is used by replacing the row-wise softmax normalization with a sparse simplex-preserving alternative. In a causal transformer, attention weights are formed from logits
01
and normalized by a map 02: 03 Standard transformers use 04; sparse variants replace this with 05-entmax (Vasylenko et al., 19 Jun 2025).
The long-context rationale is that softmax always assigns positive probability to every token, so as sequence length increases, non-informative tokens accumulate attention probability mass, leading to dispersion and representational collapse. Sparse attention mechanisms using 06-entmax can avoid these issues, due to their ability to assign exact zeros to irrelevant tokens (Vasylenko et al., 19 Jun 2025).
A central result summarized for long-context generalization is that 07-entmax can keep entropy at
08
where 09 is the number of tokens with nonzero probability, instead of softmax’s 10. This matters for tasks that depend on fixed-size patterns or local structures, including associative recall, copy/retrieval, reverse, local count, and 2-back, because sparse attention can remain focused on a small support set regardless of sequence length (Vasylenko et al., 19 Jun 2025).
The same paper introduces Adaptive-Scalable Entmax (ASEntmax),
11
with head-specific, learnable scaling. It recovers standard 12-entmax when 13, and can interpolate between sharper and smoother regimes by adapting sparsity per head and query (Vasylenko et al., 19 Jun 2025).
Positional encoding interacts strongly with 14-entmax. With ALiBi,
15
the paper’s key point is that 16-entmax converts ALiBi’s smooth recency bias into a hard attention window. With RoPE, thresholding yields multi-scale sparse attention rather than a single distance cutoff. The best reported results generally came from combining NAPE with ASEntmax (Vasylenko et al., 19 Jun 2025).
A related line uses 17-entmax in CLIP’s final visual self-attention layers as a training-free inference-time substitution for row-wise softmax. The stated motivation is that dense softmax attention spreads attention across many low-salience tokens, producing noise that obscures the fine-grained, spatially localized cues required for dense, open-vocabulary prediction. Entmax is positioned there as an implicit denoiser that zeroes contextually irrelevant dependencies while redistributing mass onto the most relevant tokens (Zohra et al., 8 Jul 2026).
6. Efficient kernels, decoding, and hierarchical routing
Because the main obstacle to practical deployment is threshold computation, several later works focus on making 18-entmax hardware-efficient. AdaSplash-2 addresses this by introducing a histogram-based initialization that reduces the number of iterations needed to compute 19 to typically 1–2. The scores are re-centered and rescaled so that the threshold search lies in 20, a coarse histogram is computed on the fly and stored in on-chip SRAM, and the resulting histogram root 21 satisfies
22
The paper states that AdaSplash-2 matches or improves per-step training time relative to FlashAttention-2 when block sparsity is moderate-to-high, e.g. 23, and that models trained with efficient 24-entmax attention match softmax baselines at short-context lengths and achieve substantial gains in long-context settings (Gonçalves et al., 16 Apr 2026).
In autoregressive decoding, exact sparsity changes the nature of the approximation problem. EntmaxKV states the key implication as
25
where 26 is the entmax support and 27 is the kept set. With softmax, every token gets positive mass, so any truncation discards probability mass; with entmax, sparse decoding is exact once support is recovered. The paper also proves the output error bound
28
and reports up to 29 speedup over full-cache softmax FlashDecoding and 30 speedup over full entmax at 1M context length (Duarte et al., 20 May 2026).
In hierarchical attention, DashAttention uses 31-entmax in Stage 1 to decide which KV chunks to keep before Stage 2 token-level softmax attention. Chunk-level routing logits
32
are transformed as
33
with support
34
The paper emphasizes that, unlike top-35, entmax gives exact zeros while keeping the support adaptive and differentiable, so the number of selected chunks can vary by query, head, and layer. It also states that DashAttention is non-dispersive and achieves comparable accuracy as full attention with 75% sparsity (Huang et al., 18 May 2026).
These systems works make a common point: 36-entmax is not only a sparsifying normalizer but also an enabling mechanism for support-aware algorithms. This suggests that once support recovery or support-aware routing becomes the main design objective, exact zeros change both algorithmic structure and attainable efficiency.
7. Practical trade-offs and recurrent points of interpretation
Across the literature, the main practical trade-off is between normalized sparsity and computational cost. Standard 37-entmax preserves the probability simplex and yields input-dependent support, but it requires threshold computation that is slower than optimized softmax. This is why alternatives such as 38-ReLU, approximate solvers, histogram-based kernels, and support-aware page selection appear repeatedly (Tezekbayev et al., 2021).
A second recurrent point is that sparsity is not uniformly beneficial. In CLIP dense prediction, the gain from attention sparsification is proportional to how much baseline attention spreads off the target class; if the baseline attention is already concentrated on the right tokens, entmax has little to fix and may even remove useful mass. Likewise, in long-context learning, fixed-temperature 39-entmax can be too sparse for some tasks and too diffuse for others, which motivates adaptive scaling in ASEntmax (Zohra et al., 8 Jul 2026).
A third point concerns differentiability relative to hard selection. The later sparse-attention systems repeatedly distinguish 40-entmax from top-41: top-42 is sparse but discontinuous and non-differentiable at selection boundaries, whereas 43-entmax remains differentiable and induces exact zeros smoothly. This makes it especially suitable for hierarchical or end-to-end trainable sparse mechanisms (Gonçalves et al., 16 Apr 2026).
A plausible implication is that 44-entmax occupies a technically distinctive middle ground. It is denser and easier to optimize than purely hard routing, but structurally more selective than softmax because it replaces dense tails with thresholded power-law support. That combination explains its continued use in machine translation output layers, long-context transformers, sparse decoding, hierarchical attention, and dense open-vocabulary vision (Vasylenko et al., 19 Jun 2025).