拆解Transformer之Embedding

Embedding 层（词嵌入层）是 Transformer / GPT 模型的第一个关键步骤，它的作用是把 「离散的 token（整数 ID）」转化为「连续的向量表示」 。我们来彻底拆解它到底做了什么。

🧩 一、输入输出关系

假设输入序列是一句话：

arduino 复制代码

"The animal didn't cross the street"

经过 tokenizer 之后，变成一串整数 token id：

python 复制代码

[ 50256,  345,  1223,  319,  407,  262,  1565 ]

Embedding 层输入这些整数，输出一个形状为：

r 复制代码

(batch_size, seq_len, embedding_dim)

的张量。例如：

scss 复制代码

(1, 7, 768)

其中 768 是 GPT-2 的 embedding 维度。

🧠 二、它内部的数据结构

Embedding 层本质上就是一个查表（lookup table）。

可以理解为：

text 复制代码

Embedding Matrix W  ∈  [vocab_size, embedding_dim]

例如，词表大小 50000，embedding 维度 768：

ini 复制代码

W.shape = (50000, 768)

每一行对应一个词的向量表示：

css 复制代码

W[0] → <pad>
W[1] → <unk>
W[2] → the
W[3] → animal
...

当输入 token id [3, 4, 7] 时，Embedding 层做的操作其实就是：

python 复制代码

output = W[[3, 4, 7]]

也就是从矩阵中选取对应的行。

这一步是 O(1) 查表操作，不涉及矩阵乘法，因此非常快。

⚙️ 三、数学形式（与线性层的关系）

线性层（nn.Linear) 做的是矩阵乘法：

ini 复制代码

y = x @ W^T + b

Embedding 层其实相当于：

scss 复制代码

one_hot(x) @ W

其中 one_hot(x) 是一个稀疏矩阵，只有一个元素为 1。

因为大多数时候 x 是一个整数索引（不是 one-hot），所以 PyTorch 的 nn.Embedding 实际上就是优化后的查表操作。

📐 四、Embedding 的作用与意义

降维与连续化

把「离散符号」转为「连续向量」，便于模型通过矩阵运算处理。
语义空间映射

在训练中，模型会 自动学习到语义相近的词在向量空间中更接近。

例如：
复制代码
```
king - man + woman ≈ queen
```
与输出层共享权重（weight tying） GPT-2 等模型通常将输入的 embedding 矩阵与输出 softmax 层的权重共享：
python 复制代码
```
self.head.weight = self.tok_emb.weight
```
这行代码让输出层 head 与输入层 tok_emb 指向同一块权重内存区域。在反向传播时：
- embedding 权重既被输入端更新；
- 又通过输出 softmax 的梯度共同更新。
- 两处梯度会累加在同一个参数上。这样参数更少，且理论上提升泛化。
从数学上理解就是：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> logits i = h t ⋅ E [ i ] T \text{logits}_i = h_t \cdot E[i]^T </math>logitsi=ht⋅E[i]T

即模型对每个词的打分 = 当前上下文表示 <math xmlns="http://www.w3.org/1998/Math/MathML"> h t h_t </math>ht 与该词 embedding 的点积。

🧭 五、在 GPT 中的 embedding 流程

以 GPT 为例：

python 复制代码

# 假设输入 batch (B,T)
idx = torch.tensor([[12, 45, 6, 89]])

tok_emb = self.tok_emb(idx)         # (B,T,C)
pos_emb = self.pos_emb[:, :T, :]    # (1,T,C)
x = tok_emb + pos_emb               # 加上位置编码

其中：

tok_emb：每个 token 的词向量。
pos_emb：每个位置的向量（表示顺序信息）。
两者相加 → 得到最终的输入表示。

🧮 六、直观例子

python 复制代码

import torch
import torch.nn as nn

# 模拟一个小词表
vocab_size = 10
embed_dim = 4

embedding = nn.Embedding(vocab_size, embed_dim)
print("Embedding matrix:\n", embedding.weight)

# 输入一批 token id
token_ids = torch.tensor([[2, 5, 3]])
print("Token IDs:", token_ids)

out = embedding(token_ids)
print("Embedding Output:\n", out)
print("Shape:", out.shape)

输出类似：

⚙️ 七、Embedding 参数的初始化

在实现上，nn.Embedding(vocab_size, hidden_size) 的参数其实就是一个形状为 (vocab_size, hidden_size) 的矩阵。

每一行代表一个 token 的向量。
这些向量的初始值不是固定的语义值，而是从某种分布中随机初始化出来的。
模型进行训练时通过反向传播，会逐渐把这些随机向量"拉"到合适的位置，使相似词的向量更近。

常见策略包括：

初始化方式	公式/区间	说明
Xavier Uniform	𝑈(−√(6/(in+out)), √(6/(in+out)))	PyTorch 默认用这种方式来初始化线性层、Embedding 等，保证前后层方差一致。
Normal(0, 0.02)	均值 0，标准差 0.02	GPT-2 原论文和 OpenAI 实现中使用此法，与 BERT 一致。
Uniform(-0.1, 0.1)	早期模型常用	较简单但效果略差。

不同实现会略有差异，但核心原则是一致的：让模型在初始阶段保持数值稳定，不爆炸也不消失。

初始化的目的就是给每个 token 一个"可学习的随机起点 "，然后让模型通过梯度下降慢慢学习语言结构。

🔍 八、小结

步骤	操作	输出形状	备注
输入 token id	`[3, 4, 7]`	`(3,)`	离散整数
查表 `W[id]`	`(3, embed_dim)`	连续向量
加上位置嵌入	`(3, embed_dim)`	引入序列顺序信息
输出送入 Transformer Block	`(B, T, C)`	开始自注意力