Softmax function

 ( Computational Neuroscience ) 

Formally, instead of considering the arg max as a function with categorical output $1,\; \backslash dots,\; n$ (corresponding to the index), consider the arg max function with one-hot representation of the output (assuming there is a unique maximum arg): $$\backslash operatorname\{arg\backslash ,max\}(z\_1,\backslash ,\; \backslash dots,\backslash ,\; z\_n)\; =\; (y\_1,\backslash ,\; \backslash dots,\backslash ,\; y\_n)\; =\; (0,\backslash ,\; \backslash dots,\backslash ,\; 0,\backslash ,\; 1,\backslash ,\; 0,\backslash ,\; \backslash dots,\backslash ,\; 0),$$ where the output coordinate $y\_i\; =\; 1$ if and only if $i$ is the arg max of $(z\_1,\; \backslash dots,\; z\_n)$, meaning $z\_i$ is the unique maximum value of $(z\_1,\backslash ,\; \backslash dots,\backslash ,\; z\_n)$. For example, in this encoding $\backslash operatorname\{arg\backslash ,max\}(1,\; 5,\; 10)\; =\; (0,\; 0,\; 1),$ since the third argument is the maximum.

This can be generalized to multiple arg max values (multiple equal $z\_i$ being the maximum) by dividing the 1 between all max args; formally where is the number of arguments assuming the maximum. For example, $\backslash operatorname\{arg\backslash ,max\}(1,\backslash ,\; 5,\backslash ,\; 5)\; =\; (0,\backslash ,\; 1/2,\backslash ,\; 1/2),$ since the second and third argument are both the maximum. In case all arguments are equal, this is simply $\backslash operatorname\{arg\backslash ,max\}(z,\; \backslash dots,\; z)\; =\; (1/n,\; \backslash dots,\; 1/n).$ Points with multiple arg max values are singular points (or singularities, and form the singular set) – these are the points where arg max is discontinuous (with a jump discontinuity) – while points with a single arg max are known as non-singular or regular points.

With the last expression given in the introduction, softargmax is now a smooth approximation of arg max: as , softargmax converges to arg max. There are various notions of convergence of a function; softargmax converges to arg max pointwise, meaning for each fixed input as , $\backslash sigma\_\backslash beta(\backslash mathbf\{z\})\; \backslash to\; \backslash operatorname\{arg\backslash ,max\}(\backslash mathbf\{z\}).$ However, softargmax does not converge uniformly to arg max, meaning intuitively that different points converge at different rates, and may converge arbitrarily slowly. In fact, softargmax is continuous, but arg max is not continuous at the singular set where two coordinates are equal, while the uniform limit of continuous functions is continuous. The reason it fails to converge uniformly is that for inputs where two coordinates are almost equal (and one is the maximum), the arg max is the index of one or the other, so a small change in input yields a large change in output. For example, $\backslash sigma\_\backslash beta(1,\backslash ,\; 1.0001)\; \backslash to\; (0,\; 1),$ but $\backslash sigma\_\backslash beta(1,\backslash ,\; 0.9999)\; \backslash to\; (1,\backslash ,\; 0),$ and $\backslash sigma\_\backslash beta(1,\backslash ,\; 1)\; =\; 1/2$ for all inputs: the closer the points are to the singular set $(x,\; x)$, the slower they converge. However, softargmax does converge compactly on the non-singular set.

Conversely, as , softargmax converges to arg min in the same way, where here the singular set is points with two arg min values. In the language of tropical analysis, the softmax is a deformation or "quantization" of arg max and arg min, corresponding to using the log semiring instead of the max-plus semiring (respectively min-plus semiring), and recovering the arg max or arg min by taking the limit is called "tropicalization" or "dequantization".

It is also the case that, for any fixed , if one input is much larger than the others relative to the temperature, $T\; =\; 1/\backslash beta$, the output is approximately the arg max. For example, a difference of 10 is large relative to a temperature of 1: $$\backslash sigma(0,\backslash ,\; 10)\; :=\; \backslash sigma\_1(0,\backslash ,\; 10)\; =\; \backslash left(1/\backslash left(1\; +\; e^\{10\}\backslash right),\backslash ,\; e^\{10\}/\backslash left(1\; +\; e^\{10\}\backslash right)\backslash right)\; \backslash approx\; (0.00005,\backslash ,\; 0.99995)$$ However, if the difference is small relative to the temperature, the value is not close to the arg max. For example, a difference of 10 is small relative to a temperature of 100: $$\backslash sigma\_\{1/100\}(0,\backslash ,\; 10)\; =\; \backslash left(1/\backslash left(1\; +\; e^\{1/10\}\backslash right),\backslash ,\; e^\{1/10\}/\backslash left(1\; +\; e^\{1/10\}\backslash right)\backslash right)\; \backslash approx\; (0.475,\backslash ,\; 0.525).$$ As , temperature goes to zero, $T\; =\; 1/\backslash beta\; \backslash to\; 0$, so eventually all differences become large (relative to a shrinking temperature), which gives another interpretation for the limit behavior.

