Why use cross-entropy instead of MSE for classification?

An analysis from a gradient perspective

When it comes to training machine learning models for classification tasks, a common question arises: Why do we prefer using cross-entropy loss instead of mean squared error (MSE)? While both are loss functions used to measure how well a model’s predictions match the actual data, cross-entropy is often favored for classification problems. Let’s explore why this is the case, particularly from the perspective of gradients and how they affect training.

Understanding the Basics: Classification and Loss Functions

Classification tasks involve categorizing inputs into different classes. For example, a model might predict whether an image contains a cat, a dog, or a bird. To evaluate how well the model performs, we need a loss function, which is a mathematical way of measuring the difference between the model’s predictions and the actual labels.

Two common loss functions are:

1. Mean Squared Error (MSE):

MSE is often used in regression tasks, where the model predicts a continuous value. The formula for MSE is:

\( L_{\text{MSE}} = \frac{1}{n} \sum_{i=1}^{n} (y_i - \hat{y}_i)^2 \)

Here,  y_i  represents the actual value, and  y_i with a hat is the predicted value. We sum up the squared differences for all predictions and then take the average. The goal is to minimize this loss.

2. Cross-Entropy Loss:

Cross-entropy loss is widely used for classification tasks, particularly when using the softmax function. The formula for cross-entropy loss for a single example is:

\( L_{\text{CE}} = -\sum_{i=1}^{C} y_i \log(p_i) \)

• C  is the number of classes.

• y_i  is the actual label (1 if it is the correct class, 0 otherwise).

• p_i  is the predicted probability of the class.

How Cross-Entropy and MSE Differ: A Gradient Perspective

The Role of Gradients in Training

To train a neural network, we use an algorithm called gradient descent, which adjusts the model’s parameters to minimize the loss function. The gradient is like a “slope” that tells the model which direction to move its parameters to reduce the loss.

When the gradient is large, the model updates its parameters quickly. If the gradient is small, the updates are slow. Therefore, the size of the gradient is crucial for efficient training.

Gradients in MSE Loss

Let’s look at what happens when we use MSE in a classification task. Suppose we have a neural network with softmax output. Softmax is a function that converts the model’s raw outputs (called logits) into probabilities:

\( p_i = \frac{e^{z_i}}{\sum_{j=1}^{C} e^{z_j}} \)

Here:

• z_i  is the logit for class  i .

• p_i  is the probability assigned to class  i .

If we calculate the gradient of the MSE loss with respect to the logits, we get:

\( \frac{\partial L_{\text{MSE}}}{\partial z_i} = (p_i - y_i) \cdot p_i \cdot (1 - p_i) \)

This formula shows that the gradient is influenced by three factors:

1. (p_i - y_i)  – the difference between the predicted probability and the actual label.

2. p_i  – the predicted probability itself.

3. (1 - p_i)  – one minus the predicted probability.

The problem here is that when  p_i  is either very small (close to 0) or very large (close to 1), the term  p_i (1 - p_i)  becomes very small. This makes the overall gradient tiny, which can slow down training significantly. This issue is called the vanishing gradient problem, and it becomes more pronounced in deeper networks.

Gradients in Cross-Entropy Loss

Now, let’s consider the gradient of cross-entropy loss:

\( \frac{\partial L_{\text{CE}}}{\partial z_i} = p_i - y_i \)

This gradient is much simpler and does not involve additional terms like  p_i(1 - p_i) . As a result:

• The gradient remains reasonably large, even when  p_i  is close to 0 or 1.

• This prevents the vanishing gradient problem and allows the model to update its parameters more effectively.

Why Cross-Entropy Is Better for Classification

In classification tasks, the model’s output is usually a probability distribution across different classes (using softmax). Cross-entropy loss directly measures the difference between the predicted probabilities and the actual labels in a way that encourages the model to become more confident in its predictions. It also ensures that the gradients are stable and effective for training, avoiding issues like vanishing gradients.

Visualizing the Difference: An Example

Imagine a simple classification problem with two classes: A and B. Suppose the model initially predicts 0.1 for class A and 0.9 for class B, but the actual label is A (which means it should predict 1.0 for A and 0.0 for B).

Using MSE:

• The gradient becomes small because of the additional terms  p_i(1 - p_i) .

• The model updates its parameters slowly, and it might take longer to correct the error.

Using Cross-Entropy:

• The gradient depends only on  p_i - y_i , so it remains relatively large.

• The model adjusts quickly to make class A’s probability higher, leading to faster and more efficient learning.

Conclusion

Cross-entropy loss is preferred over MSE in classification tasks because it provides gradients that are more effective for learning. It avoids the vanishing gradient problem, which allows the model to train faster and perform better.