Token Embedding + Positional Embedding

模型第一件事是把离散的 token id 翻译成连续向量。字符 ‘e’ 的 id=4，但 4 这个数字本身没有”语义”——4 和 5 之间的距离不比 4 和 26 更近。所以用一张可学习查找表（embedding matrix）查表：

# wte: (vocab_size, d_model) = (27, 16)，每行一个字符的"含义向量"
# wpe: (block_size, d_model) = (8, 16)，每行一个位置的"位置向量"
x = wte[ids] + wpe[positions]   # (T, 16)

wte[4] 取出 ‘e’ 的 16 维向量。这两张表都是模型参数（4000 多个数字里的一部分），由训练学到。

为什么用查表而不是 one-hot

one-hot 把 ‘e’ 表示为 27 维向量 $[0,0,0,1,0,\dots]$。它有两个致命问题：

1. 维度爆炸：词表到几万时（GPT-2 是 50257），one-hot 几万维，全连接层参数爆炸。
2. 任意两个不同 token 距离相同：one-hot 任意两向量点积为 0，模型分不出”e 和 f 比较像” vs “e 和 z 完全不像”。嵌入向量是稠密的，训练后语义相近的 token 在 16 维空间里距离也近——这是 word2vec (Mikolov 2013) 发现的著名性质（如 $\vec{king} - \vec{man} + \vec{woman} \approx \vec{queen}$）。

为什么位置用加法而不是拼接

拼接 concat(wte[4], wpe[0]) 会让维度翻倍（32 维），且模型必须从头学”前一半是字符、后一半是位置”的分工。相加保持 16 维，且让位置和字符信息纠缠在同一组分量里——后续注意力能从这种纠缠里”解码”出两者（线性层 $\mathrm{softmax}(QK^\top/\sqrt d)V$ 对加法可分）。

更深的原因来自原始 Transformer 的正弦位置编码：相加等价于在频域里调制，类似 AM 广播把信号叠加到载波上。学习式位置嵌入（microgpt 用的）继承了这个加法约定。

位置编码的必要性：置换不变性

注意力本身是集合操作——对输入顺序不敏感（精确地说，self-attention 对输入做置换等变：重排输入序列，输出也跟着重排，但每位置的值不变）。如果没有位置信息，“abc” 和 “cba” 在模型看来完全一样（每个字符的查询结果只取决于”是哪些字符”，不取决于顺序）。位置嵌入打破这种对称性，让模型知道”e 在第 0 位”和”e 在第 3 位”是不同的。

易错点

• wpe 的最大长度限制了上下文窗口。microgpt 的 block_size=8，wpe 只有 8 行，序列长度超过 8 就越界——这就是为什么 GPT 有”上下文窗口”限制（GPT-4 是 8K/128K）。位置编码必须为该长度预先准备好。
• wte 和输出层 lm_head 可以共享权重（weight tying）。nanoGPT 可选 tie_word_embeddings：让 lm_head = wte.T，节省参数且效果常更好。microgpt 不共享（独立 lm_head），见 B16。

参考

• Word embeddings (维基百科)
• Attention Is All You Need (Vaswani 2017) — 位置编码
• nanoGPT (Karpathy)

Transformer Block：注意力 + MLP 的对称结构

Transformer 由若干相同 block 堆叠（GPT-4 几十层，microgpt 1 层，结构完全一致）。每个 block 内部是对称的两阶段：

x ─→ RMSNorm ─→ Attention ─→(+)─→ RMSNorm ─→ MLP ─→(+)─→ 输出
└──────────────────────────┘ 残差     └──────────────────┘ 残差

def block(x):
    x = x + attn(rmsnorm(x))     # 子层 1：注意力 + 残差
    x = x + mlp(rmsnorm(x))      # 子层 2：MLP + 残差
    return x

为什么是 attention + MLP 这对组合

这两个子层功能正交，缺一不可：

• Attention：跨 token 交互。位置 $i$ 的向量融合了上下文相关信息（“哪些位置和我相关”）。
• MLP：token 内部特征变换。每个位置独立地把 16 维特征做非线性重组（“这个位置的特征该怎么重新组合”）。

只堆 attention 没法做复杂特征变换（attention 本身是线性的点积 + softmax，表达力有限）；只堆 MLP 没法跨位置通信。两者交替：attention 收集上下文 → MLP 加工成更抽象的特征 → 下一层 attention 再收集……这种”通信-计算”的交替堆叠是 Transformer 取代 RNN/CNN 的核心结构创新。

