Implement Multi-Head Self-Attention

Multi-Head Self-Attention Implementation

Implement a MultiHeadAttention class using PyTorch. The class should map an input tensor x of shape (batch_size, seq_len, embed_dim) to an output of the same shape. You must perform the following steps:

1. Linear Projections: Apply three separate linear layers to transform the input into Queries (Q), Keys (K), and Values (V).
2. Head Splitting: Reshape the Q, K, and V tensors into (batch_size, num_heads, seq_len, head_dim), where head_dim = embed_dim / num_heads.
3. Scaled Dot-Product Attention: Compute Attention(Q, K, V) = softmax((Q @ K.transpose(-2, -1)) / sqrt(head_dim)) @ V.
4. Concatenation and Projection: Concatenate the heads back together and pass them through a final linear output layer.

Function Signature

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        # Initialize your linear layers and parameters here
        pass

    def forward(self, x):
        # x shape: (batch_size, seq_len, embed_dim)
        # returns: (batch_size, seq_len, embed_dim)
        pass

Constraints

• embed_dim must be divisible by num_heads.
• Use torch.softmax for the attention scores.

Example

mha = MultiHeadAttention(embed_dim=512, num_heads=8)
x = torch.randn(32, 10, 512) # batch=32, seq_len=10, dim=512
output = mha(x)
# output.shape should be (32, 10, 512)

Edge Cases

• Consider scenarios where seq_len is 1.
• Ensure the implementation handles arbitrary batch_size.