Multi-head + decoder block 结构

阶段: B2 · attention/decoder/KV + judge 校准（Day 11-20） 标签: #multi-head #decoder-block #pre-LN #residual #FFN

今日导引（由浅入深）

Day 11 把单头 causal attention 跑通并用单测证明了「未来不影响过去」。但单头只能学一种「关系」。今天 Day 12 沿能力曲线上爬一格：把 d 维投影切成 h 个头各自做 attention 再拼回（multi-head），并把它包进一个完整 decoder block——causal multi-head attention + 残差 + LayerNorm + FFN + 残差。这就是 GPT/Llama/Qwen/DeepSeek 一层 Transformer 的真身，堆 N 层就是整个模型骨架。它紧接昨天，因为多头复用昨天那块 attention() 核；它通向明天 Day 13 的「数值对齐」（验证我们这块手写 block 算得对）和 Day 14 的 KV cache（在这块 block 上省重算）。今天的「最小可判定产出」：block 的 forward 输出 shape 必须对齐输入 [seq, d]，且整块仍保持 causal——扰动最后一个 token 只改它自己的输出，不改前面的。

1. 机理精读

Multi-head 为什么要切头：单头 attention 用一组 Q/K/V 投影，只能在一个子空间里捕捉一种相关模式。Multi-head 把 d 维线性投影后切成 h 个 d/h 维的子空间，每个头独立做一遍 scaled dot-product attention，再 concat 回 d 维、过一个输出投影 Wo。直觉是：不同头可以同时学「指代消解」「语法依赖」「位置邻接」等不同关系，并行而不互相挤占。代价是每个头维度变小（d/h），所以 √d 缩放在头内用的是 d/h 而非全维 d。

要点：multi-head 不增加总计算量（h 个 d/h 维头 ≈ 一个 d 维头的算力），但增加了「表达多样性」——h 组独立的低维注意力比一组高维注意力能并行表征更多种关系。这是「分而治之」的容量再分配，不是堆算力。Wo（输出投影）的作用是把 concat 后各头的信息重新混合、对齐回残差流的 d 维空间。

decoder block 的五件套：一层 decoder-only block = (1) causal multi-head attention →(2) 残差连接 →(3) LayerNorm →(4) FFN（两层线性夹一个激活，中间维 hidden 通常 4d）→(5) 残差。残差让梯度有「高速公路」直达深层，训练 100+ 层才不退化；LayerNorm 把每个 token 向量归一化到零均值单位方差，稳定数值尺度；FFN 提供逐位置的非线性变换容量（attention 负责「混 token」，FFN 负责「在每个 token 内做计算」）。

逐位置（position-wise）是 FFN 的关键词：FFN 对每个 token 独立地做相同的两层变换，不跨 token 混信息（混信息是 attention 的活）。所以一层 block 的分工很清晰——attention 横向混 token、FFN 纵向算特征。FFN 的中间维 hidden（常 4d）提供了远超 d 的非线性容量，是模型「存知识」的主战场（近年可解释性研究发现大量事实记忆落在 FFN 层）。

pre-LN vs post-LN：原始 Transformer（2017）是 post-LN——LayerNorm 放在残差之后（LN(x + Sublayer(x))）。后来发现深层 post-LN 训练不稳、要 warmup。pre-LN（x + Sublayer(LN(x))）把 LayerNorm 挪到子层输入处，残差路径全程不被归一化、梯度更干净，已成 2026 主流（Llama 4 / Qwen3 / DeepSeek-V3 均用 pre-LN，且常用 RMSNorm 这一变体）。今天教学实现走 pre-LN。

残差为什么是「梯度高速公路」：y = x + Sublayer(LN(x)) 对 x 求导含一个恒等项 ∂y/∂x ⊇ I，反向传播时梯度有一条不衰减的直通路径，即使 Sublayer 的雅可比很小，深层梯度也不会指数消失。这是 ResNet（2015）的核心思想被 Transformer 直接继承的部分，也是能堆几十上百层的前提。

