ASEntmax: Adaptive-Scalable Entmax

• ASEntmax is an attention mechanism that extends α-entmax by learning a data- and length-adaptive temperature scale to bridge sparse and dense attention.
• It leverages a closed-form Jacobian for efficient gradient computation and robust training even with extremely long sequences.
• Empirical results demonstrate ASEntmax’s superior out-of-distribution generalization and pattern-focused performance compared to softmax and fixed-α entmax.

Adaptive-Scalable Entmax (ASEntmax) is an attention mechanism for transformer-based architectures that extends α-entmax by introducing a learnable, data- and length-adaptive temperature scale. This parametric scaling allows attention distributions to interpolate between sparse (pattern-focused) and dense (softmax-like) regimes, thereby overcoming the dispersion and representational collapse observed when traditional softmax is applied to very long contexts. ASEntmax has demonstrated strong out-of-distribution generalization on a suite of synthetic algorithmic tasks, highlighting its ability to retain precise focus even as sequence lengths grow by orders of magnitude (Vasylenko et al., 19 Jun 2025).

1. Mathematical Formulation

ASEntmax generalizes the α-entmax transformation by learning a temperature (scale) parameter as a function of sequence length. Given attention logits $z \in \mathbb{R}^n$, α-entmax is defined as the solution to a Tsallis-entropy-regularized maximization problem:

$$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$

where $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$. For $\alpha > 1$, the mapping has the closed form:

$$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$

with $\tau(z)$ chosen so the distribution sums to 1.

ASEntmax introduces a per-head, per-query scale $s$, parameterized as:

$$s = \delta + \beta(\log n)^{\gamma}, \quad \delta, \beta, \gamma \in \mathbb{R}$$

with $\beta \geq 0$ (enforced via softplus), and $\gamma$ typically constrained with tanh. Then, the ASEntmax transformation becomes:

$$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$0

or equivalently,

$$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$1

The scalar $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$2 is computed so that the vector $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$3 is properly normalized, using a sorting-based root-finder with complexity $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$4 or $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$5 for standard α-entmax implementations.

2. Gradient Computation and Optimization

Gradient backpropagation through ASEntmax leverages the closed-form Jacobian for α-entmax. For a scalar loss $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$6:

• The gradient with respect to the logits $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$7 is

$$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$8

where $$\operatorname{aentmax}(z) = \arg\max_{p\in\Delta_n} \left\{ p^\top z + H_\alpha(p) \right\},\quad H_\alpha(p) = \frac{1}{\alpha(\alpha-1)} \sum_{i=1}^n (p_i - p_i^\alpha),$$9, for $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$0, is constructed over the active support $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$1:

$\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$2

and zero otherwise.

• Gradients with respect to parameters $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$3 and $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$4 are computed recursively via the scaling function’s derivatives, with

$\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$5

and for $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$6 an analogous term involving $\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$7.

In practice, only components corresponding to the nonzero support in the sparse distribution need to be computed, ensuring efficient backpropagation and avoiding redundancy.

3. Behavioral Regimes and Theoretical Properties

ASEntmax's temperature scaling provides explicit interpolation between sparse and dense attention:

• High scale ($\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$8): Sharpens the logits, increasing sparsity; yields output distributions approaching sparsemax (α=2) or even one-hot argmax in the extremal case.
• Low scale ($\Delta_n = \{ p \geq 0, \sum_{i} p_i = 1 \}$9): Collapses the logits; the solution approaches $\alpha > 1$0, i.e., a maximally dense uniform distribution over all tokens.

Theoretical results establish that, for a “two-level” logit structure (with $\alpha > 1$1 tokens at $\alpha > 1$2 and the rest at $\alpha > 1$3), α-entmax concentrates probability exactly and exclusively on the $\alpha > 1$4 relevant tokens with $\alpha > 1$5 assignments, provided that $\alpha > 1$6, independent of total sequence length $\alpha > 1$7. In contrast, softmax with a fixed temperature always disperses the distribution as $\alpha > 1$8 grows, requiring the temperature to scale with $\alpha > 1$9 to maintain any fixed concentration, which ASEntmax achieves automatically by selecting the scale as a learnable function of $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$0 [(Vasylenko et al., 19 Jun 2025), Remark 3.1]. This property ensures non-dispersive and pattern-focused attention distributions even for very long contexts.

4. Implementation and Practical Considerations

Parameterization employs:

