从零实现Transformer：第 3 部分 - 多头注意力分数维度布局（Dimension Layout of Attention Scores in Multi-Head Attention）

flyfish

[B, H, Q, K]

名称
按顺序：

1. B = Batch Size
   批次维度 / 批量维度
2. H = Head Num / Num Heads
   多头维度 / 注意力头数维度
3. Q = Query Sequence Length
   查询序列长度
4. K = Key Sequence Length
   键序列长度

Q：Query 向量总个数 = 查询侧的序列长度
K：Key 向量总个数 = 键侧的序列长度

再来一遍
$B$：batch_size，一批有多少条样本
$H$：num_heads，多头注意力有多少个独立注意力头
$Q$：Query 向量的总数量（Query 序列长度）
$K$：Key 向量的总数量（Key 序列长度）

注意力分数，是每一个 Query向量和每一个 Key向量两两算匹配相似度
两两组合，天然形成一张：$Q$ 行 × $K$ 列 的分数表

矩阵乘法 $Q=K^T$
设单头简化看（先不管H、B）：

$\mathbb{R}$ = 全体实数包含：整数、小数、正数、负数、浮点数（神经网络里所有数值都属于实数）

1. 每个输入序列，映射得到：
   $Q\in {\mathbb{R}}^{Q\times d}$ ：有 $Q$ 个Query向量，每个向量维度 $d$
   $K\in {\mathbb{R}}^{K\times d}$ ：有 $K$ 个Key向量，每个向量维度 $d$

格式：
$${\mathbb{R}}^{\mathbf{行数}\times \mathbf{列数}}$$
含义：所有 m行、d列 的实数矩阵 构成的空间

2. Key 最后两维转置
   $$K^T\in {\mathbb{R}}^{d\times K}$$

3. 矩阵相乘：
   $$score=Q\cdot K^T$$
   $${\mathbb{R}}^{Q\times d}\times {\mathbb{R}}^{d\times K}⟹{\mathbb{R}}^{Q\times K}$$

这就是为什么最后一定是 $[Q,K]$ 二维矩阵，是线性代数矩阵乘法硬性决定的行列数。

结果行数 = Query向量个数 $Q$
结果列数 = Key向量个数 $K$

套上 batch 和多头，就变成完整四维：
$$[B,H,Q,d]@[B,H,d,K]⟹[\mathbf{B},\mathbf{H},\mathbf{Q},\mathbf{K}]$$

import torch

# 假设：单条样本、单头
Q_len = 4   # 一共有 4 个 Query向量
K_len = 4   # 一共有 4 个 Key向量
dim = 8     # 每个向量的特征维度

# 构造 Q、K
Q = torch.randn(1, 1, Q_len, dim)   # [B=1, H=1, Q=4, d=8]
K = torch.randn(1, 1, K_len, dim)   # [B=1, H=1, K=4, d=8]

# K 最后两维转置
K_T = K.transpose(-2, -1)           # [1, 1, 8, 4]

# 两两相似度打分
score = Q @ K_T                     # 矩阵相乘

print("Q shape:", Q.shape)
print("K_T shape:", K_T.shape)
print("相似度分数矩阵 shape [B,H,Q,K]:", score.shape)

输出

Q shape: torch.Size([1, 1, 4, 8])
K_T shape: torch.Size([1, 1, 8, 4])
相似度分数矩阵 shape [B,H,Q,K]: torch.Size([1, 1, 4, 4])

Decoder 自注意力：自己的序列既当 Query、也当 Key

$$Q=K=序列长度S$$所以分数矩阵就是 $[B,H,S,S]$

举例子

1 符号

设：
一个序列里一共有 $n$ 个 token，这里的 $n$就是后面的$Q_{num}、K_{num}$都是指：向量的个数 = 序列长度
经过线性变换后，每个 token 对应 1 个向量
每个向量的特征维度 = $d$

2 Query 举例

$\mathbf{Q}\in {\mathbb{R}}^{\mathbf{n}\times d}$
含义：一共有 $\mathbf{n}$ 个 Query 向量，每个向量长度是 $d$

这里的例子
序列长度（token个数） $\mathbf{n}=4$
每个向量维度 $\mathbf{d}=3$

那：
$$\mathbf{Q}\in {\mathbb{R}}^{4\times 3}$$

