Softmax Regression for Multiclass

Logistic regression elegantly solves binary classification, but most real-world problems—from recognizing handwritten digits to categorizing news articles—involve more than two classes. Softmax regression, also known as multinomial logistic regression, is the direct and natural extension of logistic regression to handle multiple, mutually exclusive classes. It provides a probabilistic framework that outputs a well-calibrated probability distribution over all possible classes, making it a foundational algorithm in machine learning and a critical component of modern neural networks.

From Logistic to Softmax: Extending the Probability Model

To understand softmax regression, recall that binary logistic regression uses the sigmoid function to squash a linear score $z$ into a probability between 0 and 1. The probability for the positive class is $P(y=1|\mathbf{x})=\frac{1}{1+e^{-z}}=\sigma (z)$. For $K$ classes, we need a function that takes $K$ linear scores (one per class) and outputs $K$ probabilities that sum to 1.

This is achieved by the softmax function. For a given input feature vector $\mathbf{x}$, we first compute a score $z_k$ for each class $k$. This score is typically a linear combination: $z_k={\mathbf{w}}_k^T\mathbf{x}+b_k$, where ${\mathbf{w}}_k$ and $b_k$ are the weight vector and bias for class $k$. The softmax function then converts these scores into a probability distribution:

$$P(y=k|\mathbf{x})=\frac{e^{z_k}}{{\sum }_{j=1}^K{e}^{z_j}}=\frac{e^{{\mathbf{w}}_k^T\mathbf{x}+b_k}}{{\sum }_{j=1}^K{e}^{{\mathbf{w}}_j^T\mathbf{x}+b_j}}$$

The softmax function is a generalization of the sigmoid. It exponentiates each score (making it positive) and then normalizes by the sum of all exponentiated scores. A key property is that it is sensitive to differences in scores: the class with the highest score will be assigned a probability close to 1, while others receive probabilities close to 0. This makes it ideal for classification.

The Cross-Entropy Loss Function for Multiclass Targets

Training a softmax regression model requires a loss function that measures how well the predicted probability distribution matches the true, one-hot encoded label. The standard loss function is multiclass cross-entropy loss, also called log loss.

For a single training example $({\mathbf{x}}^{(i)},y^{(i)})$, where $y^{(i)}$ is the true class index, the loss is the negative log probability assigned to the correct class:

$$L^{(i)}=-\log P(y=y^{(i)}|{\mathbf{x}}^{(i)})$$

If we represent the true label as a one-hot vector ${\mathbf{y}}^{(i)}$ (a vector of all zeros except a 1 at the index of the true class), the loss for a single example can be written compactly as:

$$L^{(i)}=-\sum _{k=1}^K{y}_k^{(i)}\log (p_k^{(i)})$$

where $p_k^{(i)}=P(y=k|{\mathbf{x}}^{(i)})$ is the softmax output for class $k$. For a full dataset of $m$ examples, the average loss, or cost, is:

$$J(\mathbf{W},\mathbf{b})=-\frac{1}{m}\sum _{i=1}^m\sum _{k=1}^K{y}_k^{(i)}\log \left(p_k^{(i)}\right)$$

This objective function is convex for linear softmax regression, meaning gradient descent is guaranteed to find the global minimum given appropriate hyperparameters.

Optimizing with Gradient Descent

To learn the optimal weight matrix $\mathbf{W}$ (which stacks all ${\mathbf{w}}_k^T$) and bias vector $\mathbf{b}$, we use gradient descent. The update rule for any parameter $\theta$ is: $\theta :=\theta -\alpha \frac{\partial J}{\partial \theta }$, where $\alpha$ is the learning rate.

The gradient of the loss with respect to the score $z_k$ has a remarkably simple and elegant form in softmax regression:

$$\frac{\partial L^{(i)}}{\partial z_k}=p_k^{(i)}-y_k^{(i)}$$

This gradient is the difference between the predicted probability and the true label for class $k$. For the correct class, where $y_k=1$, the gradient is negative $(p_k-1)$, pushing the model to increase the score for that class. For incorrect classes $(y_k=0)$, the gradient is positive $(p_k)$, pushing the model to decrease those scores. This gradient is then propagated back through the linear model to update the weights and biases. This efficient computation is a primary reason for the softmax-cross-entropy pairing's popularity in deep learning.

Comparison with One-vs-Rest Logistic Regression

An alternative approach to multiclass classification is the one-vs-rest (OvR) strategy, which involves training $K$ separate binary logistic regression models. Each model $k$ is trained to distinguish class $k$ from all other classes combined.

