transformer模型维度详解

某驾校的程序员

665人浏览 · 2026-02-03 21:07:07

某驾校的程序员 · 2026-02-03 21:07:07 发布

一、词嵌入

一个简单的例子：我爱深度学习 -> 我 | 爱 | 深度 | 学习分词之后得到四个token

假设词映射表：

"我"    →  15
"爱"    →  256
"深度"  →  1024
"学习"  →  4096
这句话就变成：[15, 256, 1024, 4096]

什么是词表大小（vocab size）？

词表大小 = 模型认识多少个"词"

什么是词嵌入维度(embedding_dim 或者 d_model) ?

"cat" → [0.12, -0.87, 0.33, ..., 0.05] （长度 = embedding_dim）
一个词用多少维的向量表示，词嵌入维度就等于多少

Embedding 层本质是一个矩阵（可学习参数，能反向传播，能计算梯度）：

Embedding Matrix: (vocab_size, embedding_dim)
= (10000, 512)

词嵌入矩阵：

embedding = nn.Embedding(10000, 512)

我们根据[15, 256, 1024, 4096] 这个4维token 查询词嵌入矩阵，做的就是取第 15，256，1024，4096 行，得到一个 4 * 512 的一个输入矩阵。到此我们便实现了从一句话 -> 一个矩阵的映射 [token个数，词嵌入维度]。在训练中，我们通常会多个batch一起训练，就会增加一个新的维度变成：[batch, token个数，词嵌入维度]

以上部分就对应transformer 架构中的下面这个部分，值得一提的是，再transformer中encoder和decorder中使用的词嵌入举证是同一个矩阵，这里inputs 输入的是 "我爱深度学习" ,outputs的输入就是 "I love deep learning"

二、位置编码

如果我们直接将上面的输出喂给self-attention中，在self-attention 中，我们对不同位置的token的输出结果是一样的。模型不能理解到不同位置词性所表达的意思。比如说：

我爱你
你爱我

模型不能分出上面两句话的差异。因此我么们需要对位置进行编码。所以需要把「位置信息」编码成向量，加到词向量里。论文中给出的位置编码公式如下：

其中 pos表示的是第几个token,即词在句子中的位置。i 是维度索引，embedding 维度是 d_model。这里举一个简单的例子：假设d_model = 8，每个维度的固定信息如下，接下来只需要把pos信息加到分子中去就行了。

