Token 与位置嵌入:从整数到向量

Token 与位置嵌入:从整数到向量

本节摘要:id 是整数,模型要向量。两者之间坐两张查表,位置那张的选择决定模型能学什么。模型对 token id 的首次接触是 token 嵌入矩阵的行查表——矩阵每词表 id 一行、每模型维一列,查表返回的向量被模型其余部分当作 id 的含义,反向传播更新前向用到的行。token id 本身无序,模型需第二信号告诉它位置一与位置十七不同。两种主导选择:学习的位置嵌入(第二张查表,每位置一行)与固定正弦位置嵌入(无参数学公式)。本节构建两者,与 token 嵌入合成注意力块的输入,并对比它们在长度泛化与参数量上的差异。

对应原课程:Phase 19 · Lesson 32 · token-positional-embeddings(原英文 phases/19-capstone-projects/32-token-positional-embeddings/docs/en.md)。本节属「从零构建 GPT」赛道第三节。

学习目标

阅读完本节,你应当能够:

1. 构建把词表 id 映到稠密向量的 token 嵌入查表。
2. 构建按位置索引的学习位置嵌入查表。
3. 构建按位置索引、无参的固定正弦位置嵌入。
4. 把 token 与位置嵌入合成 transformer 块的单个输入。
5. 对比学习与正弦嵌入在长度泛化与参数量上的差异。

一、问题与直觉

位置嵌入的选择有后果。学习表是参数,受模型训练的最大上下文长约束;正弦表理论上无参,公式延展到任何位置,但本节的 SinusoidalPositionalEmbedding 在 max_context_length 处预计算固定表,forward 过界抛错——两模块在此都强制最大上下文长。模型即便表够大可索引,过训练长度仍可能挣扎。

合成是求和而非拼接——求和让 D 在网络中恒定,让模型在每层逐特征决定是 token 含义还是位置主导。

二、从零实现

token 嵌入矩阵:形状 (V, D) 的参数张量,V 是词表大小,PyTorch 暴露为 nn.Embedding(V, D)。初始化从小高斯抽(传统均值零、std 约 0.02),精确初始化不如跨运行一致重要。前向是单次索引:PyTorch 把 (B, T) int64 id 经 gather 映成 (B, T, D) float,反向只把梯度累积进前向触到的行。微妙细节:token 嵌入与模型末端输出投影常共享权重(weight tying),此时每次反向经输出侧触每行;本节两者分模块,但同矩阵可扮两角。

学习位置嵌入:第二个 nn.Embedding,形状 (max_context_length, D),按位置 id 0..T-1 查,前向把位置向量沿批维广播。缺点:模型只训到位置 T-1 时,查位置 T 行不存在;用此方案的生产 decoder-only 模型把最大上下文长烤进架构,拒处理更长输入。

正弦位置嵌入:位置到向量的函数,位置 p 与特征 i 产:

angle = p / (10000 ** (2 * (i // 2) / D))
emb[p, 2k]     = sin(angle)
emb[p, 2k + 1] = cos(angle)
​

函数无参,每位置向量唯一。波长跨特征维几何变化,低维编粗位置、高维编细位置。选 sin 与 cos 共同带来的性质:位置 p+k 的向量是位置 p 向量的线性函数——给注意力层一条学相对位置偏移的捷径,无需单独参数表达「回看五个 token」。本节构造时一次算全正弦表,前向时索引。

三、合成与对比

输入管线三步:读 token id、查 token 向量、加位置向量、返回和。求和步的广播把 (T, D) 位置张量沿批维复制,PyTorch 因位置张量 unsqueeze 后形 (1, T, D) 自动处理。

本节在同输入上跑两变体,打两个诊断。参数量:学习变体加 max_context_length * D 参数,正弦加零。相邻位置余弦相似度:正弦因函数连续有平滑可预测衰减;学习变体初始化时因行独立抽取近随机相似度,训练后通常发展出类似平滑结构,但它得从数据发现那结构。

四、可复用产物

main.py 定义三模块:TokenEmbedding(包 nn.Embedding(V, D))、LearnedPositionalEmbedding(包 nn.Embedding(L, D))、SinusoidalPositionalEmbedding(预计算表,暴露为 buffer)。EmbeddingComposer 把 token 与位置嵌入系一起。demo 打形状、参数量、邻位置相似度诊断。code/tests/test_embeddings.py 钉形状、广播行为、参数量、正弦公式。把模型维 D 从 64 调 32,看正弦波长带变化。

五、框架对比

学习位置嵌入(GPT-2 风格)简单但硬绑训练长度;正弦(原 Transformer 论文)无参且相对位置友好,但绝对位置信号弱。现代生产 transformer 走 RoPE(旋转位置编码)或 AliBi——它们遵循同样形状契约((B, T, D) 向量的位置依赖变换),但在注意力的 QK 投影步应用,而非输入。本节不建 RoPE/AliBi,下一节建注意力块,可选扩展是把旋转折进 QK 投影。本节不训练嵌入——训练需损失,损失需模型输出,需注意力与 LM 头,那是后两节。

六、练习

1. 参数量对比:断言学习变体比正弦多 max_context_length * D 参数。
2. 邻位置相似度:打印两变体在相邻位置的余弦相似度,确认正弦平滑、学习初始化近随机。
3. 过界行为:构造长度超 max_context_length 的输入,确认两模块都抛错。
4. D 变化:D 从 64 调 32,确认正弦波长带变化(低维粗、高维细)。
5. weight tying:把 token 嵌入与输出投影设为同矩阵,确认反向触每行。

本节要点回顾

1. 两张查表:token 嵌入(V, D)+ 位置嵌入(L, D),合成是求和(保 D 恒定)。
2. 学习 vs 正弦:学习硬绑训练长度、有参;正弦无参、相对位置友好但绝对信号弱。
3. 正弦性质:位置 p+k 是 p 的线性函数,给相对位置学习的捷径。
4. 过界都抛:两模块在此都强制最大上下文长,即便表够大模型过训练长仍挣扎。
5. token 嵌入反向只触用过的行:除非 weight tying,否则输出侧触每行。
6. 现代走 RoPE/AliBi:在 QK 投影步应用位置变换,而非输入,本节为下一节铺路。

下一节,我们建「多头自注意力」——一次线性投影、三个视图、H 个并行头、一个因果掩码。