$$P(y=j\backslash mid\; \backslash mathbf\{x\})\; =\; \backslash frac\{e^\{\backslash mathbf\{x\}^\backslash mathsf\{T\}\backslash mathbf\{w\}\_j\}\}\{\backslash sum\_\{k=1\}^K\; e^\{\backslash mathbf\{x\}^\backslash mathsf\{T\}\backslash mathbf\{w\}\_k\}\}$$

This can be seen as the composition of linear functions $\backslash mathbf\{x\}\; \backslash mapsto\; \backslash mathbf\{x\}^\backslash mathsf\{T\}\backslash mathbf\{w\}\_1,\; \backslash ldots,\; \backslash mathbf\{x\}\; \backslash mapsto\; \backslash mathbf\{x\}^\backslash mathsf\{T\}\backslash mathbf\{w\}\_K$ and the softmax function (where $\backslash mathbf\{x\}^\backslash mathsf\{T\}\backslash mathbf\{w\}$ denotes the inner product of $\backslash mathbf\{x\}$ and $\backslash mathbf\{w\}$). The operation is equivalent to applying a linear operator defined by $\backslash mathbf\{w\}$ to tuples $\backslash mathbf\{x\}$, thus transforming the original, probably highly-dimensional, input to vectors in a -dimensional space $\backslash mathbb\{R\}^K$.

Since the function maps a tuple and a specific index $i$ to a real value, the derivative needs to take the index into account:

$$\backslash frac\{\backslash partial\}\{\backslash partial\; q\_k\}\backslash sigma(\backslash textbf\{q\},\; i)\; =\; \backslash sigma(\backslash textbf\{q\},\; i)(\backslash delta\_\{ik\}\; -\; \backslash sigma(\backslash textbf\{q\},\; k)).$$

This expression is symmetrical in the indexes $i,\; k$ and thus may also be expressed as

$$\backslash frac\{\backslash partial\}\{\backslash partial\; q\_k\}\backslash sigma(\backslash textbf\{q\},\; i)\; =\; \backslash sigma(\backslash textbf\{q\},\; k)(\backslash delta\_\{ik\}\; -\; \backslash sigma(\backslash textbf\{q\},\; i)).$$

Here, the Kronecker delta is used for simplicity (cf. the derivative of a sigmoid function, being expressed via the function itself).

To ensure stable numerical computations subtracting the maximum value from the input tuple is common. This approach, while not altering the output or the derivative theoretically, enhances stability by directly controlling the maximum exponent value computed.

If the function is scaled with the parameter $\backslash beta$, then these expressions must be multiplied by $\backslash beta$.

See multinomial logit for a probability model which uses the softmax activation function.

where the action value $q\_t(a)$ corresponds to the expected reward of following action a and $\backslash tau$ is called a temperature parameter (in allusion to statistical mechanics). For high temperatures ($\backslash tau\; \backslash to\; \backslash infty$), all actions have nearly the same probability and the lower the temperature, the more expected rewards affect the probability. For a low temperature ($\backslash tau\; \backslash to\; 0^+$), the probability of the action with the highest expected reward tends to 1.

What's more, the gradient descent backpropagation method for training such a neural network involves calculating the softmax for every training example, and the number of training examples can also become large. The computational effort for the softmax became a major limiting factor in the development of larger neural language models, motivating various remedies to reduce training times.

Approaches that reorganize the softmax layer for more efficient calculation include the hierarchical softmax and the differentiated softmax. The hierarchical softmax (introduced by Morin and Yoshua Bengio in 2005) uses a binary tree structure where the outcomes (vocabulary words) are the leaves and the intermediate nodes are suitably selected "classes" of outcomes, forming . The desired probability (softmax value) of a leaf (outcome) can then be calculated as the product of the probabilities of all nodes on the path from the root to that leaf. Ideally, when the tree is balanced, this would reduce the computational complexity from $O(K)$ to $O(\backslash log\_2\; K)$. In practice, results depend on choosing a good strategy for clustering the outcomes into classes. A Huffman tree was used for this in Google's word2vec models (introduced in 2013) to achieve scalability.

A second kind of remedies is based on approximating the softmax (during training) with modified loss functions that avoid the calculation of the full normalization factor. These include methods that restrict the normalization sum to a sample of outcomes (e.g. Importance Sampling, Target Sampling).

