α‑entmax on the simplex

Project score vectors onto the probability simplex — or compare transformations as surfaces over the input logits.

1.00
1 · softmax1.52 · sparsemax∞ · argmax
fewermore
avg. support
% with a zero

Drag to rotate in 3‑D · toggle the simplex to reveal the mapping

Panel A
Panel B
α1.50

Two transformations side by side, as a function of two input logits \((z_1,z_2)\), with the third logit fixed at \(0\); the height \(p_0\) is the probability of that fixed key. Softmax is a smooth ramp that never reaches 0 or 1; α‑entmax has a sharp cliff and a flat‑zero region; top‑k softmax has a non‑differentiable fold. Drag a 3‑D plot to rotate both.

What is the probability simplex?

An attention distribution over \(n\) keys is a vector \(\boldsymbol p\) whose entries are non‑negative and sum to one. The set of all such vectors is the probability simplex:

For \(n=3\) it is exactly the triangle in the Simplex tab, with the three corners being the one‑hot vectors \(\boldsymbol e_1,\boldsymbol e_2,\boldsymbol e_3\).

From scores to a distribution

A transformation turns a score vector \(\boldsymbol z\in\mathbb R^n\) into a point \(\boldsymbol p\in\triangle^n\). The standard choice in transformers is softmax, which is dense — every entry is strictly positive, so it always lands in the interior:

Mathematical definition of α‑entmax

The \(\alpha\)-entmax transformation of a score vector \(\boldsymbol z \in \mathbb R^n\) is defined as a regularized projection onto the simplex:

where \(H_\alpha(\boldsymbol p)\) is the Tsallis(\(\alpha\)) entropy. The closed form for \(\alpha\text{-entmax}\) with \(\alpha>1\) is a thresholded power:

Here \(\tau(\boldsymbol z)\) is chosen so that \(\boldsymbol p^\star\) sums to 1. The dial \(\alpha\) interpolates a whole family: \(\alpha=1\) is softmax (dense, interior); \(\alpha=2\) is sparsemax, the plain Euclidean projection (zeros, on the edges); \(\alpha\to\infty\) is argmax (one‑hot, on the corners). The sweet spot for attention is in between — e.g. \(\alpha=1.5\).

Ben Peters, Vlad Niculae, and André F. T. Martins. 2019. Sparse Sequence‑to‑Sequence Models. ACL 2019. arXiv:1905.05702

← α‑entmax viz hub  ·  α‑entmax → AdaSplash explainer