残差流（residual stream）视角：可解释性研究把贯穿所有层的那条 x 叫「残差流」——每个子层（attention/FFN）从残差流读信息、算出增量、再写回残差流（加法）。于是整个网络像一条共享的「信息总线」，attention 负责「跨位置搬运」、FFN 负责「就地加工」，层层累加。这个视角在 B15-B18 的可解释性/agent 调试里会反复用到——今天先建立「block = 读残差流 → 子层计算 → 写回残差流」的心智模型。

decoder-only 的边界：GPT 系是 decoder-only——没有 encoder、没有 cross-attention，只有一种 causal self-attention 堆叠。这与原始论文的 encoder-decoder（翻译用，含 cross-attention）不同。现代生成式 LLM 几乎全是 decoder-only，因为「下一 token 预测」这一个目标足够通用。参考 Karpathy nanoGPT README (2023-2024)。

2. 推导 / 手算 / 代码走读

本仓的 multi-head + decoder block 已实现在 src/agent/transformer/transformer.ts（seed 规划的独立文件 decoderBlock.ts 最终合并进这份统一实现，函数同名 decoderBlock，机制完整存在；独立文件未单建）。真实符号走读：

1. splitHeads(x, heads)（第 90-95 行）：把 [seq][d] 按列切成 h 个 [seq][d/h]；开头 if (d % heads !== 0) throw 防止维度不整除导致静默错位——这是个好工程习惯。
2. concatHeads(parts)（第 97-100 行）：用 parts.flatMap(p => p[i]) 把每个头第 i 行横向拼回 [seq][d]。
3. multiHeadCausalAttention(x, p)（第 102-108 行）：splitHeads(matmul(x, p.Wq), heads) 等先做 Q/K/V 投影再切头；逐头调 attention(qh[h], kh[h], vh[h], true)（复用 Day 11 那块、causal=true）；concatHeads 后乘输出投影 p.Wo。
4. decoderBlock(x, p)（第 111-116 行，pre-LN）：
   • const a = multiHeadCausalAttention(layerNorm(x, p.g1, p.b1), p)——先 LN 再 attention（pre-LN 标志）。
   • const x1 = x.map((row, i) => row.map((v, j) => v + a[i][j]))——attention 残差。
   • const f = linear(relu(linear(layerNorm(x1, ...), p.W1, p.bff1)), p.W2, p.bff2)——LN → FFN（linear→relu→linear，中间维 hidden）。
   • return x1.map(...v + f...)——FFN 残差，输出 [seq][d]。
5. randomParams(d, heads, hidden, rng)（第 119-141 行）：注入式 RNG（mulberry32）生成确定性小随机权重，供测试/demo——保证可复现。权重幅度用 (rng()-0.5)*0.2（小初始化避免一开始就饱和），LayerNorm 的 g1/g2 初始为 1、b1/b2 为 0（恒等初始化），FFN bias 全 0——都是合理的初始化惯例。
6. 内部 linear(x, W, b)（第 55-57 行）= matmul(x, W) 后逐行加 bias；relu(x)（第 59-61 行）= max(0, v) 逐元素——FFN 的两段线性夹 ReLU 就由这两个私有函数拼出。

这套实现共约 142 行、零依赖、纯 number[][]，刻意不引张量库——目的是「机制即代码、代码即可测」，而非性能。它正是 seed「implement attention, don't just note it」（文件头注释原话）这条纪律的落地。

shape 推演（与 seed 的 [4,8] 例对齐），逐阶段追踪一个 [seq=4, d=8], heads=2, hidden=16 的前向：

全程输出 shape = 输入 [4,8]，中间 attention 为 [h, seq, seq]=[2,4,4]——这正是 seed 要断言的产出。shape 守恒（输入=输出）是堆 N 层的前提：第 k 层输出能直接喂第 k+1 层。

顺手算这一层的参数量（d=8, hidden=16，看 attention 与 FFN 的占比）：