---

Numerical algorithms

The standard softmax is numerically unstable because of large exponentiations. The safe softmax method calculates instead$$\backslash sigma(\backslash mathbf\{z\})\_i\; =\; \backslash frac\{e^\{\backslash beta\; (z\_i-m)\}\}\{\backslash sum\_\{j=1\}^K\; e^\{\backslash beta\; (z\_j-m)\}\}$$where $m\; =\; \backslash max\_i\; z\_i$ is the largest factor involved. Subtracting by it guarantees that the exponentiations result in at most 1.

The attention mechanism in Transformers takes three arguments: a "query vector" $q$, a list of "key vectors" $k\_1,\; \backslash dots,\; k\_N$, and a list of "value vectors" $v\_1,\; \backslash dots,\; v\_N$, and outputs a softmax-weighted sum over value vectors:$$o\; =\; \backslash sum\_\{i=1\}^N\; \backslash frac\{e^\{q^T\; k\_i\; -\; m\}\}\{\backslash sum\_\{j=1\}^N\; e^\{q^T\; k\_j\; -\; m\}\}\; v\_i$$The standard softmax method involves several loops over the inputs, which would be bottlenecked by memory bandwidth. The FlashAttention method is a communication-avoiding algorithm that fuses these operations into a single loop, increasing the arithmetic intensity. It is an online algorithm that computes the following quantities:and returns $o\_N/l\_N$. In practice, FlashAttention operates over multiple queries and keys per loop iteration, in a similar way as blocked matrix multiplication. If backpropagation is needed, then the output vectors and the intermediate arrays $m\_1,,\; l\_1,$ are cached, and during the backward pass, attention matrices are rematerialized from these, making it a form of gradient checkpointing.

---

Mathematical properties

Geometrically the softmax function maps the Euclidean space $\backslash mathbb\{R\}^K$ to the boundary of the standard simplex, cutting the dimension by one (the range is a $(K\; -\; 1)$-dimensional simplex in $K$-dimensional space), due to the linear constraint that all output sum to 1 meaning it lies on a hyperplane.

Along the main diagonal $(x,\backslash ,\; x,\backslash ,\; \backslash dots,\backslash ,\; x),$ softmax is just the uniform distribution on outputs, $(1/n,\; \backslash dots,\; 1/n)$: equal scores yield equal probabilities.

More generally, softmax is invariant under translation by the same value in each coordinate: adding $\backslash mathbf\{c\}\; =\; (c,\backslash ,\; \backslash dots,\backslash ,\; c)$ to the inputs $\backslash mathbf\{z\}$ yields $\backslash sigma(\backslash mathbf\{z\}\; +\; \backslash mathbf\{c\})\; =\; \backslash sigma(\backslash mathbf\{z\})$, because it multiplies each exponent by the same factor, $e^c$ (because $e^\{z\_i\; +\; c\}\; =\; e^\{z\_i\}\; \backslash cdot\; e^c$), so the ratios do not change: $$\backslash sigma(\backslash mathbf\{z\}\; +\; \backslash mathbf\{c\})\_j\; =\; \backslash frac\{e^\{z\_j\; +\; c\}\}\{\backslash sum\_\{k=1\}^K\; e^\{z\_k\; +\; c\}\}\; =\; \backslash frac\{e^\{z\_j\}\; \backslash cdot\; e^c\}\{\backslash sum\_\{k=1\}^K\; e^\{z\_k\}\; \backslash cdot\; e^c\}\; =\; \backslash sigma(\backslash mathbf\{z\})\_j.$$

Geometrically, softmax is constant along diagonals: this is the dimension that is eliminated, and corresponds to the softmax output being independent of a translation in the input scores (a choice of 0 score). One can normalize input scores by assuming that the sum is zero (subtract the average: $\backslash mathbf\{c\}$ where $c\; =\; \backslash frac\{1\}\{n\}\; \backslash sum\; z\_i$), and then the softmax takes the hyperplane of points that sum to zero, $\backslash sum\; z\_i\; =\; 0$, to the open simplex of positive values that sum to 1$\backslash sum\; \backslash sigma(\backslash mathbf\{z\})\_i\; =\; 1$, analogously to how the exponent takes 0 to 1, $e^0\; =\; 1$ and is positive.

By contrast, softmax is not invariant under scaling. For instance, $\backslash sigma\backslash bigl((0,\backslash ,\; 1)\backslash bigr)\; =\; \backslash bigl(1/(1\; +\; e),\backslash ,\; e/(1\; +\; e)\backslash bigr)$ but $\backslash sigma\backslash bigl((0,\; 2)\backslash bigr)\; =\; \backslash bigl(1/\backslash left(1\; +\; e^2\backslash right),\backslash ,\; e^2/\backslash left(1\; +\; e^2\backslash right)\backslash bigr).$

