返回 MLE 工程师 思维导图
中文·English
💻 MLE 工程师ID: mle-coding-self-attention

手撕 Multi-Head Self-Attention

Live Coding: Multi-Head Self-Attention
🎯核心定义
手撕多头自注意力算子 (Live Coding: Multi-Head Self-Attention in Pure Numpy / PyTorch) 是机器学习工程师 (MLE) 白板编码与现场代码考核中最高频、最核心的压轴题目;考察从零编写标准的 Scaled Dot-Product Multi-Head Attention:1) 线性投影:将输入张量 XRB×L×dmodelX \in \mathbb{R}^{B \times L \times d_{\text{model}}} 分别通过线性权重映射为 Q,K,VQ, K, V;2) 维度变换与切分:将张量 Reshape 并 Transpose 为 (B,H,L,dk)(B, H, L, d_k)(其中 dk=dmodel/Hd_k = d_{\text{model}} / H);3) 点积注意力得分:计算 S=QKTdkS = \frac{Q K^T}{\sqrt{d_k}} 并根据需要施加因果掩码 (Causal Mask / Upper Triangular Mask with -\infty);4) 稳定 Softmax 与加权求和:Attention(Q,K,V)=softmax(S)V\text{Attention}(Q, K, V) = \text{softmax}(S) V;5) 多头拼接与最终线性投影输出。
💡使用场景
顶级大厂 MLE 面试现场 20 分钟手撕代码考核、Transformer 核心算子自研与调试。
解决的核心痛点
考察候选人是否真正理解多头机制的张量维度流转 (Tensor Shape Transitions)、广播机制、缩放因子 dk\sqrt{d_k} 的数学根因以及掩码加法的数值稳定性,而非单纯调用黑盒库。
🎯5 个高频面试考点 (Exam Points)
1
现场写出包含 Batch 维度的 Scaled Dot-Product Attention 核心 Pure Numpy 代码,并标注每一步的 Tensor Shape?
2
数学推导:为什么必须除以缩放因子 dk\sqrt{d_k}(证明当 qi,kiN(0,1)q_i, k_i \sim \mathcal{N}(0, 1) 独立同分布时,内积点积的方差为 dkd_k)?
3
手写因果自回归掩码 (Causal Mask: `np.triu(np.ones(...), k=1)`) 的生成与为什么掩码值要填充为 `-1e9` 或 `-np.inf`?
4
对比 `torch.einsum('bhqd,bhkd->bhqk', q, k)` 与标准 `torch.matmul(q, k.transpose(-2, -1))` 的可读性与底层算子执行?
5
FlashAttention 核心优化原理手撕阐述:如何通过 Tiling 分块计算与 Online Softmax 消除中间注意力矩阵 O(N2)O(N^2) 的 HBM 读写?
📖 关联深度指南:📄 mle-coding-and-algo-prep
更新于 2026-08-14
🎯
检验攻克程度:针对「手撕 Multi-Head Self-Attention」专属刷题排雷
做单选排雷题、推导选项机制,答错自动收录进专属错题本。
🚀 开始本考点专项刷题
上一个知识点模型压缩、剪枝与量化基础下一个知识点手写 Softmax 防溢出与 Cross-Entropy

🔗 更多 MLE 工程师 知识点卡片

偏差-方差权衡与过拟合诊断常见损失函数选型与梯度特性优化器收敛性与动量选型准则模型集成 Stacking 与 Blending