The softmax function, also known as softargmax or normalized exponential function, is a mathematical function that converts a vector of real numbers into a probability distribution. It is a generalization of the logistic function to multiple dimensions and is widely used in machine learning, particularly in deep learning for multi-class classification problems.
Definition
Given an input vector \( \mathbf{z} = (z_1, \dots, z_K) \in \mathbb{R}^K \) with \( K > 1 \), the softmax function \( \sigma: \mathbb{R}^K \to (0,1)^K \) is defined for each component \( i \) as:
\[ \sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} \]
The output is a vector of positive values that sum to 1, making it interpretable as a categorical probability distribution. The exponential operation amplifies differences in the input, so larger input values yield disproportionately larger probabilities. For example, applying softmax to the vector \( (1, 2, 8) \) yields approximately \( (0.001, 0.002, 0.997) \), assigning nearly all probability mass to the largest element.
Properties
The softmax function has several important properties:
- Output range: Each component lies strictly between 0 and 1.
- Sum to one: The sum of all output components equals 1.
- Order preservation: If \( z_i > z_j \), then \( \sigma(\mathbf{z})_i > \sigma(\mathbf{z})_j \).
- Invariance to constant shifts: Adding a constant \( c \) to all inputs leaves the output unchanged, since \( e^{z_i + c} / \sum_j e^{z_j + c} = e^{z_i} / \sum_j e^{z_j} \). This property is often used for numerical stability by subtracting the maximum input value before computation.
Temperature Parameter
A variant of the softmax function introduces a temperature parameter \( \beta \) (or its inverse \( T = 1/\beta \)):
\[ \sigma(\mathbf{z})_i = \frac{e^{\beta z_i}}{\sum_{j=1}^{K} e^{\beta z_j}} \]
When \( \beta > 1 \), the distribution becomes more concentrated ("sharper") around the maximum input, while \( 0 < \beta < 1 \) produces a more uniform distribution. This parameter is commonly used in top-k sampling and temperature scaling for generative AI models to control output randomness.
Applications
Softmax is a core component in many neural network architectures:
- Classification: As the final activation layer, it converts raw logits into class probabilities, enabling training with cross-entropy loss.
- Attention mechanisms: In Transformer (architecture) models, softmax is used to compute attention weights over sequence positions, as seen in multi-head attention.
- Reinforcement learning: It converts action preferences into a stochastic policy.
- Mixture models: It is used to compute mixture weights in models like Gaussian mixture models.
Numerical Stability
Direct computation of the exponential can lead to overflow for large inputs. A common remedy is to subtract the maximum input value before exponentiation:
\[ \sigma(\mathbf{z})_i = \frac{e^{z_i - \max(\mathbf{z})}}{\sum_{j=1}^{K} e^{z_j - \max(\mathbf{z})}} \]
This transformation does not change the result due to the shift invariance property, but it ensures that the largest exponent is zero, preventing overflow.
Related Concepts
The softmax function is closely related to the logistic function, which is its special case for \( K = 2 \). It also appears in the context of loss functions like cross-entropy, and in beam search decoding for sequence generation. In large language model inference, softmax is applied to logits to obtain token probabilities, often with temperature adjustments to balance creativity and coherence.