nn.Embedding

nn.Embedding 把稀疏离散 ID 映射为稠密可学习向量，本质是一张带梯度的查找表：理解它的参数语义、padding 处理与 weight tying 惯例，就掌握了序列模型的第一层。

在 torch.nn 的模块家族里，nn.Embedding 是最不像”层”的层——没有非线性，也不在相邻位置间运算，forward 只做一件事：按下标从矩阵里取行。它与 图像Embedding 这类概念笔记分工：那边讲 embedding 作为向量表示的空间性质，本篇只讲这一层的 API 面与工程惯例。

查找表的本质：索引替代乘法

它唯一的参数是 self.weight，shape (num_embeddings, embedding_dim)，第 $i$ 行就是 ID $i$ 的向量表示，随训练更新。数学上 forward 等价于 one-hot 乘权重：长度为 $V$ 的 one-hot 行向量乘 $(V,D)$ 矩阵，恰好选出第 $i$ 行。但没人真去乘——词表 $V$ 动辄数万，矩阵乘要做 $O(V\cdot D)$ 次浮点运算，one-hot 矩阵本身也巨大而近乎全零；直接索引只是一次内存 gather，读出 $D$ 个数。效率来源就是用索引跳过乘法，这也是离散特征一律走 embedding 而非 one-hot 输入的原因。

形状契约与 padding_idx

输入必须是 LongTensor——整数才能当索引；形状任意，最常见 (B, L) 的 ID 序列，输出在末尾拼上 embedding_dim，得到 (B, L, D)。

变长序列拼 batch 必须补齐长度，而 padding 位若不加处理，会作为正常词参与训练。padding_idx 指定某个 ID（惯例取 0）为 padding：该行初始化为零、每次 forward 后仍固定为零向量，且不计梯度——模型学不到补齐符的语义，attention 侧再配 mask 即可彻底屏蔽。padding ID 的取值属于 词表设计与裁剪 的一部分，须与 tokenizer 约定一致。

import torch
import torch.nn as nn
 
vocab, dim, pad = 1000, 64, 0
emb = nn.Embedding(vocab, dim, padding_idx=pad)
 
ids = torch.tensor([[5, 17, 23, 0], [8, 42, 0, 0]])  # (B, L) LongTensor，0 为 padding
out = emb(ids)                                       # (B, L, D)
assert out.shape == (2, 4, dim)
assert out[0, 3].abs().sum() == 0                    # padding 位恒为零向量

同一签名里还有三个参数一句话定位：max_norm 在每次取出时把超范数的向量重归一化，防个别 embedding 爆炸；scale_grad_by_freq 按 batch 内词频倒数缩放梯度，把求和变求平均；sparse=True 让反向只产生被查过行的稀疏梯度——大词表下稠密梯度是整张 $(V,D)$ 矩阵，稀疏更新须换 SparseAdam 这类支持稀疏梯度的优化器。

weight tying 与冻结

embedding 学到的向量不只用于输入：语言模型还要从表示空间回到词表，标准做法是让输出投影共享同一矩阵——输入端把 ID 映入表示空间，输出端用转置方向把表示映射回 logits，即 weight tying。实现上是把 nn.Linear 的 weight 直接指到 emb.weight，省下整份 $(V,D)$ 参数，对大词表模型收益显著，代价是输入输出表示被绑定。

import torch.nn as nn
 
class TinyLM(nn.Module):
    def __init__(self, vocab, dim):
        super().__init__()
        self.emb = nn.Embedding(vocab, dim, padding_idx=0)
        self.head = nn.Linear(dim, vocab, bias=False)
        self.head.weight = self.emb.weight  # weight tying：输入输出共享
 
    def forward(self, ids):
        return self.head(self.emb(ids))     # (B, L, vocab)，展平后接 [[nn.CrossEntropyLoss]]

预训练向量则有两种消费方式：emb.weight.requires_grad_(False) 冻结——向量不更新，但仍注册为参数，随 .to(device) 与 state_dict() 迁移，下游层照常训练；彻底只读的场景可改用 buffer 存向量，不进优化器、不占梯度图。落到实践，词表规模、padding 与 mask 是否配套、要不要 tying，这三个决定在建模期就要定下。