Softmax Classifier Optimization via Gradient Descent

Softmax Classifier Optimization

The softmax classifier (also known as multinomial logistic regression) is a fundamental algorithm for multi-class classification problems. It generalizes binary logistic regression to handle K > 2 mutually exclusive classes by computing a probability distribution over all possible outcomes.

Core Mathematical Formulation

Given a dataset with n samples and d features, where each sample belongs to one of K classes, the softmax classifier learns a weight matrix W of dimensions (d × K). For an input feature vector x, the model computes class probabilities using the softmax function:

$$P(y = k | x) = \frac{e^{x^T w_k}}{\sum_{j=1}^{K} e^{x^T w_j}}$$

where w_k is the k-th column of the weight matrix, representing the learned parameters for class k.

Cross-Entropy Loss Function

The objective is to minimize the cross-entropy loss (also called negative log-likelihood), which measures the dissimilarity between predicted probabilities and true labels:

$$L(W) = -\frac{1}{n} \sum_{i=1}^{n} \log P(y_i | x_i)$$

For a true label y_i = k, this becomes -log(P(y = k | x_i)), penalizing low probability assignments to the correct class.

Gradient Descent Optimization

To minimize the loss, we compute the gradient with respect to the weights and iteratively update them:

$$W \leftarrow W - \alpha \cdot abla_W L$$

where α is the learning rate. The gradient has an elegant form: for each class k, it equals X^T (P - Y) / n, where P is the matrix of predicted probabilities and Y is the one-hot encoded label matrix.

Implementation Details

Your implementation should:

1. Initialize weights to zeros with shape (n_features, n_classes)
2. Forward pass: Compute softmax probabilities for all samples
3. Compute loss: Calculate cross-entropy loss for the current weights
4. Backward pass: Compute gradients with respect to weights
5. Update weights: Apply gradient descent step
6. Record loss: Store the loss value at each iteration
7. Return results: Final weight matrix and loss trajectory, both rounded to 4 decimal places

Numerical Stability Consideration

When implementing softmax, subtract the maximum value in each row of the logits before exponentiation to prevent numerical overflow: $$\text{softmax}(z)_k = \frac{e^{z_k - \max(z)}}{\sum_j e^{z_j - \max(z)}}$$

This mathematically equivalent formulation prevents inf values from appearing during computation.