The Road to Probability Distributions: A Survey of Softmax and Its Alternatives
Whether in basic classification tasks or in the now-ubiquitous attention mechanism, constructing a probability distribution is a key step. Concretely, this means converting an arbitrary $n$-dimensional vector into a discrete probability distribution over $n$ elements. As we all know, the standard answer to this problem is Softmax, which takes the form of exponential normalization; it's relatively simple and intuitive, and it comes with many nice properties, which is why it has become the "default choice" in most scenarios.
That said, Softmax also has some shortcomings in certain scenarios—for instance, it isn't sparse enough and can never output an exact zero—which is why many alternatives have emerged. In this post, we'll briefly summarize the relevant properties of Softmax, and then survey and compare some of its alternatives.
Recap of Softmax
First, let's introduce some general notation: $\boldsymbol{x} = (x_1,x_2,\cdots,x_n)\in\mathbb{R}^n$ is the $n$-dimensional vector that needs to be converted into a probability distribution; its components can be positive or negative, with no fixed bounds. Let $\Delta^{n-1}$ denote the set of all discrete probability distributions over $n$ elements, i.e.
\begin{equation}\Delta^{n-1} = \left\{\boldsymbol{p}=(p_1,p_2,\cdots,p_n)\left|\, p_1,p_2,\cdots,p_n\geq 0,\sum_{i=1}^n p_i = 1\right.\right\}\end{equation}
The reason we write $n-1$ rather than $n$ is that the constraint $\sum\limits_{i=1}^n p_i = 1$ defines a $n-1$-dimensional subplane within the $n$-dimensional space, and once we add the constraint $p_i\geq 0$, the set $(p_1,p_2,\cdots,p_n)$ is only a subset of that plane—so its actual dimensionality is only $n-1$.more
With this notation in hand, the theme of this post can be simply stated as an exploration of the mapping $\mathbb{R}^n\mapsto\Delta^{n-1}$, where $\boldsymbol{x}\in\mathbb{R}^n$ is what we conventionally call logits, or scores.
Basic definition
The definition of Softmax is simple:
\begin{equation}p_i = softmax(\boldsymbol{x})_i = \frac{e^{x_i}}{\sum\limits_{j=1}^n e^{x_j}}\end{equation}
There are far too many origins and interpretations of Softmax—energy models, statistical mechanics, or simply as a smooth approximation of $\text{argmax}$—so it's hard to pin down its earliest source, and we won't attempt to do so here. Often we also introduce a temperature parameter, i.e. we consider $softmax(\boldsymbol{x}/\tau)$, but $\tau$ itself can be absorbed into the definition of $\boldsymbol{x}$, so here we won't separate out the $\tau$ parameter specially.
The denominator of Softmax is usually denoted $Z(\boldsymbol{x})$, and its logarithm is exactly the $\text{logsumexp}$ operation built into most deep learning frameworks, which is a smooth approximation of $\max$:
\begin{align}\log Z(\boldsymbol{x}) = \log \sum\limits_{j=1}^n e^{x_j} = \text{logsumexp}(\boldsymbol{x})\\ \lim_{\tau\to 0^+} \tau\,\text{logsumexp}(\boldsymbol{x}/\tau) = \max(\boldsymbol{x})\end{align}
When $\tau$ is taken to be $1$, we can write $\text{logsumexp}(\boldsymbol{x}) \approx \max(\boldsymbol{x})$, and the larger the variance of $\boldsymbol{x}$, the better the approximation. For further discussion see Seeking a Smooth Maximum Function.
Two properties
Besides converting an arbitrary vector into a probability distribution, Softmax also satisfies two properties:
\begin{align}&{\color{red}{monotonicity}}:\quad p_i > p_j \Leftrightarrow x_i > x_j,\quad p_i = p_j \Leftrightarrow x_i = x_j \\[5pt] &{\color{red}{invariance}}:\quad softmax(\boldsymbol{x}) = softmax(\boldsymbol{x} + c),\,\,\forall c\in\mathbb{R} \end{align}
Monotonicity means that Softmax is order-preserving: the maximum/minimum of $\boldsymbol{x}$ corresponds to the maximum/minimum of $\boldsymbol{p}$. Invariance means that if every component of $\boldsymbol{x}$ is shifted by the same constant, the result of Softmax doesn't change—this is the same property that $\text{argmax}$ has, i.e. we likewise have $\text{argmax}(\boldsymbol{x}) = \text{argmax}(\boldsymbol{x} + c)$.
Based on these two properties, we can conclude that Softmax is in fact a smooth approximation of $\text{argmax}$ (more precisely, a smooth approximation of $\text{onehot}(\text{argmax}(\cdot))$); more specifically, we have
\begin{equation}\lim_{\tau\to 0^+} softmax(\boldsymbol{x}/\tau) = \text{onehot}(\text{argmax}(\boldsymbol{x}))\end{equation}
This is presumably where the name "Softmax" comes from. Note that we shouldn't confuse things here: Softmax is a smooth approximation of $\text{argmax}$, not of $\max$—the smooth approximation of $\max$ is instead $\text{logsumexp}$.
Gradient computation
For deep learning, one of the important ways to understand the properties of a function is to understand its gradient. For Softmax, we computed this in Viewing Attention's Scale Operation Through Gradient Maximization:
\begin{equation}\frac{\partial p_i}{\partial x_j} = p_i\delta_{i,j} - p_i p_j = \left\{\begin{aligned} p_i - p_i^2,&\quad i=j\\ - p_i p_j,&\quad i\neq j \end{aligned}\right.\end{equation}
The matrix formed by arranging these values is also known as the Jacobian matrix of Softmax, and its L1 norm has a simple form:
\begin{equation}\frac{1}{2}\left\Vert\frac{\partial \boldsymbol{p}}{\partial \boldsymbol{x}}\right\Vert_1=\frac{1}{2}\sum_{i,j}\left|\frac{\partial p_i}{\partial x_j}\right|=\frac{1}{2}\sum_i (p_i - p_i^2) + \frac{1}{2}\sum_{i\neq j} p_i p_j = 1 - \sum_i p_i^2\end{equation}
When $\boldsymbol{p}$ is a one-hot distribution, the above expression equals zero, which means that the closer the Softmax output is to one-hot, the more severe the vanishing gradient phenomenon becomes. So, at least during initialization, we should not initialize Softmax to be close to one-hot. The rightmost side of the expression above also connects to the concept of Rényi entropy, which is similar to the familiar Shannon entropy.
Reference implementation
A direct implementation of Softmax is simple: just take $\exp$ and normalize. The reference code in Numpy is:
def softmax(x):
y = np.exp(x)
return y / y.sum()
However, if $\boldsymbol{x}$ contains large components, computing $\exp$ can easily overflow. So we usually make use of Softmax's invariance property, first subtracting the maximum of all components from each component, and then computing Softmax—this ensures that every component fed into $\exp$ is at most 0, guaranteeing no overflow:
def softmax(x):
y = np.exp(x - x.max())
return y / y.sum()
Loss functions
One of the main uses of constructing a probability distribution is for the output of single-label multi-class classification tasks. That is, suppose we have an $n$-class classification task, and $\boldsymbol{x}$ is the model's output; we hope to predict the probability of each class via $\boldsymbol{p}=softmax(\boldsymbol{x})$. To train this model, we need a loss function. Assuming the target class is $t$, a common choice is the cross-entropy loss:
\begin{equation}\mathcal{L}_t = - \log p_t = - \log softmax(\boldsymbol{x})_t\end{equation}
We can compute its gradient:
\begin{equation}-\frac{\partial \log p_t}{\partial x_j} = p_j - \delta_{t,j} = \left\{\begin{aligned} p_t - 1,&\quad j=t\\ p_j,&\quad j\neq t \end{aligned}\right.\end{equation}
Note that $t$ is given, so what $\delta_{t,j}$ actually expresses is the target distribution $\text{onehot(t)}$, and the full $p_j$ is exactly $\boldsymbol{p}$ itself. So the above expression can be written more intuitively as:
\begin{equation}-\frac{\partial \log p_t}{\partial \boldsymbol{x}} = \boldsymbol{p} - \text{onehot(t)}\label{eq:softmax-ce-grad}\end{equation}
In other words, its gradient is exactly the difference between the target distribution and the predicted distribution. As long as the two are not equal, the gradient will keep existing, and optimization can continue—this is an advantage of cross-entropy. Of course, in some cases this is also a drawback, because Softmax only produces a true one-hot output when $\tau\to 0^+$, and under normal circumstances one-hot never actually occurs. This means optimization never fully stops, which could lead to over-optimization—this is also the motivation behind some of the alternatives discussed below.
Besides cross-entropy, there are other losses one could use, such as $-p_t$, which can be understood as the negative of a smooth approximation of accuracy. However, it can suffer from vanishing gradients, so its optimization efficiency is generally inferior to cross-entropy, and it's usually only suitable for fine-tuning rather than training from scratch. For more discussion, see How to Train Your Accuracy?.
Softmax variants
Having introduced Softmax, let's now summarize some of the Softmax-related variants that have been discussed on this blog before, such as Margin Softmax, Taylor Softmax, and Sparse Softmax. These are all derivatives built on top of Softmax, each focusing on a different aspect of improvement—loss functions, sparsity, long-tailed behavior, and so on.
Margin Softmax
First, let's introduce a series of Softmax variants originating from face recognition, which can collectively be called Margin Softmax. They were later also applied to sentence embedding training in NLP. This blog previously discussed one such variant, AM-Softmax, in A Sentence Similarity Model Based on GRU and AM-Softmax, and later gave a more general discussion in From the Triangle Inequality to Margin Softmax.
Although Margin Softmax bears the name "Softmax," it is in fact more of an improvement to the loss function. Taking AM-Softmax as an example, it has two distinguishing features. First, logits are constructed in the form of $\cos$, i.e. of the form $\boldsymbol{x} = [\cos(\boldsymbol{z},\boldsymbol{c}_1),\cos(\boldsymbol{z},\boldsymbol{c}_2),\cdots,\cos(\boldsymbol{z},\boldsymbol{c}_n)]/\tau$, in which case the temperature parameter $\tau$ is essential, because the plain value range of $\cos$ is $[-1,1]$, which is not enough to pull apart the differences between class probabilities. Second, rather than simply using $-\log p_t$ as the loss, it strengthens the requirement:
\begin{equation}\mathcal{L} = - \log \frac{e^{[\cos(\boldsymbol{z},\boldsymbol{c}_t)-m]/\tau}}{e^{[\cos(\boldsymbol{z},\boldsymbol{c}_t)-m]/\tau} + \sum_{j\neq t} e^{\cos(\boldsymbol{z},\boldsymbol{c}_j)/\tau}}\end{equation}
Intuitively, whereas ordinary cross-entropy just wants $x_t$ to be the largest among all the components of $\boldsymbol{x}$, AM-Softmax not only wants $x_t$ to be the largest, but also wants it to exceed the second-largest component by at least $m/\tau$, where $m/\tau$ is called the margin.
Why add this extra requirement on the target class? This comes from the application scenario. As mentioned, Margin Softmax originates from face recognition; when transplanted to NLP, it can be used for semantic retrieval—meaning that although the application scenario is retrieval, the training approach is classification. If the model is trained purely with the cross-entropy loss of a classification task, the features it encodes are not guaranteed to satisfy retrieval requirements well, so a margin needs to be added to make the features more compact. For a more detailed discussion, see From the Triangle Inequality to Margin Softmax, or consult the relevant papers.
Taylor Softmax
Next, let's introduce Taylor Softmax, discussed in The Even-Order Taylor Expansion of exp(x) at x=0 Is Always Positive, which makes use of an interesting property of the Taylor expansion of $\exp(x)$:
For any real number $x$ and even number $k$, we always have $f_k(x)\triangleq\sum\limits_{m=0}^k \frac{x^m}{m!} > 0$, i.e. the even-order Taylor expansion of $e^x$ at $x=0$ is always positive.
Using this always-positive property, we can construct a Softmax variant (where $k > 0$ is any even number):
\begin{equation}taylor\text{-}softmax(\boldsymbol{x}, k)_i = \frac{f_k(x_i)}{\sum\limits_{j=1}^n f_k(x_j)}\end{equation}
Since it is built from the Taylor expansion of $\exp$, Taylor Softmax approximates ordinary Softmax to some extent within a certain range, and in some scenarios Taylor Softmax can be used as a drop-in replacement for Softmax. So what characterizes Taylor Softmax? The answer is that it has a heavier tail—since Taylor Softmax normalizes a polynomial function, which decays more slowly than an exponential function, it tends to assign higher probability to tail classes than Softmax does, which may help alleviate the overconfidence phenomenon of Softmax.
The most recent application of Taylor Softmax is to replace the Softmax inside attention, reducing the original quadratic complexity to linear complexity. For the relevant theoretical derivation, see The Path to Upgrading Transformers: 5. Linear Attention as an Infinite-Dimensional Case. The latest practical realization of this idea is a model called Based, which uses $e^x\approx 1+x+x^2/2$ to linearize attention, claiming to be more efficient than attention and better-performing than Mamba. For more details, see the blog posts Zoology (Blogpost 2): Simple, Input-Dependent, and Sub-Quadratic Sequence Mixers and BASED: Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff.
Sparse Softmax
Sparse Softmax is a simple sparse variant of Softmax that the author proposed while participating in the 2020 CAIL (China AI and Law Challenge). It was first published in the post SPACES: An "Extract-then-Generate" Long-Text Summarization Approach (CAIL Summary), and later, with additional experiments, was written up as a short paper, Sparse-softmax: A Simpler and Faster Alternative Softmax Transformation.
As we know, in text generation we commonly use either deterministic beam search decoding, or stochastic Top-K/Top-P sampling. What these algorithms have in common is that they only keep a handful of tokens with the highest predicted probability for traversal or sampling, which is equivalent to treating the probability of the remaining tokens as zero. However, if training directly uses Softmax to construct the probability distribution, there's no way to get an exact zero, which creates an inconsistency between training and inference. Sparse Softmax aims to address this inconsistency. The idea is simple: during training, we also zero out the probability of tokens outside the top $k$:
$$\begin{array}{c|c|c} \hline & Softmax & Sparse\text{ }Softmax \\ \hline \text{basic definition} & p_i = \frac{e^{x_i}}{\sum\limits_{j=1}^n e^{x_j}} & p_i=\left\{\begin{aligned}&\frac{e^{x_i}}{\sum\limits_{j\in\Omega_k} e^{x_j}},\,i\in\Omega_k\\ &\quad 0,\,i\not\in\Omega_k\end{aligned}\right.\\ \hline \text{loss function} & \log\left(\sum\limits_{i=1}^n e^{x_i}\right) - x_t & \log\left(\sum\limits_{i\in\Omega_k} e^{x_i}\right) - x_t\\ \hline \end{array}$$
where $\Omega_k$ is the set of original indices of the top $k$ elements after sorting $x_1,x_2,\cdots,x_n$ in descending order. In short, the same truncation operation used at inference time is also applied during training. The selection of $\Omega_k$ can also follow the Top-$p$ scheme used in nucleus sampling, depending on the specific requirements. It should be noted, however, that Sparse Softmax forcibly truncates the probability of the remaining part, which means that these logits cannot receive gradients through backpropagation. As a result, Sparse Softmax is less efficient to train than ordinary Softmax, so it is generally suitable only for fine-tuning scenarios, rather than training from scratch.
Perturb Max
In this section we introduce a new way to construct probability distributions, which we call Perturb Max here. It generalizes Gumbel Max, and was first introduced on this site in the post The Construction of Discrete Probability Distributions from a Reparameterization Perspective; it is also discussed in the paper EXACT: How to Train Your Accuracy. As for any earlier origins, I haven't looked further into it.
Rethinking the Problem
First, let's recall that constructing a mapping $\mathbb{R}^n\mapsto\Delta^{n-1}$ is not hard at all: as long as $f(x)$ is a mapping from $\mathbb{R}\mapsto \mathbb{R}^*$ (from the reals to the non-negative reals), such as $x^2$, then simply letting
\begin{equation}p_i = \frac{f(x_i)}{\sum\limits_{j=1}^n f(x_j)}\end{equation}
gives a mapping satisfying the condition. What if we also want to add "monotonicity" from the "two properties"? That's not hard either — we just need $\mathbb{R}\mapsto \mathbb{R}^*$ to be a monotonically increasing function, and there are plenty of such functions, e.g. $\text{sigmoid}(x)$. But what if we further require "invariance"? Can we still casually write down a $\mathbb{R}^n\mapsto\Delta^{n-1}$ mapping that satisfies invariance? (I certainly can't.)
Some readers might wonder: why insist on preserving monotonicity and invariance at all? Indeed, purely from the perspective of fitting a probability distribution, neither property seems strictly necessary — it's all "brute force wins": as long as the model is large enough, there's nothing it can't fit. But from the perspective of a "Softmax replacement," we want the newly defined probability distribution to likewise serve as a smooth approximation of $\text{argmax}$, and for that we need to preserve as many of the same properties of $\text{argmax}$ as possible. This is the main reason we want to keep monotonicity and invariance.
Gumbel Max
Perturb Max is constructed by generalizing Gumbel Max. Readers unfamiliar with Gumbel Max may first check out Random Musings on Reparameterization: From the Normal Distribution to Gumbel Softmax. Briefly, Gumbel Max is the observation that:
\begin{equation}P[\text{argmax}(\boldsymbol{x}+\boldsymbol{\varepsilon}) = i] = softmax(\boldsymbol{x})_i,\quad \boldsymbol{\varepsilon}\sim Gumbel\text{ }Noise\end{equation}
How should we understand this result? First, $\boldsymbol{\varepsilon}\sim Gumbel\text{ }Noise$ here means that each component of $\boldsymbol{\varepsilon}$ is independently sampled from a Gumbel distribution; next, we know that given the vector $\boldsymbol{x}$, $\text{argmax}(\boldsymbol{x})$ would ordinarily be a deterministic result, but after adding random noise $\boldsymbol{\varepsilon}$, the result of $\text{argmax}(\boldsymbol{x}+\boldsymbol{\varepsilon})$ also becomes random, so each $i$ has its own probability; finally, Gumbel Max tells us that if the added noise is Gumbel noise, then the probability of $i$ occurring is exactly $softmax(\boldsymbol{x})_i$.
The most direct use of Gumbel Max is that it provides a way to sample from the distribution $softmax(\boldsymbol{x})$, though of course if all we want is sampling there are simpler methods — no need to "use a sledgehammer to crack a nut." The real value of Gumbel Max lies in "reparameterization": it shifts the randomness of the problem from the parameterized discrete distribution $\boldsymbol{x}$ onto the parameter-free $\boldsymbol{\varepsilon}$. Combined with the fact that Softmax is a smooth approximation of $\text{argmax}$, we get that $softmax(\boldsymbol{x} + \boldsymbol{\varepsilon})$ is a smooth approximation of Gumbel Max — this is Gumbel Softmax, a common technique for training models with "learnable parameters inside a discrete sampling module."
General Noise
Perturb Max derives directly from Gumbel Max: since Softmax can be derived from the Gumbel distribution, wouldn't replacing the Gumbel distribution with a general distribution, say a normal distribution, let us derive a new form of probability distribution? That is, we directly define
\begin{equation}p_i = P[\text{argmax}(\boldsymbol{x}+\boldsymbol{\varepsilon}) = i],\quad \varepsilon_1,\varepsilon_2,\cdots,\varepsilon_n\sim p(\varepsilon)\end{equation}
Repeating the derivation of Gumbel Max, we obtain
\begin{equation}p_i = \int_{-\infty}^{\infty} p(\varepsilon_i)\left[\prod_{j\neq i} \Phi(x_i - x_j + \varepsilon_i)\right]d\varepsilon_i = \mathbb{E}_{\varepsilon}\left[\prod_{j\neq i} \Phi(x_i - x_j + \varepsilon)\right]\end{equation}
where $\Phi(\varepsilon)$ is the cumulative distribution function of $p(\varepsilon)$. For a general distribution, even one as simple as the standard normal, this expression has no analytical solution in general, so it can only be estimated numerically. To get a deterministic computation, we can perform uniform sampling via the inverse CDF: first uniformly draw $t$ from $[0,1]$, then obtain $\varepsilon$ by solving $t=\Phi(\varepsilon)$.
From the definition of Perturb Max, or from the final form of $p_i$, we can assert that Perturb Max satisfies monotonicity and invariance; we won't spell out the details here. So in what scenarios does it play a unique role? Honestly, I'm not sure. EXACT: How to Train Your Accuracy uses it to construct a new probability distribution and optimize a smooth approximation of accuracy, but my own experiments showed no particularly notable effect. My personal feeling is that it might show special value in certain scenarios that require reparameterization.
Sparsemax
Next up is a probability mapping called Sparsemax, from the 2016 paper From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification. Like the Sparse Softmax I proposed myself, it is a modification aimed at sparsity, but the authors' motivation was to provide better interpretability in attention. Unlike Sparse Softmax, which directly and forcibly truncates to the top-$k$ components, Sparsemax provides a more adaptive way of constructing a sparse-type probability distribution.
Basic Definition
The original paper defines Sparsemax as the solution to the following optimization problem:
\begin{equation}sparsemax(\boldsymbol{x}) = \mathop{\text{argmin}}\limits_{\boldsymbol{p}\in\Delta^{n-1}}\Vert \boldsymbol{p} - \boldsymbol{x}\Vert^2\label{eq:sparsemax-opt}\end{equation}
The exact solution's expression can be derived via the method of Lagrange multipliers. However, this approach is not very intuitive, and it doesn't easily reveal the connection to Softmax. Below I offer what I think is a clearer way to introduce it, of my own devising.
First, notice that Softmax can be equivalently written as
\begin{equation}\boldsymbol{p} = softmax(\boldsymbol{x}) = \exp(\boldsymbol{x} - \lambda(\boldsymbol{x}))\label{eq:sparsemax-softmax}\end{equation}
where $\lambda(\boldsymbol{x})$ is the constant that makes the components of $\boldsymbol{p}$ sum to 1; for Softmax we can solve for $\lambda(\boldsymbol{x})=\log\sum\limits_i e^{x_i}$ explicitly.
Then, in the Taylor Softmax section we noted that the even-order Taylor expansion of $\exp(x)$ is always positive, so it can be used to build Softmax variants. But what about odd orders? For instance $\exp(x)\approx 1 + x$ is not always non-negative, but we can force it to be non-negative by adding $\text{relu}$, i.e., $\exp(x)\approx \text{relu}(1 + x)$. Substituting this approximation for $\exp$ in equation $\eqref{eq:sparsemax-softmax}$ gives us Sparsemax:
\begin{equation}\boldsymbol{p} = sparsemax(\boldsymbol{x}) = \text{relu}(1+\boldsymbol{x} - \lambda(\boldsymbol{x}))\end{equation}
where $\lambda(\boldsymbol{x})$ is again the constant making the components of $\boldsymbol{p}$ sum to 1, and the constant $1$ can also be absorbed into $\lambda(\boldsymbol{x})$, so the above is equivalent to
\begin{equation}\boldsymbol{p} = sparsemax(\boldsymbol{x}) = \text{relu}(\boldsymbol{x} - \lambda(\boldsymbol{x}))\end{equation}
Solution Algorithm
So far, Sparsemax remains a purely formal definition, since how to concretely compute $\lambda(\boldsymbol{x})$ is still unclear — that is the topic of this section. Even so, from the definition alone it's not hard to see that Sparsemax satisfies monotonicity and invariance; readers who aren't fully convinced can try proving it themselves.
Now let's turn to computing $\lambda(\boldsymbol{x})$. Without loss of generality, assume the components of $\boldsymbol{x}$ are already sorted in descending order, i.e., $x_1\geq x_2\geq \cdots\geq x_n$. Suppose for the moment that we already know $x_k\geq \lambda(\boldsymbol{x})\geq x_{k+1}$; then clearly
\begin{equation}sparsemax(\boldsymbol{x}) = [x_1 - \lambda(\boldsymbol{x}),\cdots,x_k - \lambda(\boldsymbol{x}),0,\cdots,0]\end{equation}
By the definition of $\lambda(\boldsymbol{x})$, we have
\begin{equation}\sum_{i=1}^k [x_i - \lambda(\boldsymbol{x})] = 1\quad\Rightarrow\quad 1 + k\lambda(\boldsymbol{x}) = \sum_{i=1}^k x_i\end{equation}
which lets us solve for $\lambda(\boldsymbol{x})$. Of course, we don't know $x_k\geq \lambda(\boldsymbol{x})\geq x_{k+1}$ in advance, but we can enumerate over $k=1,2,\cdots,n$, compute $\lambda_k(\boldsymbol{x})$ via the above formula for each, and take the $\lambda_k(\boldsymbol{x})$ satisfying $x_k\geq \lambda_k(\boldsymbol{x})\geq x_{k+1}$ — equivalently, find the largest $k$ satisfying $x_k\geq \lambda_k(\boldsymbol{x})$, and then return the corresponding $\lambda_k(\boldsymbol{x})$.
Reference implementation:
def sparsemax(x):
x_sort = np.sort(x)[::-1]
x_lamb = (np.cumsum(x_sort) - 1) / np.arange(1, len(x) + 1)
lamb = x_lamb[(x_sort >= x_lamb).argmin() - 1]
return np.maximum(x - lamb, 0)
Gradient Computation
For convenience, let's introduce the notation
\begin{equation}\Omega(\boldsymbol{x}) = \big\{k\big|x_k > \lambda(\boldsymbol{x})\big\}\end{equation}
Then we can write
\begin{equation}\boldsymbol{p} = sparsemax(\boldsymbol{x}) = \left\{\begin{aligned} &x_i - \frac{1}{|\Omega(\boldsymbol{x})|}\left(-1 + \sum_{j\in\Omega(\boldsymbol{x})}x_j\right),\quad &i\in \Omega(\boldsymbol{x})\\ &0,\quad &i \not\in \Omega(\boldsymbol{x}) \end{aligned}\right.\end{equation}
From this equivalent form we can see that, just like Sparse Softmax, Sparsemax also only has gradients for a subset of classes. We can directly compute the Jacobian matrix:
\begin{equation}\frac{\partial p_i}{\partial x_j} = \left\{\begin{aligned} &1 - \frac{1}{|\Omega(\boldsymbol{x})|},\quad &i,j\in \Omega(\boldsymbol{x}),i=j\\[5pt] &- \frac{1}{|\Omega(\boldsymbol{x})|},\quad &i,j\in \Omega(\boldsymbol{x}),i\neq j\\[5pt] &0,\quad &i \not\in \Omega(\boldsymbol{x})\text{ or }j \not\in \Omega(\boldsymbol{x}) \end{aligned}\right.\end{equation}
This shows that for classes within $\Omega(\boldsymbol{x})$, Sparsemax doesn't suffer from vanishing gradients, since the gradient there is a constant. But the overall magnitude of the gradient depends on the number of elements in $\Omega(\boldsymbol{x})$ — the fewer there are, the sparser it is, meaning the gradient is also sparser.
Loss Function
Finally, let's discuss the loss function to use when Sparsemax serves as a classification output. The intuitive idea would be to use cross-entropy $-\log p_t$, just as with Softmax. But since Sparsemax's output can be exactly zero, to avoid a $\log 0$ error we'd need to add $\epsilon$ to every component, giving the final cross-entropy form $-\log\frac{p_t + \epsilon}{1 + n\epsilon}$. However, this is ugly, and moreover it is not a convex function, so it's not an ideal choice.
In fact, the reason cross-entropy works well with Softmax is precisely that its gradient has the form $\eqref{eq:softmax-ce-grad}$. So for Sparsemax, let's likewise assume the loss function's gradient is $\boldsymbol{p} - \text{onehot(t)}$, and work backward to figure out what the loss function should look like:
\begin{equation}\frac{\partial \mathcal{L}_t}{\partial \boldsymbol{x}} = \boldsymbol{p} - \text{onehot(t)}\quad\Rightarrow\quad \mathcal{L}_t = \frac{1}{2} - x_t + \sum_{i\in\Omega(\boldsymbol{x})}\frac{1}{2}\left(x_i^2 - \lambda^2(\boldsymbol{x})\right)\end{equation}
Verifying from right to left is fairly easy; deriving from left to right may be a bit trickier, but not overly so — some patient trial and error should get you there. The constant $\frac{1}{2}$ at the front is there to guarantee the loss function's non-negativity. We can check an extreme case: suppose training has converged perfectly, so that $\boldsymbol{p}$ is also one-hot; then $x_t\to\infty$ and $\lambda(\boldsymbol{x}) = x_t - 1$, so
\begin{equation}- x_t + \sum_{i\in\Omega(\boldsymbol{x})}\frac{1}{2}\left(x_i^2 - \lambda^2(\boldsymbol{x})\right) = -x_t + \frac{1}{2}x_t^2 - \frac{1}{2}(x_t - 1)^2 = -\frac{1}{2}\end{equation}
hence the need to add the constant $\frac{1}{2}$.
Entmax-α
Entmax-$\alpha$ is a generalization of Sparsemax, motivated by the fact that Sparsemax tends to be overly sparse, which can lower learning efficiency and hurt final performance. Entmax-$\alpha$ therefore introduces the parameter $\alpha$, providing a smooth transition from Softmax ($\alpha=1$) to Sparsemax ($\alpha=2$). Entmax-$\alpha$ comes from the paper Sparse Sequence-to-Sequence Models, by the same author as Sparsemax, Andre F. T. Martins, who has done a great deal of work on sparse Softmax and sparse attention. Interested readers can check out his homepage for related work.
Basic Definition
Like Sparsemax, the original paper defines Entmax-$\alpha$ as the solution to an optimization problem similar to $\eqref{eq:sparsemax-opt}$, but this definition involves the concept of Tsallis entropy (which is also where the "Ent" in Entmax comes from), and solving it requires the method of Lagrange multipliers, which is relatively complex. We won't use that approach for the introduction here.
Our introduction is likewise based on the approximation $\exp(x)\approx \text{relu}(1 + x)$ from the previous section. For Softmax and Sparsemax, we have
\begin{align}&{\color{red}{Softmax}}:\quad &\exp(\boldsymbol{x} - \lambda(\boldsymbol{x})) \\[5pt] &{\color{red}{Sparsemax}}:\quad &\text{relu}(1+\boldsymbol{x} - \lambda(\boldsymbol{x})) \end{align}
The underlying reason Sparsemax is too sparse can also be understood as the $\exp(x)\approx \text{relu}(1 + x)$ approximation not being accurate enough. We can evolve from it a higher-precision approximation
\begin{equation}\exp(x) = \exp(\beta x / \beta) = \exp^{1/\beta}(\beta x)\approx \text{relu}^{1/\beta}(1 + \beta x)\end{equation}
As long as $0 \leq \beta < 1$, the rightmost term is a better approximation than $\text{relu}(1 + x)$ (think about why). Using this new approximation, we can construct
\begin{equation}{\color{red}{Entmax\text{-}\alpha}}:\quad \text{relu}^{1/\beta}(1+\beta\boldsymbol{x} - \lambda(\boldsymbol{x}))\end{equation}
Here $\alpha = \beta + 1$ is used to align with the notation of the original paper; in fact, using $\beta$ would be more concise. Likewise, the constant $1$ can be absorbed into the definition of $\lambda(\boldsymbol{x})$, so the final definition simplifies to
\begin{equation}Entmax_{\alpha}(\boldsymbol{x}) = \text{relu}^{1/\beta}(\beta\boldsymbol{x} - \lambda(\boldsymbol{x}))\end{equation}
Solution Algorithm
For general $\beta$, solving for $\lambda(\boldsymbol{x})$ is fairly troublesome, and usually can only be done via bisection.
First, let $\boldsymbol{z}=\beta\boldsymbol{x}$, and assume without loss of generality that $z_1\geq z_2\geq \cdots \geq z_n$. We can then observe that Entmax-$\alpha$ satisfies monotonicity and invariance, so using invariance we can, without loss of generality, set $z_1 = 1$ (if not, simply subtract $z_1 - 1$ from each $z_i$). Now we can check that when $\lambda=0$, the sum of all components of $\text{relu}^{1/\beta}(\beta\boldsymbol{x} - \lambda)$ is greater than or equal to 1, and when $\lambda=1$, the sum of all components of $\text{relu}^{1/\beta}(\beta\boldsymbol{x} - \lambda)$ equals 0. So the value of $\lambda(\boldsymbol{x})$ that makes the components sum to 1 must lie within $[0,1)$, and we can then use bisection to progressively converge on the optimal $\lambda(\boldsymbol{x})$.
For certain special values of $\beta$, we can get an algorithm for the exact solution. Sparsemax corresponds to $\beta=1$, whose solution process we already gave above; another example admitting an analytical solution is $\beta=1/2$, which is also the main case of interest in the original paper — indeed, if unspecified, "Entmax" by default refers to Entmax-1.5. Following the same approach as with Sparsemax, suppose we already know $z_k\geq \lambda(\boldsymbol{x})\geq z_{k+1}$; then
\begin{equation}\sum_{i=1}^k [z_i - \lambda(\boldsymbol{x})]^2 = 1\end{equation}
This is just a quadratic equation in $\lambda(\boldsymbol{x})$, which can be solved to give
\begin{equation}\lambda(\boldsymbol{x}) = \mu_k - \sqrt{\frac{1}{k} - \sigma_k^2},\quad \mu_k = \frac{1}{k}\sum_{i=1}^k z_i,\quad\sigma_k^2 = \frac{1}{k}\left(\sum_{i=1}^k z_i^2\right) - \mu_k^2\end{equation}
When we don't know $x_k\geq \lambda(\boldsymbol{x})\geq x_{k+1}$ in advance, we can enumerate over $k=1,2,\cdots,n$, compute $\lambda_k(\boldsymbol{x})$ via the above formula for each, and take the $\lambda_k(\boldsymbol{x})$ satisfying $x_k\geq \lambda_k(\boldsymbol{x})\geq x_{k+1}$ — though note this is no longer equivalent to finding the largest $k$ satisfying $x_k\geq \lambda_k(\boldsymbol{x})$.
Full reference implementation:
def entmat(x):
x_sort = np.sort(x / 2)[::-1]
k = np.arange(1, len(x) + 1)
x_mu = np.cumsum(x_sort) / k
x_sigma2 = np.cumsum(x_sort**2) / k - x_mu**2
x_lamb = x_mu - np.sqrt(np.maximum(1. / k - x_sigma2, 0))
x_sort_shift = np.pad(x_sort[1:], (0, 1), constant_values=-np.inf)
lamb = x_lamb[(x_sort > x_lamb) & (x_lamb > x_sort_shift)]
return np.maximum(x / 2 - lamb, 0)**2
Other topics
The gradient of Entmax-$\alpha$ is broadly similar to that of Sparsemax, so we won't go into detail here — readers can either work through the derivation themselves or consult the original paper. As for the loss function, one can likewise work backward from the gradient $\frac{\partial \mathcal{L}_t}{\partial \boldsymbol{x}} = \boldsymbol{p} - \text{onehot(t)}$ to reconstruct the loss, but the resulting form is somewhat complicated; readers interested in the details can refer to the original papers Sparse Sequence-to-Sequence Models and Learning with Fenchel-Young Losses.
That said, in my view it's simpler and more general to define the loss function directly via the $\text{stop_gradient}$ operator, which avoids the whole business of finding an antiderivative:
\begin{equation}\mathcal{L}_t = (\boldsymbol{p} - \text{onehot(t)})\cdot \text{stop_gradient}(\boldsymbol{x})\end{equation}
Here $\,\cdot\,$ denotes a vector inner product. With this definition, the gradient of the resulting loss is exactly $\boldsymbol{p} - \text{onehot(t)}$, but note that only the gradient of this loss is meaningful — the value of the loss itself carries no useful information. For example, it can be positive or negative, and smaller is not necessarily better. So if you want to track training progress or evaluate performance, you'll need to set up a separate metric (such as cross-entropy or accuracy).
Summary
This post has given a brief review and organization of Softmax and some of its substitutes, covering the definitions and properties of Softmax, Margin Softmax, Taylor Softmax, Sparse Softmax, Perturb Max, Sparsemax, and Entmax-$\alpha$.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.