残差让堆叠可行

每层 x + f(x) 的形式让深度网络可训练（详见 B14）。GPT-3 有 96 层，若没有残差，反向传播的梯度穿过 96 层线性 + softmax 早就消失了。残差提供一条”无障碍通道”，梯度能直接传到底层。

归一化保持数值稳定

每子层前 rmsnorm(x) 把向量幅度归一化，防止深堆叠下数值爆炸/消失。microgpt 用 pre-norm（归一化在残差内部），比原始 Transformer 的 post-norm 更稳定。

为什么所有层结构相同

不是”必须相同”，而是”相同已足够”。层与层之间参数独立（每层自己的 attn、mlp 权重），但结构一致。这种 uniform 设计有两个好处：

1. 超参数少：只需指定层数 $L$、维度 $d$、头数 $h$ 等少量超参。
2. 可扩展：从 microgpt (1 层) 到 GPT-4 (上百层) 用同一份代码，只是改 $L$。

不同层学到的功能不同（浅层学语法、深层学语义），但这是训练自适应的结果，不需要人为设计差异。

易错点

• microgpt 只有 1 层，所以”堆叠”概念在这个项目里看不到，但代码结构和大模型一致——把 block 函数调用 $L$ 次即可。
• pre-norm 与 post-norm 不可混用。原始 Transformer 论文用 post-norm，现代大模型（GPT-2、LLaMA）用 pre-norm。读不同代码时注意区分，否则复现行为不一致。

参考

• Attention Is All You Need (Vaswani et al., 2017)
• On Layer Normalization in the Transformer Architecture (Xiong et al., 2020) — pre-norm vs post-norm

查询/键/值：信息检索的数据库隐喻

注意力本质是一个软查询（soft lookup）：每个位置都问”我该从前面哪些位置取多少信息”。三个矩阵把同一个输入向量投到三个不同子空间：

Q = x @ wq   # (T, 16) @ (16, 16) → (T, 16)  查询：我要找什么
K = x @ wk   # 键：我能提供什么
V = x @ wv   # 值：你看上我，我给你什么

三个矩阵 wq/wk/wv 是模型参数，训练学到。

数据库视角（类比）

数据库	注意力
Query SELECT ... WHERE key = ?	$Q$ 和每个 $K$ 算相似度
返回完全匹配的行的 value	返回所有行 value 的加权和
离散、硬匹配	连续、软匹配（按相似度加权）

软匹配让模型可以”半信半疑”地从多个位置各取一部分——比如预测 “joh” 后的字符时，‘n’ 的位置贡献大，但 ‘o’ 的位置也可能贡献一点（提供”这是元音-辅音交替”的线索）。

为什么 Q、K、V 用不同矩阵

直接用 $Q=K=V=x$ 也能跑（叫 “linear attention” 的退化形式），但表达力大减。用三个独立投影矩阵 wq/wk/wv 让模型学到三种不同的”视角”：

• wq 把 $x$ 投成”我要找的查询条件”。
• wk 把 $x$ 投成”我对外广告的标签”。
• wv 把 $x$ 投成”如果有人选我，我真正交付的内容”。

同一个 ‘n’ 字符，作为查询时它的”问题”（“还有别的辅音吗？“）和作为键时它的”标签”（“我是辅音 n”）是不同的投影——分开训练才能学到。

Q/K/V 从同一个 x 出发：self-attention

注意 Q、K、V 都从同一个序列 $x$ 算出来，这是 self-attention（自注意力）。区别于 cross-attention（Q 来自 decoder、K/V 来自 encoder，用于翻译模型）。self-attention 让序列内部任意两位置直接交互，是 Transformer 取代 RNN 的关键——RNN 要逐步传递信息，self-attention 一步到位（任意两位置 $O(1)$ 跳跃）。

维度推导

设序列长 $T$、模型维度 $d_{\text{model}}=16$：

$X \in \mathbb{R}^{T \times 16}, \quad W_Q, W_K, W_V \in \mathbb{R}^{16 \times 16}$

$Q = X W_Q,\quad K = X W_K,\quad V = X W_V \;\in \mathbb{R}^{T \times 16}$

每个位置 $i$ 的查询向量 $q_i \in \mathbb{R}^{16}$、每个位置 $j$ 的键 $k_j \in \mathbb{R}^{16}$、值 $v_j \in \mathbb{R}^{16}$。