The standard logistic function is the special case for a 1-dimensional axis in 2-dimensional space, say the x-axis in the plane. One variable is fixed at 0 (say $z\_2\; =\; 0$), so $e^0\; =\; 1$, and the other variable can vary, denote it $z\_1\; =\; x$, so $e^\{z\_1\}/\backslash sum\_\{k=1\}^2\; e^\{z\_k\}\; =\; e^x/\backslash left(e^x\; +\; 1\backslash right),$ the standard logistic function, and $e^\{z\_2\}/\backslash sum\_\{k=1\}^2\; e^\{z\_k\}\; =\; 1/\backslash left(e^x\; +\; 1\backslash right),$ its complement (meaning they add up to 1). The 1-dimensional input could alternatively be expressed as the line $(x/2,\backslash ,\; -x/2)$, with outputs $e^\{x/2\}/\backslash left(e^\{x/2\}\; +\; e^\{-x/2\}\backslash right)\; =\; e^x/\backslash left(e^x\; +\; 1\backslash right)$ and $e^\{-x/2\}/\backslash left(e^\{x/2\}\; +\; e^\{-x/2\}\backslash right)\; =\; 1/\backslash left(e^x\; +\; 1\backslash right).$

---

Gradients

The softmax function is also the gradient of the LogSumExp function:$$\backslash frac\{\backslash partial\}\{\backslash partial\; z\_i\}\; \backslash operatorname\{LSE\}(\backslash mathbf\{z\})\; =\; \backslash frac\{\backslash exp\; z\_i\}\{\backslash sum\_\{j=1\}^\{K\}\; \backslash exp\; z\_j\}\; =\; \backslash sigma(\backslash mathbf\{z\})\_i,\; \backslash quad\; \backslash text\{\; for\; \}\; i\; =\; 1,\; \backslash dotsc\; ,\; K,\; \backslash quad\; \backslash mathbf\{z\}\; =\; (z\_1,\backslash ,\; \backslash dotsc,\backslash ,\; z\_K)\; \backslash in\backslash R^K,$$where the LogSumExp function is defined as $\backslash operatorname\{LSE\}(z\_1,\backslash ,\; \backslash dots,\backslash ,\; z\_n)\; =\; \backslash log\backslash left(\backslash exp(z\_1)\; +\; \backslash cdots\; +\; \backslash exp(z\_n)\backslash right)$.

The gradient of softmax is thus $\backslash partial\_\{z\_j\}\; \backslash sigma\_i\; =\; \backslash sigma\_i\; (\backslash delta\_\{ij\}\; -\; \backslash sigma\_j)$.

---

History

The softmax function was used in statistical mechanics as the Boltzmann distribution in the foundational paper , formalized and popularized in the influential textbook .

The use of the softmax in decision theory is credited to R. Duncan Luce, who used the axiom of independence of irrelevant alternatives in rational choice theory to deduce the softmax in Luce's choice axiom for relative preferences.

In machine learning, the term "softmax" is credited to John S. Bridle in two 1989 conference papers, : and :

---

Example

With an input of , the softmax is approximately . The output has most of its weight where the "4" was in the original input. This is what the function is normally used for: to highlight the largest values and suppress values which are significantly below the maximum value. But note: a change of temperature changes the output. When the temperature is multiplied by 10, the inputs are effectively and the softmax is approximately . This shows that high temperatures de-emphasize the maximum value.

Computation of this example using Python code:

>>> import numpy as np >>> z = np.array(1.0,) >>> beta = 1.0 >>> np.exp(beta * z) / np.sum(np.exp(beta * z)) array([0.02364054, 0.06426166, 0.1746813, 0.474833, 0.02364054,

      0.06426166, 0.1746813])
     

---

Alternatives

The softmax function generates probability predictions densely distributed over its support. Other functions like sparsemax or α-entmax can be used when sparse probability predictions are desired."Speeding Up Entmax" /ref> Also the Gumbel-softmax reparametrization trick can be used when sampling from a discrete-discrete distribution needs to be mimicked in a differentiable manner.

---

See also

• Softplus
• Multinomial logistic regression
• Dirichlet distribution – an alternative way to sample categorical distributions
• Partition function
• Exponential tilting – a generalization of Softmax to more general probability distributions

---

Notes

Categories: Computational Neuroscience, Logistic Regression, Artificial Neural Networks, Functions And Mappings, Articles With Example Python Code, Exponentials, Articles With Example Julia Code, Articles With Example R Code