Mask the padding before the softmax

Attention scores for padded positions have to be removed before the softmax, and the way you remove them is to set them to negative infinity, so that exp sends them to zero.

scores   [[1.0,  2.0,  3.0],
          [4.0,  5.0,  6.0]]

mask     [[F, F, T],          T means "this position is padding"
          [F, T, T]]

result   [[1.0,  2.0, -inf],
          [4.0, -inf, -inf]]

Three ways to write it

# per element, in Python
for r in range(rows):
    for c in range(cols):
        if mask[r, c]:
            out[r, c] = float("-inf")

# one op
scores.masked_fill(mask, float("-inf"))

# one op, and the form to know
torch.where(mask, float("-inf"), scores)

The loop dispatches several kernels per element. The other two dispatch one each, and read as a sentence.

masked_fill or where

masked_fill(mask, value) replaces where the mask is true and keeps everything else. One tensor, one scalar.

where(condition, a, b) picks elementwise between two tensors, either of which may be a scalar. Strictly more general, and the one to reach for when both branches are tensors:

torch.where(mask, other_scores, scores)

Both return a new tensor. The trailing-underscore masked_fill_ writes in place, which matters inside a layer and is a correctness hazard everywhere else, for the reasons the previous problem covered.

Why negative infinity and not a large negative number

-1e9 is the version you will see in older code. It works until the dtype is float16, where -1e9 is not representable and becomes -inf anyway, or until a row is entirely padding and the softmax over all -1e9 gives a uniform distribution over nothing. Use -inf and handle the all-padded row explicitly if it can occur.

Your task

def mask_scores(scores: torch.Tensor, mask: torch.Tensor) -> torch.Tensor

Return a new tensor with -inf wherever mask is true, in at most two dispatched operations. scores must be unchanged.