Implement Causal Masking for Attention

Transformers

Medium

Apply an autoregressive mask to attention scores so position $i$ can attend only to positions $j \leq i$ . For every score entry, return

M_{i,j} = \begin{cases} S_{i,j} & j \leq i \\ m & j > i \end{cases}

Here, $S$ is the input score tensor and $m$ is mask_value. The final two dimensions are square attention matrices, while any preceding dimensions represent batches or heads. Return a masked NumPy array with the same shape without modifying the input.

Loading visualization...

Examples

Input: scores = [[1, 2, 3], [4, 5, 6], [7, 8, 9]], mask_value = -1000000000.0

Output: [[1.0, -1000000000.0, -1000000000.0], [4.0, 5.0, -1000000000.0], [7.0, 8.0, 9.0]]

Explanation: Entries above the main diagonal represent future positions and are replaced by mask_value.

Input: scores = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], mask_value = -1000000000.0

Output: [[1.0, -1000000000.0, -1000000000.0, -1000000000.0], [5.0, 6.0, -1000000000.0, -1000000000.0], [9.0, 10.0, 11.0, -1000000000.0], [13.0, 14.0, 15.0, 16.0]]

Hint 1

np.triu(np.ones((T, T), dtype=bool), k=1) marks future positions.

Hint 2

A two-dimensional mask broadcasts across any leading batch and head dimensions.

Requirements

Support score tensors with shape $(T,T)$ or $(\ldots,T,T)$
Replace only entries above the main diagonal
Return a floating-point NumPy array with the same shape
Do not modify the input

Constraints

The final two dimensions are equal
Use NumPy only

Try Similar Problems

Positional Encoding Softmax Function Gelu Linear Layer Forward Dropout Training

Prev Next