Softmax function

The softmax function takes as input a vector z of K real numbers, and normalizes it into a probability distribution consisting of K probabilities proportional to the exponentials of the input numbers. That is, prior to applying softmax, some vector components could be negative, or greater than one; and might not sum to 1; but after applying softmax, each component will be in the interval ${\displaystyle (0,1)}$, and the components will add up to 1, so that they can be interpreted as probabilities. Furthermore, the larger input components will correspond to larger probabilities.

The standard (unit) softmax function ${\displaystyle \sigma$ :\mathbb {R} ^{K}\to (0,1)^{K}} is defined when ${\displaystyle K\geq 1}$ by the formula

$${\displaystyle \sigma (\mathbf {z} )_{i}={\frac {e^{z_{i}}}{\sum _{j=1}^{K}e^{z_{j}}}}\ \ {\text{ for }}i=1,\dotsc ,K{\text{ and }}\mathbf {z} =(z_{1},\dotsc ,z_{K})\in \mathbb {R} ^{K}.}$$

In simple words, it applies the standard exponential function to each element ${\displaystyle z_{i}}$ of the input vector ${\displaystyle \mathbf {z} }$ and normalizes these values by dividing by the sum of all these exponentials; this normalization ensures that the sum of the components of the output vector ${\displaystyle \sigma (\mathbf {z} )}$ is 1.

Instead of e, a different base b > 0 can be used. If 0 < b < 1, smaller input components will result in larger output probabilities, and decreasing the value of b will create probability distributions that are more concentrated around the positions of the smallest input values. Conversely, if b > 1, larger input components will result in larger output probabilities, and increasing the value of b will create probability distributions that are more concentrated around the positions of the largest input values. Writing ${\displaystyle b=e^{\beta }}$ or ${\displaystyle b=e^{-\beta }}$[lower-alpha 1] (for real β)[lower-alpha 2] yields the expressions:[lower-alpha 3]

$${\displaystyle \sigma (\mathbf {z} )_{i}={\frac {e^{\beta z_{i}}}{\sum _{j=1}^{K}e^{\beta z_{j}}}}{\text{ or }}\sigma (\mathbf {z} )_{i}={\frac {e^{-\beta z_{i}}}{\sum _{j=1}^{K}e^{-\beta z_{j}}}}{\text{ for }}i=1,\dotsc ,K.}$$

In some fields, the base is fixed, corresponding to a fixed scale,[lower-alpha 4] while in others the parameter β is varied.

The softmax function is used in various multiclass classification methods, such as multinomial logistic regression (also known as softmax regression)[2]^{: 206–209} , multiclass linear discriminant analysis, naive Bayes classifiers, and artificial neural networks.[6] Specifically, in multinomial logistic regression and linear discriminant analysis, the input to the function is the result of K distinct linear functions, and the predicted probability for the jth class given a sample vector x and a weighting vector w is:

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

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

Neural networks

The standard softmax function is often used in the final layer of a neural network-based classifier. Such networks are commonly trained under a log loss (or cross-entropy) regime, giving a non-linear variant of multinomial logistic regression.

Since the function maps a vector and a specific index ${\displaystyle i}$ to a real value, the derivative needs to take the index into account:

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

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

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

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

In order to achieve stable numerical computations of the derivative, one often subtracts a constant from the input vector. In theory, this does not change the output, and neither the derivative. But it is more stable as it can control explicitly the largest value computed in each exponent.

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

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

Reinforcement learning

In the field of reinforcement learning, a softmax function can be used to convert values into action probabilities. The function commonly used is:[7]

$${\displaystyle P_{t}(a)={\frac {\exp(q_{t}(a)/\tau )}{\sum _{i=1}^{n}\exp(q_{t}(i)/\tau )}}{\text{,}}}$$

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

In neural network applications, the number K of possible outcomes is often large, e.g. in case of neural language models that predict the most likely outcome out of a vocabulary which might contain millions of possible words.[8] This can make the calculations for the softmax layer (i.e. the matrix multiplications to determine the ${\displaystyle z_{i}}$, followed by the application of the softmax function itself) computationally expensive.[8][9] 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.[8][9]

