Tiny LLM 手算：从汉字到下一个 token 的完整矩阵计算

0. 先定一个能手算的小模型

真实 LLM 有几万到几十万词表、几千到上万维 hidden state、几十到上百层 Transformer。为了看清底层计算，我们把它压缩成一个玩具模型。它不追求真实能力，只追求每个数都能看见。

词表<BOS>、我、爱、猫、狗、。、<EOS>，共 7 个 token。

模型维度每个 token 变成 2 维向量，所以矩阵都能写在页面上。

任务输入 <BOS> 我爱，预测下一个 token 应该是“猫”。

重要简化：为了让你能手算，我们省略了 LayerNorm、MLP、RoPE、多头拆分、多层堆叠和 KV cache。真实模型把这里的同一套计算重复很多层、很多头、很多维。

1. 词表：汉字先变成 token id

Tokenizer 的作用是把字符串拆成模型认识的整数。这里用最简单的字符级 tokenizer。

token id	token	含义
0	`<BOS>`	句子开始
1	我	汉字 token
2	爱	汉字 token
3	猫	汉字 token
4	狗	汉字 token
5	。	标点
6	`<EOS>`	句子结束

比如训练文本是：

我爱猫。

加入起止符后：

<BOS> 我爱猫。 <EOS> → [0, 1, 2, 3, 5, 6]

next-token 训练会把它右移一位，输入和目标如下：

位置	模型看到的 token	模型要预测的下一个 token
0	`<BOS>`	我
1	我	爱
2	爱	猫
3	猫	。
4	。	`<EOS>`

2. token id 变成 one-hot 矩阵

我们拿推理时的上下文 <BOS> 我爱 来算下一字。token id 是：

[0, 1, 2]

把每个 id 展成长度为 7 的 one-hot 向量：

X_ids = [0,1,2]
OneHot = 1 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 0

每一行表示一个位置，每一列对应词表里的一个 token。one-hot 本身还不是语义向量，它只是“开关”。

3. embedding：查表得到 2 维向量

Embedding 矩阵可以理解成“每个 token 的坐标表”。我们的 embedding 表 E 是 7×2：

token	id	embedding
`<BOS>`	0	[0, 0]
我	1	[1, 0]
爱	2	[0, 1]
猫	3	[1, 1]
狗	4	[1, -1]
。	5	[-1, 0]
`<EOS>`	6	[0, -1]

矩阵乘法就是查表：

TokenEmbedding = OneHot × E
= 0 0 1 0 0 1

再加位置向量。这里位置 0、1、2 的位置编码设为：

P = 0.0 0.0 0.1 0.0 0.0 0.1

最终进入 Transformer 的输入矩阵：

X = TokenEmbedding + P = 0.0 0.0 1.1 0.0 0.0 1.1

4. attention：当前位置“爱”回看前文

我们只算最后一个位置，也就是模型看到 <BOS> 我爱 后，要预测下一个 token。为了手算，令：

W_Q = W_K = W_V = I

所以 Q、K、V 都等于 X：

Q = K = V = 0.0 0.0 1.1 0.0 0.0 1.1

最后位置的 query 是：

q₃ = [0.0, 1.1]

它和所有允许看的 key 做点积。decoder 使用 causal mask，所以位置 2 可以看 0、1、2，不能看未来。

score_i = q₃ · k_i / √2
score = [0/√2, 0/√2, 1.21/√2] = [0.000, 0.000, 0.856]

softmax 后得到 attention 权重：

softmax([0.000, 0.000, 0.856]) = [0.230, 0.230, 0.540]

这表示：预测“爱”后面的字时，模型大约 23% 看 BOS，23% 看“我”，54% 看“爱”。然后对 V 加权求和：

c₃ = 0.230·[0,0] + 0.230·[1.1,0] + 0.540·[0,1.1]
c₃ = [0.253, 0.594]

再加残差连接：

h₃ = x₃ + c₃ = [0,1.1] + [0.253,0.594] = [0.253, 1.694]

5. 输出头：hidden vector 变成 7 个 logits

现在模型已经把上下文压成了一个 2 维向量：

h₃ = [0.253, 1.694]

输出头是一个 2×7 矩阵，每一列对应一个 token。它把 hidden vector 投影成词表里每个 token 的分数。

W_U = -1.0 0.6 -0.5 0.8 1.0 -0.8 -1.0 -1.0 -0.5 0.6 1.2 -0.4 0.3 0.5

计算：

logits = h₃ × W_U
= [-1.947, -0.695, 0.890, 2.235, -0.425, 0.306, 0.594]

id	token	logit	直觉
0	`<BOS>`	-1.947	几乎不可能再开始
1	我	-0.695	不太像
2	爱	0.890	有一定可能重复
3	猫	2.235	最高分
4	狗	-0.425	可能性较低
5	。	0.306	有些可能
6	`<EOS>`	0.594	有些可能

6. softmax：logits 变成概率分布

Logit 不是概率。softmax 把任意实数分数变成总和为 1 的概率。

p_i = exp(logit_i) / Σ_j exp(logit_j)

