MiniMind 学习笔记之 05 从 768 维到下一个 token

上承〔04 · 注意力之外〕。前四篇把模型内部拆完了：文字进来变成向量，向量经过 8 层积木，每层做注意力、FFN、残差、归一化。但积木的输出还是 768 维的向量。模型怎么把这个向量变成下一个 token？这一篇只讲 MiniMind 实际用到的部分：输出层、Softmax、温度、Top-k、Top-p，以及生成循环。

一、从 768 维到 6400 个分数

1.1 最后一层输出的是什么

输入"今天天气"，经过 embedding、8 层积木，最后一层输出一个张量，形状是：

(batch, seq_len, 768)

batch 是批次大小，seq_len 是序列长度，768 是 hidden_size。

要预测下一个 token，只关心最后一个位置的向量。因为因果掩码的存在，最后一个位置已经看过了前面所有 token 的信息。它的 768 维向量，就是模型对"接下来该是什么"的全部理解。

把这个向量取出来：

hidden_states = hidden_states[:, -1, :]   # (batch, 768)

1.2 lm_head：从 768 到 6400

拿到 768 维向量后，过一个线性层：

self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

MiniMind 的 hidden_size 是 768，vocab_size 是 6400。所以 lm_head 的形状是 768 × 6400。

输入 768 维，输出 6400 维。这 6400 个数，对应词表里的 6400 个 token，每个数是一个分数。

比如：

"好" → 8.3
"啊" → 6.1
"，" → 5.7
"冷" → 4.2
...

这些分数叫 logits。可正可负，不是概率，也没有归一化。分数越高，模型认为这个 token 越可能是下一个。

1.3 权重绑定：lm_head 和 embed_tokens 共享矩阵

第一篇算参数时提到过一件事：MiniMind 的 embedding 矩阵被存了两份，一份给输入，一份给输出。这就是权重绑定（weight tying）。

embed_tokens 的形状是 6400 × 768，负责把 token id 变成向量。 lm_head 的形状是 768 × 6400，负责把向量变回 token 分数。

这两个矩阵，形状互为转置，功能也对称------一个从词表到向量，一个从向量回词表。所以可以共享同一组参数。

源码里：

self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
self.lm_head.weight = self.embed_tokens.weight   # 共享

共享的好处有两个：

省参数。 6400 × 768 = 4,915,200 个数，共享之后只存一份。

语义一致。 输入和输出用同一套向量表示。同一个 token，进来的向量和出去的向量在同一个空间里，方向一致。这种对称性对训练有帮助。

第一篇对账时，多出来的那 4,915,200 个数，就是这一块。计算参数时算了两遍，存文件时写了两份，实际模型里只有一份。

二、Softmax：从分数到概率

2.1 logits 不是概率

6400 个 logits，可正可负，没有范围限制。它们不是概率------不保证为正，加起来也不等于 1。

要从中选一个 token，得先把它们变成概率分布。

2.2 Softmax 的做法

Softmax 的公式：
$$softmax(zi)=\; ezi\; \sum jezj\; \backslash text\{softmax\}(z\_i)\; =\; \backslash frac\{e^\{z\_i\}\}\{\backslash sum\_j\; e^\{z\_j\}\}$$ softmax(zi)=∑jezjezi

两步：

第一步：取 e 的幂。 每个 logit 变成 $ezie^\{z\_i\}$ ezi。这一步保证所有数都为正，而且原来大的数变得更大，原来小的数变得更小。

第二步：归一化。 除以总和。所有数加起来等于 1。

举个例子。假设词表只有 4 个 token，logits 是：

好 → 2.0
啊 → 1.0
， → 0.5
冷 → -1.0

取 e 的幂：

e^2.0 ≈ 7.39
e^1.0 ≈ 2.72
e^0.5 ≈ 1.65
e^-1.0 ≈ 0.37

总和 = 7.39 + 2.72 + 1.65 + 0.37 = 12.13

归一化：

好 → 7.39 / 12.13 ≈ 0.61
啊 → 2.72 / 12.13 ≈ 0.22
， → 1.65 / 12.13 ≈ 0.14
冷 → 0.37 / 12.13 ≈ 0.03

现在这 4 个数都是正的，加起来等于 1。这是一个概率分布。模型认为"好"的概率是 61%，"啊"是 22%，以此类推。

2.3 温度：调整分布的陡峭程度

Softmax 里可以加一个温度参数 $TT$ T：
$$softmax(zi/T)=\; e\; zi/T\; \sum je\; zj/T\; \backslash text\{softmax\}(z\_i\; /\; T)\; =\; \backslash frac\{e^\{z\_i\; /\; T\}\}\{\backslash sum\_j\; e^\{z\_j\; /\; T\}\}$$ softmax(zi/T)=∑jezj/Tezi/T

温度做的事：在取 e 的幂之前，先把 logits 除以 T。

• T = 1：不缩放，原始分布。
• T < 1：logits 被放大，分布变得更陡。高分 token 的概率更高，低分 token 的概率更低。T 很小的时候，接近"永远选最高分"。
• T > 1：logits 被缩小，分布变得更平。各 token 的概率差距缩小，随机性增加。T 很大的时候，接近均匀分布。

