“手撕”：不依赖任何现成的库（比如nn.MultiheadAttention），让你在白板或者在线代码编辑器里把代码完整敲出来。
面试官想通过这个题目考察你两点：
你是否真正理解了 Transformer 的核心底层数学逻辑。
你对深度学习框架中张量维度的变换（如 view/reshape、transpose、维度广播）是否烂熟于心。这往往是大家最容易写错的地方。
import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() # 确保特征维度可以被头数整除 assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads # 每个头的维度 # 定义 Q, K, V 的线性变换矩阵 self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) # 定义最后的输出线性变换矩阵 self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): # x 的维度通常是: (batch_size, seq_len, d_model) batch_size, seq_len, _ = x.size() # 1. 线性变换，并重塑维度来切分多头 # 维度变化: (batch_size, seq_len, d_model) -> (batch_size, seq_len, num_heads, d_k) # 然后交换维度 1和2 变成: (batch_size, num_heads, seq_len, d_k) Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 2. 计算 Scaled Dot-Product Attention # K 转置的最后两个维度，为了满足矩阵乘法: (..., seq_len, d_k) x (..., d_k, seq_len) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 3. 处理 Mask (可选，但面试时提一下会加分) if mask is not None: # 将 mask 中为 0 的位置替换为一个很小的负数，这样 softmax 之后就会变成 0 scores = scores.masked_fill(mask == 0, -1e9) # 4. Softmax 与 V 相乘 attention_weights = torch.softmax(scores, dim=-1) # 维度: (batch_size, num_heads, seq_len, seq_len) @ (batch_size, num_heads, seq_len, d_k) # 结果维度: (batch_size, num_heads, seq_len, d_k) context = torch.matmul(attention_weights, V) # 5. 拼接多头的结果 # 把维度换回来: (batch_size, seq_len, num_heads, d_k) # contiguous() 是为了保证内存连续，才能调用 view context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) # 6. 最后的线性映射 output = self.W_o(context) return output
)