Softmax for Multi-Class Classification

Introduction

In our last lesson, we successfully navigated the forward pass, propagating an input through the network's layers to generate a vector of raw output values, which we call logits. This is a significant step, but it leaves us with a critical question: how do we interpret these raw, unbounded numbers, especially in the context of classification? If our goal is to assign an input to one of several categories (e.g., classifying a handwritten digit as 0 through 9), we need a principled way to convert these logits into a set of probabilities.

This lesson introduces the softmax function, the standard mechanism in neural networks for this very purpose. We will define the function, explore its use in multi-class classification, and derive it from first principles. In line with your goal of bridging classical and modern concepts, we'll see how softmax elegantly generalizes the sigmoid function used in binary logistic regression to the multi-class setting. We will also investigate its key mathematical properties, which are foundational for understanding how classification networks are trained.

From Binary to Multi-Class Classification

Let's begin by revisiting a simpler case: binary classification. Here, we typically have a single output neuron producing a logit $z$. We apply the sigmoid function, $\sigma (z)=1/(1+e^{-z})$, to map this logit to a probability $p\in [0,1]$. If $p$ is the probability of class 1, the probability of class 0 is simply $1-p$.

Now, consider a problem with $K>2$ classes. A single output neuron is no longer sufficient.

Question: What is the most direct way to adapt a network's output layer to handle $K$ classes? And what two properties must the final output vector satisfy to be considered a valid probability distribution?

Answer

The most straightforward approach is to have $K$ output neurons, one for each class. Each neuron $k$ produces its own logit, $z_k$. The final output of the network is a vector of these logits, $\mathbf{z} = [z_1, z_2, ..., z_K]^T$.

For this output vector to represent a valid probability distribution, it must satisfy two conditions:

1. Each element must be a valid probability: $0\le p_k\le 1$ for all $k$.
2. The probabilities for all classes must sum to one: ${\sum }_{k=1}^K{p}_k=1$.

Our challenge, then, is to find a function that takes the vector of arbitrary real-valued logits $\mathbf{z}$ and transforms it into a probability vector $\mathbf{p}$ that meets these two conditions. This is precisely the role of the softmax function.

The image below provides a great schematic for this conceptual jump. On the top, you see binary logistic regression using a sigmoid function for one-of-two classes. On the bottom, you see softmax regression, which handles one-of-K classes.

Deriving the Softmax Function from First Principles

Let's build the softmax function from the ground up. We start with our vector of logits, $\mathbf{z}=[z_1,z_2,...,z_K{]}^T$, where each $z_i$ can be any real number.

A simple, naive attempt to create probabilities might be to normalize by the sum: $p_i=z_i/{\sum }_j{z}_j$. This satisfies the "sum-to-one" property. However, what is the critical flaw in this approach, given that logits can be negative? And what is a common mathematical operation that can map any real number to a strictly positive one?

The flaw is that if any $z_i$ is negative, we would get a negative "probability," which is nonsensical. To solve this, we need to ensure all our scores are positive before normalization. The exponential function, $f(x)=e^x$, is a perfect candidate for this. It maps any real number to a positive number and has the added benefit of preserving the order of the original logits (if $z_i>z_j$, then $e^{z_i}>e^{z_j}$).

By first exponentiating each logit and then normalizing by the sum of these exponentiated values, we arrive at the softmax function:

$$\sigma (\mathbf{z}{)}_i=p_i=\frac{e^{z_i}}{{\sum }_{j=1}^K{e}^{z_j}}\text{for }i=1,\dots ,K$$

Each resulting $p_i$ is guaranteed to be in $(0,1)$, and their sum is guaranteed to be 1, fulfilling our requirements for a probability distribution. The name "softmax" comes from the fact that the exponential function strongly amplifies the largest logit, making its corresponding probability close to 1, while suppressing the others. It's a "soft," or differentiable, approximation of the non-differentiable argmax function.

The following article provides a clear, narrative-driven derivation of this exact logic.

Properties and Visual Intuition

To truly understand softmax, it's helpful to visualize its behavior and examine its mathematical properties. The function has several key characteristics that are critical for both its use in modeling and its role in training.

Let's formalize two of these properties.

1. Shift Invariance

As shown in the video, if we add a constant $c$ to every logit, the output of the softmax function remains unchanged. That is, $\text{softmax}(\mathbf{z})=\text{softmax}(\mathbf{z}+c)$.

Let's prove this algebraically. Let $z_i^{\prime }=z_i+c$. The new softmax output for class $i$ is:

$$p_i^{\prime }=\frac{e^{z_i^{\prime }}}{{\sum }_{j=1}^K{e}^{z_j^{\prime }}}=\frac{e^{z_i+c}}{{\sum }_{j=1}^K{e}^{z_j+c}}=\frac{e^{z_i}{e}^c}{{\sum }_{j=1}^K(e^{z_j}{e}^c)}=\frac{e^{z_i}{e}^c}{e^c{\sum }_{j=1}^K{e}^{z_j}}$$