偶数维度所在的行是sin函数，奇数维度所在行是cos函数（从0开始），我们来看一下第一列的输出维度，由于pos = 0,所以带入得到第一个token的位置编码信息如下：pos=0位置对应的编码向量是0，1交替的序列，（ PE(1）里面的分母只是取了一个大概的近似）

我们再回到编码公式，可以发现分母：10000^(2i / d_model) 不同维度使用不同“频率”的正弦波。为什么要用「不同频率」的正弦和余弦？低维度 → 变化快（高频），适合表示：相邻词、短距离关系；高维度 → 变化慢（低频），适合表示：长距离依赖

这样我们就得到了一个位置编码矩阵，他的维度和输入的词向量是完全一样的，因此位置向量和词向量能够直接相加。那么直接相加不会破坏词向量本身的信息吗?我的猜想是模型的庞大能够支撑模型自动识别出这一操作。（就一黑盒，模型能学到，可以把模型一个智力MAX的大脑？）自此我们的输出维度依然是：[batch, token个数，词嵌入维度]，只是在上一步的基础上加上了位置信息。

self.pos_emb = nn.Embedding.from_pretrained(
                get_sinusoid_encoding_table(src_len+1, d_model),
                freeze=True)

三、多头注意力机制

我们做出以下前提假设：

Batch size: B = 2

序列长度（token个数）: seq_len = 4

Transformer d_model = 8（为了展示，真实常是 512）

多头数量: h = 2

每个 head 的维度: d_k = d_model / h = 4

假设输入为x则：

X.shape = (B, seq_len, d_model) = (2, 4, 8)

运算过程如下：

1、生成QKV矩阵

其中 W_Q,W_K,W_V的维度是：(d_model, d_model) = (8, 8)，X只有后两个维度和QKV矩阵进行计算。X：（2，4，8）;计算得到QKV矩阵的维度为：（B，seq_len,d_model） = (2 , 4 , 8),随后对d_model维度进行拆分（拆成多个head），每个head的维度是 d_k = d_model / h = 4，维度变化如下：

(B, seq_len, d_model) → (B, seq_len, h, d_k)
(2, 4, 8) → (2, 4, 2, 4)

然后 transpose，让 batch 和 head 在前面

(B, seq_len, h, d_k) → (B, h, seq_len, d_k)
(2, 4, 2, 4) → (2, 2, 4, 4)

现在每个 head 都独立处理自己的 Q, K, V

2、计算每个 head 的注意力矩阵

计算公式：

2.1 计算Q @ Kᵀ -> Attention_scores

Q，K 维度信息如下：

Q: (B, h, seq_len, d_k) = (2, 2, 4, 4)

Kᵀ: (B, h, d_k, seq_len) = (2, 2, 4, 4)

两者相乘之后的维度：Attention_scores.shape = (B, h, seq_len, seq_len) = (2, 2, 4, 4)，这样每一个head得到一个（4，4）的注意力矩阵，表示每个词对句子中所有词的注意力分布

2.2、缩放 + softmax

缩放便是公式中的除以 $\sqrt{d_k}$ , 两个相同n维度的服从均值0方差1的矩阵相乘，得到的新矩阵的方差是n,因此除以 $\sqrt{d_k}$ 是为了缩放点积，使方差保持在 1 左右，从而softmax 输出更平滑,梯度不会消失,模型训练更稳定(方差固定，不随维度增长而膨胀)。

该层输出维度为：(B, h, seq_len, seq_len) = (2, 2, 4, 4)

2.3、加权求和 V

关注下面两个矩阵的维度：

Attention_scores: (B, h, seq_len, seq_len) = (2, 2, 4, 4)

V: (B, h, seq_len, d_k) = (2, 2, 4, 4)

做矩阵乘法（最后两个维度）

(B, h, seq_len, seq_len) @ (B, h, seq_len, d_k) → (B, h, seq_len, d_k)
(2, 2, 4, 4) @ (2, 2, 4, 4) → (2, 2, 4, 4) 即每个head的维度为（seq_len, d_k）= （4，4）

随后通过矩阵转置(B, h, seq_len, d_k) - > (B, seq_len, h, d_k) = (2, 4, 2, 4),在使用reshape合并多头：concat_heads: (B, seq_len, h*d_k) = (2, 4, 8)。随后我们加上层归一化和残差连接。

这样我们就得到了下面这部分的输出：

四、通过最终线性层W_O:(d_model, d_model) = (8, 8)

output = concat_heads @ W_O

output.shape = (B, seq_len, d_model) = (2, 4, 8),最终输出维度与输入一致

到此，我们讲解完了encoder部分的维度变化

五、Masked Mutil-Head Attention

用一个例子的讲一下masked的作用：输入（真实序列）: 我爱深度学习，当decoder在生成的时候，生成第1个词（”我“）的时候，不能看到任何后面的词；第 2 个词（“爱”） → 只能看到 “我”；第 3 个词（“深度”） → 只能看到 “我爱”。

那么什么是Masked Multi-Head Attention？ Masked Multi-Head Attention = Multi-Head Self-Attention + 上三角 Mask 。Attention 结构 完全一样，但是在 softmax 之前，对“未来位置”强制设为不可见。在 Decoder Self-Attention 中：scores.shape = (B, h, seq_len, seq_len)。例如 seq_len = 4，单个 head 的 score 矩阵是：

        0     1     2     3
      ------------------------
0 |    ✓     ✗     ✗     ✗
1 |    ✓     ✓     ✗     ✗
2 |    ✓     ✓     ✓     ✗
3 |    ✓     ✓     ✓     ✓

mask =
[
 [ 0,  -∞,  -∞, -∞ ],
 [ 0,   0,  -∞, -∞ ],
 [ 0,   0,   0, -∞ ],
 [ 0,   0,   0,  0 ]
]

0允许看到，-∞不能看到。scores = scores + mask

那么score矩阵能够看到什么？某个head的score矩阵：scores[i][j] = token i 对 token j 的“关注分数”当没有加mask矩阵的时候，假设scores[2] = [2.1, 1.3, 0.7, 4.9] 加上mask：masked_scores[2] = [2.1, 1.3, 0.7, -∞]再经过softmax([2.1, 1.3, 0.7, -∞]) → [0.57, 0.26, 0.17, 0.00]
被masked的这一项经过soft-max之后就变成了0。