Neural Mastery
← Practice
Numerically Stable Softmax and Cross-Entropyhard

Numerically Stable Softmax and Cross-Entropy

hard Math for ML 15-20 min

NumPy allowed · plain Python earns 10 bonus points

Linear algebra step 1/4

A vector transforms

Matrix multiplication is the first computation behind every model.
Open lesson →
01vector
02matrix
03transform
04result
[2, −1] × [[1, 2], [−.5, 1]] = [0, −2]

Computed matrix × vector transformation

Compute the softmax probability vector and negative log likelihood loss without floating-point overflow.

Related Skills

Task

Implement stable_softmax_and_ce(logits, target_idx): subtract m = max(logits) to calculate p_i = exp(z_i - m) / sum_j exp(z_j - m). Compute cross entropy loss = -log(max(p_target, 1e-15)). Return {"probs": list[float], "loss": float} rounded to 4 decimals.

Function Signature

stable_softmax_and_ce(logits: list[float], target_idx: int) -> dict[str, Any]

Examples

Example 1: Small Logits Binary
Input: logits = [2, 1] target_idx = 0
Output: {'probs': [0.7311, 0.2689], 'loss': 0.3133}
Explanation: Binary classification softmax.
Example 2: Large Overflowing Logits
Input: logits = [1000, 1001, 999] target_idx = 1
Output: {'probs': [0.2447, 0.6652, 0.09], 'loss': 0.4076}
Explanation: Logits exceeding exp(709) without overflow.

Constraints

  • 0 <= target_idx < len(logits).
  • Logits can be arbitrarily large or small.
Python3Saved ✓
def stable_softmax_and_ce(logits: list[float], target_idx: int) -> dict[str, Any]:
"""Compute numerically stable softmax probabilities and cross-entropy loss.
Subtract max(logits) before exponentiating. Return {'probs': list[float], 'loss': float}."""
# Your implementation here
pass

Case 1: Small Logits Binary
Input: logits = [2, 1] target_idx = 0
Expected: {'probs': [0.7311, 0.2689], 'loss': 0.3133}
Case 2: Large Overflowing Logits
Input: logits = [1000, 1001, 999] target_idx = 1
Expected: {'probs': [0.2447, 0.6652, 0.09], 'loss': 0.4076}