写成数字矩阵
$$\mathbf{Q}=\left[\begin{array}{c}{\mathbf{q}}_1\\ {\mathbf{q}}_2\\ {\mathbf{q}}_3\\ {\mathbf{q}}_4\end{array}\right]=\left[\begin{array}{ccc}1.1& 2.2& 3.3\\ 4.4& 5.5& 6.6\\ 7.7& 8.8& 9.9\\ 1.2& 2.3& 3.4\end{array}\right]$$

解释：
一共 4 行 → 对应 4 个 Query 向量（4个token各生成一个query）
一共 3 列 → 每一个 Query 向量本身是 3维
${\mathbf{q}}_1$ 是第1个token的query向量，${\mathbf{q}}_2$ 第2个……

3 Key 一模一样套格式

$\mathbf{K}\in {\mathbb{R}}^{\mathbf{n}\times d}$
含义：一共有 $\mathbf{n}$ 个 Key 向量，每个向量长度是 $d$

同样例子：$n=4,d=3$
$$\mathbf{K}\in {\mathbb{R}}^{4\times 3}$$

$$\mathbf{K}=\left[\begin{array}{c}{\mathbf{k}}_1\\ {\mathbf{k}}_2\\ {\mathbf{k}}_3\\ {\mathbf{k}}_4\end{array}\right]=\left[\begin{array}{ccc}0.1& 0.2& 0.3\\ 0.4& 0.5& 0.6\\ 0.7& 0.8& 0.9\\ 1.0& 1.1& 1.2\end{array}\right]$$

4行 = 4个Key向量
3列 = 每个Key向量是3维

4 最后一定是 $[Q_{num},K_{num}]$ 矩阵

注意力公式：
$$\mathbf{Score}=\mathbf{Q}\cdot {\mathbf{K}}^{\mathrm{T}}$$

1. $\mathbf{Q}:{\mathbb{R}}^{4\times 3}$
2. 对 K 转置：${\mathbf{K}}^{\mathrm{T}}:{\mathbb{R}}^{3\times 4}$

矩阵乘法规则：
$$(4\times 3)\times (3\times 4)⟹4\times 4$$

结果 4行 = 原来 Query 的数量（4个query）
结果 4列 = 原来 Key 的数量（4个key）
矩阵里每一个元素：某一个Query 和 某一个Key 的相似度打分

所以打分矩阵就是：
$$\mathbf{Score}\in {\mathbb{R}}^{{\mathbf{Q}}_{\mathbf{num}}\times {\mathbf{K}}_{\mathbf{num}}}$$

5 带Batch、多头

批量 + 多头后：
$B$：批次大小
$H$：注意力头数
$Q_{num}=4,d=3$

$$\begin{array}{rl}\mathbf{Q}& \in {\mathbb{R}}^{B,H,4,3}\\ \mathbf{K}& \in {\mathbb{R}}^{B,H,4,3}\end{array}$$

相乘：
$$Q@K.transpose(-2,-1)\Rightarrow {\mathbb{R}}^{B,H,4,4}$$

序列有 4 个 token
Query 序列长度 = 4
自动就有 4 个 Query 向量
四维第三维就是 4

Key 同理：
Key 序列长度 = 4 → 4 个 Key 向量 → 四维第四维就是 4

相乘后：
$$[B,H,4,3]@[B,H,3,4]\Rightarrow [\mathbf{B},\mathbf{H},4,4]$$
最后两维：[Query序列长, Key序列长] 或者说 [Query Sequence Length, Key Sequence Length]

import torch

# n=4个向量，每个向量d=3维
Q = torch.tensor([
    [1.1, 2.2, 3.3],
    [4.4, 5.5, 6.6],
    [7.7, 8.8, 9.9],
    [1.2, 2.3, 3.4]
])

K = torch.tensor([
    [0.1, 0.2, 0.3],
    [0.4, 0.5, 0.6],
    [0.7, 0.8, 0.9],
    [1.0, 1.1, 1.2]
])

print("Q shape:", Q.shape)  # [4, 3]
print("K shape:", K.shape)  # [4, 3]

score = Q @ K.t()
print("相似度分数矩阵 shape:", score.shape)  # 输出 [4, 4]

在Masked Multi-Head Attention 中 [B, H, Q, K] 也就是mask 的维度 [batch, num_heads, seq_q, seq_k]