Softmax and Categorical Cross-Entropy

1. The Softmax Function: Probabilities for Multiple Classes

Think of the softmax function as the multi-class generalization of the sigmoid function. While sigmoid takes a single score and squashes it to a probability between 0 and 1 (ideal for binary choices), softmax takes a vector of scores and squashes it into a vector of probabilities for K classes. The key properties of the output are:

• Each probability is between 0 and 1.
• The sum of all probabilities in the vector is exactly 1.

This gives us a true probability distribution over all possible classes. Mathematically, for a vector of raw scores (logits) $\mathbf{z}=[z_1,z_2,...,z_K]$, the softmax function is defined for each element $i$ as:

$\text{softmax}(\mathbf{z}{)}_i=\frac{e^{z_i}}{{\sum }_{j=1}^K{e}^{z_j}}$

To see this in action, the following short video provides a clear conceptual overview.

Implementing a Numerically Stable Softmax

A direct translation of the formula into NumPy is straightforward: np.exp(x) / np.sum(np.exp(x)). However, this naive approach has a critical flaw. The np.exp() function can easily overflow if the input numbers (logits) are large, resulting in inf and nan values, which will break your entire training process.

Given your background in CSE, you'll appreciate that we need a more robust, numerically stable implementation. The solution is surprisingly simple but elegant. The following article explains the problem and the fix.

As the article explains, we can subtract a constant from every logit before exponentiating without changing the result. By choosing this constant to be max(z), we ensure the largest input to np.exp() is 0, completely preventing overflow.

Let's implement this stable version.

import numpy as np

def softmax(z):
    """
    Computes the softmax function for a vector or a batch of vectors.
    
    Args:
        z: A numpy array of shape (n_samples, n_classes) or (n_classes,).
           These are the raw output scores (logits) from a linear layer.
           
    Returns:
        A numpy array of the same shape as z, containing the probabilities.
    """
    # Ensure z is at least 2D for consistent axis handling
    if z.ndim == 1:
        z = z.reshape(1, -1)
        
    # Subtract the max for numerical stability. `keepdims=True` is crucial
    # for broadcasting to work correctly.
    z_stable = z - np.max(z, axis=1, keepdims=True)
    
    # Exponentiate
    exp_z = np.exp(z_stable)
    
    # Normalize to get probabilities
    probabilities = exp_z / np.sum(exp_z, axis=1, keepdims=True)
    
    # If original input was 1D, flatten the output
    if probabilities.shape[0] == 1 and z.ndim == 1:
        return probabilities.flatten()
        
    return probabilities

# Example usage:
logits = np.array([2.0, 1.0, 0.1])
probs = softmax(logits)
print(f"Logits: {logits}")
print(f"Probabilities: {probs}")
print(f"Sum of probabilities: {np.sum(probs)}")

# Example with a large logit that would overflow a naive implementation
large_logits = np.array([1000, 1010, 990])
stable_probs = softmax(large_logits)
print(f"\nLarge Logits: {large_logits}")
print(f"Stable Probabilities: {stable_probs}")
print(f"Sum of stable probabilities: {np.sum(stable_probs)}")

Notice the use of axis=1 and keepdims=True. This makes our function robust enough to handle a batch of inputs (a 2D array where each row is a different sample) correctly, which will be vital when we build our full training loop.

2. Categorical Cross-Entropy Loss

Now that our model can produce a probability distribution using softmax, we need a loss function to measure how "good" that distribution is compared to the true label. For multi-class classification, the standard is categorical cross-entropy.

The general formula is:
$L=-{\sum }_{i=1}^K{y}_i\log ({\hat{y}}_i)$

Where:

• $K$ is the number of classes.
• ${\hat{y}}_i$ is the model's predicted probability for class $i$.
• $y_i$ is the true label for class $i$. For this calculation, we represent the true label as a one-hot encoded vector. For a 3-class problem where the correct class is 1, the one-hot vector is [0, 1, 0].

Because the true label vector $y$ is all zeros except for a single 1 at the correct class index $c$, the big summation simplifies to just one term:
$L=-1\cdot \log ({\hat{y}}_c)=-\log ({\hat{y}}_c)$
The loss is simply the negative logarithm of the probability assigned to the correct class. This is intuitive:

• If the model is confident and correct (e.g., ${\hat{y}}_c=0.99$), the loss is small: $-\log (0.99)\approx 0.01$.
• If the model is confident and wrong (e.g., ${\hat{y}}_c=0.01$), the loss is large: $-\log (0.01)\approx 4.6$.

