Implement Softmax and Cross-Entropy Loss

Problem

Implement the softmax activation function and cross-entropy loss for multi-class classification.

Softmax: Converts logits (raw scores) into probabilities that sum to 1. Formula: softmax(x_i) = exp(x_i) / sum(exp(x_j)) for all j

Cross-Entropy Loss: Measures difference between predicted probabilities and true labels. Formula: -sum(y_true * log(y_pred))

Your implementation should:

1. Handle numerical stability (prevent overflow in exp)
2. Work with batched inputs (multiple samples)
3. Support both one-hot encoded and integer labels

Constraints

• Batch size: 1 ≤ batch_size ≤ 1000
• Number of classes: 2 ≤ num_classes ≤ 1000
• Logits are float values (can be negative)
• Labels are either one-hot vectors or integers

Examples

Example 1

Input:

# Single sample, 3 classes
logits = [2.0, 1.0, 0.1]

Output:

probabilities = [0.659, 0.242, 0.099]
# Sum = 1.0

Explanation: Softmax converts scores to probabilities

Example 2

Input:

# Cross-entropy loss
y_pred = [0.7, 0.2, 0.1]  # Predicted probabilities
y_true = [1, 0, 0]         # True label is class 0

Output:

loss = -log(0.7) = 0.357

Explanation: Loss is negative log probability of correct class

Function Signature

def softmax(logits: np.ndarray) -> np.ndarray:
    """
    Compute softmax probabilities.

    Args:
        logits: Raw scores, shape (batch_size, num_classes) or (num_classes,)

    Returns:
        probabilities: Same shape as logits, values sum to 1 along last axis
    """
    pass

def cross_entropy_loss(y_pred: np.ndarray, y_true: np.ndarray) -> float:
    """
    Compute cross-entropy loss.

    Args:
        y_pred: Predicted probabilities, shape (batch_size, num_classes)
        y_true: True labels, either:
                - One-hot encoded: shape (batch_size, num_classes)
                - Integer labels: shape (batch_size,)

    Returns:
        loss: Average cross-entropy loss (scalar)
    """
    pass

Estimated Time

20-30 minutes

Tags

neural-networks classification loss-functions numpy easy