易错点

• Q/K/V 不是”三种 token”，而是同一 token 的三种投影。每个位置同时拥有自己的 Q、K、V——既是查询者又是被查询者。
• wq/wk/wv 在多头里会按头拆开（B13）。microgpt 是 4 头 × 4 维 = 16 维，每个 wq 实际是 16×16 但被 reshape 成 (4 头, 4 维) 分头用。

参考

• Attention Is All You Need (Vaswani et al., 2017)
• The Illustrated Transformer (Jay Alammar)

Scaled Dot-Product Attention

完整的注意力运算：

$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V$

分三步：① 算 logits $QK^\top$；② 缩放 $/\sqrt{d_k}$；③ softmax 后对 $V$ 加权求和。

第一步：QKᵀ 算相似度矩阵

$S = Q K^\top \in \mathbb{R}^{T \times T}, \quad S_{ij} = q_i \cdot k_j$

$S_{ij}$ 是位置 $i$ 的查询 $q_i$ 与位置 $j$ 的键 $k_j$ 的点积。点积越大，两者越”匹配”（在投影后的子空间里方向越一致）。这个 $T\times T$ 矩阵就是 attention logits。

scores = Q @ K.transpose(-2,-1)   # (T, d) @ (d, T) → (T, T)

第二步：为什么除以 √d_k

点积 $q_i \cdot k_j = \sum_{l=1}^{d_k} q_{il} k_{jl}$。若 $q, k$ 的分量是均值 0、方差 1 的独立随机变量，点积的方差正比于 $d_k$（$\mathrm{Var}(\sum) = d_k \cdot \mathrm{Var}(q \cdot k) \approx d_k$）。

当 $d_k$ 大（microgpt 是 4，GPT-2 是 64），点积数值会很大。softmax 对大数值敏感：最大的那个 logit 经过 $\exp$ 会指数膨胀，把概率几乎全压在一个位置，梯度趋近 0（饱和区）——训练停滞。除以 $\sqrt{d_k}$ 把方差拉回 1 附近，softmax 工作在线性区：

$\mathrm{Var}\!\left(\frac{q \cdot k}{\sqrt{d_k}}\right) \approx \frac{d_k}{d_k} = 1$

第三步：softmax 归一化

$P_{ij} = \frac{\exp(S_{ij}/\sqrt{d_k})}{\sum_{j'} \exp(S_{ij'}/\sqrt{d_k})}$

每行 $P_{i\cdot}$ 是一个概率分布（和为 1），表示位置 $i$ 对各位置的”关注权重”。

attn = (scores / sqrt(d_k)).softmax(dim=-1)   # (T, T)，每行和为 1
out  = attn @ V                                 # (T, T) @ (T, d) → (T, d)

因果掩码：只能看前面

GPT 是自回归模型（autoregressive），预测第 $i$ 个字符时不能偷看第 $i, i+1, \dots$ 的真实值。所以在算 attention 前，把 $S_{ij}$ 中 $j > i$ 的位置（未来）设为 $-\infty$：

tril = torch.tril(torch.ones(T, T))              # 下三角
scores = scores.masked_fill(tril==0, float('-inf'))

softmax 后这些位置的权重变成 $\exp(-\infty)/Z = 0$，位置 $i$ 只能看 $0..i$。这就是 causal self-attention。

输出 = 加权融合的 V

$\mathrm{out}_i = \sum_{j=0}^{i} P_{ij}\, v_j$

位置 $i$ 的输出是所有”前文” value 的加权和，权重由 Q-K 相似度决定。这就是”当前位置融合了相关位置信息”的数学实现。

易错点

• softmax(dim=-1) 沿最后一个维度（即 $K$ 维度）归一化，不是沿 $Q$。意思是”每个 query 对所有 key 的权重和为 1”。若搞错 dim，模型完全错乱。
• 除以 $\sqrt{d_k}$ 不是 $\sqrt{d_{\text{model}}}$。多头里 $d_k$ 是每头的维度（microgpt 是 4），不是总模型维度（16）。nanoGPT 实现里用 K.shape[-1]**0.5 自动取对。

参考

• Attention Is All You Need — 3.2.1 Scaled Dot-Product Attention
• Softmax function (维基百科)
• The Annotated Transformer (Harvard NLP)

Multi-Head Attention：并行子空间

