Xây dựng Lớp Multi-Head Attention (MHA) Đầy Đủ trong PyTorch

Mục tiêu kiến thức

• Thay vì chỉ tính một đầu chú ý duy nhất, Multi-Head Attention chiếu Q, K, V thành H không gian con (subspaces) độc lập:

head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V) MHA(Q, K, V) = Concat(head₁, …, head_H) W^O

• Chiều dữ liệu:
• Đầu vào X: (B, S, D).
• Phép chiếu: W^Q, W^K, W^V ∈ R^{D × D}.
• Reshape và Transpose sang H đầu: (B, S, D) → (B, S, H, d_k) → (B, H, S, d_k) với d_k = D / H.
• Tính Scaled Dot-Product Attention trên từng đầu song song.
• Ghép đầu (Concat): (B, H, S, d_k) → (B, S, H, d_k) → (B, S, D).
• Phép chiếu đầu ra: qua lớp Linear(D, D) của W^O.

Yêu cầu

Xây dựng lớp MultiHeadAttention(torch.nn.Module):

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, embed_dim: int, num_heads: int):
        super().__init__()
        pass

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        pass

• Trả về tuple: (output, attention_weights).

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.