VuePress 站点

🎯 课程主题

词嵌入与位置编码的代码实现——Transformer 输入层的 PyTorch 实现。

📝 核心知识点

1. 词嵌入层代码实现

概念说明：使用 nn.Embedding 构建词嵌入矩阵。
关键细节：
- nn.Embedding(vocab_size, d_model) 构建 $V \times D$ 的词嵌入矩阵
- 输入 token ID → 输出词嵌入向量
- 乘以 $\sqrt{d_{model}}$ 进行缩放（论文做法）
- 缩放目的：使嵌入向量量级与位置编码、残差在同一尺度，稳定训练

2. 位置编码代码实现

概念说明：使用 sin/cos 三角函数构建位置编码矩阵。
关键细节：
- 创建全零矩阵 pe，大小为 [max_len, d_model]
- 计算位置索引 position：[0, 1, 2, ..., max_len-1]
- 计算分母项 div_term：$10000^{2i/d_{model}}$
- 偶数位用 sin，奇数位用 cos
- 增加 batch 维度：[1, max_len, d_model]
- 注册为 buffer（不参与训练，不更新）

3. 融合过程

概念说明：词嵌入向量与位置编码相加。
关键细节：
- 输入 token ID → nn.Embedding → 词嵌入向量
- 词嵌入向量 + 位置编码 → 融合输出
- 可选 Dropout 层
- 输出包含语义信息和位置信息

🧮 核心公式与推导

词嵌入缩放：

$$\text{Embedding}(x) = \text{nn.Embedding}(x) \times \sqrt{d_{model}}$$

位置编码：

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

$$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

融合：

$$\text{Input} = \text{Embedding}(x) \times \sqrt{d_{model}} + PE$$

🏗️ 模型架构与数据流向

💻 代码实战

import torch
import torch.nn as nn
import math

class TokenEmbedding(nn.Module):
    """词嵌入层"""
    def __init__(self, vocab_size, d_model):
        """
        vocab_size: V, 词表大小
        d_model: D, 嵌入维度 (如512)
        """
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model
    
    def forward(self, x):
        """
        x: [batch, seq_len] token IDs
        return: [batch, seq_len, d_model]
        """
        # 乘以sqrt(d_model)进行缩放, 使量级与位置编码一致
        return self.embedding(x) * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    """位置编码层"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        """
        d_model: D, 嵌入维度
        max_len: 最大序列长度
        dropout: dropout比例
        """
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 创建位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        
        # 位置索引 [max_len, 1]
        position = torch.arange(0, max_len).unsqueeze(1).float()
        
        # 分母项: 10000^(2i/d_model) → 使用exp简化计算
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model)
        )
        
        # 偶数位用sin, 奇数位用cos
        pe[:, 0::2] = torch.sin(position * div_term)  # [max_len, d_model/2]
        pe[:, 1::2] = torch.cos(position * div_term)  # [max_len, d_model/2]
        
        # 增加 batch 维度 [1, max_len, d_model]
        pe = pe.unsqueeze(0)
        
        # 注册为buffer, 不参与训练 (位置编码是固定的)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """
        x: [batch, seq_len, d_model] 词嵌入向量
        return: [batch, seq_len, d_model] 融合后的向量
        """
        # 词嵌入 + 位置编码 (自动广播)
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)


# 完整的输入层
class TransformerInput(nn.Module):
    """Transformer输入层 = 词嵌入 + 位置编码"""
    def __init__(self, vocab_size, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.token_embedding = TokenEmbedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_len, dropout)
    
    def forward(self, x):
        """
        x: [batch, seq_len] token IDs
        return: [batch, seq_len, d_model]
        """
        # 1. 词嵌入 (含缩放)
        embedding = self.token_embedding(x)
        
        # 2. 加位置编码
        output = self.positional_encoding(embedding)
        
        return output


# 使用示例
vocab_size = 32000
d_model = 512
batch_size = 2
seq_len = 10

input_layer = TransformerInput(vocab_size, d_model)

# 模拟输入: token IDs
token_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
output = input_layer(token_ids)
print(f"输入: {token_ids.shape}")   # [2, 10]
print(f"输出: {output.shape}")      # [2, 10, 512]

⚠️ 常见问题与避坑指南

词嵌入需乘以 $\sqrt{d_{model}}$ 缩放，否则量级与位置编码不匹配
位置编码用 register_buffer 注册，不参与训练（非可学习参数）
位置编码的 pe 矩阵只需计算一次，通过切片适配不同序列长度
div_term 使用 exp 计算而非直接幂运算，数值更稳定

💡 个人总结与延伸

词嵌入与位置编码是 Transformer 的输入层，将离散的 token ID 转换为连续的语义向量并注入位置信息。代码实现与理论公式完全对应。现代大模型中，位置编码有多种变体（可学习位置编码、RoPE 旋转位置编码等），但输入层的基本结构不变。