AI大模型入门到实战系列（五）上下文嵌入向量（contextualized embedding）

上下文嵌入向量（contextualized embedding）

- 示例：处理句子 "I love AI"
- - [步骤1: 初始化设置](#步骤1: 初始化设置)
  - [步骤2: 输入嵌入和位置编码](#步骤2: 输入嵌入和位置编码)
  - [步骤3: 生成Q、K、V矩阵](#步骤3: 生成Q、K、V矩阵)
  - [步骤4: 注意力计算（详细过程）](#步骤4: 注意力计算（详细过程）)
  - [步骤5: 合并多头输出](#步骤5: 合并多头输出)
  - [步骤6: 残差连接和层归一化](#步骤6: 残差连接和层归一化)
  - [步骤7: 前馈神经网络](#步骤7: 前馈神经网络)
  - [步骤8: 第二个残差连接和层归一化](#步骤8: 第二个残差连接和层归一化)
  - [步骤9: 完整流程总结](#步骤9: 完整流程总结)
- 可视化流程

这节我们展示从输入嵌入向量(input embedding)到上下文嵌入向量(contextualized embedding)的完整过程。我们将用一个简化版的Transformer层来处理一个短句。

示例：处理句子 "I love AI"

步骤1: 初始化设置

python 复制代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# 设置超参数
d_model = 12       # 嵌入维度（通常为768，这里简化）
n_heads = 3        # 注意力头数（通常为12）
d_k = d_model // n_heads  # 每个头的维度 = 12/3 = 4
seq_len = 3        # 句子长度：I, love, AI
batch_size = 1     # 批大小
vocab_size = 1000  # 词汇表大小

查看值变化

python 复制代码

print(f"模型配置: d_model={d_model}, n_heads={n_heads}, d_k={d_k}")
print(f"输入序列: 'I love AI' (长度={seq_len})")

输出

将模型配置为: d_model=12, n_heads=3, d_k=4

输入序列: 'I love AI' (长度=3)

步骤2: 输入嵌入和位置编码

python 复制代码

class EmbeddingsWithPosition(nn.Module):
    """嵌入层 + 位置编码"""
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(1000, d_model)  # 最大位置
        
    def forward(self, token_ids):
	    # token_ids: [batch, seq_len] = [1, 3]
	    # 示例：token_ids = tensor([[10, 20, 30]])
	    
	    # ========== 步骤1：生成位置索引 ==========
	    positions = torch.arange(token_ids.size(1)).unsqueeze(0)
	    # token_ids.size(1) = 3（序列长度）
	    # torch.arange(3) = tensor([0, 1, 2])
	    # .unsqueeze(0) 增加batch维度 => tensor([[0, 1, 2]])
	    
	    positions = positions.expand_as(token_ids)
	    # positions: [[0, 1, 2]]
	    # token_ids: [[10, 20, 30]] 形状相同
	    # expand_as复制到batch维度 => 仍然是 [[0, 1, 2]]
	    
	    # ========== 步骤2：获取词嵌入 ==========
	    token_embeds = self.token_embedding(token_ids)
	    # token_embeds.shape = [1, 3, 12]
	    # 过程：从词嵌入矩阵中查找第10、20、30行的向量
	    
	    # ========== 步骤3：获取位置嵌入 ==========
	    pos_embeds = self.position_embedding(positions)
	    # pos_embeds.shape = [1, 3, 12]
	    # 过程：从位置嵌入矩阵中查找第0、1、2行的向量
	    
	    # ========== 步骤4：相加 ==========
	    return token_embeds + pos_embeds
	    # 形状：仍为 [1, 3, 12]
	       

# 创建嵌入层
embeddings = EmbeddingsWithPosition(vocab_size, d_model)

# 假设token IDs: I=10, love=20, AI=30
input_ids = torch.tensor([[10, 20, 30]])  # [1, 3]

# 获得输入嵌入
input_embeddings = embeddings(input_ids)  # [1, 3, 12]

查看值变化

python 复制代码

print("步骤1: 输入嵌入 + 位置编码")
print(f"输入形状: {input_embeddings.shape}")
print(f"输入值示例（第一个token）:")
print(input_embeddings[0, 0].detach().numpy().round(3))
print()

输出

输入嵌入 + 位置编码

输入形状: torch.Size([1, 3, 12])

输入值示例（第一个token）:
$1.252 0.676 -0.077 -0.821 -0.958 -1.387 1.118 0.18 -0.108 -2.223 0.748 0.784$

nn.Embedding(vocab_size, d_model) 就是创建了一个 (vocab_size, d_model) 的可训练矩阵，专门用于将整数索引映射为稠密向量。

与 nn.Linear 的区别： nn.Linear 是做矩阵乘法 + 偏置，而 nn.Embedding 是纯粹的索引查找。不过，从数学上看，用一个one-hot向量乘以嵌入矩阵，结果就是索引查找。因此，Embedding层可以看作是一个没有偏置、输入为one-hot向量的Linear层的高效实现。

self.token_embedding(token_ids) 本质上就是在查找或索引每个token_id对应的嵌入向量。

这里需要解释下self.position_embedding = nn.Embedding(1000, d_model)位置不是固定的嘛，比如0,1,2，为什么还要用训练矩阵？

位置编号（0,1,2,...）只告诉你"第几个"，但可训练的位置嵌入告诉你"这个位置通常有什么作用"。类比：

位置编号 = 你的学号（2023001，2023002，...）→ 只是个编号

位置嵌入 = 你的座位特点（前排学霸区、后排休息区、靠窗风景位）→ 有实际意义

为什么需要学习：

不同任务：情感分析中句首重要，机器翻译中每个位置都重要

不同语言：中文动词靠前，英文动词位置灵活

不同长度：第10个字在短文中是结尾，在长文中是开头

模型自己发现规律：让数据告诉模型什么位置重要

self.position_embedding 和self.token_embedding的列数（维度）必须一致，但行数（数量）可以不同

步骤3: 生成Q、K、V矩阵

python 复制代码

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        # Q、K、V的线性变换
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)
        
    def generate_qkv(self, x):
        """生成Q、K、V矩阵"""
        batch_size, seq_len, _ = x.shape
        
        # 线性变换得到Q、K、V
        Q = self.W_Q(x)  # [batch, seq_len, d_model]
        K = self.W_K(x)  # [batch, seq_len, d_model]
        V = self.W_V(x)  # [batch, seq_len, d_model]
        
        # 重塑为多头格式 [batch, seq_len, n_heads, d_k]
        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_k)
        K = K.view(batch_size, seq_len, self.n_heads, self.d_k)
        V = V.view(batch_size, seq_len, self.n_heads, self.d_k)
        
        return Q, K, V

# 创建注意力层
attention = MultiHeadAttention(d_model, n_heads)

# 生成Q、K、V
Q, K, V = attention.generate_qkv(input_embeddings)

查看值变化

python 复制代码

print("生成Q、K、V矩阵")
print(f"Q形状: {Q.shape}")  # [1, 3, 3, 4]
print(f"K形状: {K.shape}")
print(f"V形状: {V.shape}")
print()

# 查看第一个token的第一个头的Q、K、V值
print("第一个token（'I'）在第一个头的表示:")
print(f"Q (查询向量): {Q[0, 0, 0].detach().numpy().round(3)}")
print(f"K (键向量):   {K[0, 0, 0].detach().numpy().round(3)}")
print(f"V (值向量):   {V[0, 0, 0].detach().numpy().round(3)}")
print()

输出

生成Q、K、V矩阵

Q形状: torch.Size([1, 3, 3, 4])

K形状: torch.Size([1, 3, 3, 4])

V形状: torch.Size([1, 3, 3, 4])

第一个token（'I'）在第一个头的表示:

Q (查询向量): [ 1.16 0.105 -1.654 -1.009]

K (键向量): [ 1.78 -0.608 1.395 2.054]

V (值向量): [-0.301 0.126 -1.675 0.474]

步骤4: 注意力计算（详细过程）

python 复制代码

def compute_attention(Q, K, V, mask=None):
    """计算缩放点积注意力"""
    batch_size, seq_len, n_heads, d_k = Q.shape
    
    # 1. 计算Q和K的点积
    # 维度: [batch, n_heads, seq_len, d_k] × [batch, n_heads, d_k, seq_len]
    # 结果: [batch, n_heads, seq_len, seq_len]
    K_t = K.transpose(2, 3)  # 转置最后两个维度
    scores = torch.matmul(Q, K_t) / math.sqrt(d_k)
    
    # 2. 应用softmax得到注意力权重
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    attention_weights = F.softmax(scores, dim=-1)
    
    # 3. 用注意力权重加权V
    # [batch, n_heads, seq_len, seq_len] × [batch, n_heads, seq_len, d_k]
    # 结果: [batch, n_heads, seq_len, d_k]
    weighted_V = torch.matmul(attention_weights, V)
    
    return weighted_V, attention_weights

# 调整维度顺序: [batch, seq_len, n_heads, d_k] -> [batch, n_heads, seq_len, d_k]
Q_perm = Q.permute(0, 2, 1, 3)
K_perm = K.permute(0, 2, 1, 3)
V_perm = V.permute(0, 2, 1, 3)

# 计算注意力
weighted_V, attention_weights = compute_attention(Q_perm, K_perm, V_perm)

查看值变化

python 复制代码

print(" 注意力计算")
print(f"注意力输出形状: {weighted_V.shape}")
print(f"注意力权重形状: {attention_weights.shape}")

print("\n注意力权重矩阵（第一个头）:")
print("行 = 查询token (关注者)")
print("列 = 键token (被关注者)")
print(attention_weights[0, 0].detach().numpy().round(3))

print("\n解释: 每个token如何关注其他token")
for i in range(seq_len):
    token_name = ["I", "love", "AI"][i]
    weights = attention_weights[0, 0, i].detach().numpy()
    print(f"'{token_name}' 的关注分布: {weights.round(3)}")
print()

输出

注意力计算

注意力输出形状: torch.Size([1, 3, 3, 4])

注意力权重形状: torch.Size([1, 3, 3, 3])

注意力权重矩阵（第一个头）:

行 = 查询token (关注者)

列 = 键token (被关注者)
$\[0.245 0.644 0.111$ $0.391 0.275 0.334$ $0.412 0.281 0.307\]$
解释: 每个token如何关注其他token

'I' 的关注分布: [0.245 0.644 0.111]

'love' 的关注分布: [0.391 0.275 0.334]

'AI' 的关注分布: [0.412 0.281 0.307]

步骤5: 合并多头输出

python 复制代码

def combine_heads(weighted_V, batch_size, seq_len, d_model):
    """合并多个注意力头的输出"""
    # 当前形状: [batch, n_heads, seq_len, d_k]
    # 目标形状: [batch, seq_len, d_model]
    
    # 1. 转置回 [batch, seq_len, n_heads, d_k]
    weighted_V = weighted_V.permute(0, 2, 1, 3).contiguous()
    
    # 2. 合并最后两个维度
    combined = weighted_V.view(batch_size, seq_len, d_model)
    
    return combined
# 合并多头
multi_head_output = combine_heads(weighted_V, batch_size, seq_len, d_model)

查看值变化

python 复制代码

print(" 合并多头输出")
print(f"合并后形状: {multi_head_output.shape}")
print(f"第一个token合并后表示: {multi_head_output[0, 0].detach().numpy().round(3)}")
print()

# 通过输出线性层
attention_output = attention.W_O(multi_head_output)
print(" 注意力子层输出线性变换")
print(f"注意力子层输出形状: {attention_output.shape}")
print(f"第一个token输出: {attention_output[0, 0].detach().numpy().round(3)}")
print()

输出

合并多头输出

合并后形状: torch.Size([1, 3, 12])

第一个token合并后表示: [-0.26 -0.52 -0.104 0.536 -0.23 0.92 -0.151 0.235 -0.375 -0.657 0.072 0.791]

注意力子层输出线性变换

注意力子层输出形状: torch.Size([1, 3, 12])

第一个token输出: [ 0.014 0.172 0.23 0.118 -0.232 0.185 -0.274 0.207 -0.297 -0.075 0.644 -0.492]

步骤6: 残差连接和层归一化

python 复制代码

class LayerNorm(nn.Module):
    """简化版层归一化"""
    def __init__(self, features, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))
        self.eps = eps
        
    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

# 残差连接
residual = input_embeddings + attention_output

# 层归一化
layer_norm1 = LayerNorm(d_model)
norm_output1 = layer_norm1(residual)

查看值变化

python 复制代码

print(" 残差连接 + 层归一化")
print(f"残差连接: 输入嵌入 + 注意力输出")
print(f"归一化后形状: {norm_output1.shape}")
print(f"第一个token归一化后: {norm_output1[0, 0].detach().numpy().round(3)}")
print()

输出

残差连接 + 层归一化
残差连接: 输入嵌入 + 注意力输出

归一化后形状: torch.Size([1, 3, 12])

第一个token归一化后: [ 2.259 -0.111 -0.147 0.104 -0.085 -0.475 -0.026 0.058 -1.686 0.425

0.97 -1.286]

步骤7: 前馈神经网络

python 复制代码

class FeedForwardNetwork(nn.Module):
    """前馈神经网络"""
    def __init__(self, d_model, d_ff=48):  # d_ff通常是d_model的4倍
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.activation = nn.ReLU()
        
    def forward(self, x):
        return self.linear2(self.activation(self.linear1(x)))

# 前馈网络
ffn = FeedForwardNetwork(d_model, d_ff=48)
ffn_output = ffn(norm_output1)

查看值变化

python 复制代码

print("前馈神经网络")
print(f"FFN输出形状: {ffn_output.shape}")
print(f"FFN内部维度: d_model={d_model} → d_ff=48 → d_model={d_model}")
print(f"第一个token FFN输出: {ffn_output[0, 0].detach().numpy().round(3)}")
print()

输出

前馈神经网络

FFN输出形状: torch.Size([1, 3, 12])

FFN内部维度: d_model=12 → d_ff=48 → d_model=12

第一个token FFN输出: [-0.285 -0.161 -0.256 0.26 -0.139 -0.164 0.041 0.24 0.065 -0.145 -0.022 -0.098]

步骤8: 第二个残差连接和层归一化

python 复制代码

# 第二次残差连接
final_residual = norm_output1 + ffn_output

# 第二次层归一化
layer_norm2 = LayerNorm(d_model)
context_embeddings = layer_norm2(final_residual)

查看值变化

python 复制代码

print(" 最终残差连接 + 层归一化")
print("→ 得到上下文嵌入向量!")
print(f"上下文嵌入向量形状: {context_embeddings.shape}")
print(f"第一个token的上下文嵌入: {context_embeddings[0, 0].detach().numpy().round(3)}")
print()

输出

最终残差连接 + 层归一化

→ 得到上下文嵌入向量!

上下文嵌入向量形状: torch.Size([1, 3, 12])

第一个token的上下文嵌入: [ 2.097 -0.224 -0.359 0.433 -0.175 -0.602 0.073 0.365 -1.619 0.346 1.038 -1.373]

步骤9: 完整流程总结

python 复制代码

class TransformerBlock(nn.Module):
    """完整的Transformer块"""
    def __init__(self, d_model, n_heads, d_ff=None):
        super().__init__()
        if d_ff is None:
            d_ff = d_model * 4
            
        self.attention = MultiHeadAttention(d_model, n_heads)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.ffn = FeedForwardNetwork(d_model, d_ff)
        
    def forward(self, x):
        """完整的前向传播"""
        print("=" * 60)
        print("Transformer块完整流程:")
        print("=" * 60)
        
        # 1. 输入
        print(f"1. 输入嵌入: {x.shape}")
        
        # 2. 注意力机制
        Q, K, V = self.attention.generate_qkv(x)
        print(f"2. 生成Q/K/V: {Q.shape}")
        
        # 3. 计算注意力
        Q_perm = Q.permute(0, 2, 1, 3)
        K_perm = K.permute(0, 2, 1, 3)
        V_perm = V.permute(0, 2, 1, 3)
        weighted_V, attn_weights = compute_attention(Q_perm, K_perm, V_perm)
        print(f"3. 注意力计算: {weighted_V.shape}")
        
        # 4. 合并多头
        batch_size, seq_len = x.shape[:2]
        multi_head = combine_heads(weighted_V, batch_size, seq_len, self.attention.d_model)
        attention_output = self.attention.W_O(multi_head)
        print(f"4. 注意力子层输出: {attention_output.shape}")
        
        # 5. 残差连接 + 层归一化
        residual1 = x + attention_output
        norm1_output = self.norm1(residual1)
        print(f"5. 第一个残差连接+层归一化: {norm1_output.shape}")
        
        # 6. 前馈网络
        ffn_output = self.ffn(norm1_output)
        print(f"6. 前馈网络: {ffn_output.shape}")
        
        # 7. 最终输出
        final_residual = norm1_output + ffn_output
        context_embeddings = self.norm2(final_residual)
        print(f"7. 最终输出（上下文嵌入）: {context_embeddings.shape}")
        
        return context_embeddings

# 运行完整流程
print("\n" + "=" * 60)
print("完整流程演示")
print("=" * 60)

transformer_block = TransformerBlock(d_model, n_heads)
context_embeddings = transformer_block(input_embeddings)

print("\n" + "=" * 60)
print("最终结果对比")
print("=" * 60)

print("\n输入嵌入（无上下文）:")
print(input_embeddings[0, 0].detach().numpy().round(3))

print("\n上下文嵌入（有关联信息）:")
print(context_embeddings[0, 0].detach().numpy().round(3))

print("\nV矩阵（注意力机制的输入）:")
print(V[0, 0, 0].detach().numpy().round(3))

print("\n结论: V ≠ 上下文嵌入向量")
print("- V: 是注意力计算中的'值'矩阵")
print("- 上下文嵌入: 经过完整Transformer层处理后的结果")

可视化流程

复制代码

完整Transformer层流程:

输入嵌入 [1, 3, 12]
    ↓
┌─────────────────────────────────┐
│  生成Q、K、V                     │
│  - 线性变换                     │
│  - 分割多头: [1, 3, 3, 4]      │
└─────────────────────────────────┘
    ↓
┌─────────────────────────────────┐
│  注意力计算                     │
│  Q·K^T / √d_k → softmax → 权重  │
│  权重 × V → 加权输出            │
└─────────────────────────────────┘
    ↓
┌─────────────────────────────────┐
│  合并多头输出                   │
│  [1, 3, 3, 4] → [1, 3, 12]     │
│  线性变换                       │
└─────────────────────────────────┘
    ↓
┌─────────────────────────────────┐
│  残差连接 + 层归一化            │
│  output = norm(x + attention(x))│
└─────────────────────────────────┘
    ↓
┌─────────────────────────────────┐
│  前馈神经网络                  │
│  d_model → d_ff → d_model      │
└─────────────────────────────────┘
    ↓
┌─────────────────────────────────┐
│  残差连接 + 层归一化            │
│  output = norm(x + FFN(x))     │
└─────────────────────────────────┘
    ↓
上下文嵌入向量 [1, 3, 12]