• $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$1 produced via a small linear layer plus softplus ($$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$2),
• $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$3 via a linear layer and tanh (bounded within a finite range),
• $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$4 either fixed (commonly $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$5) or freely learned.

Root-finding for $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$6 utilizes the robust $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$7 sorting+cumsum kernel for $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$8-entmax, stable for $$p_i = \left[(\alpha-1) z_i - \tau(z)\right]_+^{1/(\alpha-1)}, \quad \sum_i p_i = 1,$$9 and sequence lengths up to millions.

Computational complexity:

Method	Complexity	Additional Remarks
Softmax	$\tau(z)$0	Dense exp+sum
α-entmax	$\tau(z)$1	Often $\tau(z)$2 with selection
ASEntmax	$\tau(z)$3	Only $\tau(z)$4 cost per query for $\tau(z)$5

GPU-efficient sparse attention is possible via custom kernels (e.g., AdaSplash by Gonçalves et al. 2025). ASEntmax adds minimal per-query compute compared to α-entmax.

Compared to fixed-α entmax and softmax:

• Fixed-α entmax may become over-sparse as $\tau(z)$6 increases, resulting in under-representation for moderately dense tasks (e.g. Copy).
• Softmax with learned global temperature improves upon naive approaches but still fails to maintain focus as $\tau(z)$7 diverges.
• ASEntmax assigns per-head, per-query scaling, customizing the sparsity schedule to both data and length.

5. Interaction with Positional Encodings

The interplay between positional encoding schemes and sparse attention is pronounced in ASEntmax-equipped models:

• NoPE (no positional encoding): With softmax, inevitably develops a “first-token” bias as transformer depth increases (cf. Wu et al. 2025), which ASEntmax can disrupt via attention pruning. Pure NoPE, however, disables recency cues.
• ALiBi (linear distance bias): With ASEntmax, induces a hard window: beyond a certain head-specific cutoff $\tau(z)$8, logits fall below the threshold and are thus pruned entirely ($\tau(z)$9). The cutoff is determined by the maximal/minimal logits and the linear bias slope.
• RoPE (rotary): Generates frequency-dependent sparse attention windows. High-frequency heads focus local, low-frequency heads attend globally, and periodic thresholding produces overlapping sparse windows.
• NAPE (“NoPE + ALiBi” hybrid): Splits attention heads such that half use ALiBi, enforcing recency bias, while half use NoPE, enabling long-range, content-driven focus. NAPE yields the highest long-context generalization in tandem with ASEntmax.

6. Empirical Results and Performance

On synthetic algorithmic tasks (Copy, Reverse, Sort, Max Retrieval, 2Back, Local Count, MQMTAR), with models trained on lengths up to 64 and tested to 65,000, ASEntmax combined with NAPE achieves the highest out-of-distribution accuracy. Typical transformer sizes are 2–6 decoder layers, 8–16 heads, 2–10 million parameters, trained on 10–50 million instances.

Select results:

Task	Test Length (vs. Train)	ID Accuracy	ASEntmax Accuracy	Softmax Accuracy	Sparsemax Accuracy
MQMTAR	256×	100%	96.4%	80.2%	66.8%
MQMTAR	1024×	—	76.7%	—	—
Reverse	4×	100%	96.4%	0%	—
Copy	64×	—	86.6%	0%	—
Sort	4×	—	79.7%	0%	—
Max Retrieval	$s$0	—	$s$1	$s$2	—

These findings corroborate the hypothesis that (a) sparsity prevents attention mass dispersion; (b) adaptive scaling via $s$3 counteracts oversparsification in extreme lengths; (c) hybrid positional encodings optimize representational fidelity across tasks.

7. Significance and Research Context

ASEntmax is positioned as a robust, theoretically justified extension of α-entmax. By providing learnable, length-adaptive scaling, it overcomes the representational collapse inherent to softmax and the over-sparsification of fixed-parameter sparse attention in long contexts. Its per-head, per-query adaptivity supports optimal allocation of attention mass, preserving pattern focus hundreds to thousands of times beyond training lengths. When paired with the NAPE positional encoding scheme, ASEntmax achieves state-of-the-art out-of-distribution generalization in tasks that stress fixed-size pattern retrieval and manipulation in very long sequences (Vasylenko et al., 19 Jun 2025).

A plausible implication is that the principles embodied in ASEntmax may generalize to new architectures and application domains requiring flexible attention sparsity and robust non-dispersive representations over growing input lengths.