softmax

The softmax function (also known as normalized exponential) is a generalization and extension of the sigmoid to move from binary classification to multi-class classification.

It takes a vector of $K$ real numbers (logits) and normalizes it into a probability distribution consisting of $K$ probabilities proportional to the exponentials of the input numbers.

For an input vector $\mathbf{z} = (z_1, z_2, \dots, z_K) \in \mathbb{R}^K$, the Softmax function $\sigma(\mathbf{z})_i$ for a specific component $i$ is defined as:

Properties

#

• Range: Each output probability \(\sigma(\mathbf{z})_i \in (0, 1)\).
• Sum to One: The sum of all outputs equals exactly one: $\sum_{i=1}^{K} \sigma(\mathbf{z})_i = 1$.
• Monotonicity: Larger inputs map to larger probabilities.
• When applied to a binary classification task $K=2$, the Softmax probability for the first class simplifies mathematically to the standard Sigmoid function: $$\sigma(\mathbf{z})_1 = \frac{e^{z_1}}{e^{z_1} + e^{z_2}} = \frac{1}{1 + e^{-(z_1 - z_2)}} = \text{Sigmoid}(z_1 - z_2)$$

Machine Learning Applications

#

Multi-Class Classification

#

Softmax is almost universally used as the activation function in the final layer of a neural network when performing mutually exclusive multi-class classification tasks.

Cross-Entropy Loss

#

In training pipelines, Softmax is typically paired with the Cross-Entropy loss function.
Together, they create a smooth gradient that heavily penalizes confident but incorrect predictions, accelerating backpropagation convergence.

Implementation

#

Given a vector of logits, Softmax can be implemented following this algorithm:

• find the max value and subtract it from all the values in the vector obtaining the shifted values
• calculate the exponential of the shifted values
• divide by the total sum of the exponential of the shifted values to find probabilities

In numpy:

logits = np.array([[2., 1., 0.], [-2., 0., 2.]])

def softmaxNumpy(logits):
	#operations target columns, so axis=-1
	row_max = logits.max(axis=-1, keepdims=True)
	shifted = logits - row_max
	exponentials = np.exp(shifted)
	denominators = exponentials.sum(axis=-1, keepdims=True)
	probabilities = exponentials / denominators
	
	return probabilities

In PyTorch:

logits = torch.tensor([[2., 1., 0.], [-2., 0., 2.]])

def softmaxPyTorch(logits)
	#operations target columns, so dim=-1
	row_max = logits.amax(dim=-1, keepdims=True)
	shifted = logits - row_max
	exponentials = torch.exp(shifted)
	denominators = exponentials.sum(dim=-1, keepdims=True)
	probabilities = exponentials / denominators
	
	return probabilities