The $e^c$ terms in the numerator and denominator cancel out, leaving us with the original softmax expression.

$$p_i^{\prime }=\frac{e^{z_i}}{{\sum }_{j=1}^K{e}^{z_j}}=p_i$$

This property is not just a mathematical curiosity. From a systems engineering perspective, it's vital for numerical stability. If the logits contain very large values, calculating $e^{z_i}$ can lead to overflow. By subtracting the maximum logit from all logits before applying the softmax ($c=-\max (\mathbf{z})$), we shift the values so the largest is 0, preventing overflow while yielding the exact same result.

2. Equivalence to Sigmoid for K=2

Softmax is a true generalization of the sigmoid. Let's demonstrate that for two classes ($K=2$), softmax reduces to the sigmoid function. Our logits are $[z_1,z_2]$.

The softmax probability for class 1 is:

$$p_1=\frac{e^{z_1}}{e^{z_1}+e^{z_2}}$$

We can divide the numerator and denominator by $e^{z_1}$:

$$p_1=\frac{1}{1+e^{z_2}/e^{z_1}}=\frac{1}{1+e^{z_2-z_1}}$$

This is equivalent to the sigmoid function applied to the difference in logits, $\sigma (z_1-z_2)$. This perfectly illustrates the "Classical to Modern Bridge": the multi-class softmax gracefully collapses into the familiar binary sigmoid.

For a formal write-up of this property and others, the following resource is excellent.

The Derivative of Softmax: A Glimpse into Backpropagation

For our next lesson, we will need to compute the gradient of the loss function with respect to the network's parameters. A key step in that process involves finding the derivative of the softmax output with respect to its input logits. This is the Jacobian matrix of the softmax function, where each entry is $\frac{\partial p_i}{\partial z_j}$.

Let's derive this, as it reveals the dynamics of how a change in one logit affects all probabilities. We must consider two cases.

Case 1: $i=j$ (The derivative of a softmax output with respect to its own logit)

We use the quotient rule for $\frac{\partial p_i}{\partial z_i}$ where $p_i=\frac{e^{z_i}}{{\sum }_k{e}^{z_k}}$.

$$\frac{\partial p_i}{\partial z_i}=\frac{(e^{z_i})({\sum }_k{e}^{z_k})-(e^{z_i})(e^{z_i})}{({\sum }_k{e}^{z_k}{)}^2}$$

Splitting the fraction:

$$=\frac{e^{z_i}}{{\sum }_k{e}^{z_k}}-{\left(\frac{e^{z_i}}{{\sum }_k{e}^{z_k}}\right)}^2=p_i-p_i^2=p_i(1-p_i)$$

This is identical to the derivative of the sigmoid function, which makes sense given their relationship.

Case 2: $i\ne j$ (The derivative of a softmax output with respect to another logit)

$$\frac{\partial p_i}{\partial z_j}=\frac{(0)({\sum }_k{e}^{z_k})-(e^{z_i})(e^{z_j})}{({\sum }_k{e}^{z_k}{)}^2}$$

Rearranging the terms:

$$=-\left(\frac{e^{z_i}}{{\sum }_k{e}^{z_k}}\right)\left(\frac{e^{z_j}}{{\sum }_k{e}^{z_k}}\right)=-p_i{p}_j$$

This elegant result shows that increasing the logit $z_j$ has a negative effect on all other probabilities $p_i$, proportional to the magnitude of both probabilities. This interaction is the core of the competitive nature of multi-class classification. The mathematical derivation is also covered clearly in the following video.

Conclusion

In this lesson, we have formally defined the softmax function and established its critical role in multi-class classification. We have seen that it is not an arbitrary choice, but a principled function that can be derived from the fundamental need to convert unbounded logits into a valid probability distribution.

Key Takeaways:

• Purpose: The softmax function transforms a vector of $K$ real-valued logits into a $K$-dimensional vector of probabilities that sum to 1.
• Mechanism: It works by exponentiating its inputs to ensure positivity and then normalizing them.
• Bridge Concept: It is a direct generalization of the sigmoid function; for $K=2$, softmax is equivalent to the sigmoid.
• Properties: It is invariant to shifts in the input (crucial for numerical stability) and its derivative has a simple, elegant form ($p_i(1-p_i)$ or $-p_i{p}_j$) that captures the competitive relationship between class probabilities.

Preview of the Next Lesson:

Now that our network can produce a vector of probabilities $\mathbf{p}$, we can compare this prediction to the true, one-hot encoded label $\mathbf{y}$. To do this, we need a loss function. For multi-class classification, the standard is the cross-entropy loss. In our next lesson, we will derive this loss function and then combine its derivative with the softmax derivative we found today to compute the gradient of the loss with respect to the network's logits. This will yield the famous p - y gradient, the starting point for backpropagation in nearly all modern classification architectures.