Numerically Stable Softmaxeasy

Numerically Stable Softmax

easy Neural Networks by Hand 15 min

NumPy allowed · plain Python earns 10 bonus points

Transform unnormalized log-odds (logits) into a valid probability distribution without floating point overflow.

Related Skills

Task

Implement softmax_stable(logits): for each row z in logits, compute exp(z_i - max(z)) / sum(exp(z_j - max(z))). Subtracting the row maximum ensures that the largest exponent evaluated is exp(0) = 1, preventing float overflow. Round every output probability to 4 decimal places.

Function Signature

softmax_stable(logits: list[list[float]]) -> list[list[float]]

Examples

Example 1: Standard small logits
Input: logits = [[1, 2, 3]]
Output: [[0.09, 0.2447, 0.6652]]

Example 2: Large logits requiring stabilization
Input: logits = [[1000, 1001, 1002]]
Output: [[0.09, 0.2447, 0.6652]]

Example 3: Batch of 2 rows
Input: logits = [[0, 0], [2, -2]]
Output: [[0.5, 0.5], [0.982, 0.018]]

Constraints

1 <= N <= 100
1 <= C <= 1000
Logits are real numbers

Python3Saved ✓

def softmax_stable(logits: list[list[float]]) -> list[list[float]]:

"""Compute row-wise softmax probabilities with max subtraction for numerical stability.

Args:

logits: 2D matrix of shape (N, C).

Returns:

2D probability matrix of shape (N, C) rounded to 4 decimals.

"""

pass

Case 1: Standard small logits

Input: logits = [[1, 2, 3]]

Expected: [[0.09, 0.2447, 0.6652]]

Case 2: Large logits requiring stabilization

Input: logits = [[1000, 1001, 1002]]

Expected: [[0.09, 0.2447, 0.6652]]

Case 3: Batch of 2 rows

Input: logits = [[0, 0], [2, -2]]

Expected: [[0.5, 0.5], [0.982, 0.018]]