单头注意力只能学到一种”匹配模式”。多头让模型在同一层里同时跑多个独立的注意力，每个在更小的子空间里学不同的关系：

# microgpt: d_model=16, n_heads=4, head_dim=4
Q = x @ wq   # (T, 16)
Q = Q.view(T, 4, 4).transpose(0,1)   # (4 头, T, 4 维)
# 每头独立做 scaled dot-product attention
heads = [attention(Q[h], K[h], V[h]) for h in range(4)]   # 每头输出 (T, 4)
out = torch.cat(heads, dim=-1) @ wo   # 拼回 (T, 16) 再投影

维度分配

16 维拆成 4 头 × 4 维：

$\underbrace{d_{\text{model}}}_{16} = \underbrace{n_{\text{heads}}}_{4} \times \underbrace{d_k}_{4}$

总计算量与单头 16 维几乎相同（都是 $T^2 \cdot d_{\text{model}}$ 量级），但拆成 4 个独立的”视角”，每个 4 维子空间里学自己的 Q-K 匹配模式。

为什么多头有效

单头在一个 16 维空间里只能学一种相似度度量（一个 $W_Q W_K^\top$ 的”问题-标签”配对）。多头让模型有 4 套独立的 $W_Q^{(h)} W_K^{(h)\top}$，每套学不同的关系：

• 头 1 可能学”看前一个字符”（bigram 模式）
• 头 2 可能学”看句子开头”（位置模式）
• 头 3 可能学”看相同字符”（重复检测）
• 头 4 可能学”看元音”（语音/语法模式）

这些模式在不同子空间里互不打架。最后 concat 拼起来再用 wo 线性混合，让模型决定怎么综合多视角信息。

维度推导

每头参数量 $W_Q^{(h)} \in \mathbb{R}^{16 \times 4}$，4 头共 $4 \times 16 \times 4 = 256$ 参数——和单头 $W_Q \in \mathbb{R}^{16\times 16} = 256$ 完全一样。多头不增加参数量，只是把它们重新组织：单头是一个 16→16 矩阵，多头是 4 个 16→4 矩阵（拼接后等价于一个 block-diagonal 结构）。

数学上，多头注意力等价于把 $W_Q, W_K, W_V$ 限制为分块对角形式（每个头的 query/key/value 在自己的 4 维子空间里不与其他头交互），再用 $W_O$ 做一次跨头混合。

wo 的作用：跨头混合

单纯 concat 只是堆叠 4 头的输出，没有信息交流。wo（16×16 矩阵）做线性变换，让头 1 的信息可以影响头 2 对应的维度——这是”多头协同”的关键。如果省掉 wo（直接用 concat 结果），各头的预测能力被局限在自己的子空间，效果显著下降。

易错点

• d_model % n_heads 必须整除。microgpt 是 16/4=4 整除；如果 d_model=17, n_heads=4 就无法均分，工程上会报错。GPT-2 (768/12=64)、GPT-3 (12288/96=128) 都满足。
• head_dim 不是 d_model。$d_k = d_{\text{model}}/n_{\text{heads}} = 4$，B12 里除以 $\sqrt{d_k}=\sqrt{4}=2$ 而不是 $\sqrt{16}=4$。读代码时 scores / math.sqrt(K.size(-1)) 自动取每头维度。

参考

• Attention Is All You Need — 3.2.2 Multi-Head Attention
• Multi-head attention (Illustrated Transformer)

残差与归一化：让深层可训练

每个 Transformer block 的结构（pre-norm 变体，microgpt 用这个）：

x = x + attn(rmsnorm(x))      # 残差 + 注意力（输入先归一化）
x = x + mlp(rmsnorm(x))       # 残差 + MLP

两个细节：残差连接 x + f(x) 和 RMSNorm 归一化。

残差连接（Residual Connection）

形式 $y = x + f(x)$，“短路”绕过 $f$。两层意义：

1. 前向：信号可以”绕过”$f$ 直接传到后面，即使 $f$ 学到全 0（网络层退化），整体仍是恒等映射 $y=x$。深层网络里这是关键——某层学不到东西不会破坏整体。
2. 反向：梯度 $\partial y/\partial x = 1 + \partial f/\partial x$。即使 $\partial f/\partial x$ 很小（梯度消失），那个 +1 保证梯度能一路传回，深层不退化。

残差由 He et al. (2015) 在 ResNet 提出，解决”网络越深反而越差”的悖论。Transformer 几乎所有层都用它。

