模型内部 · 反嵌入与采样

反嵌入与采样

LM head · 与 WE 共享 · 真实运行

图中数值
真实 GPT-2
dmodel
768 画出 8 个
词表
50,257
LM head
WEᵀ，共享权重

全部步骤

  1. 最后的 LayerNorm

    只有最后一个位置预测下一个词元。它的向量在经过全部 12 个块之后，再通过最后的 LayerNorm（lnf）。这些数值来自 GPT-2 small 的真实运行。

    [1 × 768]

  2. Logits

    这个向量与嵌入矩阵的每一行做点积：反嵌入矩阵 WU（即 LM head）在 GPT-2 中就是 WEᵀ。于是每个词元得到一个 logit。图中为最可能的 6 个词元画出了 768 维中的 8 维；每个 logit 实际上是全部 768 维之和。

    [1×768]·[768×50,257] · 77M FLOPs

  3. 温度

    在 softmax 之前，logits 先除以温度 T。只有它们之间的差距才重要：较低的 T 拉大差距，让分布更尖锐；较高的 T 让分布更平坦。

    z / T

  4. Softmax

    softmax 把缩放后的 logits 变成覆盖全部 50,257 个词元、总和为 1 的概率。T 低时前几名拿走大部分概率；提高 T 则把概率挪向长尾。

    p = softmax(z / T)

  5. 采样

    由一种策略选出下一个词元：贪心取概率最高的；采样把概率首尾相接排开，由一个随机数 u 选中一个；top-k 或 top-p 则先砍掉长尾。生成会一直重复，直到出现 <|endoftext|> 或达到长度上限。

    采样 · 贪心 · top-k · top-p

代码

x = self.transformer.ln_f(x)
logits = self.lm_head(x[:, [-1], :])      # lm_head.weight is wte.weight (tied)
logits = logits[:, -1, :] / temperature
if top_k is not None:
    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
    logits[logits < v[:, [-1]]] = -float("Inf")
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)   # greedy: probs.argmax(-1)
idx = torch.cat((idx, idx_next), dim=1)              # append, run again

延伸阅读