注意力机制（Self-Attention）

直觉理解

> 类比合同审查：当你分析一份合同中"甲方"这个词时，你需要回头找到合同开头对甲方的定义。注意力机制就是模型自动学会"回头找相关信息"的能力。

每个Token在处理时，会对掩码允许的位置分配"注意力权重"，权重越高表示关联越强。

---

数学公式

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

三个矩阵的含义：

**Q（Query）**：当前词在"问"：谁和我相关？

**K（Key）**：每个词在"广播"：我有什么特征？

**V（Value）**：每个词在"提供"：如果你关注我，获取这些信息

除以 $\sqrt{d_k}$：防止点积过大导致 softmax 饱和（梯度消失）

---

nanoGPT 实现


def forward(self, x):
    B, T, C = x.size()  # batch, sequence length, embedding dim
    
    # 计算 Q, K, V（一次线性变换得到三个）
    q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
    
    # 拆分为多头
    k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
    q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
    v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
    
    # SDPA接口：具体内核由输入、设备及后端条件决定
    y = torch.nn.functional.scaled_dot_product_attention(
        q, k, v,
        attn_mask=None,
        dropout_p=self.dropout if self.training else 0,
        is_causal=True  # 关键：因果掩码
    )
    
    # 重新拼接多头输出
    y = y.transpose(1, 2).contiguous().view(B, T, C)
    y = self.resid_dropout(self.c_proj(y))
    return y

---

多头注意力（Multi-Head Attention）

不只做一次注意力，而是并行做多次（GPT-2 默认 12 头），各头可学习不同的表示；以下语义分工仅为示意，不是逐头实测结论。


Head 1: 关注语法关系（主谓宾）
Head 2: 关注指代关系（代词 → 名词）
Head 3: 关注语义相似性
...
Head 12: 关注某种领域特定模式

> 每头的维度 = n_embd / n_head = 768 / 12 = 64 维

---

因果掩码（Causal Mask）

语言模型训练时，位置 $t$ 不能"看到"位置 $t+1$ 及之后的 Token（否则预测变得没有意义）。


位置 1 可以看到: [1]
位置 2 可以看到: [1, 2]
位置 3 可以看到: [1, 2, 3]
...

掩码矩阵（下三角）：
1 0 0 0
1 1 0 0
1 1 1 0
1 1 1 1

---

Flash Attention

nanoGPT 注释："flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0"

Flash Attention 是一种 IO 感知的注意力实现，避免将大型注意力矩阵写入显存，具体速度收益随硬件、精度和输入形状变化，本文未做本机性能测试。

详见：02-技术/Flash-Attention优化

---

计算复杂度

标准注意力：$O(T^2 \cdot d)$，其中 $T$ 是序列长度。

这就是为什么长上下文（长合同、长文档）处理成本高——在固定维度的标准密集注意力部分，序列长度翻倍带来约四倍该部分计算量，不能直接推成整个应用耗时。

---

来源

[nanoGPT model.py](https://github.com/karpathy/nanoGPT/blob/master/model.py)

Vaswani et al., "Attention Is All You Need" (2017)

本轮核查来源与限制（2026-09-17）

https://raw.githubusercontent.com/karpathy/nanoGPT/master/model.py

https://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html

https://arxiv.org/abs/2205.14135

只修本轮明确列出的陈述，不将旧页整篇升级为已验证。themis_only为本地保存和禁止外发的兼容标记，不代表本地模型推理。

⚖ Mike Wu Law

注意力机制