Neural Networks & LLM Components · Medium

Masked Softmax

Normalize only valid logits while assigning masked positions zero probability.

Task

Implement masked_softmax(logits: list[float], mask: list[bool]) -> list[float]. Apply numerically stable softmax only to positions where mask is True, while returning zero probability at every masked position.

Requirements

Logits and mask have equal length; True means the position is valid.
Compute the maximum and normalization sum using only valid logits.
Return zero at every masked position without changing input order.
Return all zeros when every position is masked, and an empty result for empty input.

Example

logits = [1.0, 2.0, 3.0]
mask = [True, False, True]

masked_softmax(logits, mask)
# [0.11920292, 0.0, 0.88079708]