Implement Multi-Head Attention from Scratch in NumPy

Multi-Head Attention Implementation

Implement the MultiHeadAttention class using NumPy. Your implementation should perform the core attention mechanism used in Transformer architectures.

Requirements

1. Linear Projections: Given input X of shape (batch_size, seq_len, d_model), project into Queries, Keys, and Values using weight matrices Wq, Wk, Wv of shape (d_model, d_model).
2. Split Heads: Reshape the projected matrices to (batch_size, num_heads, seq_len, d_k) where d_k = d_model / num_heads.
3. Scaled Dot-Product Attention: Compute Attention(Q, K, V) = softmax((Q @ K.T) / sqrt(d_k)) @ V.
4. Causal Masking: Optionally apply a mask to prevent tokens from attending to future positions (i.e., setting values in the upper triangle of the attention matrix to negative infinity before softmax).
5. Output: Concatenate the results from heads and apply a final linear projection Wo.

Function Signature

import numpy as np

class MultiHeadAttention:
    def __init__(self, d_model, num_heads):
        # Initialize weight matrices here
        pass

    def forward(self, X, mask=None):
        # Perform the multi-head attention operation
        # Returns the output tensor of shape (batch_size, seq_len, d_model)
        pass

Constraints

• d_model will always be divisible by num_heads.
• Use a numerically stable softmax implementation: exp(x - max(x)).

Example

Input: X (1, 3, 8), num_heads = 2. Internal d_k = 4. The attention scores for each head will be (1, 2, 3, 3) shaped tensors before being concatenated back to (1, 3, 8).

Edge Cases

• mask=None: Standard self-attention.
• mask=True: Apply a causal (look-ahead) mask.
• Single-head vs. Multi-head scenarios.