代入上面的 7 个 logits：

id	token	logit	exp(logit)	probability
0	`<BOS>`	-1.947	0.143	0.0088
1	我	-0.695	0.499	0.0307
2	爱	0.890	2.435	0.1499
3	猫	2.235	9.346	0.5754
4	狗	-0.425	0.654	0.0403
5	。	0.306	1.358	0.0836
6	`<EOS>`	0.594	1.812	0.1115

于是：

P(next token = 猫 | <BOS> 我爱) = 57.54%

7. 训练：模型怎样知道“猫”应该更高

训练时，真实下一个 token 是“猫”，也就是 id=3。交叉熵损失是：

Loss = -log P(猫) = -log(0.5754) = 0.5528

对 logits 的梯度非常简单：

∂L/∂logits = p - one_hot(猫)
= [0.0088, 0.0307, 0.1499, -0.4246, 0.0403, 0.0836, 0.1115]

含义很直观：正确答案“猫”的梯度是负数，梯度下降会把它的 logit 往上推；其他 token 的梯度是正数，会被往下压。

为了手算清楚，这里只更新输出头 W_U。真实训练会把梯度继续反传到 attention、embedding 和所有层。

∂L/∂W_U = h₃ᵀ · (p - y)
= 0.0022 0.0078 0.0379 -0.1074 0.0102 0.0211 0.0282 0.0149 0.0520 0.2539 -0.7193 0.0682 0.1416 0.1888

学习率设为 0.1，更新：

W_U_new = W_U - 0.1 · ∂L/∂W_U

更新后的输出头：

W_U_new = -1.0002 0.5992 -0.5038 0.8107 0.9990 -0.8021 -1.0028 -1.0015 -0.5052 0.5746 1.2719 -0.4068 0.2858 0.4811

注意“猫”这一列从 [0.8, 1.2] 变成了 [0.8107, 1.2719]，更贴近当前 hidden vector 的方向，所以它下一次会更容易被选中。

8. 更新后再算一次：猫的概率升高

用同一个 hidden vector 乘新的输出头：

logits_new = [-1.950, -0.704, 0.846, 2.360, -0.436, 0.281, 0.561]

softmax 后：

token	更新前概率	更新后概率
`<BOS>`	0.0088	0.0082
我	0.0307	0.0286
爱	0.1499	0.1348
猫	0.5754	0.6128
狗	0.0403	0.0374
。	0.0836	0.0767
`<EOS>`	0.1115	0.1014

训练的本质：不是模型“理解了猫”，而是大量样本不断调整矩阵，让正确 token 在对应上下文下的 logit 越来越高，错误 token 的 logit 越来越低。

9. 推理采样：概率分布怎样变成一个汉字

推理时不会再计算 loss，也不会更新权重。它只做 forward，得到概率分布，然后选下一个 token。

最简单是 argmax：选概率最大的 token。

argmax(probability) = id 3 = 猫

如果用随机采样，假设随机数 u = 0.42。按词表顺序累加概率：

token	概率	累计概率	是否命中 u=0.42
`<BOS>`	0.0082	0.0082	否
我	0.0286	0.0368	否
爱	0.1348	0.1716	否
猫	0.6128	0.7844	是

所以采样得到：

sampled_token_id = 3 → tokenizer.decode([3]) = "猫"

10. 生成不是一次，而是循环

得到“猫”以后，模型把它追加到上下文后面：

旧上下文：<BOS> 我爱
新上下文：<BOS> 我爱猫

然后重复同样流程：tokenize → embedding → attention → logits → softmax → sample。下一步模型可能生成“。”，再下一步生成 <EOS>，于是得到完整文本：

我爱猫。

真实服务会用 KV cache 保存过去 token 的 K/V，避免每次从头算所有历史。这不改变数学结果，只是减少重复计算。

11. 把手算映射回真实大模型

本文 tiny 模型	真实大模型	本质是否相同
7 个 token	几十万 token 词表	相同，都是 id 和词表映射
2 维 embedding	几千到上万维 hidden state	相同，都是查表得到向量
1 个 attention head	几十层、多头、GQA/MQA、RoPE	相同，都是 QKᵀ 得分、softmax、加权 V
只更新输出头	反传更新所有可训练参数	梯度链路更长，但交叉熵和梯度方向相同
一次生成“猫”	循环生成长文本、工具调用、代码	相同，都是每次预测一个或一组 token

最底层的直觉：模型不是直接“吐字”。它每一步都先把上下文压成一个向量，再用这个向量和词表中每个 token 的输出方向做匹配，得到分数，softmax 成概率，最后采样出 token id，再由 tokenizer decode 成文字。

来源与说明

本文承接你给出的 AI Infra 学习指南，把其中“tokenizer、预训练 next-token loss、推理 KV cache”等概念向下展开为一个可手算实例。所有矩阵、权重和数字都是为教学构造的 toy example，不代表任何真实模型权重。

AI Infra 学习指南：从预训练到后训练，再到推理上线
核心机制：decoder-only Transformer、causal self-attention、softmax、cross entropy、gradient descent。