从零实现 Self-Attention

注意力是一张查找表，每个词都在问"谁对我重要？"——并学会答案。

类型： Build 语言： Python 前置要求： 阶段 3（深度学习核心）、阶段 5 第 10 课（序列到序列） 预计时间： ~90 分钟

学习目标

• 只用 NumPy 从零实现缩放点积 self-attention，包括 query/key/value 投影和 softmax 加权求和
• 构建一个 multi-head attention 层，拆分多个头、并行计算注意力、再拼接结果
• 追踪注意力矩阵如何捕捉 token 之间的关系，并解释为什么除以 sqrt(d_k) 能防止 softmax 饱和
• 应用因果掩码，把双向注意力转成自回归（解码器风格）注意力

问题背景

RNN 一次处理一个 token。等你走到第 50 个 token，来自第 1 个 token 的信息已经被压缩了 50 步。长程依赖被挤进一个固定大小的隐藏状态——这个瓶颈靠再多 LSTM 门控也没法彻底解决。

2014 年的 Bahdanau 注意力论文给出了修法：让解码器回头看每一个编码器位置，决定哪些对当前这一步重要。但它还是焊在 RNN 上的。2017 年的《Attention Is All You Need》问了一个更尖锐的问题：如果注意力是唯一的机制呢？没有循环，没有卷积，只有注意力。

self-attention 让序列里的每个位置在一个并行步骤里关注其他每一个位置。这正是 transformer 又快、又能扩展、又称霸的原因。

核心概念

数据库查找的类比

把注意力想成一次软性的数据库查找：

传统数据库：
  Query: "capital of France"  -->  精确匹配  -->  "Paris"

注意力：
  Query: "capital of France"  -->  和所有 key 算相似度  -->  对所有 value 加权混合

每个 token 生成三个向量：

• Query (Q)："我在找什么？"
• Key (K)："我包含什么？"
• Value (V)："如果被选中，我提供什么信息？"

query 和所有 key 的点积产生注意力分数。分数高意味着"这个 key 匹配我的 query"。这些分数给 value 加权。输出是 value 的加权和。

Q、K、V 的计算

每个 token 嵌入都经过三个学到的权重矩阵投影：

输入嵌入（n 个 token 的序列，每个 d 维）：

  X = [x1, x2, x3, ..., xn]       shape: (n, d)

三个权重矩阵：

  Wq  shape: (d, dk)
  Wk  shape: (d, dk)
  Wv  shape: (d, dv)

投影：

  Q = X @ Wq    shape: (n, dk)      每个 token 的 query
  K = X @ Wk    shape: (n, dk)      每个 token 的 key
  V = X @ Wv    shape: (n, dv)      每个 token 的 value

对单个 token，直观地看：

             Wq
  x_i ------[*]------> q_i    "我在找什么？"
       |
       |     Wk
       +----[*]------> k_i    "我包含什么？"
       |
       |     Wv
       +----[*]------> v_i    "我提供什么？"

注意力矩阵

一旦你有了所有 token 的 Q、K、V，注意力分数就构成一个矩阵：

Scores = Q @ K^T    shape: (n, n)

              k1    k2    k3    k4    k5
        +-----+-----+-----+-----+-----+
   q1   | 2.1 | 0.3 | 0.1 | 0.8 | 0.2 |   <- q1 对每个 key 关注多少
        +-----+-----+-----+-----+-----+
   q2   | 0.4 | 1.9 | 0.7 | 0.1 | 0.3 |
        +-----+-----+-----+-----+-----+
   q3   | 0.2 | 0.6 | 2.3 | 0.5 | 0.1 |
        +-----+-----+-----+-----+-----+
   q4   | 0.9 | 0.1 | 0.4 | 1.7 | 0.6 |
        +-----+-----+-----+-----+-----+
   q5   | 0.1 | 0.3 | 0.2 | 0.5 | 2.0 |
        +-----+-----+-----+-----+-----+

每一行：一个 token 对整个序列的注意力

看着一个 query 逐个扫过所有 key：每一行给每个 token 打分，softmax 把分数变成权重，context 向量就是 value 的加权融合。

attention-matrix

为什么要缩放？

点积随维度 dk 增大。如果 dk = 64，点积可能落在几十的量级，把 softmax 推进梯度消失的区域。修法是：除以 sqrt(dk)。

Scaled scores = (Q @ K^T) / sqrt(dk)

这让数值保持在 softmax 能产生有用梯度的范围里。

Softmax 把分数变成权重

softmax 把原始分数转成每一行上的一个概率分布：

q1 的原始分数：   [2.1, 0.3, 0.1, 0.8, 0.2]
                            |
                         softmax
                            |
注意力权重：      [0.52, 0.09, 0.07, 0.14, 0.08]   (和约等于 1.0)

现在每个 token 都有一组权重，说明它该对其他每个 token 关注多少。

value 的加权和

每个 token 的最终输出是所有 value 向量的加权和：

output_i = sum( attention_weight[i][j] * v_j  for all j )

对 token 1：
  output_1 = 0.52 * v1 + 0.09 * v2 + 0.07 * v3 + 0.14 * v4 + 0.08 * v5

完整流水线

flowchart LR
  X["X (输入)"] --> Q["Q = X · Wq"]
  X --> K["K = X · Wk"]
  X --> V["V = X · Wv"]
  Q --> S["Q · Kᵀ / √dk"]
  K --> S
  S --> SM["softmax"]
  SM --> WS["加权求和"]
  V --> WS
  WS --> O["输出"]

一行公式：

Attention(Q, K, V) = softmax( Q @ K^T / sqrt(dk) ) @ V

softmax-attention-scaling

动手构建

