2、PyTorch中的Embedding层：nn.Embedding模块详解、参数初始化、前向传播与反向传播

好，咱们进入实战环节。上一章我们把词嵌入的理论基础捋了一遍，这一章直接上手PyTorch。说实话，nn.Embedding这个模块看起来简单，但坑不少。我刚开始用的时候，以为就是个查表操作，结果在梯度传播上栽过跟头。今天咱们把它彻底讲透。

2.1 nn.Embedding模块的核心机制

先看最基础的东西。nn.Embedding本质上是一个可训练的查找表。你给它一个索引，它返回对应的向量。就这么简单？嗯，但背后的参数初始化、梯度流动，才是关键。

import torch
import torch.nn as nn

# 创建一个Embedding层：词汇表大小10，每个词映射到5维向量
embedding = nn.Embedding(num_embeddings=10, embedding_dim=5)

# 输入：一批索引（假设batch_size=3）
input_indices = torch.tensor([1, 5, 8])
output = embedding(input_indices)
print(output.shape)  # torch.Size([3, 5])

你看，输入是三个整数，输出是三个5维向量。这就是前向传播的全部。但注意，num_embeddings是词汇表大小，embedding_dim是向量维度。这两个参数决定了权重矩阵的形状：[num_embeddings, embedding_dim]。

2.2 参数初始化：别小看这一步

默认情况下，nn.Embedding使用标准正态分布初始化。但我个人习惯，会根据任务调整。为什么？因为初始化直接影响收敛速度，甚至模型效果。

我在做推荐系统项目时遇到过一个问题：冷启动词向量训练不稳定。后来发现是初始化范围太大，导致梯度爆炸。改成均匀分布小范围初始化后，问题解决了。

# 自定义初始化
def init_embedding_weights(embedding_layer, init_range=0.1):
    nn.init.uniform_(embedding_layer.weight, -init_range, init_range)

# 或者使用xavier初始化（适合与线性层配合）
nn.init.xavier_uniform_(embedding_layer.weight)

# 也可以加载预训练词向量
pretrained_vectors = torch.randn(10, 5)  # 假设已有
embedding.weight.data.copy_(pretrained_vectors)

2.3 前向传播：查表操作的背后

前向传播其实就是一个索引+拷贝的过程。但这里有个细节：输入可以是任意形状的张量，输出会保留输入的形状，只是最后一维变成embedding_dim。

# 输入可以是多维的
input_2d = torch.tensor([[1, 2], [3, 4]])  # shape: [2, 2]
output_2d = embedding(input_2d)
print(output_2d.shape)  # torch.Size([2, 2, 5])

为什么会这样？因为PyTorch把每个索引独立映射成向量，然后保持其他维度不变。这在处理序列数据时特别有用——你直接输入[batch_size, seq_len]，输出就是[batch_size, seq_len, embedding_dim]。

嗯，这里要注意：索引不能越界。如果词汇表大小是10，你传索引10进去，会报错。所以padding_idx参数很重要，它指定一个特殊索引（通常是0），对应的向量不参与梯度更新。

# 带padding的Embedding
embedding_pad = nn.Embedding(10, 5, padding_idx=0)
input_with_pad = torch.tensor([0, 1, 2])  # 0是padding
output_pad = embedding_pad(input_with_pad)
# 索引0对应的向量始终为0，且梯度为0

2.4 反向传播：梯度如何流动？

这是很多人忽略的地方。反向传播时，Embedding层的梯度计算很简单：只有被索引到的行才有梯度。但问题来了——如果同一个索引在batch中出现多次，梯度会累加。

举个例子：

# 同一个索引出现两次
input_dup = torch.tensor([1, 1, 2])
output_dup = embedding(input_dup)
loss = output_dup.sum()
loss.backward()

# 查看梯度
print(embedding.weight.grad[1])  # 梯度是2（因为索引1出现了两次）
print(embedding.weight.grad[2])  # 梯度是1
print(embedding.weight.grad[3])  # 梯度是0（没被索引到）

我曾经在训练语言模型时，发现某些词的向量更新特别快，而另一些几乎不动。排查了半天，才发现是数据分布不均——高频词被频繁索引，梯度累加导致更新幅度大；低频词几乎没梯度。这就是Embedding层的梯度稀疏性问题。

2.5 实战技巧：权重共享与梯度优化

说到权重共享，最经典的就是Transformer中的输入Embedding与输出Embedding共享权重。为什么这么做？因为输入和输出的词汇表是同一个，共享权重可以减少参数量，同时让模型学到更一致的表示。

class SharedEmbeddingModel(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        # 输出层共享权重
        self.out_proj = nn.Linear(embed_dim, vocab_size, bias=False)
        self.out_proj.weight = self.embedding.weight  # 共享
        
    def forward(self, x):
        x = self.embedding(x)
        return self.out_proj(x)

注意：共享权重后，梯度会同时更新Embedding层和输出层的权重。但因为是同一个权重矩阵，梯度会累加。这在实践中效果不错，但如果你发现训练不稳定，可以尝试给梯度加个缩放因子。

2.6 总结：记住这三点

好了，这一章的内容就这些。我帮你梳理一下核心要点：

1. nn.Embedding就是查表——输入索引，输出向量。权重矩阵的每一行对应一个词的向量。
2. 梯度只更新被索引的行——高频词更新快，低频词更新慢。注意数据平衡。
3. 权重共享是常用技巧——减少参数量，提升表示一致性。但要注意梯度累加问题。

下一章我们会深入Embedding层的梯度优化实战，包括如何解决梯度稀疏、如何加速训练。到时候我会分享一个我在NLP竞赛中用的优化方案，保证让你眼前一亮。