用刚才的例子，T = 0.5：

logits / 0.5：
好 → 4.0
啊 → 2.0
， → 1.0
冷 → -2.0

取 e 的幂：

e^4.0 ≈ 54.6
e^2.0 ≈ 7.39
e^1.0 ≈ 2.72
e^-2.0 ≈ 0.14

总和 ≈ 64.85

归一化：

好 → 0.84
啊 → 0.11
， → 0.04
冷 → 0.002

"好"的概率从 61% 涨到 84%。分布更尖了。

T = 2.0：

logits / 2.0：
好 → 1.0
啊 → 0.5
， → 0.25
冷 → -0.5

取 e 的幂：

e^1.0 ≈ 2.72
e^0.5 ≈ 1.65
e^0.25 ≈ 1.28
e^-0.5 ≈ 0.61

总和 ≈ 6.26

归一化：

好 → 0.43
啊 → 0.26
， → 0.20
冷 → 0.10

"好"的概率从 61% 降到 43%。分布更平了。

MiniMind 的 generate 函数里，temperature 是一个可调参数。默认值通常是 0.7 到 1.0 之间。

三、采样：从概率里选一个 token

3.1 贪心：永远选最高分

最简单的方式：选概率最高的那个。

next_token = torch.argmax(probs)

优点：确定性强，每次输出一样。 缺点：太死板，容易生成重复、无聊的文本。

3.2 随机采样：按概率抽

把概率分布当作一个转盘，按比例抽。

next_token = torch.multinomial(probs, num_samples=1)

"好"的概率是 0.61，那它被抽中的概率就是 61%。"啊"是 0.22，被抽中的概率就是 22%。

优点：多样，每次输出可能不同。 缺点：可能抽到低概率的奇怪 token，生成质量不稳定。

3.3 Top-k：只在最高的 k 个里抽

先选出概率最高的 k 个 token，把剩下的概率归零，再重新归一化，然后在这 k 个里随机抽。

比如 k = 3，词表里 6400 个 token，只保留概率最高的 3 个，其他 6397 个的概率设成 0。

top_k_probs, top_k_indices = torch.topk(probs, k=3)
top_k_probs = top_k_probs / top_k_probs.sum()   # 重新归一化
next_token = top_k_indices[torch.multinomial(top_k_probs, 1)]

优点：过滤掉了长尾里的垃圾 token，生成质量更稳。 缺点：k 固定，遇到分布很尖的情况（只有 1 个明显高分），k=3 还是保留了 2 个不该保留的；遇到分布很平的情况，k=3 又砍掉了太多。

3.4 Top-p：按累计概率截断

Top-p，也叫 nucleus sampling。不固定 k，而是固定一个累计概率 p。

把 token 按概率从高到低排。从最高开始，一个一个累加概率，直到累计值超过 p。只保留这些 token，后面的全砍掉。

比如 p = 0.9，概率分布是：

好 → 0.45  累计 0.45
啊 → 0.25  累计 0.70
， → 0.15  累计 0.85
冷 → 0.08  累计 0.93  ← 超过 0.9，保留
热 → 0.04  累计 0.97  ← 砍掉
...

保留"好、啊、，，冷"这 4 个，砍掉后面的。然后在这 4 个里随机抽。

sorted_probs, sorted_indices = torch.sort(probs, descending=True)
cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
sorted_indices_to_remove = cumulative_probs > p
sorted_probs[sorted_indices_to_remove] = 0
sorted_probs = sorted_probs / sorted_probs.sum()
next_token = sorted_indices[torch.multinomial(sorted_probs, 1)]

Top-p 的好处：动态调整保留的 token 数量。

• 分布很尖的时候，前 1 个 token 就占了 90% 的概率，只保留 1 个。
• 分布很平的时候，可能要保留几十个才凑够 90%。

比 Top-k 更灵活。

3.5 MiniMind 用哪个

MiniMind 的 generate 函数里，同时支持 temperature、top-k、top-p。实际用的时候，通常组合起来：

# 先按温度缩放
logits = logits / temperature
# 再 Top-k 过滤
# 再 Top-p 过滤
# 最后采样

具体的默认值要看 MiniMind 的配置。一般来说，temperature 默认 0.7 到 1.0，top-k 默认 50 左右，top-p 默认 0.9 左右。

四、生成循环：一个 token 一个 token 往外蹦

现在把整个推理流程串起来。

输入"今天天气"，目标是生成后续的 token。

第一步：前向传播。

"今天天气"经过分词、embedding、8 层积木，得到最后一层的输出。取最后一个位置的 768 维向量。

第二步：算分数。

过 lm_head，得到 6400 个 logits。

第三步：Softmax。

按温度缩放 logits，做 Softmax，得到 6400 个概率。

第四步：采样。

按 top-k、top-p 过滤，从剩余的概率里抽一个 token。假设抽到了"真"。

第五步：拼接。

把"真"拼到输入后面，变成"今天天气真"。

第六步：重复。

