Cross-Entropy Loss from Logitseasy

Cross-Entropy Loss from Logits

easy Inside an LLM 15 min

NumPy allowed · plain Python earns 10 bonus points

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.

Open lesson →

01tokens

02Q K V

03scores

04attention

modelslearnrelationships

models32%

learn38%

relationships30%

Q · Kᵀ / √d scores normalized with softmax

Compute the cross-entropy loss for a single-token prediction from raw logits and true target index.

Task

Implement cross_entropy_loss(logits, target_idx). Apply numerically stable softmax by subtracting max(logits), then return -log(p_target) rounded to 4 decimal places. Raise ValueError if logits is empty, or IndexError if target_idx is out of bounds.

Function Signature

cross_entropy_loss(logits: list[float], target_idx: int) -> float

Examples

Example 1: confident correct target
Input: logits = [10, 1, 0.5]
target_idx = 0
Output: 0.0002

Example 2: uncertain target
Input: logits = [2, 2, 2]
target_idx = 1
Output: 1.0986

Example 3: confident wrong target
Input: logits = [10, 1, 0.5]
target_idx = 1
Output: 9.0002

Constraints

`len(logits) > 0`
`0 <= target_idx < len(logits)`
Round result to 4 decimal places.

Python3Saved ✓

def cross_entropy_loss(logits: list[float], target_idx: int) -> float:

# Your implementation here

pass

Case 1: confident correct target

Input: logits = [10, 1, 0.5]
target_idx = 0

Expected: 0.0002

Case 2: uncertain target

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

Expected: 1.0986

Case 3: confident wrong target

Input: logits = [10, 1, 0.5]
target_idx = 1

Expected: 9.0002