Approaches that reorganize the softmax layer for more efficient calculation include the hierarchical softmax and the differentiated softmax.[8] The hierarchical softmax (introduced by Morin and 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 latent variables.[9][10] 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.[9] Ideally, when the tree is balanced, this would reduce the computational complexity from ${\displaystyle O(K)}$ to ${\displaystyle O(\log _{2}K)}$.[10] In practice, results depend on choosing a good strategy for clustering the outcomes into classes.[9][10] A Huffman tree was used for this in Google's word2vec models (introduced in 2013) to achieve scalability.[8]

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.[8] These include methods that restrict the normalization sum to a sample of outcomes (e.g. Importance Sampling, Target Sampling).[8][9]

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

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

More generally, softmax is invariant under translation by the same value in each coordinate: adding ${\displaystyle \mathbf {c} =(c,\,\dots ,\,c)}$ to the inputs ${\displaystyle \mathbf {z} }$ yields ${\displaystyle \sigma (\mathbf {z} +\mathbf {c} )=\sigma (\mathbf {z} )}$, because it multiplies each exponent by the same factor, ${\displaystyle e^{c}}$ (because ${\displaystyle e^{z_{i}+c}=e^{z_{i}}\cdot e^{c}}$), so the ratios do not change:

${\displaystyle \sigma (\mathbf {z} +\mathbf {c} )_{j}={\frac {e^{z_{j}+c}}{\sum _{k=1}^{K}e^{z_{k}+c}}}={\frac {e^{z_{j}}\cdot e^{c}}{\sum _{k=1}^{K}e^{z_{k}}\cdot e^{c}}}=\sigma (\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: ${\displaystyle \mathbf {c} }$ where ${\textstyle c={\frac {1}{n}}\sum z_{i}}$), and then the softmax takes the hyperplane of points that sum to zero, ${\textstyle \sum z_{i}=0}$, to the open simplex of positive values that sum to 1${\textstyle \sum \sigma (\mathbf {z} )_{i}=1}$, analogously to how the exponent takes 0 to 1, ${\displaystyle e^{0}=1}$ and is positive.

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

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

The softmax function is also the gradient of the LogSumExp function, a smooth maximum:

${\displaystyle {\frac {\partial }{\partial z_{i}}}\operatorname {LSE} (\mathbf {z} )={\frac {\exp z_{i}}{\sum _{j=1}^{K}\exp z_{j}}}=\sigma (\mathbf {z} )_{i},\quad {\text{ for }}i=1,\dotsc ,K,\quad \mathbf {z} =(z_{1},\,\dotsc ,\,z_{K})\in \mathbb {R} ^{K},}$

where the LogSumExp function is defined as ${\displaystyle \operatorname {LSE} (z_{1},\,\dots ,\,z_{n})=\log \left(\exp(z_{1})+\cdots +\exp(z_{n})\right)}$.

The softmax function was used in statistical mechanics as the Boltzmann distribution in the foundational paper Boltzmann (1868),[11] formalized and popularized in the influential textbook Gibbs (1902).[12]

The use of the softmax in decision theory is credited to Luce (1959) harvtxt error: no target: CITEREFLuce1959 (help),[13]^: 1 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, Bridle (1990a):[13]^: 1 and Bridle (1990b):[3]

We are concerned with feed-forward non-linear networks (multi-layer perceptrons, or MLPs) with multiple outputs. We wish to treat the outputs of the network as probabilities of alternatives (e.g. pattern classes), conditioned on the inputs. We look for appropriate output non-linearities and for appropriate criteria for adaptation of the parameters of the network (e.g. weights). We explain two modifications: probability scoring, which is an alternative to squared error minimisation, and a normalised exponential (softmax) multi-input generalisation of the logistic non-linearity.[14]^: 227

For any input, the outputs must all be positive and they must sum to unity. ...

Given a set of unconstrained values, ${\displaystyle V_{j}(x)}$, we can ensure both conditions by using a Normalised Exponential transformation:

${\displaystyle Q_{j}(x)=\left.e^{V_{j}(x)}\right/\sum _{k}e^{V_{k}(x)}}$

This transformation can be considered a multi-input generalisation of the logistic, operating on the whole output layer. It preserves the rank order of its input values, and is a differentiable generalisation of the 'winner-take-all' operation of picking the maximum value. For this reason we like to refer to it as softmax.[15]^: 213

If we take an input of [1, 2, 3, 4, 1, 2, 3], the softmax of that is [0.024, 0.064, 0.175, 0.475, 0.024, 0.064, 0.175]. 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: softmax is not scale invariant, so if the input were [0.1, 0.2, 0.3, 0.4, 0.1, 0.2, 0.3] (which sums to 1.6) the softmax would be [0.125, 0.138, 0.153, 0.169, 0.125, 0.138, 0.153]. This shows that for values between 0 and 1 softmax, in fact, de-emphasizes the maximum value (note that 0.169 is not only less than 0.475, it is also less than the initial proportion of 0.4/1.6=0.25).

Computation of this example using Python code:

>>> import numpy as np
>>> a = [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0]
>>> np.exp(a) / np.sum(np.exp(a)) 
array([0.02364054, 0.06426166, 0.1746813, 0.474833, 0.02364054,
       0.06426166, 0.1746813])

Here is an example of Julia code:

julia> A = [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0];  # semicolon to suppress interactive output

julia> exp.(A) ./ sum(exp.(A))
7-element Array{Float64,1}:
 0.0236405
 0.0642617
 0.174681
 0.474833
 0.0236405
 0.0642617
 0.174681

Here is an example of R code:

> z <- c(1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0)
> softmax <- exp(z)/sum(exp(z))
> softmax
[1] 0.02364054 0.06426166 0.17468130 0.47483300 0.02364054 0.06426166 0.17468130

Here is an example of Elixir code:[16]

iex> t = Nx.tensor([[1, 2], [3, 4]])
iex> Nx.divide(Nx.exp(t), Nx.sum(Nx.exp(t)))

#Nx.Tensor<
  f64[2][2]
  [
    [0.03205860328008499, 0.08714431874203257],
    [0.23688281808991013, 0.6439142598879722]
  ]
>

Here is an example of Raku code:

> my @z = [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0];
> say @z.map: {exp($_)/sum(@z.map: {exp($_)})}

(0.023640543021591385 0.06426165851049616 0.17468129859572226 0.4748329997443803 0.023640543021591385 0.06426165851049616 0.17468129859572226)

• 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.