RMSNorm：LayerNorm 的简化

LayerNorm 对每个 token 的向量做去均值 + 除标准差 + 仿射：

$\text{LN}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta, \quad \mu = \frac{1}{d}\sum x_i, \quad \sigma^2 = \frac{1}{d}\sum (x_i - \mu)^2$

RMSNorm (Zhang & Sennrich 2019) 去掉减均值，只除 RMS（root mean square）：

$\text{RMSNorm}(x) = \gamma \cdot \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \epsilon}}$

只保留幅度归一化，去掉”中心化”。作者论证：LayerNorm 起作用的主要是”除以方差”（缩放不变性），减均值贡献小但占计算量。RMSNorm 在大模型（LLaMA、microgpt）广泛替代 LayerNorm，速度快 10-50% 且效果相当。

γ 是可学习的逐分量缩放（microgpt 是 16 维权重），初始化为 1。$\epsilon=10^{-5}$ 防止除零。

为什么在 f 之前归一化（pre-norm）

# Pre-norm (microgpt / GPT-2 / LLaMA)
x = x + attn(rmsnorm(x))

# Post-norm (原始 Transformer 论文)
x = rmsnorm(x + attn(x))

Post-norm 把归一化放在残差外，梯度流要穿过 norm 层，深层容易不稳定（原始论文要 warmup 学习率才能训）。Pre-norm 让残差主路”裸奔”，归一化只作用于 $f$ 的输入，梯度可以无障碍地沿残差通道传到底——这是为什么现代大模型都改用 pre-norm，能直接训上百层。

归一化的方向：沿特征维而不是 token 维

rmsnorm 沿最后一个维度（特征维 d_model=16）做，不是沿序列维 T。意思是每个 token 独立归一化自己的 16 个分量，token 之间互不影响。这与 BatchNorm（沿 batch 维归一化）正交——LayerNorm/RMSNorm 不依赖 batch，所以推理（batch=1）和训练行为一致，适合自回归生成。

易错点

• RMSNorm 的 γ 是逐分量的，不是标量。16 维向量有 16 个独立的 $\gamma$，让模型对每个特征维度做不同缩放。误写成标量会失去灵活性。
• 残差连接要求维度匹配。x + f(x) 要求 $f$ 输入输出同维（都是 16）。若中间层改维度（如 16→64→16，见 B15 MLP），残差必须加在 16 维那一层不能跨 64 维。
• epsilon 的位置在根号内：$\sqrt{\text{RMS}^2 + \epsilon}$ 不是 $\text{RMS} + \epsilon$。前者在 RMS 接近 0 时主导项是 $\sqrt\epsilon$，仍能稳定除；后者主导是 $\epsilon$ 本身，量级太小。

参考

• Deep Residual Learning (He et al., 2015)
• Root Mean Square Layer Normalization (Zhang & Sennrich, 2019)
• Layer Normalization (Ba et al., 2016)

Position-wise Feed-Forward Network

每个 Transformer block 在注意力之后还有一个 MLP，对每个位置独立做：

# d_model=16, d_ff=64（4× 升维）
h = relu(x @ w1)   # (T, 16) @ (16, 64) → (T, 64)
out = h @ w2       # (T, 64) @ (64, 16) → (T, 16)

两层全连接，中间夹 ReLU，升维 4 倍再压回。

为什么升维到 4×

16→64→16 看似”先扩张后收缩”的浪费，但这是有意的：

• 升维提供更大容量：64 维空间里模型可以学更复杂的非线性变换（键值记忆、特征组合）。
• ReLU 在高维空间更有效：低维空间 ReLU 的”半边激活”很快把一半神经元关掉，剩余表达力不足。高维下大量神经元中总有一些对当前输入有响应，保持信息流。

GPT-2/3 也用 4× 比例（d_model=768 → d_ff=3072），是经验上的甜点。LLaMA 等用 SwiGLU 时会调整成约 8/3×（因为有门控分支，参数重平衡）。

ReLU：非线性从哪来

$\text{ReLU}(x) = \max(0, x)$

没有非线性激活，两层全连接 $\mathrm{out} = (x W_1) W_2 = x(W_1 W_2) = x W$ 等价于一个单层线性变换——无论堆多少层，整体仍是线性的，学不到复杂函数。ReLU 的”折线”形状提供分段线性近似能力：足够多的 ReLU 神经元可以逼近任意连续函数（万能逼近定理）。

