手写 Multi-Head Attention。

核心回答

MHA先将输入分别投影为Q、K、V，拆成多个头；每个头计算softmax(QKᵀ/√d_k)V，再拼接各头并通过输出投影。下面用PyTorch实现batch-first自注意力，包含布尔注意力掩码和key padding掩码；为突出机制，省略dropout、KV缓存及残差/归一化。

PyTorch实现

输入x形状为[B,T,D]，D必须能被头数H整除。这里两类掩码都约定True表示禁止关注：attn_mask为[T,T]，key_padding_mask为[B,T]。每个query至少保留一个可见key，全遮蔽行明确报错。

import math
import torch
from torch import nn

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        if d_model <= 0 or num_heads <= 0 or d_model % num_heads:
            raise ValueError("d_model必须为正且能被正num_heads整除")
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, attn_mask=None, key_padding_mask=None):
        if x.ndim != 3 or x.shape[-1] != self.d_model or x.shape[1] == 0:
            raise ValueError("x应为非空序列[B,T,D]")
        batch, length, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        def split(t):
            return t.reshape(batch, length, self.num_heads,
                             self.head_dim).transpose(1, 2)
        q, k, v = split(q), split(k), split(v)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        blocked = torch.zeros((batch, 1, length, length),
                              dtype=torch.bool, device=x.device)
        if attn_mask is not None:
            if attn_mask.dtype != torch.bool or attn_mask.shape != (length, length):
                raise ValueError("attn_mask须为布尔[T,T]")
            blocked = blocked | attn_mask.to(x.device)[None, None]
        if key_padding_mask is not None:
            if (key_padding_mask.dtype != torch.bool or
                    key_padding_mask.shape != (batch, length)):
                raise ValueError("key_padding_mask须为布尔[B,T]")
            blocked = blocked | key_padding_mask.to(x.device)[:, None, None, :]
        if blocked.all(dim=-1).any():
            raise ValueError("每个query至少需要一个可见key")
        weights = scores.masked_fill(blocked, float("-inf")).softmax(dim=-1)
        heads = weights @ v
        merged = heads.transpose(1, 2).contiguous().reshape(batch, length, self.d_model)
        return self.out(merged), weights

维度与容易写错的位置

投影后Q/K/V各为[B,T,D]，分头变成[B,H,T,D/H]，分数和权重为[B,H,T,T]。softmax沿最后的key维做；加权V后回到[B,H,T,D/H]，再还原[B,T,D]，最后仍需输出投影W_O。

例如B=2、T=5、D=12、H=3，每头维度4，缩放除以2，分数形状是[2,3,5,5]。拼头时先交换回T和H的次序，再整理内存，不能直接把原[B,H,T,d]按[B,T,D]解释。

因果掩码把主对角线上方设为True，禁止读取未来。key padding掩码屏蔽的是被关注的key列，不会自动把padding query的输出清零；训练时还需相应损失掩码。不同PyTorch API对布尔掩码的约定可能不同，不能机械复用。

复杂度与扩展

对长度T、隐藏维度D，线性投影约O(BTD²)，注意力约O(BT²D)；显式权重占O(BHT²)空间。交叉注意力改为从一组输入产生Q、另一组输入产生K/V，并允许query与key长度不同。生产中通常使用已验证的融合实现；这段代码用于展示机制，不代表最快内核。

对应教程

参考资料