While OvR is simpler to implement with existing binary classifiers, softmax regression has distinct advantages:

• Coherent Probabilities: Softmax produces a single, normalized probability distribution over all classes. OvR produces $K$ independent probability scores, which may not sum to 1, requiring post-hoc normalization if interpreted as a joint distribution.
• Unified Decision Boundary: Softmax regression is trained jointly on all classes. This leads to a single, consistent set of weight vectors that define the decision boundaries between all classes simultaneously. In contrast, OvR trains classifiers independently, which can lead to inconsistencies in regions where multiple classifiers predict a positive class.
• Theoretical Foundation: Softmax is derived from a single probabilistic model (multinomial distribution), making it more statistically principled for mutually exclusive classes.

In practice, softmax is generally preferred when the classes are mutually exclusive. OvR can be useful when classes are not mutually exclusive or when you need to leverage highly optimized binary classifiers.

Decision Boundary Geometry in Multiclass Settings

Understanding the geometry of the decision boundaries helps visualize how softmax regression partitions the feature space. The model predicts the class $k$ with the highest estimated probability $P(y=k|\mathbf{x})$, which is equivalent to predicting the class with the highest score $z_k={\mathbf{w}}_k^T\mathbf{x}+b_k$.

The decision boundary between any two classes, say class $i$ and class $j$, is the set of points $\mathbf{x}$ where $P(y=i|\mathbf{x})=P(y=j|\mathbf{x})$. Given the softmax formula, this equality simplifies to:

$${\mathbf{w}}_i^T\mathbf{x}+b_i={\mathbf{w}}_j^T\mathbf{x}+b_j$$

This can be rewritten as $({\mathbf{w}}_i-{\mathbf{w}}_j{)}^T\mathbf{x}+(b_i-b_j)=0$, which is the equation of a linear boundary (a line in 2D, a plane in 3D, or a hyperplane in higher dimensions). Crucially, softmax regression produces piecewise linear decision boundaries. The feature space is divided into $K$ convex regions (where region $k$ contains all points for which class $k$ has the highest score), each separated from the others by linear boundaries. This linearity is a direct result of using linear score functions; non-linear boundaries require feature transformations (like polynomial features) or moving to neural networks.

Common Pitfalls

1. Numerical Instability with Large Scores: Direct computation of $e^{z_k}$ can overflow (result in inf) if any score $z_k$ is very large. The standard remedy is to use a normalization trick: subtract the maximum score, $c=\max (z_k)$, from all scores before exponentiating. This shifts the values without changing the output probabilities:

$P(y=k|\mathbf{x})=\frac{e^{z_k-c}}{{\sum }_{j=1}^K{e}^{z_j-c}}$. This ensures the largest exponentiated value is $e^0=1$, preventing overflow.

1. Interpreting Outputs as Confidence Calibration: While softmax outputs a probability distribution, these values are not always well-calibrated measures of confidence, especially for complex models. A probability of 0.9 does not guarantee a 90% chance of being correct. The probabilities tend to be overconfident, particularly in deep neural networks. Techniques like temperature scaling or using Bayesian methods are needed for true calibration.

1. Assuming Mutually Exclusive Classes: Softmax regression is designed for single-label, multiclass classification where each example belongs to exactly one class. Applying it to multi-label problems (where an example can have multiple true labels) will yield misleading results. For multi-label tasks, you would train independent sigmoid output units for each class instead.

1. Ignoring Class Imbalance: The standard cross-entropy loss treats each example equally. If your dataset has severe class imbalance (e.g., 1000 examples of class A and 10 of class B), the model may become biased toward the majority class. Mitigation strategies include resampling the dataset, weighting the loss function inversely by class frequency, or using focal loss.

Summary

• Softmax regression is the fundamental extension of logistic regression to $K>2$ mutually exclusive classes, outputting a full probability distribution via the softmax function.
• It is trained by minimizing the multiclass cross-entropy loss, which measures the discrepancy between the predicted probability distribution and the true one-hot encoded label.
• Optimization via gradient descent is efficient due to the simple gradient form: $\frac{\partial L}{\partial z_k}=p_k-y_k$.
• Compared to the one-vs-rest strategy, softmax provides a single, coherent probabilistic model and unified decision boundaries, making it the preferred method for mutually exclusive classes.
• The model creates piecewise linear decision boundaries in the feature space, with linear separators between each pair of classes defined by the learned weight vectors.