第10节——手搓Token Embedding代码

10.Token Embedding

• Token Embedding 层其实是 Encoder、也是整个 Transformer 最基础，最前端的模块
• 比前面的 Attention 简单很多，但非常关键

---

10.1 Token Embedding 是什么？

• 之前我们说过，Transformer 的输入真正喂进 Encoder 的，不是 token id，而是：
  X_embed: (B, L, d_model)
  也就是每个 token 被映射成一个 d_model 维的向量。
• 最开始我们拿到的通常是这样的 token 映射：

  src = torch.tensor([
    [5, 9, 2, 7, 0],
    [4, 3, 8, 0, 0],
  ])  # (B=2, L=5)
  

• 这里的数字是 token id（词表中的索引），比如：

  5 → "I"
  9 → "love"
  2 → "you"
  0 → <PAD>
  

• Token Embedding 的任务就是：
  把这些离散的 id → 映射成连续的向量表示。
  公式：
  Embedding:{0,1,2,…,V−1}→Rdmodel\text{Embedding}: \{0,1,2,\dots,V-1\} \to \mathbb{R}^{d_{\text{model}}}Embedding:{0,1,2,…,V−1}→Rdmodel​

---

10.2 nn.Embedding 在数学上等价于什么？

• 代码：

  embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model)
  

  • vocab_size:词汇表大小，原始文本数据集中所有唯一字符的数量
• PyTorch 会内部创建一个矩阵：

  E: (vocab_size, d_model)
  

  • 第 i 行：就是 id 为 i 的 token 对应的向量表示
  • 查表方式：给一个 id，它就返回 E[id] 向量（词向量）
• 本质上就是“查表”，但 E 的每一行是可训练参数，会在反向传播中被更新（理解词在高维空间中的向量表示）。

---

10.3 手搓一个 TokenEmbedding 模块

• 写一个小类，把词表大小：vocab_size、模型维度：d_model、pad_id：padding 的 id（比如 0）都封装进去。

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, pad_id=0):
        """
        vocab_size: 词表大小（总共有多少个 token id）
        d_model:    每个 token 映射后的向量维度（要和 Transformer 的 d_model 对齐）
        pad_id:     用于 padding 的 token id（通常是 0）
        """
        super().__init__()
        
        # nn.Embedding 本质上就是一个 (vocab_size, d_model) 的查表矩阵 E
        # 第 i 行是 token i 的向量表示，会在训练中被更新
        self.embedding = nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=d_model,
            padding_idx=pad_id,  # 这个 id 的向量在训练时不会被更新
        )
    
    def forward(self, x):
        """
        x: (B, L)，里面是 token id（int 类型，如 0,1,2,...）
        返回:
            embed: (B, L, d_model)，每个 token 对应一个 d_model 维向量
        """
        return self.embedding(x)

# 先运行上面的类再运行测试代码
# 假设：
vocab_size = 10   # 词表里有 0~9 这 10 个 token
d_model = 8       # 每个 token 映射成 8 维向量
PAD = 0

embed_layer = TokenEmbedding(vocab_size, d_model, pad_id=PAD)

# 模拟一个 batch 的 token id 输入：
src = torch.tensor([
    [1, 2, 3, 0, 0],   # 句子1（后两个是 PAD）
    [4, 5, 0, 0, 0],   # 句子2
])  # (B=2, L=5)

emb = embed_layer(src)
print("src shape:", src.shape)   # (2, 5)
print("emb shape:", emb.shape)   # (2, 5, 8)
print("emb[0, 0]:", emb[0, 0])   # 第一句第一个 token（1） 的向量

src shape: torch.Size([2, 5])
emb shape: torch.Size([2, 5, 8])
emb[0, 0]: tensor([ 0.7941, -0.2266, -0.9386, -0.9746, -0.4068,  0.4354,  0.0848, -1.6910],
       grad_fn=<SelectBackward0>)

• 同一个 id（比如 1）在不同位置得到的 embedding 是一模一样的（共享参数）（就是查字典了）