Causal Multi-Head Self-Attention

Write the attention layer of a GPT-style model as an nn.Module.

CausalSelfAttention(dim, n_heads) holds two layers:

• self.qkv = nn.Linear(dim, 3 * dim), which makes queries, keys and values in one matrix multiply. Split its output in the order q, k, v.
• self.proj = nn.Linear(dim, dim), which mixes the heads at the end.

In forward(x, key_padding_mask=None), with x of shape (B, T, dim):

q, k, v  -> each reshaped to (B, n_heads, T, head_dim), head_dim = dim // n_heads
scores   = q @ k^T / sqrt(head_dim)                       # (B, n_heads, T, T)
position t may only attend to keys 0..t                   # causal mask
if key_padding_mask is given, padding keys get no weight  # (B, T), True = real token
out      = proj( softmax(scores) @ v, heads joined back to (B, T, dim) )

Set a score to -inf before the softmax to give it zero weight.

Function signature

class CausalSelfAttention(torch.nn.Module):
    def __init__(self, dim, n_heads): ...
    def forward(self, x, key_padding_mask=None): ...

Example

Input

T = 3, and every score in one head is equal (for example, q is all zeros)

Output

attention weights of that head:
[[1.000, 0.000, 0.000],
 [0.500, 0.500, 0.000],
 [0.333, 0.333, 0.333]]

Reasoning

The causal mask removes the keys above the diagonal. The softmax then spreads the weight evenly over the keys that are left: one key for position 0, two for position 1 and three for position 2. So the output at position 0 depends on x[0] only.

Constraints and notes

• The tests copy your weights into nn.MultiheadAttention and compare the outputs. Use the head split above: head h uses columns h * head_dim to (h + 1) * head_dim of each of q, k and v.
• Padding is on the right, so position 0 is always a real token and no row of the softmax is fully masked.
• F.scaled_dot_product_attention is allowed, but you must still combine the causal mask and the padding mask yourself.
• The tests check with autograd that no output depends on a later input.