Multi-Head Attention 维度变化

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]