Transformer Attention Mask and Heads Coding

Requirements

Implement or adapt a Transformer attention component. Your discussion and implementation should address all of the following:

1. Describe what enters self-attention and what it returns.
2. Extend single-head attention into a multi-head design.
3. Support an attention mask.
4. Preserve the correct dimensions for batches, heads, token positions, and hidden features.
5. Discuss the algorithm's computational cost and typical implementation mistakes.

Use this PyTorch interface:

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, hidden_dim: int, num_heads: int):
        ...

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        ...

Assume x has shape (batch, seq, hidden_dim) and the returned tensor has the same shape. The mask must be usable for either padding positions or causal attention.

Notes

Keep the dimension changes explicit, especially when converting projected queries, keys, and values into separate attention heads. Explain how a padding mask differs in broadcast behavior from a mask that blocks access to future positions.

Masks must affect the logits before normalization; applying one only after softmax is not acceptable.

Examples

Example 1: Unmasked attention

Input:

hidden_dim = 6
num_heads = 2
x.shape = (3, 4, 6)
attention_mask = None

Output:

output.shape = (3, 4, 6)

The batch contains three sequences of four tokens, and the output retains the input's batch, sequence, and hidden dimensions.

Example 2: Padding mask

Input:

hidden_dim = 4
num_heads = 2
x.shape = (1, 3, 4)
attention_mask.shape = (1, 1, 1, 3)
attention_mask = [1, 1, 0]

Output:

output.shape = (1, 3, 4)

The third token is unavailable as a key/value position for every query, while the result still has one sequence of length three and four hidden features.

Example 3: Causal mask

Input:

hidden_dim = 8
num_heads = 4
x.shape = (2, 5, 8)
attention_mask.shape = (1, 1, 5, 5)
attention_mask[i, j] = 1 when j <= i, otherwise 0

Output:

output.shape = (2, 5, 8)

Each position can use itself and earlier positions only, and the mask broadcasts across both batches and heads.

Constraints

• hidden_dim must be divisible by num_heads.
• The per-head width is hidden_dim // num_heads.
• Inputs use the layout (batch, seq, hidden_dim) before projection.
• Projected queries, keys, and values should be arranged so attention scores can be computed independently for every batch and head.
• The implementation must support an omitted mask as well as a supplied padding or causal mask.

Preparation

• Practice writing a compact PyTorch self-attention layer without relying on a ready-made attention module.
• Trace every shape through projection, head splitting, transposition, score calculation, normalization, head recombination, and the final projection.
• Test the implementation with both padding and causal masking.
• Be prepared to explain its time and memory requirements and to identify common shape and masking bugs.