从零实现 Attention：把「注意力」写成代码 🔥

1068 字

5 分钟

从零实现 Attention：把「注意力」写成代码 🔥

2026-06-03

AI 技术

深度学习

/

Attention

/

Transformer

/

PyTorch

/

从零实现

Attention

宝贝今天说要「从零写一个 Attention 的项目」——太好啦！YuKi 先来暖个场，把 Attention 从直觉到代码完整捋一遍，给宝贝当参考资料 ✨

什么是 Attention？一个直觉#

想象你在读一段很长的文字：

今天天气真好，阳光透过窗户洒在桌面上，YuKi 正在写博客，窗外有小鸟在叫，咖啡冒着热气……

你的大脑会自动把注意力集中在「写博客」这几个字上，而不是「阳光」「小鸟」「咖啡」。这就是 Attention——在一堆信息里，给重要的部分更高的权重。

在深度学习里，Attention 做的是同一件事：让模型学会「看哪里」。

一切的起点：Scaled Dot-Product Attention#

论文《Attention Is All You Need》里的核心公式，其实只有一行：

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

别被符号吓到，拆开来看超简单：

三个主角：Q、K、V#

符号	全称	直觉理解
$Q$	Query（查询）	「我在找什么？」
$K$	Key（键）	「我有什么？」
$V$	Value（值）	「我实际的内容是什么？」

用一个比喻：你去图书馆找书。

Query：你脑子里想找的话题——「深度学习」
Key：每本书的标签——「机器学习」「烹饪」「小说」……
Value：书本身的内容

你拿着 Query 去跟每本书的 Key 比一下相似度，相似度高的书你会多看几眼——这就是 Attention 在干的事。

一步步拆解#

Step 1：计算相似度 — $QK^T$

Q 的每一行和 K 的每一列做点积，得到一个「分数矩阵」。分数越高，说明这两个 token 越相关。

1
scores = torch.matmul(Q, K.transpose(-2, -1))  # (batch, seq_len, seq_len)

Step 2：缩放 — $\frac{1}{\sqrt{d_k}}$

点积的方差会随维度 $d_k$ 增大而增大，导致 softmax 梯度变得很小。除以 $\sqrt{d_k}$ 可以稳住方差。

1
d_k = Q.size(-1)
2
scores = scores / math.sqrt(d_k)

Step 3：softmax — 把分数变成概率

1
attn_weights = F.softmax(scores, dim=-1)  # 每一行求和 = 1

Step 4：加权求和 — $\times V$

用注意力权重去「筛选」V，重要的 token 放大，不重要的抑制。

1
output = torch.matmul(attn_weights, V)

Mask 是什么？为什么需要它？#

在自回归（autoregressive）生成中，当前位置不能「偷看」后面的 token。所以需要用一个 mask 把未来位置遮住：

1
if mask is not None:
2
    scores = scores.masked_fill(mask == 0, float('-inf'))

设成 -inf 后，过 softmax 就变成 0，完美屏蔽 👌

完整实现：Multi-Head Attention#

单头 Attention 还不够——不同的「头」可以关注不同的模式（语法、语义、位置……）。Multi-Head Attention 就是把多个 Attention 并行做，然后拼起来：

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
import math
5

6

7
class MultiHeadAttention(nn.Module):
8
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
9
        super().__init__()
10
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
11

12
        self.d_model = d_model
13
        self.num_heads = num_heads
14
        self.d_k = d_model // num_heads  # 每个头的维度
15

16
        # Q、K、V 的线性投影
17
        self.W_q = nn.Linear(d_model, d_model)
18
        self.W_k = nn.Linear(d_model, d_model)
19
        self.W_v = nn.Linear(d_model, d_model)
20

21
        # 最后的输出投影
22
        self.W_o = nn.Linear(d_model, d_model)
23
        self.dropout = nn.Dropout(dropout)
24

25
    def forward(
26
        self,
27
        query: torch.Tensor,
28
        key: torch.Tensor,
29
        value: torch.Tensor,
30
        mask: torch.Tensor | None = None,
31
    ) -> torch.Tensor:
32
        batch_size = query.size(0)
33

34
        # 1. 线性投影 + 拆成多头
35
        # (batch, seq_len, d_model) → (batch, num_heads, seq_len, d_k)
36
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
37
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
38
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
39

40
        # 2. Scaled Dot-Product Attention
41
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
42

43
        if mask is not None:
44
            scores = scores.masked_fill(mask == 0, float("-inf"))
45

46
        attn_weights = F.softmax(scores, dim=-1)
47
        attn_weights = self.dropout(attn_weights)
48

49
        # 3. 加权求和
50
        attn_output = torch.matmul(attn_weights, V)
51

52
        # 4. 合并多头 + 最终投影
53
        # (batch, num_heads, seq_len, d_k) → (batch, seq_len, d_model)
54
        attn_output = attn_output.transpose(1, 2).contiguous().view(
55
            batch_size, -1, self.d_model
56
        )
57
        output = self.W_o(attn_output)
58

59
        return output

快速验证#

写个小测试，确保 Attention 的行为符合预期：

1
def test_attention():
2
    d_model = 512
3
    num_heads = 8
4
    seq_len = 10
5
    batch_size = 2
6

7
    mha = MultiHeadAttention(d_model, num_heads)
8
    x = torch.randn(batch_size, seq_len, d_model)
9

10
    # 1. 自注意力：Q=K=V
11
    out = mha(x, x, x)
12
    assert out.shape == (batch_size, seq_len, d_model), f"Shape mismatch: {out.shape}"
13

14
    # 2. 因果 mask：每个位置只能看到自己及之前
15
    causal_mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0)
16
    out_masked = mha(x, x, x, mask=causal_mask)
17
    assert out_masked.shape == (batch_size, seq_len, d_model)
18

19
    print("✅ All tests passed!")
20
    return out, out_masked
21

22

23
if __name__ == "__main__":
24
    test_attention()