完整题目
请手动实现一个多头注意力（Multi-Head Attention）模块，要求包括QKV线性变换、缩放点积注意力、多头拼接及输出投影，并说明各部分的作用和维度变化过程。
import math
import torch
from torch import nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model: int, num_heads: int, dropout: float = 0.0):
super().__init__()
if d_model % num_heads != 0:
raise ValueError("d_model must be divisible by num_heads")
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.qkv = nn.Linear(d_model, 3 * d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
# x: [B, L, D]；mask为bool，可广播到[B, H, L, L]，True表示可见
b, l, d = x.shape
qkv = self.qkv(x).view(b, l, 3, self.num_heads, self.head_dim)
q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) # [B,H,L,Dh]
scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
if mask is not None:
mask = mask.to(torch.bool)
scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
weights = torch.softmax(scores, dim=-1)
if mask is not None:
# 避免整行均被mask时产生无意义的均匀权重
weights = weights * mask.to(weights.dtype)
weights = weights / weights.sum(-1, keepdim=True).clamp_min(torch.finfo(weights.dtype).tiny)
weights = self.dropout(weights)
context = weights @ v # [B,H,L,Dh]
context = context.transpose(1, 2).contiguous().view(b, l, d)
return self.out_proj(context), weights # [B,L,D]维度与作用
[B,L,D] -> Q/K/V [B,H,L,Dh] -> scores [B,H,L,L] -> context [B,H,L,Dh] -> [B,L,D]，其中 Dh=D/H。除以 sqrt(Dh) 是为了控制点积方差，避免softmax过早饱和。多头让不同投影子空间并行建模关系；out_proj负责跨头线性混合和恢复模型维度，它本身不增加非线性。
生产实现还要把因果mask与padding mask组合，并确认广播方向；长序列的主要瓶颈是 [L,L] 分数矩阵的计算和存储，可换用框架提供的scaled-dot-product attention或FlashAttention kernel。