第 1 步：从零写 softmax

softmax 把原始 logits 转成概率。减去最大值是为了数值稳定。

import numpy as np

def softmax(x):
    shifted = x - np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(shifted)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

logits = np.array([2.0, 1.0, 0.1])
print(f"logits:  {logits}")
print(f"softmax: {softmax(logits)}")
print(f"sum:     {softmax(logits).sum():.4f}")

第 2 步：缩放点积注意力

核心函数。接收 Q、K、V 矩阵，返回注意力输出加权重矩阵。

def scaled_dot_product_attention(Q, K, V):
    dk = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(dk)
    weights = softmax(scores)
    output = weights @ V
    return output, weights

第 3 步：带可学习投影的 self-attention 类

一个完整的 self-attention 模块，Wq、Wk、Wv 权重矩阵用类似 Xavier 的缩放初始化。

class SelfAttention:
    def __init__(self, d_model, dk, dv, seed=42):
        rng = np.random.default_rng(seed)
        scale = np.sqrt(2.0 / (d_model + dk))
        self.Wq = rng.normal(0, scale, (d_model, dk))
        self.Wk = rng.normal(0, scale, (d_model, dk))
        scale_v = np.sqrt(2.0 / (d_model + dv))
        self.Wv = rng.normal(0, scale_v, (d_model, dv))
        self.dk = dk

    def forward(self, X):
        Q = X @ self.Wq
        K = X @ self.Wk
        V = X @ self.Wv
        output, weights = scaled_dot_product_attention(Q, K, V)
        return output, weights

第 4 步：在一个句子上跑起来

为一个句子造假的嵌入，观察注意力权重。

sentence = ["The", "cat", "sat", "on", "the", "mat"]
n_tokens = len(sentence)
d_model = 8
dk = 4
dv = 4

rng = np.random.default_rng(42)
X = rng.normal(0, 1, (n_tokens, d_model))

attn = SelfAttention(d_model, dk, dv, seed=42)
output, weights = attn.forward(X)

print("Attention weights (each row: where that token looks):\n")
print(f"{'':>6}", end="")
for token in sentence:
    print(f"{token:>6}", end="")
print()

for i, token in enumerate(sentence):
    print(f"{token:>6}", end="")
    for j in range(n_tokens):
        w = weights[i][j]
        print(f"{w:6.3f}", end="")
    print()

第 5 步：用 ASCII 热力图把注意力可视化

把注意力权重映射成字符，快速看个直观效果。

def ascii_heatmap(weights, tokens, chars=" ░▒▓█"):
    n = len(tokens)
    print(f"\n{'':>6}", end="")
    for t in tokens:
        print(f"{t:>6}", end="")
    print()

    for i in range(n):
        print(f"{tokens[i]:>6}", end="")
        for j in range(n):
            level = int(weights[i][j] * (len(chars) - 1) / weights.max())
            level = min(level, len(chars) - 1)
            print(f"{'  ' + chars[level] + '   '}", end="")
        print()

ascii_heatmap(weights, sentence)

实际使用

PyTorch 的 nn.MultiheadAttention 做的正是我们刚搭的东西，外加多头拆分和输出投影：

import torch
import torch.nn as nn

d_model = 8
n_heads = 2
seq_len = 6

mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads, batch_first=True)

X_torch = torch.randn(1, seq_len, d_model)

output, attn_weights = mha(X_torch, X_torch, X_torch)

print(f"Input shape:            {X_torch.shape}")
print(f"Output shape:           {output.shape}")
print(f"Attention weight shape: {attn_weights.shape}")
print(f"\nAttn weights (averaged over heads):")
print(attn_weights[0].detach().numpy().round(3))

关键区别：multi-head attention 并行跑多个注意力函数，每个有自己的、大小为 dk = d_model / n_heads 的 Q、K、V 投影，然后拼接结果。这让模型能同时关注不同类型的关系。

拿去用

这节课产出：

• outputs/prompt-attention-explainer.md —— 一个通过数据库查找类比来讲解注意力的 prompt

练习

1. 修改 scaled_dot_product_attention，让它接受一个可选的掩码矩阵，在 softmax 之前把某些位置置为负无穷（这就是因果/解码器掩码的工作方式）
2. 从零实现 multi-head attention：把 Q、K、V 切成 n_heads 块，在每块上跑注意力，拼接，再经过一个最终权重矩阵 Wo 投影
3. 取两个长度相同但不同的句子，喂进同一个 SelfAttention 实例，对比它们的注意力模式。什么变了？什么没变？

关键术语

术语	大家嘴上怎么说	实际是什么意思
Query (Q)	"提问向量"	输入的一个可学习投影，表示这个 token 在找什么信息
Key (K)	"标签向量"	一个可学习投影，表示这个 token 包含什么信息，用来和 query 匹配
Value (V)	"内容向量"	一个可学习投影，携带真正的信息，根据注意力分数被聚合
缩放点积注意力	"注意力公式"	softmax(QK^T / sqrt(dk)) @ V —— 缩放防止高维下 softmax 饱和
Self-attention	"token 看自己也看别人"	Q、K、V 都来自同一序列的注意力，让每个位置都能关注其他每个位置
注意力权重	"关注多少"	位置上的一个概率分布，由对缩放点积做 softmax 得到
Multi-head attention	"并行注意力"	用不同投影跑多个注意力函数，再拼接结果，得到更丰富的表示

延伸阅读

• Attention Is All You Need (Vaswani et al., 2017) —— 最初的 transformer 论文
• The Illustrated Transformer (Jay Alammar) —— 对整个架构最好的图解
• The Annotated Transformer (Harvard NLP) —— 逐行 PyTorch 实现并附讲解