模型内部 · 嵌入

嵌入

词元 + 位置 · GPT-2

图中数值
玩具规模
词表
50,257
dmodel
8 GPT-2 768
WE 参数
38.6M
WP 参数
0.79M

全部步骤

  1. 独热 id

    每个词元 id 变成一个独热行:50,257 个零,在 id 处有一个 1。叠起来,提示词就是一个 5 × 50,257 的矩阵。

    [5 × 50,257]

  2. 按行查表

    独热行乘以 WE,为每个词元选出 WE 的一行:纸面上是一次 GEMM,实践中是一次查表。每一行是对应词元的一段学到的 768 维描述,用法相近的词元得到相似的行:在 GPT-2 中,Ġcat 与 Ġdog(余弦 0.55)和 Ġkitten(0.50)的距离,比与 Ġon(0.21)更近。

    GPT-2 WE [50,257 × 768] · 3860 万参数

  3. 这些行意味着什么

    WE 的每一行都是对一个词元的学到的描述。从 768 维压到 2 维后,同类的词落在一起:数字、星期和月份、家庭成员、功能词。这些都是 GPT-2 的真实行;把鼠标停在一个词上可查看与它最接近的词。

    真实 WE 行 · 768 维中的 2 维(PCA)

  4. 加上位置

    注意力把输入当作无序集合(因果掩码只给出很弱的位置线索),所以每个位置都要加上位置矩阵 WP 中属于自己的一行。第 i 个位置总是得到第 i 行。试试其他提示词:一个重复的词,以及顺序不同的同一组词。

    GPT-2 WP [1,024 × 768] · 上下文 1,024

  5. 进入残差流

    这个和就是进入第 1 块的残差流:每个词元一条宽 768 的车道,彼此还没有混合。

    h₀ [N × 768]

代码

# idx: (B, T) token ids
# wte(idx) equals F.one_hot(idx, 50257).float() @ wte.weight, without the multiply
tok_emb = self.transformer.wte(idx)          # (B, T, 768)
pos = torch.arange(0, T, device=idx.device)
pos_emb = self.transformer.wpe(pos)          # (T, 768)
x = self.transformer.drop(tok_emb + pos_emb) # the residual stream

延伸阅读