Let's watch a brief video segment that visualizes this calculation.

Implementing Categorical Cross-Entropy Loss

Similar to softmax, a naive implementation of cross-entropy can be unstable. If the model predicts a probability of 0 for the correct class, we end up with log(0), which is negative infinity. Our stable softmax function helps prevent exact zeros, but to be fully robust, it's good practice to be aware of how libraries handle this.

The same article from Jay Mody you read earlier also covers this. It introduces a log_softmax function that combines the two operations for maximum stability.

For our purposes, we will implement the loss function separately from softmax for conceptual clarity, as this makes it easier to derive the gradients in the next lesson. Our stable softmax implementation should protect us from most issues.

Here is the implementation. We will write a function that takes the predicted probabilities (from our softmax function) and the true class indices (e.g., 0, 1, 2, ...).

def categorical_cross_entropy(y_pred_proba, y_true):
    """
    Calculates the categorical cross-entropy loss for a batch of predictions.
    
    Args:
        y_pred_proba: A numpy array of shape (n_samples, n_classes) containing
                      the predicted probabilities from the softmax function.
        y_true: A numpy array of shape (n_samples,) containing the true class
                indices (e.g., 0, 1, 2, ...).
                
    Returns:
        The average loss over the batch (a single scalar value).
    """
    n_samples = y_pred_proba.shape[0]
    
    # We want to select the predicted probability for the correct class for each sample.
    # NumPy's advanced indexing allows us to do this in one vectorized operation.
    # `y_pred_proba[range(n_samples), y_true]` creates a 1D array of the
    # probabilities assigned to the true classes.
    correct_class_probabilities = y_pred_proba[range(n_samples), y_true]
    
    # Add a small epsilon for numerical stability in case a probability is exactly 0
    epsilon = 1e-15
    correct_class_probabilities = np.clip(correct_class_probabilities, epsilon, 1 - epsilon)
    
    # Calculate the log-likelihoods
    log_likelihoods = -np.log(correct_class_probabilities)
    
    # Calculate the average loss over the batch
    average_loss = np.sum(log_likelihoods) / n_samples
    
    return average_loss

# Example usage:
# Let's say we have 3 samples and 3 classes (0, 1, 2)
# The true labels are [0, 1, 2]
y_true_labels = np.array([0, 1, 2])

# Case 1: Good predictions (high probability for correct classes)
good_probas = np.array([
    [0.9, 0.05, 0.05],  # Sample 0, correct is 0
    [0.1, 0.8, 0.1],   # Sample 1, correct is 1
    [0.2, 0.2, 0.6]    # Sample 2, correct is 2
])
good_loss = categorical_cross_entropy(good_probas, y_true_labels)
print(f"Loss for good predictions: {good_loss:.4f}")

# Case 2: Bad predictions (low probability for correct classes)
bad_probas = np.array([
    [0.1, 0.4, 0.5],   # Sample 0, correct is 0
    [0.8, 0.1, 0.1],   # Sample 1, correct is 1
    [0.7, 0.2, 0.1]    # Sample 2, correct is 2
])
bad_loss = categorical_cross_entropy(bad_probas, y_true_labels)
print(f"Loss for bad predictions: {bad_loss:.4f}")

As expected, the loss for the "bad" predictions is significantly higher than for the "good" ones. This is exactly what we want from a loss function: it gives our optimization algorithm (like gradient descent) a clear signal on how to adjust the model's parameters.

Conclusion

Excellent work today! You have implemented two of the most fundamental components in modern classification models.

Key Takeaways:

• The softmax function converts raw model scores (logits) into a multi-class probability distribution, serving as the activation for the output layer.
• Implementing a numerically stable softmax by subtracting the maximum logit is crucial for building robust models that don't fail with large inputs.
• Categorical cross-entropy is the standard loss function for multi-class classification. It efficiently measures the "distance" between the predicted probability distribution and the true label.
• The loss is essentially the negative log-probability of the correct class, providing an intuitive and effective signal for model training.

Preview of the Next Lesson:
We are now perfectly positioned for the next step. In the upcoming lesson, you will implement Softmax Regression for multi-class classification. We will combine the softmax function and categorical_cross_entropy loss with your gradient descent optimizer. This will involve the final piece of the puzzle: deriving the gradient of the loss function with respect to the model's weights. Get ready to build a complete, end-to-end multi-class classifier