“手撕”不依赖任何现成的库比如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, maskNone): # 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