GPT-2 用 GELU，LLaMA 用 SwiGLU，microgpt 用最简单的 ReLU——核心作用都是引入非线性。

为什么叫 “position-wise”

MLP 对每个 token 独立做，token 之间不交互（每个位置的 16→64→16 完全独立于其他位置）。这和注意力正交：

• 注意力：跨 token 交互（“哪些位置相关”）。
• MLP：token 内部特征变换（“该位置的 16 维特征如何重组”）。

两者交替堆叠：注意力收集上下文信息 → MLP 加工成更抽象的特征 → 下一层注意力再收集……这就是 Transformer block 的设计逻辑。

MLP 的”键值记忆”假说

Geva et al. (2020) 提出：MLP 的第一层 $W_1$ 的每行可看作一个”键”（识别某种输入模式），第二层 $W_2$ 的对应列是该键激活时输出的”值”。MLP 整体是 $\sum_i \mathbb{1}[\text{key}_i \text{ matches}] \cdot \text{value}_i$ 的软版本——所以 MLP 在大模型里常被解读为事实记忆存储（如”巴黎是法国首都”存在某些 MLP 神经元里）。这解释了为什么 MLP 占模型参数的 2/3（microgpt 也如此）。

易错点

• MLP 不共享 token 间参数但共享位置间参数。同一层的 MLP 对所有位置用同一个 $W_1, W_2$（position-invariant），但每个 token 独立计算。这与卷积的”权重共享”思想一致。
• ReLU 不是 sigmoid。$\max(0,x)$ 在 $x>0$ 时梯度为 1（不会饱和），但 $x<0$ 时梯度为 0（“死 ReLU”问题）。Leaky ReLU、GELU 等变种缓解这点，但 microgpt 这种小模型用 ReLU 足够。

参考

• Attention Is All You Need — 3.3 Position-wise Feed-Forward Networks
• Transformer Feed-Forward Layers Are Key-Value Memories (Geva et al., 2020)
• Universal approximation theorem

输出层：16 维向量 → 27 个 logits

走完 Transformer block，每个位置得到一个 16 维向量 $h \in \mathbb{R}^{16}$。最后一步用 lm_head（language modeling head）把它映射到词表大小：

logits = h @ lm_head   # (T, 16) @ (16, 27) → (T, 27)

每个位置输出 27 个数（对应 26 个字母 + BOS），叫 logits（未归一化的对数概率）。最大值的位置就是模型预测的下一个字符。

logits 是”原始分数”，不是概率

logits 可以是任意实数（负无穷到正无穷），不归一化。要变成概率必须再过 softmax：

$p_i = \frac{\exp(\text{logit}_i)}{\sum_j \exp(\text{logit}_j)}$

但 microgpt 推理时不需要算 softmax 就能取 argmax——因为 softmax 是单调函数，$\arg\max_i \text{logit}_i = \arg\max_i p_i$。所以推理只做 argmax(logits)，省一次 softmax 计算。

lm_head 与 wte 的关系

lm_head 的形状是 (16, 27)，把 16 维”语义向量”投到 27 维”字符空间”。它和 embedding wte (27, 16) 形状互为转置。两者可以共享权重（weight tying）：

lm_head = wte.T   # (27,16).T = (16,27)

直觉：嵌入学的是”字符 → 语义向量”，输出层是逆向的”语义向量 → 字符”，用同一个映射合理。nanoGPT 默认不共享（独立 lm_head），microgpt 也是。共享的优点是参数少（27×16=432 个参数省下来）、有时泛化更好；缺点是限制了模型对输入和输出用不同表征的能力。

logits 的几何含义

$h \cdot \text{lm\_head}[:, i]$ 是 $h$ 与第 $i$ 个字符的”输出向量”的点积——衡量 $h$ 这个语义向量”有多像字符 $i$“。所以输出层本质是用点积度量当前隐状态与每个候选字符的相似度，点积最大的就是预测。

这与 B11 QKV 注意力里 $q_i \cdot k_j$ 的几何完全同构：都是”用点积在向量空间里找最匹配的”。

温度采样（可选）

推理时除 argmax 外还可以从 softmax 分布里采样：

$p_i = \frac{\exp(\text{logit}_i / T)}{\sum_j \exp(\text{logit}_j / T)}$

