Implement Causal Masking for Attention
Implement Causal Masking for Attention
Implement causal (autoregressive) masking for attention scores so each position can only attend to itself and past tokens, never the future. Given attention score tensors scores shaped (..., T, T), set all entries above the main diagonal (future positions) to a large negative value (e.g., -1e9) before softmax.
Mathematical Definition
In autoregressive models, token at time t must not see tokens at t+1..T−1:
zj={scoreij−∞if j≤iif j>iFunction Arguments
scores: np.ndarray- Attention scores with shape (..., T, T). Supports 2D, 3D, and 4D inputs.mask_value: float- Value for masked positions (default: -1e9)
Returns
New array (same shape) with future positions replaced by mask_value. Must not modify the input array.
Examples
Input: scores = [[1, 2, 3], [4, 5, 6], [7, 8, 9]], mask_value = -1e9
Output: [[1, -1e9, -1e9], [4, 5, -1e9], [7, 8, 9]]
Input: scores = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], mask_value = -1e9
Output: [[1, -1e9, -1e9, -1e9], [5, 6, -1e9, -1e9], [9, 10, 11, -1e9], [13, 14, 15, 16]]
Hint 1
Use np.triu(..., k=1) to create an upper-triangular boolean mask for future positions.
Hint 2
Create a single (T, T) mask and use broadcasting with scores[..., mask] for efficiency.
Requirements
- Support 2D (T, T) and higher-rank (..., T, T) tensors
- Fully vectorized (no Python loops over T)
- Use broadcasting efficiently (don't allocate giant masks per batch/head)
- Do not modify the input in-place (return a masked copy)
- Return same shape and dtype=float
Constraints
- Input tensors up to shape (32, 16, 512, 512)
- Use NumPy only
- Time limit: 200 ms
Try Similar Problems
Log in to take notes on this problem
Accepts: array
Accepts: number
Implement Causal Masking for Attention
Implement Causal Masking for Attention
Implement causal (autoregressive) masking for attention scores so each position can only attend to itself and past tokens, never the future. Given attention score tensors scores shaped (..., T, T), set all entries above the main diagonal (future positions) to a large negative value (e.g., -1e9) before softmax.
Mathematical Definition
In autoregressive models, token at time t must not see tokens at t+1..T−1:
zj={scoreij−∞if j≤iif j>iFunction Arguments
scores: np.ndarray- Attention scores with shape (..., T, T). Supports 2D, 3D, and 4D inputs.mask_value: float- Value for masked positions (default: -1e9)
Returns
New array (same shape) with future positions replaced by mask_value. Must not modify the input array.
Examples
Input: scores = [[1, 2, 3], [4, 5, 6], [7, 8, 9]], mask_value = -1e9
Output: [[1, -1e9, -1e9], [4, 5, -1e9], [7, 8, 9]]
Input: scores = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], mask_value = -1e9
Output: [[1, -1e9, -1e9, -1e9], [5, 6, -1e9, -1e9], [9, 10, 11, -1e9], [13, 14, 15, 16]]
Hint 1
Use np.triu(..., k=1) to create an upper-triangular boolean mask for future positions.
Hint 2
Create a single (T, T) mask and use broadcasting with scores[..., mask] for efficiency.
Requirements
- Support 2D (T, T) and higher-rank (..., T, T) tensors
- Fully vectorized (no Python loops over T)
- Use broadcasting efficiently (don't allocate giant masks per batch/head)
- Do not modify the input in-place (return a masked copy)
- Return same shape and dtype=float
Constraints
- Input tensors up to shape (32, 16, 512, 512)
- Use NumPy only
- Time limit: 200 ms
Try Similar Problems
Log in to take notes on this problem
Accepts: array
Accepts: number