• attention：Wq/Wk/Wv/Wo 各 d×d=64，合计 4·64 = 256。
• FFN：W1 为 d×hidden=128、W2 为 hidden×d=128，合计 256（bias 略）。
• LayerNorm：g1/b1/g2/b2 各 d=8，合计 32，量级很小。

可见 attention 与 FFN 参数量级相当；真实大模型里 hidden≈4d，FFN 参数 ≈ 2·4d² = 8d² 反而比 attention 的 4d² 更多——这也是 FFN 被视为「知识仓库」的一个旁证。

把 decoderBlock 的数据流写成伪代码，便于对照真实实现：

function decoderBlock(x, p):
    # ---- 子层 1：causal multi-head attention（pre-LN + 残差）----
    h  = layerNorm(x, p.g1, p.b1)          # pre-LN
    Q,K,V = x·Wq, x·Wk, x·Wv  (基于 h)
    qh,kh,vh = splitHeads(Q/K/V, heads)    # 切 h 个 d/h 维头
    parts = [attention(qh[i], kh[i], vh[i], causal=True) for i in heads]
    a  = concatHeads(parts)·Wo             # 拼回 d 并输出投影
    x1 = x + a                             # 残差（注意加的是原 x，非 h）

    # ---- 子层 2：FFN（pre-LN + 残差）----
    h2 = layerNorm(x1, p.g2, p.b2)
    f  = (relu(h2·W1 + b1))·W2 + b2        # 升维 hidden → ReLU → 降回 d
    return x1 + f                          # 残差

注意两处残差加的都是「子层输入」（x 和 x1），而非「LN 后的 h」——这是 pre-LN 残差的正确写法，本仓 decoderBlock（第 111-116 行）严格如此。写成 x = h + a（加 LN 后的）是常见 bug，会破坏残差流的恒等直通。

3. 今日实战

seed 计划「基于 Day 11，新建 decoderBlock.ts，扩 multi-head（split/concat），加 pre-LN 残差 + FFN，整块 <300 行，对固定输入跑一次 forward 打印各阶段 shape」。实际落地于 transformer.ts（约 142 行，远小于 300 行预算）。可复现实战：

1. 在 src/agent/__tests__/transformer.test.ts 的 describe('decoderBlock') 里，用 d=8, heads=2, seq=4, hidden=16，p = randomParams(d, heads, 16, mulberry32(3))，X = randMatrix(seq, d, mulberry32(99))。
2. 跑 decoderBlock(X, p)，断言 y.length === seq 且 y[0].length === d——shape 对齐 [4,8]。
3. causal 验证：把 X 最后一行 +1.5 得 X2，比对 y2 前 seq-1 行与 y1 逐元素相等（toBeCloseTo(...,10)），且最后一行至少一个元素差 >1e-9——证明整块仍 causal。
4. 跑 pnpm vitest src/agent/__tests__/transformer.test.ts 看 decoderBlock 全组绿。

4. 今日实测 / 产出

• 状态（按 seed）：seed 标注「待建/待跑」。规划产出为 decoderBlock.ts，forward 输出 shape 对齐输入 [4,8]，中间 attention 为 [h, seq, seq]，shape 断言入单测。
• 实际进度（诚实标注）：机制已落在统一文件 src/agent/transformer/transformer.ts 的 decoderBlock（multi-head/split-concat/pre-LN/残差/FFN 齐全），且 transformer.test.ts 的 describe('decoderBlock') 三条单测（shape 保持、causal、确定性）+ 「heads 不整除 d 报错」均已绿。seed 命名的独立文件 decoderBlock.ts 未单独创建（合并进 transformer.ts），这是诚实差异。
• 实现纯确定性、无 key 依赖；shape 断言与 causal 断言已入单测——与 seed 意图一致。
• 已就绪的单测清单（transformer.test.ts，均绿）：matmul 正确性、softmaxRows 行和=1 且单调、layerNorm 零均值单位方差 + gamma/beta、单头 causal「未来不改过去」、decoderBlock shape 保持、decoderBlock causal、确定性、randomParams 维度不整除报错。这些是后续 Day 13 数值对齐、Day 14 KV cache「加速不改数值」的验收地基。