回到第一步，用"今天天气真"作为新输入，生成下一个 token。假设抽到了"好"。

第七步：停止。

继续重复，直到抽到结束符（<|im_end|> 或 <|endoftext|>），或者达到最大长度限制。

这就是自回归生成的全过程。每一步只生成一个 token，然后把它拼到输入后面，作为下一步的输入。

4.1 KV Cache 在这里的作用

如果没有 KV Cache，每一步都要重新算前面所有 token 的 K 和 V。生成第 n 个 token 时，计算量是 O(n)。

有了 KV Cache，前面 token 的 K、V 已经缓存了，每一步只需要算新 token 的 K、V。生成第 n 个 token 时，计算量降到 O(1)。

04 篇讲过 KV Cache 的原理。这里能看到它在生成循环里的位置：每次前向传播时，把新 token 的 K、V 拼到缓存后面，前面 token 的 K、V 直接用缓存。

4.2 权重绑定在这里的作用

lm_head 和 embed_tokens 共享矩阵。这意味着：

• 输入"真"的时候，查 embed_tokens 得到 768 维向量。
• 输出的时候，过 lm_head 得到 6400 个分数。

这两个操作，用的是同一块矩阵的转置。输入和输出在同一个向量空间里，语义一致。

训练的时候，这种对称性让梯度可以同时更新输入和输出的表示，收敛更快。

五、回到第一篇的"猜下一个词"

第一篇开头说：模型做的事，是"盖住后面，猜下一个词"。

现在把这句话展开成完整的工程流程：

1. 输入一段话，分词成 token。
2. token 查 embedding，变成 768 维向量。
3. 向量过 8 层 Transformer 块，每层做注意力、FFN、残差、归一化。
4. 最后一层的输出，取最后一个位置，过 lm_head，得到 6400 个 logits。
5. Softmax 变成概率。
6. 采样选一个 token。
7. 把这个 token 拼到输入后面，重复。

第一篇说"猜词"，第二篇说"字怎么变成数字"，第三篇说"token 怎么互相看"，第四篇说"注意力之外的所有零件"，这一篇说"最后怎么变成下一个 token"。

五篇串起来，一个完整的 LLM 推理流程就走完了。

六、动手实验

实验一：看 logits 和概率

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("jingyaogong/minimind-3")
model = AutoModelForCausalLM.from_pretrained("jingyaogong/minimind-3")

text = "今天天气"
inputs = tokenizer(text, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits[:, -1, :]   # 最后一个位置的 6400 个分数

probs = torch.softmax(logits, dim=-1)

# 取 top 10
top_probs, top_indices = torch.topk(probs, k=10)
for prob, idx in zip(top_probs[0], top_indices[0]):
    token = tokenizer.decode([idx])
    print(f"{token}: {prob.item():.4f}")

跑一遍，看看"今天天气"后面，模型认为最可能的是什么。

实验二：对比不同温度

for T in [0.5, 1.0, 2.0]:
    scaled_logits = logits / T
    probs = torch.softmax(scaled_logits, dim=-1)
    top_probs, top_indices = torch.topk(probs, k=5)
    print(f"\n温度 T = {T}")
    for prob, idx in zip(top_probs[0], top_indices[0]):
        token = tokenizer.decode([idx])
        print(f"{token}: {prob.item():.4f}")

观察温度对概率分布的影响。T 越小，分布越尖；T 越大，分布越平。

七、本篇概念清单

概念	本篇交代到什么程度
lm_head	讲透：从 768 到 6400，形状、作用
权重绑定	讲透：为什么能共享、省了什么
logits	讲透：是什么、和概率的区别
Softmax	讲透：公式、两步操作
温度	讲透：怎么缩放、对分布的影响
贪心搜索	讲清：做法、优缺点
随机采样	讲清：做法、优缺点
Top-k	讲清：做法、优缺点
Top-p	讲清：做法、优缺点
生成循环	讲透：七步走完
KV Cache（呼应）	沿用 04 篇，不重复
交叉熵损失	不展开：训练时的事，本篇只讲推理

本篇要牢记的只有三个词：lm_head、Softmax、采样。

八、回到开头的问题

04 篇结尾说，积木的输出还是 768 维向量。这一篇回答了：向量怎么变成下一个 token。

过程是：过 lm_head 得到 6400 个分数，Softmax 变成概率，按温度调整、按 top-k 和 top-p 过滤，采样选一个，拼回输入，循环。

"猜下一个词"这句话，到这里才算完整拆开。

九、思考题

1. 如果把温度设成 0，会发生什么？和贪心搜索有什么区别？
2. Top-k 和 Top-p 同时用，是先 Top-k 再 Top-p，还是反过来？顺序有影响吗？
3. 训练时为什么不采样？提示：想想梯度怎么传。
4. 权重绑定让 lm_head 和 embed_tokens 共享矩阵。如果解开绑定，让它们各自独立，会发生什么？

备注：本篇的源码分析基于 MiniMind 仓库的 model/model_minimind.py 和 generate 相关代码。实验脚本建议在安装了 transformers 和 torch 的环境下运行。