温度 $T$ 控制分布锐度：$T\to 0$ 退化为 argmax（贪心），$T=1$ 是原始分布，$T>1$ 让分布更平均（更有创造性、更随机）。microgpt 训练用 argmax 评估，但生成时可以加温度增加多样性。

易错点

• logits 不是 log 概率，是 softmax 的输入。$\text{logit}_i = \log p_i + C$（$C$ 是与 $i$ 无关的常数，由归一化吸收），所以 $\text{logit}_i - \text{logit}_j = \log(p_i/p_j)$ 是对数几率（log odds）。
• lm_head 的方向：是 $(d_{\text{model}}, V)$ 不是 $(V, d_{\text{model}})$。h @ lm_head 把 16 维投到 27 维；写反成 lm_head @ h 形状不对。

参考

• Softmax function (logits → probabilities)
• Weight tying in neural language models (Press & Wolf, 2016)

Cross-Entropy Loss

模型输出 logits，真实答案是某个字符（one-hot 标签 $y$）。需要一个数字衡量”模型预测与真相差多远”——交叉熵损失：

$\mathcal{L} = -\sum_{i=1}^{V} y_i \log p_i = -\log p_{\text{true}}$

最后等号因为 $y$ 是 one-hot（只有真实字符位置为 1，其余 0）。所以交叉熵损失就是真实字符被赋予的概率的负对数。

从信息论理解

交叉熵 $H(y, p) = -\sum y_i \log p_i$ 度量”用分布 $p$ 编码来自分布 $y$ 的事件，平均需要多少比特”。模型越接近真实分布（$p \to y$），交叉熵越小；模型把概率给错地方（$p_{\text{true}} \to 0$），$-\log p_{\text{true}} \to \infty$，损失爆炸。

模型预测真实字符的概率 p_true → loss = -log(p_true)
p_true = 0.9   → loss ≈ 0.11   (自信且对)
p_true = 0.5   → loss ≈ 0.69   (不确定)
p_true = 0.01  → loss ≈ 4.6    (自信但错)
p_true = 0.001 → loss ≈ 6.9    (大错)

负对数让”低概率给真值”受重罚，符合直觉。

数值稳定的实现

直接 log(softmax(logits)) 在 logits 大时数值溢出（$\exp$ 爆炸）。实际用 log-sum-exp trick：

$\log p_i = \text{logit}_i - \log\!\left(\sum_j \exp(\text{logit}_j)\right) = \text{logit}_i - m - \log\!\left(\sum_j \exp(\text{logit}_j - m)\right)$

其中 $m = \max_j \text{logit}_j$。减去最大值让所有 $\exp$ 参数 ≤ 0，结果 ∈ (0, 1]，避免溢出。PyTorch 的 F.cross_entropy 内部就是这个：

loss = F.cross_entropy(logits.view(-1, V), targets.view(-1))
# 等价于：先 log_softmax（带 log-sum-exp trick），再 NLL

序列 loss 是位置平均

microgpt 训练时一次喂 8 个字符的序列，每个位置都预测下一个字符，共 8 个 loss：

$\mathcal{L}_{\text{batch}} = \frac{1}{T}\sum_{t=1}^{T} -\log p_{\text{true}}^{(t)}$

取平均让 loss 不依赖序列长度，便于跨 batch 比较。这个标量就是”模型本次考试的分数”，反向传播会从它出发算每个参数的梯度。

为什么不用 MSE（均方误差）

分类问题用交叉熵而不是 MSE（mean squared error），有两个理由：

1. 梯度匹配：softmax + 交叉熵的梯度是 $p - y$，简洁且与误差线性相关。softmax + MSE 的梯度带 $\frac{\partial^2}{\partial \text{logit}^2}$ 项，训练慢。
2. 概率解释：交叉熵直接对应极大似然估计（MLE）——最小化 $-\log p_{\text{true}}$ 等价于最大化真实数据的对数似然。这是统计学最原则性的目标。

易错点

• targets 是 token id 不是 one-hot。PyTorch 的 cross_entropy 内部把整数标签当作 one-hot 处理，比显式 one-hot + 计算高效。误传 one-hot 进去维度错乱。
• 忽略 padding token 的 loss。变长序列里有 padding（如 <PAD>），这些位置不该贡献 loss——实际实现里用 ignore_index 参数屏蔽。microgpt 固定长度 8，没有 padding 问题。

参考

• Cross entropy (维基百科)
• Log-sum-exp trick (numerical stability)
• Practical Guide to Cross Entropy (PyTorch docs)