multi-head 的当代变体（today→B12 的桥）：今天实现的是经典 MHA（每头一套独立 Q/K/V）。生产侧为省 KV cache 已演化出一谱系，今天先建立坐标，B12 再深入：

它们都不改 attention 数学本身，只改「KV 怎么共享/压缩」——所以今天把 MHA 写对，是理解后面所有变体的前提。

5. 常见误区 / 陷阱

1. 多头缩放用全维 d：每头 attention 的 √ 缩放应用 d/h（头内维度），用全维 d 会过度缩小 logits。
2. heads 不整除 d 却静默切：会产生错位的子空间。本仓在 splitHeads/randomParams 显式 throw，单测 randomParams(8, 3, ...) 验证报错——别省这个检查。
3. 误用 post-LN：post-LN 深层训练不稳、需 warmup，已被 pre-LN 取代。教学和复现都用 pre-LN。
4. 忘了 FFN 残差或 LN：少一条残差，深层梯度消失；少一处 LN，数值尺度漂移。五件套缺一不可。
5. concat 顺序与 Wo 不一致：拼头顺序若与训练时不同，Wo 的列对应就错了。本仓 concatHeads 用 flatMap 按头序拼，与 splitHeads 严格互逆，保持一致。
6. 把 attention 当成「有参数算子」：attention 本身无参数，参数全在 Wq/Wk/Wv/Wo 和 FFN——调容量是调这些投影/FFN 宽度，不是改 attention 公式。

6. 学习资源（每条带 YYYY-MM）

• Karpathy nanoGPT README & 《Let's build GPT》(2023-01 起，2024 持续维护)——seed 指定复看，decoder-only block 从零实现范式。
• Vaswani et al.《Attention Is All You Need》(2017-06, arXiv:1706.03762)——multi-head 与 encoder-decoder 原型。
• Xiong et al.《On Layer Normalization in the Transformer Architecture》(2020-02, arXiv:2002.04745)——pre-LN 优于 post-LN 的训练稳定性论证。
• DeepSeek-V3 Technical Report (2024-12, arXiv:2412.19437)——2026 主流骨架（pre-LN/RMSNorm + 多头变体）的当代参照。
• Shazeer《GLU Variants Improve Transformer》(2020-02, arXiv:2002.05202)——SwiGLU 取代 ReLU 的 FFN 演进依据（本日教学仍用 ReLU）。
• 本仓 src/agent/transformer/transformer.ts（AICAP-180, 2026-06）——decoderBlock/multiHeadCausalAttention/splitHeads/concatHeads 真实实现。

SOTA检查 (2026-06 更新)

• 当前主流：decoder-only + pre-LN + 残差 + FFN 仍是 2026 主流骨架（Llama 4 / Qwen3 / DeepSeek-V3 均沿用）。nanoGPT 教学价值仍在。
• 是否仍 SOTA：结构骨架是 SOTA；细节有演进——LayerNorm 多被 RMSNorm 替代、激活常用 SwiGLU 而非 ReLU、注意力侧叠 GQA/MQA/MLA 省 KV。今天的 ReLU+标准 LayerNorm 是教学最简版。
• 过时黑名单：避免照搬 post-LN（已被 pre-LN 取代，训练更稳）；不要把 encoder-decoder + cross-attention 当生成式 LLM 默认结构（现代生成模型几乎全 decoder-only）。
• 下次复查点：B2 末（Day 20）；B12 训练/推理深水区时把 ReLU→SwiGLU、LayerNorm→RMSNorm、MHA→GQA/MLA 的当代取舍补全并重验。

衔接

• 昨天：Day 11 — Scaled dot-product + causal mask（单头 attention 核，证明未来不影响过去）。
• 今天：把单头扩成 multi-head 并包成完整 decoder block（pre-LN + 残差 + FFN），输出 shape 对齐、整块 causal。
• 明天：Day 13 — Reference 对齐（用固定权重 + 数值 diff 验证这块手写 block 算得对，max abs diff < 1e-4）。