目标：用 wgpu 计算着色器（WGSL）把最耗时的算子搬到 GPU 上跑， 支持 NVIDIA 与 Intel 核显，同时保证"GPU 不可用时自动回退 CPU"。

1. 为什么选 wgpu

• 跨平台：Windows 走 DX12 / Vulkan，NVIDIA 独显和 Intel 核显都能用；
• 纯 Rust：不依赖 CUDA，也不引入深度学习框架；
• 计算着色器：WGSL 语言手写算子，和写 CPU 的 for 循环思路一致，适合教学。

2. 架构总览

Cargo.toml          gpu feature（可选依赖 wgpu / pollster）
src/gpu.rs          GPU 上下文 + 4 个 WGSL 计算入口 + 同步取回
src/tensor.rs       matmul_data()：GPU 优先，失败回退 CPU
src/main.rs         gpu::init() + demo_gpu()（第 4 个演示）

• --features gpu 开启；默认零 GPU 依赖，构建轻量。
• 初始化用 OnceLock<Option<GpuContext>>：失败静默置 None，后续自动走 CPU。

3. WGSL 计算着色器

4 个计算入口共用同一个 ShaderModule（绑定声明是 module 级的）：

参数统一走 16 字节 uniform：struct Params { p0: u32, p1: u32, p2: u32, p3: u32 } （f32 标量用 bitcast<f32> 位模式传参）。

matmul 用三维 workgroup：@workgroup_size(16,16,1)，global_invocation_id 的 x/y/z 分别对应行/列/batch。tiled 版：每个 workgroup 负责一个 16×16 输出块， 先把 A/B 的 16×16 小块载入共享内存（var<workgroup> sh_a/sh_b: array<f32,256>）， 再在片内做内积，把 K 维的全局内存读取从 K 次降到 K/16 次。

坑：sh_a 的排布必须按 (lid.x, lid.y)（线程行对应输出行），若按习惯的 (lid.y, lid.x) 会行列错乱，矩阵乘结果全错（训练 loss 卡住不下降）。 另外不能在 barrier 之前对越界线程 return——dispatch 向上取整后同一 workgroup 内控制流分歧会让 workgroupBarrier() 变成未定义行为，应"越界读补 0、写回再保护"。

4. 踩坑记录

1. WGSL 变量遮蔽：let b = params.p0 把全局 storage 数组 b 遮蔽成 u32， 再写 b[...] 报 Invalid access into expression。局部变量改名即可。
2. uniform 数组对齐：uniform 地址空间数组 stride 必须 16 字节对齐， array<u32,4> 实际占 64 字节；改用 4 个独立 u32 字段（16 字节）最省事。
3. 绑定编号：scale/relu 不用 binding 1，但声明仍是全局的；创建 bind group 时必须显式指定 binding 编号（0/2/3），不能从 0 连续排。
4. wgpu 30 API：PipelineLayoutDescriptor 无 push_constant_ranges（用 immediate_size）、bind_group_layouts 元素是 Option<_>、 PollType::Wait 是带字段 struct、get_mapped_range() 返回 Result。
5. 沙箱限制：Windows 上 GPU 驱动会写 NVIDIA DXCache、D3DSCache 等目录， 受限环境需要放行，否则进程会被杀（程序自身会先打印完结果）。

5. 运行方式

# 普通运行（无 GPU，纯 CPU）
cargo run --release -- demo

# 开启 GPU 加速（含第 4 个演示：正确性 + 性能对比）
cargo run --release --features gpu -- demo

# 训练 / 推理加 --features gpu 即自动走 GPU
cargo run --release --features gpu -- train --config config.json

实测（NVIDIA GeForce MX150 / Vulkan）：

512x512 矩阵乘：CPU 506.7ms vs GPU 57.0ms（快 8.9x）
批量矩阵乘 CPU vs GPU 最大误差 2.86e-6
逐元素算子（scale/relu/add）验证：通过

注意：MX150 上 512×512 从 naive 的 57.6ms 只微降到 57.0ms——该规模下每次调用的 固定开销（上传/调度/同步取回）已接近计算时间，共享内存的收益被抵消。大模型训练 （如 n_embd=256、block=128、batch=16）每步有几十次 GPU 矩阵乘串行同步， 固定开销会累积，低端 GPU 上仍然偏慢。教学实现以"清晰、可回退"优先，不追求极致吞吐。

6. 动手练习

1. 把 tiled 块从 16×16 改成 8×8 或 32×32，观察性能与占用率变化；
2. 把 LayerNorm 也写成 WGSL 着色器，减少 CPU↔GPU 往返；
3. 思考：为什么 GPU 矩阵乘没有"快 50x"？瓶颈在哪里（显存搬运、每次调用的同步取回）？
4. 进阶：把一次 forward/backward 的多次 dispatch 合并提交、只在最后同步一次（计算图）。

代码位置：src/train.rs（LRScheduler、train_gpt） 调用处：src/main.rs（demo_gpt 里的训练配置） 配套代码：src/optim.rs（AdamW，第 17 课）

1. 本课要搞懂的问题

1. 学习率为什么不能"一步到位"？固定 lr 训练会出什么问题？
2. warmup 是干什么的？"前 50 步 lr 从 0 线性升到 max" 在代码里怎么写的？
3. cosine decay 长什么样？为什么末期要用越来越小的学习率？
4. LRScheduler 的 lr() / step() 到底怎么算、怎么用？
5. 训练日志里的 lr 0.002947 → … → 0.000300 是怎么一步步算出来的？
6. 21 课全部学完，我手里有什么？下一步该往哪走？

2. 为什么需要学习率调度

第 6 课我们学过梯度下降：θ = θ - lr·g，学习率 lr 决定每一步走多大。那"全程用一个固定的 lr"行不行？看看两种极端：

理想的 lr 轨迹是一条"倒 U 形"：

lr
 │          ╭────────────────╮
 │         ╱                  ╲
 │        ╱                    ╲
 │  ╱────╱                      ╲
 │ ╱                             ╲
 └──────────────────────────────────▶ step
   0     warmup 结束           训练结束
         （lr 达到 max）       （lr 降到 min）

• 初期：lr 从很小开始（warmup），保护刚起步、还很脆弱的参数；
• 中期：lr 维持在最大值附近，快速下降；
• 后期：lr 逐渐变小（cosine decay），在损失曲面的"沟底"精细收敛，不会跳过最优解。

这就是本课 LRScheduler 做的事情：先 warmup 线性爬升，再 cosine 曲线衰减。

3. Warmup：先热个身

3.1 公式

warmup 阶段（step < warmup_steps），学习率从 0 线性升到最大值：

lr(step) = max_lr × (step + 1) / warmup_steps

注意分子里的 step + 1：step 从 0 开始数，step + 1 保证第 0 步的 lr 是 max_lr / warmup_steps（而不是 0）， 到第 warmup_steps - 1 步时 lr 正好是 max_lr × warmup_steps / warmup_steps = max_lr。

3.2 为什么训练初期要用小学习率

直觉：刚上冰面的新手滑冰，先小步慢走找平衡（warmup），站稳了再加速（大 lr）。

4. Cosine decay：按余弦曲线收尾

4.1 公式

warmup 结束后，学习率按半个余弦周期从 max_lr 衰减到 min_lr：

progress = (step - warmup_steps) / (total_steps - warmup_steps)     // 0 → 1
cosine   = 0.5 × (1 + cos(π × progress))                            // 1 → 0
lr(step) = min_lr + (max_lr - min_lr) × cosine

4.2 关键点

余弦曲线在两端（progress ≈ 0 和 ≈ 1）斜率平缓、在中间（progress ≈ 0.5）斜率最陡，翻译成 lr 就是：

• 刚离开 warmup 时 lr 在 max 附近"多停留一会儿"（给足中期快速下降的时间）；
• 训练后期 lr 平缓地逼近 min_lr（在沟底精细搜索，不震荡）。

为什么末期要小 lr？损失曲面在高维空间里像"沟壑"，大 lr 会在沟底来回弹跳；小 lr 才能贴着沟底慢慢走到最低点。 而且末期 loss 下降本来就慢，用小 lr 不会拖慢多少，却能让最终 loss 更漂亮。

我们的代码里 min_lr 取的是 max_lr × 0.1（见 5.3 节），即最终学习率衰减到最大值的十分之一。

5. LRScheduler 实现讲解

5.1 结构体

src/train.rs 第 20 课小节：

/// 学习率调度器：warmup + cosine decay
pub struct LRScheduler {
    warmup_steps: usize,
    total_steps: usize,
    max_lr: f32,
    min_lr: f32,
    step: usize,
}

impl LRScheduler {
    pub fn new(warmup_steps: usize, total_steps: usize, max_lr: f32, min_lr: f32) -> Self {
        LRScheduler {
            warmup_steps,
            total_steps,
            max_lr,
            min_lr,
            step: 0,
        }
    }
    // ...
}

5.2 lr()：按阶段算当前学习率

pub fn lr(&self) -> f32 {
    if self.step < self.warmup_steps {
        // 线性 warmup
        self.max_lr * (self.step as f32 + 1.0) / self.warmup_steps.max(1) as f32
    } else {
        // cosine 衰减：从 max_lr 平滑降到 min_lr
        let progress = (self.step - self.warmup_steps) as f32
            / (self.total_steps - self.warmup_steps).max(1) as f32;
        let progress = progress.min(1.0);
        let cosine = 0.5 * (1.0 + (std::f32::consts::PI * progress).cos());
        self.min_lr + (self.max_lr - self.min_lr) * cosine
    }
}

公式 ↔ 代码对照表：

注意分支边界：当 step == warmup_steps 时走 cosine 分支，此时 progress = 0、cosine = 1，lr 恰好等于 max_lr—— warmup 的终点和 cosine 的起点无缝衔接，lr 曲线没有断点。

5.3 step() 与在训练循环里的用法

pub fn step(&mut self) {
    self.step += 1;
}

step() 每训练一步调用一次，把内部计数 +1，从而推动 lr() 前进。在 train_gpt 里和 AdamW 配合：

// src/train.rs train_gpt
let mut opt = AdamW::new(cfg.max_lr, params.clone(), cfg.weight_decay);
let mut scheduler = LRScheduler::new(cfg.warmup_steps, cfg.steps, cfg.max_lr, cfg.min_lr);
// ...
for step in start_step..cfg.steps {
    // ... 前向、反向、梯度裁剪 ...
    let cur_lr = scheduler.lr(); // 先取当前步的学习率（scheduler.step() 之后会变成下一步的）
    opt.lr = cur_lr;             // 把"当前步的学习率"写进优化器（AdamW.lr 是 pub 字段）
    opt.step();                  // 优化器用这个 lr 更新参数
    opt.zero_grad();
    scheduler.step();            // 步数 +1，为下一步准备新的 lr
    // ...
}

配合逻辑一句话：调度器负责"算 lr"，优化器负责"用 lr 更新参数"，两者通过 opt.lr = scheduler.lr() 对接。

6. 训练日志解读：lr 0.002947 → 0.000300

6.1 我们的训练配置

src/main.rs 的 demo_gpt 用 TrainConfig 传参（演示 3）：

let tcfg = config::TrainConfig {
    seed: 42,
    batch_size: 8,
    steps: 600,
    max_lr: 3e-3,
    warmup_steps: 50,
    eval_every: 100,
    ..config::TrainConfig::default()
};
train::train_gpt(&model, &tokenizer, &loader, &tcfg, None, None, &mut rng);

于是 train_gpt 内部：min_lr = max_lr × 0.1 = 0.0003，调度器为 LRScheduler::new(50, 600, 0.003, 0.0003)。

6.2 日志打印的时机（避免看错）

train_gpt 里 cur_lr 在 scheduler.step() 之前取值并用于更新：

let cur_lr = scheduler.lr();   // 取"本步"的学习率（scheduler 还没 +1）
opt.lr = cur_lr;
opt.step();
opt.zero_grad();
scheduler.step();              // 之后才 +1（为下一步准备计数）
let last = step + 1 == cfg.steps;
if (step + 1) % cfg.eval_every == 0 || last {
    println!(
        "step {:>5} | lr {:.6} | train_loss {:.4}",
        step + 1,              // 日志从 1 开始数（显示 100、200、…）
        cur_lr,                // 打印的是"本步"的 lr
        loss.item()
    );
}

注意两个细节：一是日志里的 step 是 step + 1（循环变量从 0 开始，打印时 +1，显示 100、200、…）； 二是 scheduler.step() 虽然在打印之前执行，但打印的 cur_lr 是本步开头就取好的——所以日志里的 lr 就是"本步实际用于参数更新的学习率"，不存在错位（旧的实现把打印放在 scheduler.step() 之后、显示的是下一步的 lr，本轮修复后已对齐）。

6.3 三个关键值怎么来的

① 第一个打印点 0.002947（step 100）

demo 配置 eval_every = 100，所以第 100 次迭代（循环变量 step=99，scheduler 计数=99）才第一次打印：

progress = (99 - 50) / (600 - 50) = 49/550 ≈ 0.089
cosine   = 0.5 × (1 + cos(π × 0.089)) ≈ 0.9805
lr = min_lr + (max_lr - min_lr) × cosine = 0.0003 + 0.0027 × 0.9805 ≈ 0.002947   ✓

② warmup 段为什么看不到

eval_every = 100 > warmup_steps = 50，warmup 段（step 0-49）没有任何打印点，所以日志里看不到 0.00006 起步的线性爬升。把 eval_every 改成 10 就能看到：step 10/20/30/40 的 lr 依次是 0.0006 → 0.0012 → 0.0018 → 0.0024（每步增加 max_lr / warmup_steps = 0.003 / 50 = 0.00006，10 步就是 0.0006），然后平滑进入 cosine 段。

③ 最后 0.000300（step 600）

最后一步（循环变量 step=599，计数=599）progress 封顶到 1，cosine = 0.5 × (1 + cos π) = 0：

lr = min_lr + (max_lr - min_lr) × 0 = min_lr = 0.0003   ✓

6.4 真实日志序列（cargo run --release -- demo 原样输出）

这条曲线正是第 4.2 节那张表：先平、中陡、末缓。对比 loss 日志：前期 loss 快速下降（lr 大）， 后期 loss 缓慢逼近平台（lr 小、精细收敛）——这就是"调度"在训练曲线上留下的痕迹。

小技巧：把 eval_every 改成 10，就能看到 warmup 段每 10 步 +0.0006 的线性爬升， 以及 cosine 段更密的 lr 变化——对照 6.3 的公式验证每一步。

7. 全项目 21 课总结回顾

对照 docs/00-学习计划.md 的六个阶段，我们把"从零写一个 LLM"拆成了 21 课。回顾一下每阶段的成果：

阶段一：地基 —— 张量与自动微分（第 1-4 课）

里程碑：只靠标准库就实现了深度学习框架的"张量 + 自动微分"核心，cargo test 一路绿灯。

阶段二：神经网络基础（第 5-7 课）

里程碑：一个从零手写的神经网络成功收敛，证明自动微分没写错。

阶段三：分词器（第 8 课）

里程碑：文本 → token id 的通路打通，CharTokenizer 为小语料演示服务。

阶段四：Transformer 架构（第 9-12 课）

里程碑：GPT 的骨架完整成形——注意力找相关性、MLP 加工信息、归一化稳住数值、残差保梯度畅通。

阶段五：训练与推理（第 13-16 课）

里程碑：cargo run 一次跑通"训练 → 生成"，模型真的会"写"了。

阶段六：进阶优化（第 17-21 课）

一路走下来的核心收获

验收标准（来自学习计划）：cargo run 正常输出 XOR 正确率、BPE 词表、训练日志与两段生成文本； MLP 在 XOR 上收敛；小 GPT 生成通顺短文本——全部达成。

8. 后续展望：从这里出发

21 课只是"入门级 GPT"的终点，却是真实 LLM 世界的起点。按性价比排序，值得继续的方向：

8.1 数据和规模

8.2 直接能做的代码级改进

• RoPE 已是完成态：第 19 课已把它接进 MultiHeadAttention（只旋转 Q/K、KV cache 缓存"已旋转的 K"），下一步可以直接试 8.3 的长度外推技巧；
• RMSNorm：把 LayerNorm 的均值项去掉，更快且效果相当（LLaMA 在用）；
• GQA / MQA：多头共享 KV，配合 KV cache 大幅省显存；
• 更好的分词器：加大 BPE 合并次数 / 词表，或换 SentencePiece 风格。

8.3 架构前沿（了解即可）

• MoE（混合专家）：把每个 Block 的 MLP 换成多个"专家" + 一个路由器，每个 token 只激活少数专家。 参数量大幅增加但计算量可控（Mixtral、DeepSeek 系列就是代表）；
• 长度外推：RoPE 训练长度外会失效，用 NTK-aware scaling、YaRN 等方法把上下文"撑"长；
• 稀疏注意力 / Flash Attention：优化注意力本身的计算与 IO，长上下文必备。

8.4 真正的并行（从单卡到多卡 GPU）

我们现在是单线程 CPU（可选单卡 GPU）、零依赖的教学实现，真实训练要靠并行：

想实战的话，可以先把本项目移植到 candle / burn 等 Rust 深度学习框架上（它们封装了 GPU 和 autograd）， 或者用 rayon 给张量运算加上多线程——那是"工程优化"的另一个故事了。

学习这条路没有终点：从手写张量到训练小 GPT，你走完了框架背后的 90% 黑盒；剩下的 10% 是工程与规模。

9. 动手练习

1. 改 warmup 步数：把 main.rs 里的 warmup_steps 从 50 改成 0（注意 lr() 里 .max(1) 会兜底，不会除零）和 200，分别训练，观察 loss 曲线差异——warmup 太短初期是否震荡？太长是否拖慢中期？
2. 改 min_lr：把 TrainConfig::default() 里的 min_lr（3e-4）改成 3e-5 和 1.5e-3，对比最终 loss。想想 min_lr 设太高会有什么后果（后期"精细收敛"没了）。
3. 绘制 lr 曲线：写个 10 行小程序（或用 6.4 的公式在表格/脚本里算一遍），以 step 为横轴画出 600 步的 lr 曲线，确认它是"先线性爬升、再余弦衰减"的倒 U 形。
4. 对比固定 lr：把 train_gpt 改成固定 opt.lr = cfg.max_lr（不调 scheduler），训练 600 步对比 loss——体会调度带来的"同样的步数，更好的收敛"。
5. 看懂日志：把 eval_every 从 100 改成 10，重新 cargo run，观察 warmup 段的 lr 爬升——它每步增加 max_lr/warmup_steps = 0.003/50 = 0.00006，所以每 10 步增加 0.0006（0.00006 → 0.00066 → 0.00126 → …），确认和公式一致。
6. 给 21 课写总结：按第 7 节的表格，自己列一份"我掌握了什么"清单，找出最薄弱的一课回去重读——这是收尾课最重要的作业。

10. 本课总结（也是全项目的结语）

• 为什么需要调度：固定大 lr 后期震荡、固定小 lr 前期太慢，动态 lr 兼顾"稳、快、细"。
• warmup：前 warmup_steps 步 lr = max_lr × (step+1)/warmup_steps 线性爬升，保护训练初期的随机参数和 AdamW 未热身的动量。
• cosine decay：lr = min_lr + (max_lr - min_lr) × 0.5 × (1 + cos(π·progress))，从 max_lr 平滑衰减到 min_lr（本项目取 max_lr × 0.1）。
• LRScheduler：lr() 按 step 计数选分支算 lr，step() 推进一步；在 train_gpt 里通过 opt.lr = scheduler.lr() 喂给 AdamW。
• 日志解读：日志显示 step 100→600，lr 从 0.002947 一路衰减到 0.000300（min_lr）；warmup 段（step 0-49）因 eval_every=100 没有打印点；打印的是本步实际用于更新的 cur_lr，没有错位。
• 21 课全回顾：地基（张量/自动微分）→ 神经网络 → 分词器 → Transformer → 训练推理 → 现代技巧（AdamW/KV cache/RoPE/调度/GPU），cargo run 达成全部验收标准。
• 展望：更大语料与模型、RMSNorm/GQA、MoE、真正的并行（数据/张量/流水线）、长度外推——每一步都能在这个项目上继续长出来。

最后一句：一个能跑的、完全由你手写的小 GPT，就是打开大模型世界最好的钥匙。 21 课结束，学习才刚刚开始。祝玩得开心！

代码位置：src/rope.rs（rotary_pair 生产入口 + rotary 测试用 + test_rotary / test_rotary_grad_exact 测试） 配套代码：src/attention.rs（RoPE 接入 MultiHeadAttention） 配套文档：docs/11-位置编码与归一化.md

1. 本课要搞懂的问题

1. 第 11 课的正弦位置编码有什么缺陷？为什么现代 LLM（LLaMA、Qwen、Gemma……）几乎清一色用 RoPE？
2. "旋转"怎么把位置信息编码进向量？公式 x'_{2i} = x_{2i}·cosθ - x_{2i+1}·sinθ 是怎么来的？
3. 为什么说旋转是正交变换？范数不变这个性质重要在哪？
4. 为什么 RoPE 的点积只与"位置差"有关？相比正弦编码强在哪？
5. 为什么 RoPE 和 KV cache 是"天生一对"？推理时为什么只需要旋转新 token？
6. src/rope.rs 里 rotary_pair 的前向、反向具体怎么实现？test_rotary 测了什么？

2. 先回顾：正弦位置编码的两个局限

第 11 课我们实现了正弦位置编码，把位置向量加到 token embedding 上：

x = token_embedding + pos_emb

它解决了"注意力看不见位置"的问题，但有两个不足：

2.1 编码的是"绝对位置"，模型得自己推断相对关系

正弦编码给每个绝对位置 pos 一个固定向量，加到输入上。模型面对的其实是"内容 + 绝对位置"的混合表示， "A 在第 3 位、B 在第 7 位" 这个信息是间接的——模型需要自己学会"位置 7 减去位置 3 = 距离 4"。

原论文证明了 PE(pos+k) 可以表示为 PE(pos) 的线性组合（三角恒等式），即"理论上模型能学到相对距离"， 但这只是给了模型一个机会，不是保证——要靠训练把这种规律学出来。

直觉：正弦编码像给每个座位发一张写着"座位号"的号牌，模型得自己学会"6 号和 2 号的号牌之差 = 隔了 4 排"。 RoPE 的做法是直接把"隔了几排"编码进打分公式里，模型不用学这一步。

2.2 外推（extrapolation）差

我们的模型 block_size = 32（GPTConfig::tiny），训练时位置只见过 0..32。如果推理时生成更长的序列：

• 位置 32、33…… 的 pos_emb 向量虽然能算出来（正弦函数对任意 pos 都有定义），
• 但模型从没见过这种输入分布，注意力分数可能畸变，输出质量断崖式下跌。

补充：RoPE 也不是"万能外推药"。角度随位置线性增长，训练长度之外同样会失效。 业界用 NTK-aware scaling、YaRN 等技巧做长度外推，那是后话，本课先把 RoPE 本身讲透。

3. RoPE 原理：把 d 维拆成 d/2 对来旋转

RoPE 出自论文 RoFormer: Enhanced Transformer with Rotary Position Embedding（Su et al., 2021）。

3.1 二维旋转回顾

在二维平面上，把向量 (a, b) 绕原点逆时针旋转角度 θ：

a' = a·cosθ - b·sinθ
b' = a·sinθ + b·cosθ

写成矩阵就是旋转矩阵 R(θ) = [[cosθ, -sinθ], [sinθ, cosθ]]。

3.2 把 d 维向量看成 d/2 个二维向量

一个 d 维向量（d 必须为偶数），按相邻两两配对：

(x₀, x₁), (x₂, x₃), (x₄, x₅), ..., (x_{d-2}, x_{d-1})

每一对看作平面上的一个点，第 i 对 (x_{2i}, x_{2i+1}) 绕原点旋转角度：

θ_i = pos / 10000^(2i/d)      i = 0, 1, ..., d/2 - 1

旋转后：

x'_{2i}   = x_{2i}·cosθ_i - x_{2i+1}·sinθ_i
x'_{2i+1} = x_{2i}·sinθ_i + x_{2i+1}·cosθ_i

注意频率方案 10000^(2i/d) 和正弦编码一模一样（第 11 课）：i 越小频率越高（角度变化快，管近处的精细位置）， i 越大频率越低（角度变化慢，管远处的粗略位置）。这样"近距离靠高维、远距离靠低维"的分工被保留了下来。

3.3 手算一个例子

设 d = 4、位置 pos = 1，则有两对：

于是：

x'₀ = x₀·0.5403 - x₁·0.8415
x'₁ = x₀·0.8415 + x₁·0.5403
x'₂ ≈ x₂·0.99995 - x₃·0.01
x'₃ ≈ x₂·0.01 + x₃·0.99995

第一对被旋转了整整 1 弧度（幅度很大，专门区分相邻位置），第二对几乎没动（幅度极小，负责长距离）。 同一个向量，位置不同旋转角度就不同——位置信息就这样被"揉"进了向量本身。

4. 旋转是正交变换：范数不变

旋转矩阵满足 R(θ)ᵀ·R(θ) = I（转置乘自己等于单位阵），且 |det R| = 1，因此它是正交矩阵：

‖R(θ)·x‖² = (R(θ)x)ᵀ(R(θ)x) = xᵀR(θ)ᵀR(θ)x = xᵀx = ‖x‖²

也就是说旋转只改变向量的方向，不改变长度（范数）。这带来两个直接好处：

对比一下：正弦编码是"加法"，加完 ‖x + pos_emb‖ 和 ‖x‖ 一般不同，多多少少会扰动数值分布； RoPE 是"乘法"（乘以正交矩阵），只转不拉，范数严格不变。

更妙的是：正交矩阵的逆就是它的转置，而转置恰好等于负角度旋转：

R(θ)ᵀ = R(-θ)

这个性质在反向传播里极其好用——梯度回传时把角度取反再旋转一次就行（详见第 8 节）。

5. 核心性质：相对位置（点积只与位置差有关）

5.1 先看 d = 2 的情形

设查询在位置 m：q_m = R(mθ)·q；键在位置 n：k_n = R(nθ)·k。注意力分数是它们的点积：

q_m · k_n = (R(mθ)q)ᵀ · (R(nθ)k)
          = qᵀ · R(mθ)ᵀ · R(nθ) · k
          = qᵀ · R(-mθ) · R(nθ) · k      （Rᵀ = R⁻¹ = 负角度）
          = qᵀ · R((n - m)θ) · k          （旋转矩阵乘法 = 角度相加）

结果只依赖 n - m，与 m、n 各自的值无关！

5.2 推广到 d 维

d 维的旋转矩阵是块对角矩阵：

R(θ) = diag( R(θ₀), R(θ₁), ..., R(θ_{d/2-1}) )

块对角矩阵相乘等于各块分别相乘，所以 d 维下同样成立：

q_m · k_n = Σ_i  q_{2i}ᵀ · R((n-m)·θ_i) · k_{2i}    只依赖 n - m

5.3 为什么这比正弦编码强

一句话总结：正弦编码把位置塞进"表示"里，RoPE 把位置塞进"打分"里。 而注意力唯一关心位置的地方就是打分（q·k），所以 RoPE 做到了"精确投放"。

6. 与 KV cache 天然兼容：只旋转新 token

回顾第 18 课的 KV cache 推理流程（src/sample.rs 的 generate + src/attention.rs 的 KVCache）：

• 第一次：把整个 prompt 喂给模型，算出所有位置的 K/V 存入缓存；
• 之后每步：只前向最新 1 个 token，历史的 K/V 直接从缓存取。

RoPE 和这个流程是无缝衔接的：

新 token 的绝对位置 = 缓存长度 base + 它在当前窗口里的下标 j

base（已缓存的位置数）由 GPT::forward 算出来传给每层（第 18 课的设计），MultiHeadAttention::forward 里 positions 就是这么构造的：

// src/model.rs（GPT::forward）：base = 缓存长度
let base = kv_cache
    .as_ref()
    .map(|c| c.first().map(|k| k.seq_len()).unwrap_or(0))
    .unwrap_or(0);

// src/attention.rs（MultiHeadAttention::forward）：batch 内每个样本位置相同，重复 b 次
let mut positions = Vec::with_capacity(b * t);
for _ in 0..b {
    positions.extend(base..base + t);
}

于是接入 RoPE 时：

因为旋转只依赖"位置号"，而每个 token 的位置号在生成时是唯一确定的，所以历史上每个 K 只旋转一次、终身有效。 相比"每次生成都要重新取一段 pos_emb 加进输入"的加法式编码，RoPE 是"随算随旋"，不需要任何历史重算。

另一个视角：RoPE 在注意力内部旋转 Q/K，位置信息不会污染 token 的"内容表示"（残差流上的 x 始终是纯语义）， 这也让每一层的 Q/K 都能用"本层自己的位置感"去打分。

7. 前向实现讲解：cos/sin 表 + 查表旋转

src/rope.rs 把实现拆成三层：先一次性预计算 cos/sin 表，再用表查表旋转，Q/K 共用同一张表：

/// 预计算每个 (位置, 对偶下标) 的 cos/sin 表，长度 rows × (D/2)。
/// 同一批 positions 的三角只算一次：前向、反向、Q/K 复用。
fn build_cos_sin_tab(positions: &[usize], d: usize) -> (Vec<f32>, Vec<f32>) {
    let rows = positions.len();
    let mut c_tab = vec![0.0f32; rows * (d / 2)];
    let mut s_tab = vec![0.0f32; rows * (d / 2)];
    for r in 0..rows {
        let pos = positions[r] as f32;
        for i in 0..d / 2 {
            let theta = pos / 10000f32.powf((2 * i) as f32 / d as f32);
            c_tab[r * (d / 2) + i] = theta.cos();
            s_tab[r * (d / 2) + i] = theta.sin();
        }
    }
    (c_tab, s_tab)
}

/// 用现成的 cos/sin 表旋转一个张量（[rows, D]），每对元素按公式旋转。
fn rotate_with_tab(x: &Tensor, c_tab: &[f32], s_tab: &[f32]) -> Tensor {
    let (rows, d) = (x.shape[0], x.shape[1]);
    let sd = x.data.borrow();
    let mut out_data = vec![0.0f32; rows * d];
    for r in 0..rows {
        for i in 0..d / 2 {
            let (c, s) = (c_tab[r * (d / 2) + i], s_tab[r * (d / 2) + i]);
            let (a, b) = (sd[r * d + 2 * i], sd[r * d + 2 * i + 1]);
            out_data[r * d + 2 * i] = a * c - b * s;
            out_data[r * d + 2 * i + 1] = a * s + b * c;
        }
    }
    drop(sd);
    // ...（反向闭包查同一张表、按 R(θ)ᵀ 回传，见第 8 节）
    result
}

对外只暴露两个入口——rotary（旋转单个张量，仅供测试）和 rotary_pair（一次建表同时旋转 Q/K）：

pub fn rotary_pair(&self, other: &Tensor, positions: &[usize]) -> (Tensor, Tensor) {
    let d = self.shape[1];
    let (c_tab, s_tab) = build_cos_sin_tab(positions, d);
    (
        rotate_with_tab(self, &c_tab, &s_tab),
        rotate_with_tab(other, &c_tab, &s_tab),
    )
}

为什么要"先建表再查表"？注意力里 Q 和 K 的 positions 完全相同，三角函数的 cos/sin 只要算一遍； 前向算一遍、反向闭包再查一遍——相比"每次旋转现场重算三角"，批量训练能省下可观的重复计算。

逐行对照公式：

几个值得注意的设计点：

1. 三重循环的次序：建表时外层按行 r、内层按对 i，cos/sin 每个 (r, i) 只算一次；旋转时同样两层循环，但只做查表乘加，不再碰任何三角函数。
2. drop(sd)：读完输入数据后立刻释放借用，之后才创建结果张量和反向闭包——避免闭包捕获时和 self.data 的借用纠缠。
3. requires_grad 分支：如果输入不需要梯度（比如纯推理），就直接返回普通结果，不建 parents/backward，省下计算图的维护开销。这和第 18 课 KV cache 推理时"纯数据拼接、无梯度"的思路一致。
4. 反向闭包查同一张表：反向闭包为了"自包含"（只捕获 ct/st 两份表拷贝 + rg/sg 两个 Rc），直接查表取值，不再重算 theta/cos/sin。代价是闭包多持有两张表，好处是前向建的三角表被完整复用。

8. 反向实现讲解：正交矩阵的梯度

8.1 理论：梯度用"转置"回传

设 y = R(θ)·x（旋转一对）。根据链式法则：

∂L/∂x = R(θ)ᵀ · ∂L/∂y

因为旋转矩阵正交，R(θ)ᵀ = R(-θ)，所以梯度回传 = 把梯度按负角度再旋转一次。这就是方法注释里写的：

"每个 pair 的旋转矩阵正交，梯度用其转置（即负角度旋转）回传。"

展开来写，若记 (ga, gb) = (∂L/∂y_{2i}, ∂L/∂y_{2i+1})，则正确回传是：

∂L/∂x_{2i}   = ga·cosθ + gb·sinθ
∂L/∂x_{2i+1} = -ga·sinθ + gb·cosθ

验证（也可以硬算）：y₀ = a·c - b·s、y₁ = a·s + b·c， ∂L/∂a = ga·(∂y₀/∂a) + gb·(∂y₁/∂a) = ga·c + gb·s ✓ ∂L/∂b = ga·(∂y₀/∂b) + gb·(∂y₁/∂b) = -ga·s + gb·c ✓

8.2 代码实际怎么写的（已按正确形式实现）

src/rope.rs 里 rotate_with_tab 反向闭包的循环体是（c/s 直接查 ct/st 表）：

let (ga, gb) = (g[r * d + 2 * i], g[r * d + 2 * i + 1]);
// 反向 = 前向旋转矩阵的转置 R(θ)ᵀ：grad = (ga·c + gb·s, -ga·s + gb·c)
sgm[r * d + 2 * i] += ga * c + gb * s;
sgm[r * d + 2 * i + 1] += -ga * s + gb * c;

与 8.1 的结论逐行对应：第一行 ga·c + gb·s、第二行 -ga·s + gb·c，正是 R(θ)ᵀ·g。

💡 陷阱提示：反向回传必须用转置（即负角度旋转）。写代码时容易手滑把符号写成正角度旋转 (ga·c - gb·s, ga·s + gb·c)——方向错了但梯度范数不变（正交变换保范数）， 只校验范数的测试发现不了。所以测试要"既测范数、又测方向"（见第 9 节 test_rotary_grad_exact）。 本仓库初版实现确实写反过，已修复并补了逐元素断言。

对训练的影响：本课已经把 RoPE 接进了 MultiHeadAttention（src/attention.rs），训练和 KV cache 推理都用它。 GPT 已不再有 pos_emb 字段，位置信息完全由注意力内部旋转 Q/K 提供（见第 10 节）。

9. test_rotary 测试讲解

测试在 src/rope.rs 的 mod tests 里，共三个断言，分别验证三件事：

#[test]
fn test_rotary() {
    // 1. 旋转是正交变换：范数不变
    let x = Tensor::param(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![1, 6]);
    let r = x.rotary(&[3]);
    let orig_norm: f32 = x.data().iter().map(|v| v * v).sum();
    let rot_norm: f32 = r.data().iter().map(|v| v * v).sum();
    assert!((orig_norm - rot_norm).abs() < 1e-3, "范数应守恒：{} vs {}", orig_norm, rot_norm);

    // 2. pos=0 时所有角度为 0，等于恒等变换
    let x2 = Tensor::param(vec![1.0, 2.0, 3.0, 4.0], vec![1, 4]);
    let r2 = x2.rotary(&[0]);
    assert!((r2.data()[0] - 1.0).abs() < 1e-5);
    assert!((r2.data()[3] - 4.0).abs() < 1e-5);

    // 3. 梯度：sum 的梯度是单位向量，经正交矩阵回传后范数不变（= 元素数）
    let x3 = Tensor::param(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![1, 6]);
    let loss = x3.rotary(&[2]).sum();
    loss.backward();
    let g: Vec<f32> = x3.grad();
    assert!((g.iter().map(|v| v * v).sum::<f32>() - 6.0).abs() < 1e-3, "梯度范数应为 6");
}

测试 ③ 的思路很巧妙：∂sum/∂y 是"全是 1"的单位向量，只要反向实现用的是正交矩阵（无论正角度还是负角度）， 回传后梯度范数平方必然还是元素数 6。它验证了"正交性"，但正如 8.2 所说，验证不了"方向"。

仓库里还补了一个 test_rotary_grad_exact，专测方向：d=2、pos=1 时 θ₀=1 rad， 对 x=[1,2]、loss=sum，期望梯度是 R(θ)ᵀ·[1,1] = (cos1+sin1, -sin1+cos1) = (1.3818, -0.3012)， 逐元素断言。有了它，反向符号写反会立刻红：

#[test]
fn test_rotary_grad_exact() {
    let x = Tensor::param(vec![1.0, 2.0], vec![1, 2]);
    let loss = x.rotary(&[1]).sum();
    loss.backward();
    let (c, s) = (1f32.cos(), 1f32.sin());
    let (ga, gb) = (c + s, -s + c);
    assert!((x.grad()[0] - ga).abs() < 1e-5, "grad[0] = {}", x.grad()[0]);
    assert!((x.grad()[1] - gb).abs() < 1e-5, "grad[1] = {}", x.grad()[1]);
}

运行测试：

cargo test test_rotary    # 跑全部 rotary 测试

10. 如何接入注意力：对 Q/K 应用

10.1 为什么只旋转 Q 和 K，不旋转 V

注意力分数是 q·k（点积），位置信息只有在这里才"有用武之地"。而输出 out = softmax(scores)·V 是内容的加权求和， V 是"被加权的内容"，不需要带位置信息。所以标准做法是：只旋转 Q 和 K，V 保持原样。

10.2 接入点：MultiHeadAttention::forward

在 src/attention.rs 的 MultiHeadAttention::forward 里，Q/K/V 投影出来后是这样的：

let q = self.c_q.forward(x).reshape(vec![b, t, d]); // [B, T, D]
let k = self.c_k.forward(x).reshape(vec![b, t, d]);
let v = self.c_v.forward(x).reshape(vec![b, t, d]);

RoPE 有两个可选的接入时机：

我们的 rotary_pair 接口（输入 [rows, D] + 每行的 positions）两种都支持，只要最后一维是偶数即可—— GPTConfig::tiny 里 n_embd=64、head_dim=16，都满足。

10.3 实际接入方式（本仓库已接入）

src/attention.rs 里，Q/K 投影之后、KV cache 拼接之前旋转：

// 位置：Q/K 投影之后、KV cache append 之前
let mut positions = Vec::with_capacity(b * t);
for _ in 0..b {
    positions.extend(base..base + t);
}
// Q/K 用同一批 positions，一次建表同时旋转（三角函数只算一遍）
let (q, k) = q
    .reshape(vec![b * t, d])
    .rotary_pair(&k.reshape(vec![b * t, d]), &positions);
let (q, k) = (
    q.reshape(vec![b, t, d]),
    k.reshape(vec![b, t, d]),
);
// 之后照旧：cache.append(&k, &v)（缓存里存的就是旋转后的 K）

要点：

1. positions 长度必须等于行数（rotary_pair 里 assert_eq!(self.shape[0], positions.len())），这里行数是 b*t，所以 batch 内每个样本的同一列位置相同、要重复 b 次。
2. 旋转发生在 append 之前：缓存里存的是"已旋转的 K"，历史 K 直接复用，符合第 6 节的兼容性约定。
3. Q 没有缓存（推理时 Q 永远只有新 token 一个），所以 Q 总是用当前位置旋转，天然正确。
4. 由于 K/V 已旋转并缓存，推理模式（KV cache）和训练模式行为一致，不会像正弦编码那样需要区分两套位置逻辑。这一点有专门的测试守着（src/model.rs 的 test_kv_cache_matches_full_forward）。
5. rotary_pair 的反向按 R(θ)ᵀ 实现并通过 test_rotary_grad_exact 逐元素验证。

替换还是叠加？本仓库选了替换：GPT 结构体里没有 pos_emb 字段（第 11 课的正弦编码在接入 RoPE 后已删除）， 位置信息完全由注意力内部旋转 Q/K 提供。叠加方案（保留 pos_emb 再加 RoPE）一般没必要，且会稀释 RoPE 的相对位置特性。

11. 动手练习

1. 手算一对旋转：设 d = 2、pos = 1，对向量 [1.0, 0.0] 手算 rotary 的输出 （提示：θ₀ = 1/10000^0 = 1 rad，cos1 ≈ 0.5403, sin1 ≈ 0.8415），再用 x.rotary(&[1]) 验证。
2. 验证梯度方向：rotary 的反向已修复为 R(θ)ᵀ（见第 8.2 节）。自己动手把测试改成"正角度旋转"的错误版本， 跑 test_rotary 和 test_rotary_grad_exact，观察只有后者能发现错误——体会"既测范数、又测方向"的价值。
3. 对照实际接入：对照 10.3 节的代码，在 src/attention.rs 里找到 RoPE 的接入位置。 试着把 rotary_pair 改成"先 rotary 转 Q、再 rotary 转 K"（各自建表），跑 cargo test 确认结果一致——体会 rotary_pair 一次建表省下的重复计算。
4. 验证相对位置性质：写一个测试——取同一个 q，在位置 m 旋转；取同一个 k，在位置 n 旋转， 断言 q_m · k_n 只和 n - m 有关（对固定的位置差 d，换不同的 m 结果应相同）。这正是第 5 节推导的代码验证。
5. 思考外推：把 block_size 从 32 改成 64 重新训练，再让模型生成超过训练长度的文本， 观察 RoPE 的输出质量在训练长度之外如何退化，想想为什么（角度随位置线性增长，训练长度之外照样失效）。
6. 思考：为什么 θ_i = pos / 10000^(2i/d) 里的 10000 和 d 与正弦编码一模一样？ 如果把它改成 10 会怎样？（提示：所有维度频率变高，近距离区分更敏感、远距离更早"打满"一圈。）

12. 本课总结

• 正弦编码的两个局限：编码绝对位置（相对关系要模型自己学）、外推差（训练长度外没见过的输入分布）。

• RoPE 把 d 维拆成 d/2 对，第 i 对按 θ_i = pos / 10000^(2i/d) 旋转： x'_{2i} = x_{2i}·cosθ - x_{2i+1}·sinθ、x'_{2i+1} = x_{2i}·sinθ + x_{2i+1}·cosθ。

• 旋转是正交变换：R(θ)ᵀR(θ) = I，范数不变 → 数值稳定、不破坏 LayerNorm。

• 相对位置性质：q_m · k_n = qᵀR((n-m)θ)k，点积只与位置差 n - m 有关——位置被"编死"进了打分公式。

• 与 KV cache 天然兼容：新 token 按 base + j 旋转自己的 Q/K，历史 K 已旋转、直接复用。

• src/rope.rs 的实现分三层：build_cos_sin_tab 预计算 cos/sin 表 → rotate_with_tab 查表旋转 → rotary_pair 一次建表同时旋转 Q/K；反向按正交矩阵的转置 R(θ)ᵀ 回传（已修复并通过 test_rotary_grad_exact 验证）。

• test_rotary 三个断言分别验证：正交性（范数不变）、pos=0 恒等、梯度范数守恒；test_rotary_grad_exact 逐元素验证梯度方向。

• 实际接入：在 MultiHeadAttention 拆头前、KV cache append 之前对 Q/K 旋转（只转 Q/K、不转 V），已完成接入；GPT 已无 pos_emb 字段（第 11 课正弦编码被替换）。

• 下一课（第 20 课）：收尾——学习率调度（warmup + cosine decay）与整个项目的总结回顾！

代码位置：src/attention.rs（KVCache / MultiHeadAttention） 代码位置：src/model.rs（GPT::forward） 代码位置：src/sample.rs（generate） 演示入口：src/main.rs（演示 3：生成 1 / 生成 2）

1. 本课要搞懂的问题

1. 推理时为什么"历史 token 的 K/V"会被一遍遍重复计算？
2. KVCache 的数据结构长什么样？append / seq_len 各做了什么？
3. cache 模式与全量模式在 generate 里的流程有什么不同（首次 vs 之后每步）？
4. 用了缓存之后，为什么生成的概率分布和全量模式完全一样？
5. 为什么缓存模式下上下文达到 block_size 就必须停止生成？

2. 问题：为什么历史 K/V 会被重复计算

生成是逐 token的：每产生一个新 token，都要把它拼到上下文末尾，再前向一次，从输出的 logits 里采样下一个 token。

全量模式（第 16 课 generate 的 use_kv_cache=false 分支）每步都把整个上下文重新喂给模型：

第 1 步：输入 [t0]               → 前向 1 个位置
第 2 步：输入 [t0, t1]           → 前向 2 个位置
第 3 步：输入 [t0, t1, t2]       → 前向 3 个位置
...
第 T 步：输入 [t0, t1, ..., tT]  → 前向 T 个位置

关键观察：第 k 步算出来的前 k-1 个位置的 K/V，和第 k-1 步算出来的一模一样——推理时权重冻结、输入前缀相同，同一个 Linear 层（c_k、c_v）对相同输入必然给出相同输出。

既然如此，为什么要重算？直接记住不就好了？这就是 KV Cache 的动机：

注意：只有 K 和 V 需要缓存，Q 不用。因为预测"下一个 token"只关心新位置上的注意力输出，而它只需要新位置的 Q 去和所有历史位置的 K、V 做注意力。历史位置自己的注意力输出（以及它们的 Q）在生成中根本用不上。

3. KVCache 的结构

src/model.rs 里的定义：

/// KV 缓存（第 18 课）：
/// 生成第 N 个 token 时，前 N-1 个 token 的 K、V 不需要重算。
/// 把每个注意力层的 K、V 存起来，每次只算新 token 的 K、V 并拼接。
pub struct KVCache {
    k: Option<Tensor>, // [1, T, D]
    v: Option<Tensor>,
}

注意：每个注意力层各有一个 KVCache。GPT::new_kv_cache 返回 Vec<KVCache>，长度 = n_layer（本项目 2 层）：

pub fn new_kv_cache(&self) -> Vec<KVCache> {
    (0..self.cfg.n_layer).map(|_| KVCache::new()).collect()
}

3.1 append：把新 K/V 拼到缓存尾部

fn append_data(prev: &Option<Tensor>, cur: &Tensor) -> Tensor {
    match prev {
        Some(p) => {
            let mut all = p.data();
            all.extend(cur.data());
            let d = cur.shape()[2];
            Tensor::from_vec(all, vec![1, p.shape()[1] + 1, d])
        }
        None => cur.clone(),
    }
}

pub fn append(&mut self, k: &Tensor, v: &Tensor) {
    self.k = Some(Self::append_data(&self.k, k));
    self.v = Some(Self::append_data(&self.v, v));
}

以 [1, T, D] 为例，append_data 做的事：

1. 取旧缓存 p 的数据（p.data()，一维展平数组）；
2. 把新 K/V 的数据 cur.data() 拼到末尾；
3. 按 [1, 旧长度+1, D] 重新包成张量。

细节：cur 在推理模式下形状是 [1, 1, D]（只算 1 个新位置），所以长度 +1；d 从 cur.shape()[2] 取。纯数据拼接，推理时无梯度，所以没有走任何 autograd 路径。

3.2 seq_len：已缓存了多少位置

pub fn seq_len(&self) -> usize {
    self.k.as_ref().map(|t| t.shape()[1]).unwrap_or(0)
}

• 缓存为空（None）→ 0；
• 否则返回 k 张量的第 1 维大小，即已缓存的位置数。

它有两个用途（后面会看到）：一是 GPT::forward 用它算位置偏移 base；二是 generate 用它判断要不要停止。

4. MultiHeadAttention 怎么用缓存

MultiHeadAttention::forward 中与缓存相关的部分：

// 1. 投影得到 Q、K、V（Linear 输出是 2D [B*T, D]，恢复成 3D）
let q = self.c_q.forward(x).reshape(vec![b, t, d]); // [B, T, D]
let k = self.c_k.forward(x).reshape(vec![b, t, d]);
let v = self.c_v.forward(x).reshape(vec![b, t, d]);

// 2. KV cache：拼接历史的 K/V（只影响 K、V 的长度）
let (k, v) = match kv_cache {
    Some(cache) => {
        cache.append(&k, &v);
        (cache.k().unwrap(), cache.v().unwrap())
    }
    None => (k, v),
};
let t_total = k.shape()[1];

变化只有一处：K、V 变长，Q 保持 [B, T, D] 不动：

后续的拆头、注意力分数、softmax 等代码一行都不用改，因为它们是按 t_total 写的通用代码：

• 拆头时 k/v 用 t_total 做 reshape（vec![b, t_total, self.n_head, head_dim]），q 仍用 t；
• 分数 scores = q.matmul(&kt).mul_scalar(scale) 形状 [B*H, t, t_total]；
• 因果掩码 mask 是 [t, t_total]，广播相加后 softmax_last_dim()，最后 attn.matmul(&v)。

这就是 KV Cache 优雅的地方：模型代码零侵入，只把"输入 K/V 的来源"从"当场算"换成"缓存里取"。

5. GPT::forward 里的 base 偏移

推理模式下，新 token 的位置不再是"序列内的第 j 个"，而是"全局的第 base + j 个"。GPT::forward 这样处理：

// 2. 位置编码：KV cache 推理时，当前位置从缓存长度开始
let base = kv_cache
    .as_ref()
    .map(|c| c.first().map(|k| k.seq_len()).unwrap_or(0))
    .unwrap_or(0);
let mut positions = Vec::with_capacity(b * t);
for _ in 0..b {
    for j in 0..t {
        positions.push(base + j);
    }
}
...
// 3. 因果掩码：scores 形状 [B*H, T, T_total]，广播 mask [T, T_total]
let t_total = t + base;
let mut mask_data = vec![0.0f32; t * t_total];
for i in 0..t {
    for j in 0..t_total {
        if j > i + base {
            mask_data[i * t_total + j] = f32::NEG_INFINITY;
        }
    }
}

为什么掩码的下界是 base：新 token 在全局序列里的下标从 base 开始（i=0 对应全局 base），所以它能看全局 0..=base（全是缓存里的历史）+ 自己，不能看 base+1 之后（未来）。这和全量模式的因果性完全一致。

训练时 kv_cache 传 None（src/train.rs 里 model.forward(&x, b, t, None)），因为训练时权重每步都在变、历史 K/V 没有复用价值，缓存反而白占内存。

6. generate：cache 模式 vs 全量模式的流程对比

src/sample.rs 的 generate：

let block_size = model.cfg.block_size;
let mut ids = tokenizer.encode(prompt);
let mut cache = model.new_kv_cache();

for _ in 0..max_new {
    // KV cache 模式：上下文总长达到 block_size 就停（缓存无法像全量模式那样截断历史）
    if use_kv_cache && cache[0].seq_len() >= block_size {
        break;
    }
    // 只保留最近的 block_size 个 token（全量模式需要）
    let start = ids.len().saturating_sub(block_size);
    let ctx = &ids[start..];

    let logits = if use_kv_cache {
        // 首次：缓存为空，把整个 prompt 喂进去（顺便填充缓存）
        // 之后：每步只前向最新 1 个 token，历史 K/V 从缓存取
        if cache[0].seq_len() == 0 {
            model.forward(ctx, 1, ctx.len(), Some(&mut cache))
        } else {
            model.forward(&ids[ids.len() - 1..], 1, 1, Some(&mut cache))
        }
    } else {
        // 全量模式：每次把整个上下文重新算一遍（慢，但没有 cache 内存）
        model.forward(ctx, 1, ctx.len(), None)
    };

    // 取最后一个位置的 logits
    let v = model.cfg.vocab_size;
    let n = logits.numel();
    let last_row = &logits.data()[n - v..];
    let next = sample_token(last_row, temperature, top_k, top_p, rng);
    ids.push(next);
}

两种模式逐项对比：

用流程图看第一次生成（prompt = "The fox"，7 个 token，max_new=80）：

cache 模式：                             全量模式：
───────────                             ───────────
第 1 步：前向 ["The fox"(7个)]          第 1 步：前向 ["The fox"(7个)]
         ↓ 填充缓存（7 个位置）                    ↓ 结果只取最后一行，丢弃其余
         采样出第 1 个新 token
第 2 步：前向 [最新 1 个]               第 2 步：前向 ["The fox" + 1个]（8 个）
         ↓ 缓存 = 8 个位置                       ↓ 又从头算了一遍 7 个历史 K/V
第 3 步：前向 [最新 1 个]               第 3 步：前向 [9 个]
         ↓ 缓存 = 9 个位置                       ↓ 重复劳动越来越多
...                                    ...
第 26 步：前向 [最新 1 个]              第 80 步：前向 [窗口内 32 个]
         ↓ 缓存 = 32 个位置                      ↓ 80 个新 token 全部生成
         采样出第 26 个新 token
第 27 步：开头检查缓存 = 32 ≥ block_size → 停止

取 logits 的细节：generate 只取输出张量的最后一行（logits.data()[n - v..]，v = vocab_size）。全量模式算了一整段序列，但生成只需要最后一个位置的预测——前半部分的计算全部是"浪费"；缓存模式干脆只算最后一行需要的东西，正是这种浪费的反面。

7. 为什么输出分布不变

这是 KV Cache 正确性的核心论证，分三步：

1. K/V 值相同：推理时权重冻结。缓存里存的 K/V，与全量模式下同一批输入算出来的 K/V，数值逐位相同（都是同一份代码算的）。
2. 注意力计算相同：新位置的注意力输出 = softmax(Q_k·Kᵀ/√d + mask) · V。其中 K、V 是"全部历史"（缓存模式从缓存取、全量模式当场算），数值相同；Q 是新位置的投影，也相同。
3. softmax 结果相同：mask 规则一致（第 5 节已证），同一组分数经过同样的 softmax → 同样的概率分布 → 同样的采样分布。

用一句话概括：缓存只是把"这次算完就扔"的中间结果留了下来，计算路径和数值一个都没变，所以分布必然不变。

代码注释也点明了这一点（src/main.rs）：

println!("\n  （KV cache 只改计算方式、不改生成分布，两者应高度一致）");

演示里的"验证"其实是间接的：demo 用了两个不同的 prompt（"Once upon a" vs "The fox"）和同一个 rng 序列，所以两段输出文本不同是正常的。想严格验证"分布一致"，应该用相同 prompt + 相同 rng 种子分别跑 use_kv_cache=false 和 true，对比逐 token 输出是否逐位一致——这是动手练习 1。

8. 上下文达到 block_size 后停止生成

真实输出里能直接看到这个机制（第 16 课跑出来的）：

  —— 生成 2（temperature=0.8, top-k=10, top-p=0.9, 带 KV cache）——
  The fox the hidden garden, whe li

数一下："The fox" = 7 个 token，续写 the hidden garden, whe li = 26 个 token，输出共 7 + 26 = 33 个字符。逐迭代看缓存怎么涨的：

也就是说，生成到第 26 个新 token 后，下一次循环开头检查：

if use_kv_cache && cache[0].seq_len() >= block_size {
    break;
}

此时缓存（prompt 7 个 + 已前向的 25 个新 token）恰好等于 32 = block_size，直接跳出——所以 max_new=80 根本没跑完，输出戛然而止。注意最后采样的第 26 个 token 甚至没有参与前向、也没进缓存（它只是被采样并 push 进 ids，下一次循环就 break 了）。

为什么必须停？ 三个原因，都指向同一个根：

对比全量模式的生成 1：prompt "Once upon a" = 12 个 token，每步窗口都滑到最近 32 个，所以 80 个新 token 全部生成完（Once upon a to his friend, the wise old owl. One day, Red found a wold of colors and turned）。

真实 LLM 的 KV cache 比这复杂得多：支持"滑动窗口 + 丢弃最旧块"（如 Mistral 的 sliding window）、对缓存做量化压缩等。本项目的 KVCache 是最简版——只拼不丢，因此一旦填满就必须停止。把"丢了也能继续"留作动手练习 5。

9. 动手练习

1. 严格验证分布不变：在 demo_gpt 里用相同的 prompt（如都传 "The fox"）和相同的 rng 分别调 generate(..., false, ...) 与 generate(..., true, ...)，对比逐 token 输出是否一致。
2. 打印缓存形状：在 MultiHeadAttention::forward 的 cache.append 之后加一行 println!("cache seq_len = {}", cache.seq_len());，观察它从 7 一路涨到 32 的过程。
3. 把 break 条件去掉：临时注释掉 generate 里的 if use_kv_cache && cache[0].seq_len() >= block_size { break; }，运行看会发生什么——体会位置编码表 [32, D] 的下界约束。
4. 对比计算量：全量模式第 k 步前向 k 个位置、缓存模式每步只前向 1 个位置。对 block_size=32、max_new=80，估算两种模式累计前向的位置总数各是多少。
5. （进阶）给 KVCache 加"截断"：仿照全量模式的窗口滑动，给 KVCache 加一个 truncate(len) 方法（把 k.data() 裁到最近 len 个位置再包回张量），并在 generate 的缓存分支里每步调用它，让缓存模式也能像全量模式一样持续生成——对比改动前后的输出。

10. 本课总结

• 逐 token 生成时，历史位置的 K/V 每步都在被重复计算——全量模式累计 O(T²)，这是 KV Cache 要消灭的浪费

• KVCache = 每层一份的 (k, v) 张量（[1, T, D]），append 纯数据拼接、seq_len 读已缓存长度

• MultiHeadAttention 用缓存后只有 K/V 变长，Q 只算新位置，后续代码零改动；GPT::forward 用 base 修正位置编码与因果掩码

• 流程对比：首次前向整个 prompt 填缓存 → 之后每步只前向 1 个 token；全量模式则是每步重算整个窗口

• 分布不变的原因：缓存里的 K/V 与全量模式算出的数值相同，注意力、softmax 计算路径一致

• 缓存模式只拼不丢，上下文达到 block_size=32 必须停止（训练长度外是 RoPE 外推区 + 缓存无法截断历史），真实输出里生成 2 止步于 32 个 token

• 下一课：换掉正弦位置编码，用 RoPE（旋转位置编码）让位置信息融入注意力计算。

代码位置：src/optim.rs（SGD / AdamW） 代码位置：src/train.rs（train_gpt 中 AdamW 的用法） 演示入口：src/main.rs（演示 3：训练小 GPT）

1. 本课要搞懂的问题

1. 朴素 SGD 有什么致命弱点？为什么深度模型上它不够用？
2. Adam 的一阶动量 m、二阶动量 v 分别"记忆"了什么？对应公式怎么写？
3. 为什么需要偏差修正（除以 1-β^t）？前几步不修会怎样？
4. 权重衰减"解耦"（decoupled weight decay）和传统的 L2 正则化有什么区别？
5. AdamW::step 的每一行代码分别对应公式里的哪一项？

2. 先复习：SGD 的局限

第 6 课实现的 SGD 是整个故事的原点：

/// 更新一步：θ = θ - lr * g
pub fn step(&self) {
    for p in &self.params {
        let g = p.grad();
        let d = p.data();
        let updated: Vec<f32> = d.iter().zip(&g).map(|(v, g)| v - self.lr * g).collect();
        p.set_data(updated);
    }
}

公式只有一个：

θ ← θ - lr · g

它有两个致命弱点：

现实世界的梯度方向总是带噪声的（batch 采样引入的随机性、损失面本身的崎岖），只信"当前这一下"的 SGD 在深度网络上要么太慢、要么震荡发散。

3. Adam 的核心思想：一阶动量 + 二阶动量

Adam（Adaptive Moment Estimation）给每个参数维护两个额外状态，公式如下：

一阶动量（梯度的指数移动平均，记住"方向"）：
    m_t = β1 · m_{t-1} + (1 - β1) · g_t

二阶动量（梯度平方的指数移动平均，感知"陡峭程度"）：
    v_t = β2 · v_{t-1} + (1 - β2) · g_t²

参数更新：
    θ_t = θ_{t-1} - lr · m̂_t / (√v̂_t + ε)

直觉对应：

注意 m、v 都是逐参数、逐元素维护的（AdamW 里 m: Vec<Vec<f32>>、v: Vec<Vec<f32>>，和每个参数张量形状一一对应），所以"自适应步长"是精细到每个标量权重的。

为什么用指数移动平均而不是简单平均

因为只需要记住 m_{t-1} 和 v_{t-1} 两个状态就能增量更新，不需要存全部历史梯度；β1=0.9 意味着过去约 10 步的梯度主导，β2=0.999 意味着过去约 1000 步的梯度平方主导——越靠前的历史衰减得越厉害。

4. 偏差修正：除以 (1 - β^t)

问题：m_0 = 0、v_0 = 0，训练第一步 m_1 = (1-β1)·g_1，只有真实梯度的 10%（β1=0.9 时）。训练初期 m、v 被严重"低估"，直接使用会导致起步步长偏小。

修正办法：把 m_t、v_t 除以各自的累积衰减系数：

m̂_t = m_t / (1 - β1^t)
v̂_t = v_t / (1 - β2^t)

关键观察：t 越大，β^t 越接近 0，修正系数越接近 1——偏差修正只在训练初期起作用。而且 m 的 β1 小、修正结束得快，v 的 β2 接近 1、修正持续得更久。这正是代码里两个系数分开算的原因。

5. 权重衰减解耦（decoupled weight decay）

正则化思想：每步更新时，额外把参数往 0 拉一点点，防止权重过大、过拟合。

传统做法（L2 正则化）：把 λ·θ 加进损失再求导，相当于梯度变成 g + λ·θ，然后被 Adam 的"自适应步长"一通缩放——λ·θ 这个衰减项也被 1/√v̂ 缩放了，衰减强度随梯度历史变化，不受控制。

AdamW 的解耦做法：权重衰减独立于梯度，作为一个单独的减法项：

θ ← θ - lr · m̂/(√v̂ + ε)  -  lr · wd · θ
     └──── Adam 步长 ────┘   └─ 解耦的权重衰减 ─┘

在现代 LLM 训练里（GPT 系列等）几乎都用 AdamW，就是因为这个"可预期"的衰减。

6. AdamW::step 逐行讲解

完整代码（src/optim.rs）：

pub fn step(&mut self) {
    self.t += 1;
    // 偏差修正系数（训练初期 t 小，修正大）
    let bc1 = 1.0 - self.beta1.powi(self.t as i32);
    let bc2 = 1.0 - self.beta2.powi(self.t as i32);

    for (i, p) in self.params.iter().enumerate() {
        let g = p.grad();
        let d = p.data();
        let mut updated = vec![0.0f32; d.len()];
        for j in 0..d.len() {
            let gv = g[j];
            // 1. 更新动量
            self.m[i][j] = self.beta1 * self.m[i][j] + (1.0 - self.beta1) * gv;
            self.v[i][j] = self.beta2 * self.v[i][j] + (1.0 - self.beta2) * gv * gv;
            // 2. 偏差修正
            let m_hat = self.m[i][j] / bc1;
            let v_hat = self.v[i][j] / bc2;
            // 3. 更新：θ -= lr * m_hat/(√v_hat + eps) + lr * wd * θ（权重衰减解耦）
            let step = self.lr * m_hat / (v_hat.sqrt() + self.eps);
            let decay = self.lr * self.weight_decay * d[j];
            updated[j] = d[j] - step - decay;
        }
        p.set_data(updated);
    }
}

逐行对应公式：

四个值得停下来想一想的细节：

1. 为什么 v_hat.sqrt() 要加 eps：防止 v_hat ≈ 0（训练初期）时除零，数值稳定性用。eps = 1e-8 是 Adam 论文的标准值。
2. lr 是 pub 字段：src/train.rs 里每步都改它——

   opt.lr = scheduler.lr();   // 第 20 课：warmup + cosine 调度
   opt.step();
   

   调度器只负责改 lr，AdamW 的 m、v 状态跨步累积、完全不受影响。
3. 权重衰减项用的是 d[j]（更新前的旧参数）：这就是"解耦"——衰减直接作用在参数本身，而不是作用在梯度上。
4. m、v 与参数逐元素对齐：AdamW::new 里 params.iter().map(|p| vec![0.0f32; p.numel()])，每个参数张量配一个同长度的一阶/二阶动量数组，更新时按 j 同步遍历。

7. AdamW 的超参数

AdamW::new 里的默认值就是论文标准：

AdamW {
    lr,
    beta1: 0.9,
    beta2: 0.999,
    eps: 1e-8,
    weight_decay,
    t: 0,
    params,
    m,
    v,
}

本项目在 train_gpt 里这样接入：

let mut opt = AdamW::new(max_lr, params.clone(), 0.01);   // wd = 0.01
...
clip_grad_norm(&params, 1.0);   // 先裁剪梯度（防止个别大梯度冲坏 m、v 的估计）
opt.lr = scheduler.lr();        // 再设置当前学习率
opt.step();                     // 最后更新
opt.zero_grad();

顺序值得注意：梯度裁剪在 AdamW::step 之前。因为 Adam 的 m、v 是梯度历史的长期记忆，如果某步梯度爆炸没被裁掉，它会污染 m/v 很久。先裁剪、再进优化器，是 LLM 训练的标准顺序。

8. 动手练习

1. 对比 SGD vs AdamW：临时把 train_gpt 里的 AdamW::new(...) 换成 SGD::new(0.01, params.clone())（去掉 opt.lr = scheduler.lr() 那行，因为 SGD::lr 是私有字段），跑 600 步看 loss——体会 AdamW 在小语料上的收敛速度优势。
2. 手推第一步：假设某参数 θ=1.0、g=0.5、lr=0.001、wd=0.01，手算 t=1 时 m_hat、v_hat、step、decay，再在 AdamW::step 里加一行 println! 验证。
3. 看偏差修正的效果：把 bc1、bc2 改成恒为 1.0（不修正），训练对比 loss 曲线——训练初期应该明显变慢。
4. 调权重衰减：把 wd 从 0.01 改成 0.1 和 0.0，观察训练后 loss 与生成文本的差异，体会正则化的作用。
5. 思考：为什么 v 存的是 g² 而不是 |g|？如果某个维度的梯度长期是 +0.1/-0.1 交替（震荡），m 和 v 分别是什么表现？（提示：m 会被抵消趋近 0，v 会累积为正——这就是 Adam 抑制震荡的机制。）

9. 本课总结

• SGD 只有 θ -= lr·g，既没有惯性（震荡）也不能自适应步长（一刀切）

• Adam 用一阶动量 m = β1·m + (1-β1)·g 平滑方向，用二阶动量 v = β2·v + (1-β2)·g² 感知陡峭度，更新为 θ -= lr·m̂/(√v̂+ε)

• 偏差修正除以 1-β^t：只在训练初期起作用，补偿 m/v 从 0 起步的低估

• AdamW 把权重衰减解耦成独立项 θ -= lr·wd·θ，不受自适应步长干扰，衰减强度可预期

• AdamW::step 的每行代码都能在公式里找到对应项：t → 修正系数 → 逐元素动量 → 修正 → 步长/衰减 → 写回

• 下一课：推理加速神器 KV Cache——为什么生成时要缓存 K/V，怎么做到"每步只算 1 个 token"。

代码位置：src/main.rs（demo_gpt） 代码位置：src/train.rs（train_gpt / LRScheduler / clip_grad_norm） 代码位置：src/data.rs（CORPUS / DataLoader） 代码位置：src/sample.rs（generate / sample_token）

1. 本课要搞懂的问题

1. demo_gpt 从数据到生成文本，完整流程分哪几步？
2. 只有 669 个字符的小语料，训练日志里的 step / lr / loss 三列怎么读？
3. 日志里为什么看不到 warmup 段？lr 从 0.002947 一路衰减到 0.000300 是怎么来的？
4. temperature、top-k、top-p 三个参数是怎么配合采样的？
5. 为什么 loss 已经降到 0.16，模型输出的文本依然只是"像样"而不是"正确"？

2. 训练全景：demo_gpt 做了什么

src/main.rs 的演示 3（第 12-21 课）是本节的主角：

fn demo_gpt() {
    println!("=== 演示 3：训练小 GPT 并生成文本 ===");

    let mut rng = Rng::new(1234);
    let tokenizer = CharTokenizer::new(CORPUS);
    let vocab_size = tokenizer.vocab_size();
    println!("  语料 {} 字符，字符词表 {} 个", CORPUS.len(), vocab_size);

    let model = GPT::new(GPTConfig::tiny(vocab_size), &mut rng);

    // 训练（第 13、17、20 课：训练循环 + AdamW + warmup/cosine 调度）
    let loader = DataLoader::new(CORPUS, &tokenizer, model.cfg.block_size, 8);
    let tcfg = config::TrainConfig {
        seed: 42,
        batch_size: 8,
        steps: 600,
        max_lr: 3e-3,
        warmup_steps: 50,
        eval_every: 100,
        ..config::TrainConfig::default()
    };
    train::train_gpt(&model, &tokenizer, &loader, &tcfg, None, None, &mut rng);

    // 生成（无 cache）
    println!("\n  —— 生成 1（temperature=0.8, top-k=10, top-p=0.9, 无 KV cache）——");
    let out1 = generate(&model, &tokenizer, "Once upon a", 80, 0.8, 10, 0.9, false, &mut rng);
    println!("  {}", out1);

    // 生成（带 KV cache，第 18 课）
    println!("\n  —— 生成 2（temperature=0.8, top-k=10, top-p=0.9, 带 KV cache）——");
    let out2 = generate(&model, &tokenizer, "The fox", 80, 0.8, 10, 0.9, true, &mut rng);
    println!("  {}", out2);
    println!("\n  （KV cache 只改计算方式、不改生成分布，两者应高度一致）");
}

整个流程可以拆成 5 步：

注意：训练用的是字符级分词器，所以"1 个字符 = 1 个 token"，语料 669 个字符就是 669 个 token。这让后面的数字（32、80）可以直接按"字符数"理解。

3. 数据：669 字符的小语料

src/data.rs 里内置了一篇英文小故事（狐狸 Red 找金钥匙）：

pub const CORPUS: &str = "\
Once upon a time in a small village, there lived a curious little fox named Red. \
Every morning, Red would wake up early and explore the forest. ...";

训练数据是自监督的：输入 x 是一段 32 个 token 的序列，标签 y 是 x 右移一位——每个位置都预测"下一个字符是谁"，文本自己就是标签，不需要人工标注。

DataLoader::sample_batch 每次随机选 8 个起点，各截 33 个 token（前 32 个作 x，后 32 个作 y）：

pub fn sample_batch(&self, rng: &mut Rng) -> (Vec<usize>, Vec<usize>) {
    let max_start = self.tokens.len() - self.block_size - 1;
    let mut x = Vec::with_capacity(self.batch_size * self.block_size);
    let mut y = Vec::with_capacity(self.batch_size * self.block_size);
    for _ in 0..self.batch_size {
        let start = rng.choice(max_start);
        for j in 0..self.block_size {
            x.push(self.tokens[start + j]);
            y.push(self.tokens[start + j + 1]);
        }
    }
    (x, y)
}

关键点：

• 随机采样而非顺序扫描：每次 sample_batch 都在语料里随机挑起点。语料只有 669 token，但 600 步 × 8 个 batch 会反复"看到"语料的不同片段（有些片段会被重复看，有的可能一次都没被抽到）——小语料训练天然就是"背课文"。
• 返回的 x、y 都是 [B*T] = [8×32] = [256] 的展平数组，正好满足 GPT::forward(idx, b=8, t=32, None) 的输入要求（训练时 kv_cache 传 None）。

4. 超参数一览

train_gpt 的调用参数与 GPTConfig::tiny 汇总：

模型参数量：train_gpt 开头会打印一行"开始训练"（真实数字就在其中）：

开始训练：char（vocab=35）模型参数 102336 | 语料 669 tokens（训练 669 / 验证 0）| batch=8 block=32

按第 12 课的方法验证一下：词表 V=35（不是 100）时，tok_emb = 35×64 = 2240，每层 Block ≈ 49984，两层 ≈ 99968，ln_f = 128，输出头是权重绑定（复用 tok_emb 转置，无独立 lm_head 参数），总计 2240 + 99968 + 128 = 102336 ✓。约 10 万参数，CPU 上几秒就能跑完整个 demo。

5. 真实训练日志解读

运行 cargo run --release，演示 3 会打印（这是真实运行输出，不是编的）：

=== 演示 3：训练小 GPT 并生成文本 ===
  语料 669 字符，字符词表 35 个
开始训练：char（vocab=35）模型参数 102336 | 语料 669 tokens（训练 669 / 验证 0）| batch=8 block=32
step   100 | lr 0.002947 | train_loss 1.4597
step   200 | lr 0.002540 | train_loss 0.5498
step   300 | lr 0.001850 | train_loss 0.3338
step   400 | lr 0.001096 | train_loss 0.2312
step   500 | lr 0.000518 | train_loss 0.1557
step   600 | lr 0.000300 | train_loss 0.1624

5.1 三列日志分别是什么

train_gpt 里每步做 6 件事，日志打印在最后：

for step in 0..steps {
    let (x, y) = loader.sample_batch(rng);          // 1. 采样 batch
    let logits = model.forward(&x, b, t, None);     // 2. 前向
    let loss = cross_entropy_loss(&logits, &y);     //    算损失
    loss.backward();                                // 3. 反向
    clip_grad_norm(&params, 1.0);                   // 4. 梯度裁剪
    let cur_lr = scheduler.lr();                    // 5. 取当前步 lr 喂给优化器
    opt.lr = cur_lr;
    opt.step();
    opt.zero_grad();                                // 6. 清零梯度
    scheduler.step();                               //    步数 +1（为下一步准备 lr）
    let last = step + 1 == steps;
    if (step + 1) % cfg.eval_every == 0 || last {   // 每 eval_every 步（或最后一步）打印
        // lr 打印的就是本步实际用的 cur_lr，没有错位；step 打印 step + 1
        println!("step {:>5} | lr {:.6} | train_loss {:.4}", step + 1, cur_lr, loss.item());
    }
}

5.2 loss：1.46 → 0.16 说明了什么

• 第一个打印点 1.46：日志只在 step 100、200、… 打印（eval_every = 100）。随机初始化时模型对 35 个字符基本"一视同仁"，理论下界是均匀分布的交叉熵 ln(35) ≈ 3.56；训练 100 步后降到 1.46，说明已经开始学习。
• 先快后慢：step 100→300 loss 从 1.46 掉到 0.33（降了 77%），step 300→600 只从 0.33 掉到 0.16。这是训练曲线的典型形态——早期梯度大、方向明确，后期接近收敛、只能精雕细琢。
• 终点 0.16：交叉熵 0.16 意味着模型给"正确下一个字符"的平均概率约为 exp(-0.16) ≈ 0.85。对一篇 669 字符的"课文"来说，模型已经相当好地"背"下了其中的统计规律。

5.3 warmup 阶段：为什么日志里看不到

LRScheduler 的规则（src/train.rs）：

pub fn lr(&self) -> f32 {
    if self.step < self.warmup_steps {
        // 线性 warmup
        self.max_lr * (self.step as f32 + 1.0) / self.warmup_steps.max(1) as f32
    } else {
        // cosine 衰减
        let progress = (self.step - self.warmup_steps) as f32
            / (self.total_steps - self.warmup_steps).max(1) as f32;
        let progress = progress.min(1.0);
        let cosine = 0.5 * (1.0 + (std::f32::consts::PI * progress).cos());
        self.min_lr + (self.max_lr - self.min_lr) * cosine
    }
}

warmup 就是前 50 步让学习率线性爬升：

lr(step) = max_lr × (step + 1) / warmup_steps     （step < 50 时）

代入 max_lr = 0.003、warmup_steps = 50：

注意：demo 的 eval_every = 100，warmup 段（step 0-49）没有打印点，所以真实日志里看不到 0.00006 起步的爬升。 把 eval_every 改成 10，就能看到 step 10/20/30/40 的 lr = 0.0006 → 0.0012 → 0.0018 → 0.0024 （每步增加 0.003/50 = 0.00006，10 步就是 0.0006）。

为什么要 warmup？训练刚开始时参数是随机值，梯度方向噪声大、量级不可控。如果一上来就用 0.003 的大步长，很容易把参数"推飞"（loss 直接变成 NaN）。先用小步长稳住方向，再逐渐加力，是现代 LLM 训练的标准做法。

5.4 cosine 衰减：从峰值平滑降回 min_lr

第 50 步之后走 cosine 曲线，从 max_lr = 0.003 平滑降到 min_lr = 0.003 × 0.1 = 0.0003：

lr = min_lr + (max_lr - min_lr) × 0.5 × (1 + cos(π × progress))
progress = (step - 50) / (600 - 50)，超过 1 就截断到 1

验证日志里的两个数字：

• step 100：scheduler 计数 = 99，progress = (99-50)/550 ≈ 0.089，cosine ≈ 0.9805，lr = 0.0003 + 0.0027×0.9805 ≈ 0.002947 ✓
• step 600：scheduler 计数 = 599，progress = (599-50)/550 ≈ 0.998，cosine ≈ 0，lr ≈ min_lr = 0.000300 ✓

学习率全程曲线：

lr
│
0.003 ┤        ╭╮
      │       ╭╯ ╰╮
0.002 ┤      ╭╯    ╰╮
      │     ╭╯      ╰╮
0.001 ┤    ╭╯        ╰╮
      │   ╭╯          ╰╮
0.0003┤──╯             ╰────── (min_lr)
      └──┬────┬────┬────┬────→ step
         0   100  200  300  400  500  600
         └warmup(50步)┘└─── cosine 衰减 ───┘

后期的"小步慢走"是为了在 loss 接近收敛时不震荡、精细地落到更优的参数点。

6. 生成文本与采样参数

训练 600 步后调用 generate（src/sample.rs），参数 (prompt, max_new=80, temperature=0.8, top_k=10, top_p=0.9)。

sample_token 内部的 6 步采样管线：

真实生成结果（cargo run --release 原样输出）：

  —— 生成 1（temperature=0.8, top-k=10, top-p=0.9, 无 KV cache）——
  Once upon a time in a small village, there lived a curious little fox named Red. Every morn

  —— 生成 2（temperature=0.8, top-k=10, top-p=0.9, 带 KV cache）——
  The fox named all is s fend the g

（生成 2 用的是另一个 prompt "The fox"，且因缓存模式上下文达到 block_size=32 提前停止，第 18 课会专门讲；生成 1 在无缓存模式下把 80 个新字符完整生成完了。）

读这段输出：模型学会了故事的结构——"Once upon a time..." 开头、"in a small village, there lived a curious little fox named Red" 几乎完整复现语料原文、主谓宾、句号逗号。字面上"像样"，但仔细读全是毛病：生成 2 的 "named all is s fend the g" 语法不通、句子戛然而止（block_size 截断）。这就是下一节要回答的问题。

7. 为什么小模型输出只是"像样"而非"正确"

四个层面叠加，缺一不可：

用一句话总结："像样"来自学到了语料的高频统计规律；"不正确"来自语料/模型/训练都不足以学到完整语法，再加上采样本身的随机性。 想要更"正确"，方向是加大语料、加大模型、多训几步（后面第 19、20、21 课还会继续优化），但永远不可能在 669 字符上学出真正的英语——这也侧面说明了为什么现代 LLM 需要 TB 级数据和千亿参数。

8. 动手练习

1. 改种子观察差异：把 demo_gpt 里 Rng::new(1234) 改成别的数字（如 42），重新 cargo run --release。loss 曲线和生成文本都会变——思考：为什么损失曲线也会变？（提示：采样 batch 的随机起点变了）
2. 改 warmup：把 train_gpt 的 warmup_steps 从 50 改成 5 和 500，分别跑一次，对比前 100 步的 loss。体会"warmup 太短容易起飞、太长浪费步数"。
3. 改生成参数：把 generate 的 temperature 改成 0.2 和 1.5 各跑一次。观察文本变得更"死板/重复"还是更"发散/乱"。
4. 数 token：验证第 5.4 节——打印 scheduler.lr() 在 step 100、600 的计算过程，对照日志里的 0.002947 和 0.000300。
5. 思考：loss 从 1.46 降到 0.16，但为什么不能说"模型学会了英语"？模型"学会"的到底是什么？

9. 本课总结

• demo_gpt 五步走：分词 → 建模型 → 造数据 → train_gpt 训练 600 步 → generate 采样生成

• 数据是自监督的：x 是 32 个 token，y 是 x 右移一位，预测"下一个字符"

• 真实日志：loss 1.46 → 0.16，前 300 步降得最快；lr 从 0.002947 一路 cosine 衰减到 0.000300（warmup 段因 eval_every=100 没有打印点）

• 生成用 temperature=0.8 + top-k=10 + top-p=0.9：先缩放、再截断、再按概率随机抽样

• 小模型输出"像样而非正确"：语料太小、模型太小、训练不足、采样随机，四者叠加

• 下一课：换掉朴素的 SGD，给优化器装上"动量 + 自适应步长 + 权重衰减"——AdamW。

代码位置：src/sample.rs 模型接口：src/model.rs 演示入口：src/main.rs

1. 本课要搞懂的问题

1. 模型输出的 logits 是分数，怎么变成真正的文本？
2. 为什么"每次都选概率最大的 token"（argmax）效果很差？
3. temperature、top-k、top-p 分别解决了什么问题？它们在代码里怎么实现？
4. generate 生成一句文本时，每步发生了什么？KV cache 模式和全量模式有什么不同？

2. 生成的基本流程

训练完成后（第 13-14 课），模型已经学会了"给前缀，预测下一个 token"。生成文本就是把这件事反复做：

"Once upon a"
   │ 前向：model.forward(ctx, 1, T, ...)
   ▼
logits（每个 token 一个分数，未归一化）
   │ 采样：sample_token(logits, T=0.8, top_k=10, top_p=0.9, rng)
   ▼
下一个 token id
   │ 拼回上下文
   ▼
"Once upon a time" → 再前向 → 再采样 → …… 直到 max_new 个 token

关键在于中间那一步：logits 是分数，不是文本。怎么从分数里选 token，决定了生成质量。

3. 为什么不能直接 argmax

最朴素的想法：把 logits 做 softmax 变成概率，然后永远选概率最大的（这就是贪心解码 / argmax）。它的毛病很典型：

大模型的真实需求是"在合理的候选里随机一点"：既不能完全确定（呆板），也不能完全乱来（胡言乱语）。于是有了下面三个旋钮。

4. temperature：调节分布的"锐度"

softmax 的公式，引入温度 T 后：

pᵢ = exp(zᵢ / T) / Σⱼ exp(zⱼ / T)

代码（src/sample.rs sample_token 第 1 步）：

let scaled: Vec<f32> = logits
    .iter()
    .map(|&l| l / temperature.max(1e-5))   // 防止 T=0 除零
    .collect();

数值例子：logits = [1.0, 2.0, 0.5, 0.1]（词表 4 个 token）：

5. top-k：只在前 k 个里选

思想：分数排最后的那些 token 本来就是"凑数"的，干脆把它们从候选里删掉，只在前 k 个里分配概率。

代码（sample_token 第 2-3 步：先按分数从高到低排序，再截断）：

let mut items: Vec<(usize, f32)> = scaled.iter().enumerate().map(|(i, &v)| (i, v)).collect();
items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

if top_k > 0 && items.len() > top_k {
    items.truncate(top_k);   // top_k = 0 表示不启用
}

6. top-p（nucleus）：按累积概率截断

思想：不看"固定个数"，而是从高到低累加概率，直到累积概率达到 p，把后面的全部丢掉。候选集合大小随分布自动变化——分布集中时集合小，分布分散时集合大。

代码（sample_token 第 5 步，注意它作用在 softmax 之后的概率上）：

if top_p < 1.0 {
    let mut cum = 0.0;
    let mut keep = items.len();
    for (i, p) in probs.iter().enumerate() {
        cum += p;
        if cum >= top_p { keep = i + 1; break; }
    }
    items.truncate(keep);   // 丢掉尾部
    probs.truncate(keep);
    let s: f32 = probs.iter().sum();
    for p in probs.iter_mut() { *p /= s; }   // 重新归一化
}

沿用上面的例子（T=1.0，概率 [0.21, 0.57, 0.13, 0.09]，已按分数排序）：

top_p = 0.9：
  0.57  → 累计 0.57（< 0.9 继续）
  0.21  → 累计 0.78（< 0.9 继续）
  0.13  → 累计 0.91（≥ 0.9 停）→ 保留前 3 个，丢弃第 4 个

相比 top-k 的"一刀切固定个数"，top-p 更聪明：概率分布尖的时候只留 1-2 个候选，平的时候留一堆。实践中 top-k 和 top-p 经常一起开（本项目 main.rs 里就是 top_k=10, top_p=0.9 同时启用）。

7. 汇总：sample_token 的完整六步

src/sample.rs 的 sample_token 把上面所有招数串起来：

第 6 步的"轮盘赌"采样：

let mut u = rng.next_f32();          // 均匀随机数 [0, 1)
for (i, p) in probs.iter().enumerate() {
    if u < *p { return items[i].0; } // 落进第 i 段的概率 = pᵢ
    u -= p;                          // 没落进就减去这段，继续往后看
}
items.last().map(|(i, _)| *i).unwrap_or(0)

数学上：选到 token i 的概率正好等于 pᵢ。这保证了——概率高的 token 被选中的机会大，但概率低的也有机会被选中，这就是"随机性"的来源。

8. generate：一步一步生成整句文本

全量模式（无 KV cache）的流程，src/sample.rs：

let block_size = model.cfg.block_size;
let mut ids = tokenizer.encode(prompt);        // 起始文本 → id 序列
let mut cache = model.new_kv_cache();

for _ in 0..max_new {
    if use_kv_cache && cache[0].seq_len() >= block_size { break; }  // KV 模式的长度上限

    let start = ids.len().saturating_sub(block_size);   // 只保留最近 block_size 个
    let ctx = &ids[start..];

    let logits = if use_kv_cache {
        if cache[0].seq_len() == 0 {
            model.forward(ctx, 1, ctx.len(), Some(&mut cache))      // 首次：喂整个 prompt
        } else {
            model.forward(&ids[ids.len() - 1..], 1, 1, Some(&mut cache))  // 之后：只算最新 1 个
        }
    } else {
        model.forward(ctx, 1, ctx.len(), None)   // 全量：每次重算整个上下文
    };

    let v = model.cfg.vocab_size;
    let n = logits.numel();
    let last_row = &logits.data()[n - v..];      // 取最后一个位置的 logits
    let next = sample_token(last_row, temperature, top_k, top_p, rng);
    ids.push(next);                              // 拼回上下文，进入下一轮
}

tokenizer.decode(&ids)                           // id 序列 → 文本

逐步拆解：

1. 编码：tokenizer.encode(prompt) 把起始文本变成 id 序列。
2. 截断上下文：ctx = ids[ids.len() - block_size ..]，上下文超过 block_size 时只保留最近的一段——模型"记不住"更早的历史。
3. 前向拿 logits：全量模式下每次都把整个 ctx 重新算一遍（慢，但无需额外内存）。
4. 取最后一个位置：前向输出形状是 [B*T, vocab_size]，这里 B=1，最后一个位置即 data[n-v..]，它对应"基于当前全部上下文预测的下一个 token"。
5. 采样并拼接：sample_token 选出一个 id，ids.push(next)，进入下一轮。
6. 解码：循环结束后 tokenizer.decode(&ids) 把整条序列还原成文本（包含 prompt 和生成的部分）。

9. KV cache 模式的上下文限制：block_size 截断

KV cache（第 18 课详解）的思路：生成第 N 个 token 时，前 N-1 个 token 的 K/V 不需要重算，缓存起来每步只算最新的 1 个 token，大幅加速推理。但它带来一个限制，代码注释里写得很明白：

// KV cache 模式：上下文总长达到 block_size 就停（缓存无法像全量模式那样截断历史）
if use_kv_cache && cache[0].seq_len() >= block_size {
    break;
}

两种模式对"超长上下文"的处理对比：

所以 KV cache 在 max_new 还没用完时可能提前 break——这是"加速"换来的"长度上限"。main.rs 里两个生成演示（false 和 true）生成内容高度一致，正是"cache 只改计算方式、不改生成分布"的验证（提示词短时不会触发截断）。

10. 参数怎么配：一个经验表

main.rs 里的演示配置：temperature=0.8, top_k=10, top_p=0.9，是一个偏保守、够通顺的组合。

11. 动手练习

1. 把 generate 的 temperature 分别改成 0.1 和 2.0 跑一遍，观察输出从"复读机"到"胡言乱语"的变化。
2. 把 top_k 改成 1（等价于只在概率最高 token 附近贪心）再跑，观察重复现象。
3. 把 top_p 改成 1.0（关闭）但保留 top_k=10，对比输出差异。
4. 在 sample_token 的轮盘赌循环里打印每次的 u 和命中的 token，手动验证"命中概率 ≈ pᵢ"。
5. 用足够长的 prompt（超过 block_size，如 40 个字符）分别跑 use_kv_cache = false 和 true，观察 KV 模式提前停止生成的现象。

12. 本课总结

• 生成 = 前向拿 logits → 采样选 token → 拼回上下文 → 循环
• argmax 不可取：必然重复、呆板、无法回头
• 三个旋钮：temperature 调分布锐度、top-k 按个数截断、top-p 按累积概率截断，最后按概率随机采样
• sample_token 六步：缩放 → 排序 → top-k → softmax → top-p → 轮盘赌
• generate 全量模式每次重算上下文；KV cache 模式加速但受 block_size 长度限制，超限即停
• 下一步（第 16 课）：把所有零件拼起来，训练一个小 GPT 并生成有意义的长文本

代码位置：src/data.rs 训练入口：src/train.rs（第 13 课）

1. 本课要搞懂的问题

1. 语言模型的训练数据长什么样？一篇文章怎么变成一条条"题目"？
2. 什么是自监督学习？为什么训练 GPT 不需要人工标注？
3. 训练样本里的 x 和 y 是怎么配对的？为什么说 y 是 x 平移一位？
4. batch_size 和 block_size 到底是什么含义？代码里怎么体现？

2. 从文本到训练数据：三步走

训练 GPT 不需要人工标注，只需要大量纯文本。数据流水线是：

原始文本（CORPUS）
   │ ① 分词 tokenize（第 8 课：CharTokenizer.encode）
   ▼
token id 序列  [t₀, t₁, t₂, t₃, ...]
   │ ② 切窗口：每 block_size 个 token 一段
   ▼
输入窗口 x
   │ ③ 配答案：x 右移一位得到 y
   ▼
训练样本 (x, y)

本项目的语料是内置在 src/data.rs 里的一个英文小故事：

pub const CORPUS: &str = "\
Once upon a time in a small village, there lived a curious little fox named Red. \
... every adventure begins with a single step.";

3. 第一步：文本 → token id

模型只认识数字，所以先用第 8 课的 CharTokenizer（字符级分词）把文本变成 id 序列：

let tokenizer = CharTokenizer::new(CORPUS);   // 词表 = 语料中所有出现过的字符
let tokens = tokenizer.encode(text);          // "Once" -> [O, n, c, e] 的 id 序列

在本项目中每个字符就是一个 token（一个 usize），例如故事开头的 id 序列就是：

'O' 'n' 'c' 'e' ' ' 'u' 'p' 'o' 'n' ' ' 'a' ' ' 't' 'i' 'm' 'e' ' ' ...
 0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15   4  ...

4. 自监督：预测下一个 token

训练目标一句话：给一段前缀，让模型预测下一个 token 是什么。

输入：Once upon a ti
输出：m     ← 下一个字符（来自原文）

输入：Once upon a tim
输出：e     ← 下一个字符

输入：Once upon a time
输出：空格  ← 下一个字符

这为什么叫自监督？因为"标准答案"就藏在文本自己身上——预测"下一个 token"，标签就是原文里紧跟其后的那个 token，不需要任何人去标注。模型只要在海量文本上反复做这件事，就能学会语法、语义和世界知识。这也是 ChatGPT 这类大模型训练的第一阶段。

5. 第二步 + 第三步：x/y 配对（平移一位）

DataLoader::sample_batch 的核心逻辑（src/data.rs）：

pub fn sample_batch(&self, rng: &mut Rng) -> (Vec<usize>, Vec<usize>) {
    let max_start = self.tokens.len() - self.block_size - 1;
    let mut x = Vec::with_capacity(self.batch_size * self.block_size);
    let mut y = Vec::with_capacity(self.batch_size * self.block_size);
    for _ in 0..self.batch_size {
        let start = rng.choice(max_start);          // 随机选一个起点
        for j in 0..self.block_size {
            x.push(self.tokens[start + j]);         // 窗口 [start, start+T)
            y.push(self.tokens[start + j + 1]);     // 窗口 [start+1, start+T+1)
        }
    }
    (x, y)
}

关键就两行：

• x.push(tokens[start + j])：x 的第 j 个 token 是窗口里的第 j 个
• y.push(tokens[start + j + 1])：y 的第 j 个 token 是 x 的第 j 个 token 的下一个

所以 y 恰好是 x 向右平移一位，一一对应：

位置:  0   1   2   3   4   5  ...  T-1
x:   [t₀  t₁  t₂  t₃  t₄  t₅  ...  t_{T-1}]
y:   [t₁  t₂  t₃  t₄  t₅  t₆  ...  t_T    ]
       ↑   ↑   ↑
      x[0]的下一个 = y[0]，x[1]的下一个 = y[1]，……

6. 具体例子：窗口从语料中切出来

设 block_size = 16，随机起点 start = 0，故事开头 17 个字符为 Once upon a time（注意末尾还有个空格）。则这一个样本是：

x = tokens[0..16]  = O n c e ' ' u p o n ' ' a ' ' t i m
y = tokens[1..17]  = n c e ' ' u p o n ' ' a ' ' t i m e

模型在这个样本上要同时做 T=16 次"预测下一个 token"，cross_entropy_loss 对它们取平均（第 13 课）。

为什么同一段文本要重复利用？ 一个长度为 700+ 字符的语料，可以切成几百个互相重叠的窗口（起点 0、1、2、……都可以），一份语料就被"榨出"远超自身长度的训练样本。

7. batch_size 与 block_size：两个容易混淆的概念

它们合起来决定了一次前向的输入规模：

一次前向的输入张量：[B*T] 展平 → model.forward 里 reshape 成 [B, T, D]

其中 B = batch_size（样本数）、T = block_size（每个样本的 token 数）。在 train_gpt 里：

let (x, y) = loader.sample_batch(rng);
let b = batch_size;      // B = 8
let t = block_size;      // T = 32（GPTConfig::tiny 的默认值）
let logits = model.forward(&x, b, t, None);   // [B*T, vocab_size]

用一张图理解 B 和 T：

batch_size = 2, block_size = 4 时的一次采样：

         ┌───── 第 1 个样本（起点 start=0）─────┐
x = [ t0  t1  t2  t3 ]   y = [ t1  t2  t3  t4 ]
         ┌───── 第 2 个样本（起点 start=57）────┐
x = [ t57 t58 t59 t60 ]  y = [ t58 t59 t60 t61 ]
         └──────────────────────────────┘
         展平后 x 是长度 B*T = 8 的一维数组

两个起点之间没有关系——每个样本是独立从语料里随机切的，这正是"随机梯度下降"里"随机"二字的来源。

8. 一个隐藏的断言：语料不能太短

assert!(tokens.len() > block_size, "语料太短，无法切出完整序列");

sample_batch 里 max_start = tokens.len() - block_size - 1 还额外要求至少留 1 个 token 当 y 的尾巴（因为 y 比 x 多取一位）。如果语料比 block_size 还短，就切不出任何完整窗口，直接 panic 提示——这是第 1 课就强调的"出错早、出错明显"的防御性编程习惯。

9. 动手练习

1. 手动算一遍：语料有 700 个 token，block_size = 32，max_start 是多少？一共能切出多少个不同的起点？（答案：668）
2. 把 batch_size 从 8 改成 1 再训练，观察 loss 曲线的抖动幅度（batch 越小噪声越大）。
3. 把 block_size 改成 64（注意 GPTConfig 的 block_size 也要一起改）再训练，对比效果和训练速度。
4. 思考：为什么 y 要比 x 多取一个 token（tokens[start + block_size]）？如果只取 tokens[start..start+block_size] 当 y 会有什么问题？

10. 本课总结

• 数据流水线：文本 → token id → 切窗口 → 配答案
• 自监督：预测下一个 token，标签来自文本自身，无需人工标注
• 配对方式：y 是 x 右移一位，y[j] = x[j] 的下一个 token
• block_size 控制"看多长"，batch_size 控制"一次看几段"，两者相乘是一次前向的 token 总数
• sample_batch 随机选起点切重叠窗口，把一份语料复用出大量样本
• 下一步（第 15 课）：训练好模型后，如何从输出的 logits 里"采样"出通顺、不呆板的文本

代码位置：src/train.rs 演示入口：src/main.rs

1. 本课要搞懂的问题

1. 训练一个模型，代码上到底要做哪几步？顺序为什么不能乱？
2. loss.backward() 到底做了什么？为什么调用一次，所有参数的梯度就都有了？
3. 梯度爆炸是什么？clip_grad_norm 是怎么防止它的？
4. 为什么每步训练结束都要"清零梯度"？不清零会怎样？

2. 训练循环总览：六步走

训练 = 把"前向 → 算损失 → 反向 → 更新"这个动作反复执行成千上万次。 train_gpt 里每一次循环（一个 step）严格按下面的顺序执行：

顺序不能乱：前向必须发生在反向之前（没有前向就没有计算图）；更新必须发生在反向之后（没有梯度就没东西可更新）；清零必须放在最后（否则下一步的梯度会叠加在旧梯度上）。

对应代码（src/train.rs 中的 train_gpt，省略了日志打印）：

for step in 0..steps {
    // 1. 采样 batch
    let (x, y) = loader.sample_batch(rng);
    let b = batch_size;
    let t = block_size;

    // 2. 前向 + 损失
    let logits = model.forward(&x, b, t, None);
    let loss = cross_entropy_loss(&logits, &y);

    // 3. 反向
    loss.backward();

    // 4. 梯度裁剪
    clip_grad_norm(&params, 1.0);

    // 5. 更新参数（设置当前学习率）
    opt.lr = scheduler.lr();
    opt.step();

    // 6. 清零梯度
    opt.zero_grad();
    scheduler.step();
}

下面把这六步逐一拆开讲。

3. 准备阶段：参数、优化器、学习率调度器

进入循环前，train_gpt 先做了三件准备工作：

let params = model.parameters();                 // 收集模型所有参数
let mut opt = AdamW::new(max_lr, params.clone(), 0.01);   // 优化器（第 17 课详述）
let mut scheduler = LRScheduler::new(warmup_steps, steps, max_lr, max_lr * 0.1); // 学习率调度（第 20 课详述）

• parameters()：通过第 5 课实现的 Module trait，把模型里所有可训练张量（embedding、每层 Linear、LayerNorm 的权重与偏置）收进一个 Vec<Tensor>。优化器更新和梯度裁剪都靠这份清单。
• AdamW：负责"怎么更新参数"，内部为每个参数维护动量 m 和二阶动量 v（第 17 课专门讲）。
• LRScheduler：决定每一步用多大的学习率——前期 warmup 从 0 线性爬升，后期 cosine 衰减。本课只需知道它输出一个学习率、赋值给 opt.lr 即可。

4. 第一步：采样 batch

let (x, y) = loader.sample_batch(rng);

x 是一批"输入 token"，y 是"标准答案"（下一个 token）。它们怎么来的、batch_size 和 block_size 是什么意思，是第 14 课的主题，本课先把它当黑盒：每次调用就得到一批新的 (x, y)。

5. 第二步：前向 + 损失

let logits = model.forward(&x, b, t, None);          // [B*T, vocab_size]
let loss = cross_entropy_loss(&logits, &y);           // 标量

• model.forward 让数据走一遍整个 Transformer，输出 logits：每个位置"预测下一个 token"的分数（未归一化），形状 [B*T, vocab_size]。
• cross_entropy_loss（第 6 课已实现）比较预测与标准答案 y：

loss = -mean( log softmax(logits)[i, y[i]] )

模型给正确 token 的概率越高，loss 越小；loss 为 0 意味着预测完全正确。训练的目标就是把这个数压到最低。

6. 第三步：反向 —— loss.backward() 做了什么

这是自动微分（第 2 课）的核心。看 src/autograd.rs 里的实现：

pub fn backward(&self) {
    assert_eq!(self.rank(), 0, "backward() 只支持标量（0 维）输出");
    // 1. 先把输出自身的梯度设为 1（d loss / d loss = 1）
    self.grad.borrow_mut()[0] = 1.0;

    // 2. 迭代式 DFS 拓扑排序（避免递归栈溢出，见第 4 课 6.2 节）
    //    用显式栈模拟递归，得到"先子节点后父节点"的拓扑序
    let mut stack: Vec<(Tensor, usize)> = Vec::new();
    let mut visited: HashSet<usize> = HashSet::new();
    // ... 三色标记法遍历 ...

    // 3. 把拓扑序反过来：从 loss 开始，一层一层往回传梯度
    for t in order.iter().rev() {
        if let Some(b) = &t.backward { b(); }
    }
}

拆开看有三步：

1. 初始化：把 loss 自己的梯度设为 1.0。因为我们要算的是 ∂loss/∂θ，链式法则的起点就是 ∂loss/∂loss = 1。
2. 拓扑排序：前向时每个中间结果都记住了自己的 parents（谁算出了我）和 backward（如何把梯度传回给我的输入）。DFS 从 loss 出发往下走到所有叶子节点，得到"先依赖、后被依赖"的拓扑序。
3. 逆序传播：按拓扑序的反向（从 loss 往参数方向）依次调用每个节点的 backward 闭包。每个节点把自己的梯度 g 按链式法则乘上局部导数，累加到它的输入（包括参数）的梯度上。

loss ──► softmax ──► matmul ──► ... ──► Linear(W) ──► embedding
  g=1     累积到输入     累积到输入          累积到 W         累积到表

为什么是"累加"而不是"赋值"？因为一个参数会被计算图中很多地方共用（比如同一个权重矩阵被整批样本共享），每个分支的贡献都要加在一起，这才是真正的偏导数。

所以调用一次 loss.backward()，模型里所有参数（params 里每个张量）的 .grad() 就都拿到了 ∂loss/∂θ。

7. 第四步：梯度裁剪 clip_grad_norm

问题：梯度爆炸。 深层网络 + 长序列训练时，链式法则连乘会导致梯度呈指数级放大，一步更新就可能把参数"推飞"，loss 直接变成 NaN/无穷大。

对策： 在更新参数之前，先检查所有参数梯度的总范数；如果超过阈值 max_norm，就整体等比缩放，让方向不变、大小受控。

代码（src/train.rs）：

pub fn clip_grad_norm(params: &[Tensor], max_norm: f32) {
    // 1. 算总范数：所有梯度元素的平方和，再开根号
    let mut total = 0.0f32;
    for p in params {
        let g = p.grad();
        for &v in &g { total += v * v; }
    }
    let norm = total.sqrt();

    // 2. 超过阈值才裁剪：整体乘 scale = max_norm / norm
    if norm > max_norm {
        let scale = max_norm / norm;
        for p in params {
            let g = p.grad();
            let scaled: Vec<f32> = g.iter().map(|&v| v * scale).collect();
            p.grad_set(scaled);   // 覆盖写回（tensor.rs 里专为裁剪提供的 API）
        }
    }
}

对应公式：

norm = sqrt( Σ gᵢ² )
gᵢ' = gᵢ · min(1, max_norm / norm)

注意 max_norm 这里取 1.0，训练 LLM 时这是很常见的取值。裁剪不改变梯度方向，只是限制步长上限——这比粗暴地调低学习率更精准。

8. 第五步：更新参数 opt.step()

opt.lr = scheduler.lr();   // 这一步用多大的学习率，由调度器决定
opt.step();                // 按梯度更新所有参数

本课用的是 AdamW（第 17 课详述），核心更新规则（简写）：

m ← β₁·m + (1-β₁)·g           一阶动量（记住梯度方向）
v ← β₂·v + (1-β₂)·g²          二阶动量（感知坡度陡缓）
θ ← θ - lr·m̂/(√v̂ + ε) - lr·wd·θ

直觉：沿着梯度的反方向迈一小步，步子大小由学习率和"历史上梯度的大小"共同决定。与最朴素的 SGD（θ = θ - lr·g，第 6 课）相比，AdamW 对每个参数自适应地调整步长，训练更稳。

9. 第六步：清零梯度 opt.zero_grad()

opt.zero_grad();   // 内部就是对每个参数调用 p.zero_grad()

为什么必须清零？回顾第 6 节：反向传播是累加梯度。如果不清零，下一步前向+反向得到的梯度会叠加上一步的旧梯度，参数更新方向就被污染了，loss 会来回震荡甚至发散。

正确的时间线（注意第 3 步的累加发生在"同一步内、多次贡献之间"，第 6 步的清零发生在"不同步之间"）：

step 0:  梯度 = 0 → 前向 → 反向(累加) → 裁剪 → 更新 → 清零 → 梯度 = 0
step 1:  梯度 = 0 → 前向 → 反向(累加) → 裁剪 → 更新 → 清零 → 梯度 = 0
...

10. 看训练日志：loss 在下降

train_gpt 末尾每隔 eval_every 步打印一行（main.rs 里每 100 步一次）：

step     0 | lr 0.00006 | loss 3.4658
step   100 | lr 0.00138 | loss 1.8701
step   200 | lr 0.00228 | loss 1.4210
...
step   599 | lr 0.00030 | loss 0.8142

判读方法：

• loss 总体下降 → 模型在学习，六步循环工作正常；
• loss 出现 NaN/剧烈抖动 → 大概率梯度爆炸，检查 clip_grad_norm 的阈值或调低 max_lr；
• loss 降得很慢/不降 → 学习率太小，或数据/模型配置有问题（下一课会讲数据怎么来）。

11. 动手练习

1. 把 clip_grad_norm(&params, 1.0) 这一行注释掉再跑 cargo run，观察 loss 是否更容易出现剧烈波动（体会它到底在防什么）。
2. 把 max_norm 从 1.0 改成 0.01 再训练，对比收敛速度（阈值太小会"拖慢"训练）。
3. 试着删掉 opt.zero_grad()，观察 loss 曲线会发生什么，并解释原因（提示：梯度在跨步累加）。
4. 在循环里打印第 5 步更新前后某个参数 params[0].data()[0] 的变化，验证"更新确实发生在 step 之后"。

12. 本课总结

• 训练循环六步，顺序固定：采样 → 前向 → 损失 → 反向 → 裁剪 → 更新 → 清零（裁剪是防爆炸的保险，可视为第 3.5 步）
• loss.backward()：置初值 → DFS 拓扑排序 → 逆序沿链式法则累加梯度，一次调用搞定所有参数
• clip_grad_norm：总梯度范数超阈值就等比缩放，只限大小、不改方向
• opt.zero_grad()：跨步清零，防止梯度污染
• 下一步（第 14 课）：搞清楚循环第一步 sample_batch 返回的 x/y 到底是怎么从语料里构造出来的

代码位置：src/model.rs（GPTConfig / GPT / TransformerBlock） 代码位置：src/attention.rs（MultiHeadAttention / KVCache） 代码位置：src/layers.rs（Embedding / Linear / LayerNorm / gelu） 演示入口：src/main.rs（演示 3：训练小 GPT 并生成文本）

1. 本课要搞懂的问题

1. 一个完整的 GPT 由哪几大块组成？各自负责什么？
2. Transformer Block 内部的数据是怎么流动的？pre-norm 到底"pre"在哪？
3. 输入一串 token id，经过 GPT::forward 后形状怎么一步步变成 logits？
4. GPTConfig::tiny 里每个数字（64 / 4 / 2 / 32）分别代表什么？

2. GPT 架构总览

把前几课的积木全部拼起来，就是完整的 GPT：

              ┌───────────────────────────────┐
 token id     │  tok_emb: Embedding [V, D]    │  ← 每个 token 查表成向量
              └───────────────────────────────┘
              │   x = tok（位置信息不再相加：由 RoPE 在   │
              │   注意力内部旋转 Q/K 提供，见第 19 课）    │
              ┌───────────────────────────────┐
              │  blocks: N 层 Transformer      │  ← 注意力找相关性 + MLP 加工信息
              │  Block（本课第 3 节）           │
              └───────────────────────────────┘
              ┌───────────────────────────────┐
              │  ln_f: 最终 LayerNorm          │  ← 输出前再归一化一次
              └───────────────────────────────┘
              ┌───────────────────────────────┐
              │  lm_head: tok_emb 表转置       │  ← 权重绑定，映射成对每个词的打分
              └───────────────────────────────┘
                      logits [B*T, V]（预测下一个 token）

对应的 Rust 结构体（src/model.rs）：

/// 完整的 GPT 模型
pub struct GPT {
    pub cfg: GPTConfig,
    tok_emb: Embedding,          // 同时充当 lm_head（权重绑定）
    blocks: Vec<TransformerBlock>,
    ln_f: LayerNorm,
}

位置信息哪里去了？—— 第 11 课的正弦位置编码（pos_emb）在第 19 课被 RoPE 取代：不再向输入加位置向量，而是在每个注意力层内部对 Q/K 做旋转（见第 19 课与 MultiHeadAttention::forward）。所以 GPT 结构体里已经没有 pos_emb 字段了。

在 GPT::new 里把它们创建出来：

pub fn new(cfg: GPTConfig, rng: &mut Rng) -> Self {
    let n_embd = cfg.n_embd;
    let vocab_size = cfg.vocab_size;
    let blocks = (0..cfg.n_layer)
        .map(|_| TransformerBlock::new(&cfg, rng))
        .collect();
    GPT {
        cfg,
        tok_emb: Embedding::new(vocab_size, n_embd, rng),
        blocks,
        ln_f: LayerNorm::new(n_embd, 1e-5),
    }
}

注意 blocks 是 (0..cfg.n_layer).map(...).collect()——同一份代码复制 N 份，层数完全由配置决定，想加深网络只要改一个数字。

3. TransformerBlock 内部结构

3.1 pre-norm 结构

每个 Block 内部是"两个子层 + 各自残差"的 pre-norm 结构（第 11 课讲过 pre-norm 的好处）：

x ──► LN1 ──► MultiHeadAttention ──► (+残差) ──► LN2 ──► MLP(GELU) ──► (+残差) ──► 输出
      │                                ▲          │                        ▲
      └────────────────────────────────┘          └────────────────────────┘
            加回输入 x（残差连接）                         加回输入 x（残差连接）

对应的结构体与 forward：

/// Transformer Block（第 11 课）
///
/// 结构（GPT-2 风格，pre-norm）：
///   x -> LayerNorm -> Attention -> 残差 +
///   x -> LayerNorm -> MLP(GELU)  -> 残差 +
struct TransformerBlock {
    ln1: LayerNorm,
    attn: MultiHeadAttention,
    ln2: LayerNorm,
    mlp_linear1: Linear, // [D, 4D]
    mlp_linear2: Linear, // [4D, D]
}

impl TransformerBlock {
    fn forward(&self, x: &Tensor, mask: &Tensor, kv_cache: Option<&mut KVCache>) -> Tensor {
        // 注意力子层 + 残差连接
        let h = self.attn.forward(&self.ln1.forward(x), mask, kv_cache);
        let x = x.add(&h);
        // 前馈子层 + 残差连接
        let h = self.ln2.forward(&x);
        let h = gelu(&self.mlp_linear1.forward(&h));
        let h = self.mlp_linear2.forward(&h);
        x.add(&h)
    }
}

逐行拆解：

3.2 为什么 MLP 要"先升维再降维"

前馈子层 MLP(D → 4D → D) 是 Transformer 里唯一"逐位置"加工信息的部件（注意力负责"跨位置"交换信息，MLP 负责"在每个位置上"独立加工）：

MLP:  D → 4D → D
       │      │
   先放大   再压缩
  （特征空间更丰富，    （恢复原维度，方便残差相加）
   更容易学非线性）

4D 是 GPT-2 论文里的惯例比例。MLP 和注意力形成互补：注意力负责"找谁相关"，MLP 负责"想清楚该怎么表达"。

4. GPT::forward 数据流

4.1 输入与输出

pub fn forward(
    &self,
    idx: &[usize],          // [B*T] 展平的 token id
    b: usize,               // batch 大小
    t: usize,               // 序列长度
    mut kv_cache: Option<&mut Vec<KVCache>>,  // 推理缓存（第 18 课，训练时传 None）
) -> Tensor {

• 输入：idx 是一维数组，长度 b * t，里面是 token id（比如 [2, 5, 9, 3] 表示一条样本 4 个词）。
• 输出：logits，形状 [B*T, vocab_size]——每个位置对"下一个词是谁"的打分。

4.2 五步数据流

完整代码：

let d = self.cfg.n_embd;
assert_eq!(idx.len(), b * t, "输入 id 数量必须等于 b*t");

// 1. token embedding
let x = self.tok_emb.forward(idx).reshape(vec![b, t, d]);

// 2. 位置信息由 RoPE 提供（在注意力内部旋转 Q/K，见 MultiHeadAttention::forward）。
//    base = KV cache 模式下已缓存的位置数：新 token 的绝对位置 = base + 窗口内下标 j。
let base = kv_cache
    .as_ref()
    .map(|c| c.first().map(|k| k.seq_len()).unwrap_or(0))
    .unwrap_or(0);

// 3. 因果掩码：scores 形状 [B*H, T, T_total]，广播 mask [T, T_total]
let t_total = t + base;
let mut mask_data = vec![0.0f32; t * t_total];
for i in 0..t {
    for j in 0..t_total {
        if j > i + base {
            mask_data[i * t_total + j] = f32::NEG_INFINITY;
        }
    }
}
let mask = Tensor::from_vec(mask_data, vec![t, t_total]);

// 4. 逐层过 Transformer Block
let mut x = x;
for (i, block) in self.blocks.iter().enumerate() {
    let cache = kv_cache.as_mut().map(|c| &mut c[i]);
    x = block.forward(&x, &mask, cache, base);
}

// 5. 最终归一化 + 输出头（权重绑定：lm_head 复用 tok_emb.table 的转置）
let x = self.ln_f.forward(&x);
let x = x.reshape(vec![b * t, d]);
x.matmul(&self.tok_emb.table.transpose())

几个容易忽略的细节：

1. 位置信息来自 RoPE 而不是相加：第 11 课的做法是 x = tok + pos_emb（把正弦位置向量加进去）；第 19 课之后改为在注意力内部对 Q/K 做旋转（rotary_pair），GPT::forward 不再需要 pos_emb 表，只把 base 传给各层——KV cache 推理时，新 token 的绝对位置是 base + j（第 18 课）。
2. 因果掩码的构造：j > i + base 的位置设为 -inf。也就是说第 i 个 token 只能看到"它自己和它前面的"（含 KV cache 里的历史位置），未来位置在 softmax 后概率为 0——保证模型只能预测下一个词、不能偷看答案。
3. 权重绑定的输出头：tok_emb.table 是 [V, D]，它的转置 [D, V] 恰好可以把 [D] 向量打分成 [V] 个词的分数（"第 i 行 = 第 i 个词的嵌入"与当前向量做点积）。这与 GPT 的"输入输出共享词嵌入"做法一致，省掉了一份独立的 lm_head 参数（第 5.2 节参数量里会体现）。

4.3 一张形状变化总表

以 b=1, t=4, n_embd=64, vocab=100 为例：

全程只有两处形状变化：tok_emb 之后从 [B*T, D] 变成 [B, T, D]（为了层内处理），ln_f 之后从 [B, T, D] 变回 [B*T, D]（因为 Linear 把 3D 输入自动展平，输出再还原）。中间的注意力拆头/合头都在 Block 内部完成，外部看到的一直是 [B, T, D]。

5. GPTConfig::tiny 配置解读

/// 一个小配置，适合学习演示（其余字段与 Default 一致）
pub fn tiny(vocab_size: usize) -> Self {
    GPTConfig {
        vocab_size,
        ..Default::default()
    }
}

5.1 由配置推导出的关键数字

• 每个头处理多少维：head_dim = n_embd / n_head = 64 / 4 = 16（MultiHeadAttention::forward 里有断言 head_dim * n_head == d，配置必须能整除）。
• 每层的形状：注意力 4 个 Linear 都是 [64, 64]；MLP 是 [64, 256] 和 [256, 64]（4 倍升维）。
• 位置信息：RoPE 的 cos/sin 表按"位置 × 对偶下标"预计算（第 19 课），维度 block_size × (D/2) 的常数，不参与训练。

5.2 粗略参数量估算

设词表大小 = V，参数量约：

真实 GPT-2 Small 有 1.17 亿参数（n_embd=768、n_layer=12、n_head=12）。我们的 tiny 把它缩小了约 1000 倍，纯粹是为了能在普通 CPU 上几秒钟跑一步训练，把原理讲清楚。

6. 动手练习

1. 手推数据流：设 vocab=50, b=2, t=3，用 tiny 配置，写出 GPT::forward 里每一步张量的形状（从 idx 到 logits），对照第 4.3 节的表检查。
2. 改配置：自己加一个 GPTConfig::small，比如 n_embd=128, n_head=8, n_layer=4, block_size=64。注意 head_dim = 128/8 = 16 仍成立；再按 5.2 节的表估一下参数量。
3. 验证 logits 形状：在 main.rs 演示 3 里，model.forward(...) 之后加一行打印 logits.shape()，确认是 [B*T, V]。
4. 看 Block 内部分工：把 mlp_linear1 的维度改成 2 * cfg.n_embd（2 倍而不是 4 倍），训练看 loss 变化——体会 MLP 宽度对模型能力的影响。
5. 思考：我们的输出头就是 tok_emb.table 的转置（权重绑定）。为什么可以这样做？相比"独立的 lm_head"省了多少参数？（提示：[V, D] 和 [D, V] 互为转置，lm_head 的参数量 64×V+V 恰好被省掉了。）

7. 本课总结

• GPT 由 token embedding + N 层 Transformer Block + 最终 LayerNorm + 权重绑定输出头 四大部分组成（位置信息由第 19 课的 RoPE 在注意力内部提供，不再有独立的 pos_emb 表）

• Transformer Block 是 pre-norm 结构：LN → Attention → 残差，再 LN → MLP(GELU) → 残差；MLP 做 D → 4D → D 的升维降维

• GPT::forward 数据流：idx [B*T] → [B, T, D] → ... → [B*T, V]，只有进出输出头时形状变化，层内形状始终 [B, T, D]

• 因果掩码保证"只能看过去"，RoPE 保证"知道相对位置"

• 输出头权重绑定：x @ tok_emb.tableᵀ，与 GPT 的"输入输出共享词嵌入"一致

• tiny 配置：n_embd=64、n_head=4（head_dim=16）、n_layer=2、block_size=32，约 10.6 万参数，CPU 上几分钟就能跑一轮训练

• 下一课：写训练循环（前向 → 算损失 → 反向传播 → 更新参数），让这个 GPT 真的学会生成文本！

代码位置：src/model.rs（TransformerBlock 里的残差连接） 代码位置：src/layers.rs（LayerNorm） 演示入口：src/main.rs

⚠️ 历史说明：本课实现的正弦位置编码（sinusoidal_positions / pos_emb）已在第 19 课被 RoPE 取代—— 当前 src/model.rs 里不再有 pos_emb 字段，位置信息由注意力内部旋转 Q/K 提供。 本节仍是必读的：LayerNorm 与残差连接至今原样在模型里，且"位置编码要解决什么问题"是理解 RoPE 的前提。

1. 本课要搞懂的问题

1. 注意力机制天生"看不见位置"，为什么？怎么补救？
2. 正弦位置编码的公式 PE(pos, 2i) = sin(pos / 10000^(2i/d)) 到底在算什么？
3. LayerNorm 是怎么把数据"拉回"标准分布的？它和 BatchNorm 有什么区别？
4. 残差连接为什么能让我们放心地把网络加深？

2. 为什么需要位置编码

先回忆第 9-10 课的注意力公式：

scores = (Q · Kᵀ) / √d_k      # Q、K 由每个 token 的内容算出来
attn   = softmax(scores)
out    = attn · V

注意：Q、K、V 全部只由 token 的内容（embedding）计算而来。点积 q·k 衡量的是"两个 token 的内容有多相关"， 跟"它们相隔多远、谁在前谁在后"没有任何关系。

这带来一个致命问题：注意力是"排列不变"（permutation invariant）的。看一个例子：

两个句子语义完全相反，但注意力打分一模一样（因为每个词的内容没变）。如果模型只看注意力， 它永远分不清"猫追狗"和"狗追猫"。

通俗理解：注意力是个"走神的学生"，它只知道"这两个词好像有关"，但不知道"谁在左、谁在右"。 语言是顺序敏感的，所以必须把位置信息塞进去。

解决方案：给每个位置的 token embedding 加上一个位置向量：

x = token_embedding + position_embedding

位置向量就是"第 0 个位置长什么样、第 1 个位置长什么样……"的一组向量。加上之后，同一个词在不同位置就有了不同的表示。

3. 正弦位置编码

3.1 公式

原版 Transformer 论文（Attention Is All You Need）提出用不同频率的正弦波来编码位置：

PE(pos, 2i)   = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

• pos：位置（0, 1, 2, ...）
• i：embedding 里的维度下标（0, 1, 2, ..., d-1）
• d：embedding 维度（我们的模型里就是 n_embd）

位置 pos 的向量长这样：

pos_emb[pos] = [ sin(pos/1), cos(pos/1), sin(pos/10000^(2/d)), cos(pos/10000^(2/d)), ... ]
                  └── i=0 ──┘  └── i=1 ──┘      └── i=2 ──┘       └── i=3 ──┘

关键点：维度下标 i 越小，频率越高（波形越密）；i 越大，频率越低（波形越平缓）。 也就是说，编码向量的"前半段"负责记录精细的相对位置，"后半段"负责记录大致的绝对位置。

3.2 手算一个例子

设 d = 4，即每个位置只有 4 个维度（i = 0, 1, 2, 3）：

可以看到：i = 0/1 的波形变化剧烈（区分相邻位置），i = 2/3 的波形几乎不动（区分远距离）。

3.3 为什么选正弦波

3.4 对应代码

⚠️ 以下函数已从代码库中删除（被 RoPE 取代），仅供理解原理。

/// 正弦位置编码（第 11 课，已删除）
///
/// PE(pos, 2i)   = sin(pos / 10000^(2i/D))
/// PE(pos, 2i+1) = cos(pos / 10000^(2i/D))
fn sinusoidal_positions(max_len: usize, d: usize) -> Vec<f32> {
    let mut data = vec![0.0f32; max_len * d];
    for pos in 0..max_len {
        for i in 0..d {
            let freq = 10000f32.powf((2 * (i / 2)) as f32 / d as f32);
            let angle = pos as f32 / freq;
            data[pos * d + i] = if i % 2 == 0 { angle.sin() } else { angle.cos() };
        }
    }
    data
}

逐行对照公式：

在 GPT::new 里把它变成一个 [block_size, n_embd] 的张量存起来：

let pos_emb = Tensor::from_vec(
    sinusoidal_positions(cfg.block_size, cfg.n_embd),
    vec![cfg.block_size, cfg.n_embd],
);

GPT 结构体里它的类型是普通 Tensor 而不是 Tensor::param(...)：

pos_emb: Tensor, // 正弦位置编码 [block_size, D]（常数，不参与训练）

它不在 parameters() 里，反向传播时不会产生梯度——位置编码是常数，这正是正弦编码对比"可学习位置编码"的一大优点。

3.5 怎么用：加到 token embedding 上

在 GPT::forward 里，用 gather_rows 按位置取行，再和 token embedding 相加：

// 1. token embedding
let tok = self.tok_emb.forward(idx).reshape(vec![b, t, d]);
// 2. 位置编码：KV cache 推理时，当前位置从缓存长度开始
let pos_emb = self.pos_emb.gather_rows(&positions).reshape(vec![b, t, d]);
let x = tok.add(&pos_emb);

x 的每个元素 = 词的语义（token embedding）+ 词的位置（pos embedding），后续所有层都在这个"带位置信息"的表示上工作。

4. LayerNorm：层归一化

4.1 公式

LayerNorm 对最后一维（每个样本自己）做归一化，再缩放平移：

y = (x - μ) / √(σ² + ε) · γ + β

• μ：x 在最后一维上的均值
• σ²：x 在最后一维上的方差
• ε：一个极小的数（防除零，我们的代码里用 1e-5）
• γ（gamma）、β（beta）：可学习的参数，初始 γ=1、β=0

直觉：先把每个样本拉成"均值 0、方差 1"的标准分布（稳定数值范围），再让模型自己决定要不要缩放平移回去（恢复表达能力）。

4.2 对应代码逐行拆解

src/layers.rs 里的实现：

pub struct LayerNorm {
    pub gamma: Tensor, // [d] 可学习缩放
    pub beta: Tensor,  // [d] 可学习平移
    pub eps: f32,
}

impl LayerNorm {
    pub fn new(d: usize, eps: f32) -> Self {
        LayerNorm {
            gamma: Tensor::param(vec![1.0; d], vec![d]),
            beta: Tensor::param(vec![0.0; d], vec![d]),
            eps,
        }
    }

    pub fn forward(&self, x: &Tensor) -> Tensor {
        let d = x.shape()[x.rank() - 1];
        // 均值 μ：[..., 1]（保持维度，方便广播）
        let mean = x.sum_last_dim().mul_scalar(1.0 / d as f32);
        // 中心化
        let centered = x.sub(&mean);
        // 方差 σ²：中心化后平方再取均值
        let var = centered
            .mul(&centered)
            .sum_last_dim()
            .mul_scalar(1.0 / d as f32);
        // 归一化
        let norm = centered.div(&var.add_scalar(self.eps).sqrt());
        // 缩放平移（γ、β 是 [d]，广播到 [..., d]）
        norm.mul(&self.gamma).add(&self.beta)
    }
}

公式 ↔ 代码对照表：

注意一个细节：均值、方差是用 sum_last_dim 求的，所以归一化是逐样本、逐序列位置独立进行的—— 对输入 [B, T, D] 来说，就是沿着最后一维 D 归一把每个 B×T 的"行"。

4.3 LayerNorm vs BatchNorm

记忆口诀：LayerNorm 归一化"每一行"，BatchNorm 归一化"每一列"。 GPT 选 LayerNorm 还有一个重要原因：它不引入"batch 内样本之间的相互依赖"，行为可预测、实现简单（公式就 5 步）。

4.4 在模型里用在哪

TransformerBlock 里有两个（ln1 在注意力前、ln2 在 MLP 前），模型末尾还有一个 ln_f：

// TransformerBlock::new
ln1: LayerNorm::new(cfg.n_embd, 1e-5),
ln2: LayerNorm::new(cfg.n_embd, 1e-5),
// GPT::new
ln_f: LayerNorm::new(n_embd, 1e-5),

三个归一化层都在层的最前面（pre-norm 结构，下一课细讲），eps 统一用 1e-5。

5. 残差连接（Residual Connection）

5.1 公式与直觉

残差连接就是在子层的输出上加上输入本身：

y = x + F(x)

其中 F(x) 是任意子层（注意力、MLP……）。直觉上，F(x) 只负责"学一个增量/修正"，而 x 这条"高速公路"原封不动地把信息送到下一层。

5.2 为什么解决梯度消失

先回忆反向传播：梯度从 loss 一层层往回传，每过一层就要乘一次该层的导数。如果网络有 N 层，梯度大概要被乘 N 次：

∂loss/∂x₀ = ∂loss/∂xₙ · ∂xₙ/∂xₙ₋₁ · ... · ∂x₁/∂x₀

如果每层的导数都小于 1（比如常见激活函数的导数 ≤ 1），乘几十层后梯度指数级衰减，几乎变成 0——前面的层根本学不到东西，这就是"梯度消失"。

有了残差连接，情况完全不同。y = x + F(x) 对 x 求导：

∂y/∂x = 1 + ∂F/∂x

• 多出来的那个 1 是恒等映射的导数，永远在。
• 反向传播时梯度传过残差块，至少保留一份"原样拷贝"：

∂loss/∂x = ∂loss/∂y · (1 + ∂F/∂x) = ∂loss/∂y + ∂loss/∂y · ∂F/∂x

不管 F 的导数多小，梯度里始终有一个 ∂loss/∂y 直接传回上一层，不会衰减成 0。

5.3 为什么能加深网络

• 对梯度：恒等映射让深层的梯度能"直达"浅层，解决了训练困难。
• 对信息：每一层都能"无损地"跳过不需要的变换。即使某个子层学不到有用的东西，模型也可以把它当成恒等映射（F(x) ≈ 0），至少不劣化，所以我们可以放心堆很多层。

一句话：残差连接 = 梯度的高速公路 + 信息的保险丝。

5.4 对应代码

TransformerBlock::forward 里残差连接就两处 add：

fn forward(&self, x: &Tensor, mask: &Tensor, kv_cache: Option<&mut KVCache>) -> Tensor {
    // 注意力子层 + 残差连接
    let h = self.attn.forward(&self.ln1.forward(x), mask, kv_cache);
    let x = x.add(&h);   // ← 残差连接 ①：x + Attention(LN(x))
    // 前馈子层 + 残差连接
    let h = self.ln2.forward(&x);
    let h = gelu(&self.mlp_linear1.forward(&h));
    let h = self.mlp_linear2.forward(&h);
    x.add(&h)            // ← 残差连接 ②：x + MLP(LN(x))
}

5.5 和 LayerNorm 的组合：pre-norm

注意到顺序是 先 LayerNorm、再子层、最后加残差（LN → 子层 → +），这叫 pre-norm（归一化在子层前）。对比早期 Transformer 的 post-norm（子层 → + → LN）：

6. 三者如何协作

一个 Transformer Block 里，位置编码负责"把位置信息喂进来"，LayerNorm 负责"稳住数值分布"，残差连接负责"让梯度和信息畅通无阻"：

x (带位置编码的 embedding)
 │
 ├─ 残差①： x + Attention( LN1(x) )        ← LN 归一化 → 注意力找相关性 → 加回原 x
 │
 └─ 残差②： x + MLP( LN2(x) )              ← LN 归一化 → MLP 加工信息 → 加回原 x
 │
 └─ 送到下一层

7. 动手练习

1. 手算 LayerNorm：对向量 x = [1.0, 2.0, 3.0, 4.0]，手算 μ、σ²、（σ²+ε）开方后的归一化结果（ε=1e-5），再用 src/layers.rs 里的 LayerNorm::new(4, 1e-5)（γ=1、β=0 时）跑一遍对比，应该完全一致。
2. 改频率基数：把 sinusoidal_positions 里的 10000 改成 10 或 1000000，想想会怎样？（提示：10 让所有频率都变高、位置区分更"拥挤"；1000000 让波形几乎不动。跑训练对比 loss 曲线。）
3. 去掉残差：把 TransformerBlock::forward 里的 x.add(&h) 改成直接 h（两个子层都改），跑训练观察 loss 是否下降明显变慢——亲身感受梯度消失。
4. 思考：位置编码为什么是"加"而不是"拼接"（concat）？（提示：加法让每个维度只"混入"一个位置信号，且不改变维度数；拼接会让维度翻倍、参数爆炸。）
5. 思考：LayerNorm 里为什么用 sum_last_dim 求均值而不是 sum？如果 sum 把所有元素加一起会怎样？

8. 本课总结

• 注意力是位置盲的，必须把位置信息加进输入 → 正弦位置编码 PE(pos,2i)=sin(pos/10000^(2i/d))

• 正弦编码用不同频率的波形区分位置，是常数（不参与训练），实现就一个双重循环

• LayerNorm 对最后一维做 y=(x-μ)/√(σ²+ε)·γ+β，用 sum_last_dim 求均值方差，γ/β 可学习

• LayerNorm 与 BatchNorm 的核心区别：归一行 vs 归一系列，NLP 用前者

• 残差连接 y = x + F(x) 让梯度里永远保留一份 ∂loss/∂y，解决梯度消失、支持深网络

• pre-norm（LN → 子层 → 残差）是 GPT 系列的标准结构

• 下一课：把这些积木全部拼起来，组成完整的 GPT 模型！

• （本课的正弦位置编码在第 19 课被 RoPE 取代，见 docs/19-RoPE旋转位置编码.md；LayerNorm 与残差连接保留至今。）

代码位置：src/attention.rs（MultiHeadAttention 结构 + forward） 底层算子：src/tensor.rs（reshape / permute / matmul / softmax_last_dim）

1. 本课要搞懂的问题

1. 为什么"一个头"不够？多个头到底多学到了什么？
2. 拆头为什么是 reshape → permute → reshape 三步，不能一步到位？
3. 合并头之后为什么还要一个输出投影 c_proj？
4. MultiHeadAttention::forward 的 8 步，每一步形状怎么变？

2. 动机：一个头只能学一种"相关性"

第 9 课的注意力，一个查询位置对所有键算一组权重——它只能表达一种"谁和谁相关"的模式。

但文本里的相关性是多种多样的：

• 相邻词之间的语法依存（"红 的 苹果"）
• 指代关系（"它" 指谁？）
• 全局语义（某个主题词反复出现）

一个头顾不过来。多头的思路：把 D 维空间切成 H 份，每份独立算注意力，让不同头学到不同模式，最后再合并。

数学形式：

MultiHead(Q, K, V) = Concat(head_1, ..., head_H) · W_O
head_i = Attention(Q·W_i^Q, K·W_i^K, V·W_i^V)

3. 结构一览

struct MultiHeadAttention {
    c_q: Linear,      // [D, D]
    c_k: Linear,      // [D, D]
    c_v: Linear,      // [D, D]
    c_proj: Linear,   // [D, D] 输出投影
    n_head: usize,
}

数据流：

x → c_q / c_k / c_v 投影 → 拆头 → 每头各自算注意力 → 合并 → c_proj → 输出

n_head 来自配置 GPTConfig::tiny：n_embd = 64，n_head = 4 → head_dim = 64 / 4 = 16。

4. 拆头：reshape + permute + reshape

// 3. 拆头：[B, T, D] -> [B*H, T, head_dim]
//    （先 reshape 出 H 维，再 permute 把 H 提到第 2 维）
let q = q
    .reshape(vec![b, t, self.n_head, head_dim])   // [B,T,D] -> [B,T,H,head_dim]
    .permute(&[0, 2, 1, 3])                       // -> [B,H,T,head_dim]
    .reshape(vec![b * self.n_head, t, head_dim]); // -> [B*H,T,head_dim]

为什么不能直接 reshape 成 [B*H, T, head_dim]？

关键在内存布局。以 D=8、H=2、head_dim=4、单个 token 为例，x 最后一维的数据排列是：

[ 头0的4维 | 头1的4维 ]     ← 头 0 在前、头 1 在后（行优先存储）

• reshape([B, T, H, head_dim])：只是重新解释数据，头 0 / 头 1 依然按块排列 ✓
• 但如果跳过 permute 直接 reshape 成 [B*H, T, head_dim]，B 维会按 b0-t0-h0, b0-t0-h1, b0-t1-h0, ... 的顺序切片——同一个 token 的多个头被混进 T 维，完全错乱 ✗

所以必须先 permute(&[0, 2, 1, 3]) 把"头"维提到第 2 维（H 排在 T 前面），让每个 (b, h) 的数据在内存里连成一块，再 reshape 成 [B*H, T, head_dim]，才能保证每个"批"恰好是一个头。

permute 和 reshape 都不搬动数据，只改变"维度解释"；三步合起来的效果 = "把 (b, h) 提成独立的批"。 这正是第 1 课讲的行优先存储的实战应用。

K、V 的拆头

K、V 的拆法完全一样（注意这里的 T 变成了 t_total——带 KV cache 时更长，第 18 课讲解）：

let k = k
    .reshape(vec![b, t_total, self.n_head, head_dim])
    .permute(&[0, 2, 1, 3])
    .reshape(vec![b * self.n_head, t_total, head_dim]);
// v 同理

为什么"先整体投影再拆头"是等价的？

c_q 是 [D, D] 线性层。对头 h 来说，它取输出的第 h 块（head_dim 维），这正好是输入 x 经过 W_Q 第 h 块列的作用——数学上等价于每个头一个独立的投影 W_h^Q（GPT-2 的实现也是这个思路，只是把投影合并写在一起）。

5. 并行注意力：一个批量 matmul 算完所有头

拆头后 [B*H, T, head_dim] 恰好是 matmul 支持的 3D 批量形状，第 9 课的公式一行不多、一行不少：

// 4. 注意力分数：scores = Q·Kᵀ / √d_k
let scale = 1.0 / (head_dim as f32).sqrt();
let kt = k.permute(&[0, 2, 1]);                  // [B*H, head_dim, T_total]
let scores = q.matmul(&kt).mul_scalar(scale);    // [B*H, T, T_total]

// 5. 因果掩码：把"未来位置"变成 -inf，softmax 后概率为 0
let scores = scores.add(mask);

// 6. softmax 得到注意力权重，加权求和
let attn = scores.softmax_last_dim();            // [B*H, T, T_total]
let out = attn.matmul(&v);                       // [B*H, T, head_dim]

H 个头互不干扰，只是被"压扁"进了批维——这是拆头最大的好处：一次批量运算同时算完 H 个独立的注意力。

6. 合并头：拆头的逆操作

// 7. 合并头回 [B, T, D]
let out = out
    .reshape(vec![b, self.n_head, t, head_dim]) // [B*H,T,head_dim] -> [B,H,T,head_dim]
    .permute(&[0, 2, 1, 3])                     // -> [B,T,H,head_dim]
    .reshape(vec![b, t, d]);                    // -> [B,T,D]（H×head_dim 拼回 D）

注意合并的顺序与拆头镜像对称：

拼回后，每个 (b, t) 的 D 维 = [头0的16维, 头1的16维, 头2的16维, 头3的16维]，与拆头前的布局完全一致——因为 matmul 的批维上每个 (b, h) 都是独立的，头之间从不串扰。

7. 输出投影 c_proj：让模型决定"怎么混"

// 8. 输出投影
self.c_proj.forward(&out)

• 合并头只是"物理拼接"，各头之间还没有交互
• c_proj 是一个 [D, D] 线性层：学习如何把 H 个头的信息融合（线性混合、重组）
• 至此一个完整的多头注意力子层结束，输出形状 [B, T, D] 与输入一致，方便后面接残差连接（第 11 课）

数学上：Concat(head_1..head_H) · W_O，W_O 就是 c_proj 的权重。

8. forward 全流程：8 步形状对照表

以 GPTConfig::tiny（D=64、H=4、head_dim=16）为例，B / T 视输入而定（训练演示时 B=8、T=block_size=32）：

完整的 forward 骨架（略去 KV cache 分支，逻辑与源码一致）：

fn forward(&self, x: &Tensor, mask: &Tensor, kv_cache: Option<&mut KVCache>) -> Tensor {
    let (b, t, d) = (x.shape()[0], x.shape()[1], x.shape()[2]);
    let head_dim = d / self.n_head;

    // ① 投影得到 Q、K、V
    let q = self.c_q.forward(x).reshape(vec![b, t, d]);
    let k = self.c_k.forward(x).reshape(vec![b, t, d]);
    let v = self.c_v.forward(x).reshape(vec![b, t, d]);

    // ③ 拆头
    let q = q.reshape(vec![b, t, self.n_head, head_dim])
             .permute(&[0, 2, 1, 3])
             .reshape(vec![b * self.n_head, t, head_dim]);
    let k = k.reshape(vec![b, t_total, self.n_head, head_dim])
             .permute(&[0, 2, 1, 3])
             .reshape(vec![b * self.n_head, t_total, head_dim]);
    let v = v.reshape(vec![b, t_total, self.n_head, head_dim])
             .permute(&[0, 2, 1, 3])
             .reshape(vec![b * self.n_head, t_total, head_dim]);

    // ④⑤⑥ 注意力（第 9 课）
    let scale = 1.0 / (head_dim as f32).sqrt();
    let kt = k.permute(&[0, 2, 1]);
    let scores = q.matmul(&kt).mul_scalar(scale).add(mask);
    let attn = scores.softmax_last_dim();
    let out = attn.matmul(&v);

    // ⑦ 合并头，⑧ 输出投影
    let out = out.reshape(vec![b, self.n_head, t, head_dim])
                 .permute(&[0, 2, 1, 3])
                 .reshape(vec![b, t, d]);
    self.c_proj.forward(&out)
}

9. 参数：训练时谁在被更新？

Module for MultiHeadAttention 负责收集全部可学习参数：

impl Module for MultiHeadAttention {
    fn parameters(&self) -> Vec<Tensor> {
        let mut ps = self.c_q.parameters();   // weight [64,64] + bias [64]
        ps.extend(self.c_k.parameters());
        ps.extend(self.c_v.parameters());
        ps.extend(self.c_proj.parameters());  // 4 个 Linear = 8 个参数张量
        ps
    }
}

前向算出的误差通过反向传播（第 2 课）一路传回这 8 个张量，4 个投影层的权重随训练不断调整——"学什么相关性、怎么融合"全都由数据决定。

forward 第 ② 步的 KV cache 只是推理时缓存历史 K/V 的加速手段，不改变计算结果，第 18 课专门讲解。

10. 运行与测试

cargo test   # tensor 的 test_permute / test_matmul_3d 验证拆头、批量 matmul 所需算子
cargo run    # 演示 3 训练小 GPT：4 头注意力参与每一轮前向/反向

11. 动手练习

1. 把 GPTConfig::tiny 的 n_head 从 4 改成 8（head_dim = 8）重新训练，观察 loss 是否变化（代码有 assert_eq!(head_dim * self.n_head, d) 保证必须能整除）。
2. 手动构造一个 [2, 3, 8] 的 Tensor，按 reshape([2,3,2,4]) → permute([0,2,1,3]) → reshape([4,3,4]) 手写一遍，画出每个元素的归属，验证第 4 节的内存布局分析。
3. 修改代码：去掉第 ⑧ 步 c_proj，直接返回合并后的 out，重新训练对比 loss（思考：表达能力损失在哪里）。
4. 阅读 forward 中 kv_cache 分支的 cache.append，解释为什么推理时它能省掉"历史 token 的重复计算"（第 18 课预告）。

12. 本课总结

• 多头 = 把 D 切成 H 份，H 个注意力并行，学不同的相关性模式
• 拆头三步曲：reshape（切块）→ permute（H 提前）→ reshape（压平 B·H），顺序不能乱
• 合并是拆头的镜像；c_proj 负责融合各头信息
• 一次批量 matmul 同时算所有头，高效且简洁
• 下一步（第 11 课）：位置编码 + LayerNorm + 残差连接，拼出完整的 Transformer Block！

代码位置：src/attention.rs（MultiHeadAttention::forward 中的投影 / 打分 / softmax / 掩码部分） 底层算子：src/tensor.rs（softmax_last_dim / matmul / permute / mul_scalar） 掩码构造：src/model.rs（GPT::forward）

1. 本课要搞懂的问题

1. 为什么一个 token 需要"看"别的 token？注意力到底在算什么？
2. Query / Key / Value 是什么？为什么不能直接用 x 打分？
3. 为什么要除以 √d_k？softmax 在干什么？
4. 为什么 GPT 只能看过去？-inf 掩码怎么实现？

2. 从词向量到"上下文向量"

第 12 课会把每个 token 嵌入成一个向量，但 Embedding 查表得到的向量只代表这个 token 自己：它不知道周围发生了什么。"bank" 到底是"河岸"还是"银行"，单看它自己无法判断，必须结合上下文。

注意力要做的事：让每个位置根据内容，从其他位置"收集"信息。输出不再是"我自己的向量"，而是"我 + 我关注的别人的向量的加权混合"。

3. 检索的视角：Query / Key / Value

注意力最直观的类比是图书馆检索：

流程：拿 query 去和所有 key 比对相似度 → 得到权重 → 按权重把 value 加权求和。

我们的模型里 Q、K、V 都来自同一个输入 x（每个位置既当"提问者"又当"被看者"），这叫自注意力（Self-Attention）。

4. Q / K / V 投影：三个角色各学各的表示

直接用 x 当 Q、K、V 行不行？可以，但不好——三个角色应该学不同的表示。所以每个角色先过一层可学习的线性投影（Linear: y = xW + b）：

// 1. 投影得到 Q、K、V
let q = self.c_q.forward(x).reshape(vec![b, t, d]); // [B, T, D]
let k = self.c_k.forward(x).reshape(vec![b, t, d]);
let v = self.c_v.forward(x).reshape(vec![b, t, d]);

• c_q / c_k / c_v 是三个独立的 Linear::new(D, D)（MultiHeadAttention 结构体里的字段），各有自己的权重和偏置，在训练中学习
• Linear::forward 对 3D 输入会自动展平成 [B*T, D] 计算再还原 3D；这里的 reshape(vec![b, t, d]) 保证形状明确是 [B, T, D]
• D = n_embd。模型配置 GPTConfig::tiny：n_embd = 64，n_head = 4 → head_dim = 16（第 10 课讲怎么拆头；本课先假设只有一个头，公式层面 H = 1）

5. 打分：scores = Q·Kᵀ / √d_k

query 和 key 的相似度用点积衡量：q_i · k_j 越大越相似。一次性算所有两两相似度，就是矩阵乘法：

S = Q · Kᵀ / √d_k

• Q: [T, d_k]，K: [T, d_k] → Kᵀ: [d_k, T] → S: [T, T]
• S[i][j] = (q_i · k_j) / √d_k：第 i 个查询对第 j 个键的"原始分数"

对应代码：

// 4. 注意力分数：scores = Q·Kᵀ / √d_k
let scale = 1.0 / (head_dim as f32).sqrt();
let kt = k.permute(&[0, 2, 1]);                     // K 转置：[B*H, head_dim, T_total]
let scores = q.matmul(&kt).mul_scalar(scale);       // [B*H, T, T_total]

• permute(&[0, 2, 1])：把 K 的 [T, head_dim] 两维交换，得到数学上的 Kᵀ
• q.matmul(&kt)：3D 批量矩阵乘法（第 1 课学的 matmul，这里按"批"逐矩阵相乘）
• mul_scalar(scale)：每个元素乘以 1/√d_k

5.1 手算例子（d_k = 2，两个 token）

Q = [ [1, 0],          K = [ [1, 0],          V = [ [1, 0],
      [0, 1] ]               [0, 1] ]               [0, 1] ]

Q·Kᵀ = [ [1, 0],
         [0, 1] ]

除以 √2 后：S = [ [0.707, 0    ],
                 [0,     0.707] ]

softmax 按行归一化（见第 7 节）后，第 1 行 ≈ [0.670, 0.330]，第 2 行 ≈ [0.330, 0.670]——即"每个位置主要关注自己"（对角线），很符合直觉。

6. 为什么要除以 √d_k

假设 q、k 每个维度独立、均值 0、方差 1。两个 d_k 维向量的点积：

Var(q·k) = Var(Σ q_i·k_i) = d_k          →  标准差 = √d_k

• 除以 √d_k 把方差拉回 1，softmax 保持"温和"的分布
• 所以这个公式叫 Scaled Dot-Product Attention——缩放点积注意力

7. softmax 归一化：把分数变成权重

原始分数可正可负、量纲不定，不能直接当权重。softmax 把每一行变成非负且和为 1 的概率分布：

a_ij = exp(S_ij) / Σ_j' exp(S_ij')

含义：第 i 个查询给第 j 个键的"注意力权重"，一行加起来等于 1。

数值稳定技巧

exp 对大的正数会爆炸（如 exp(1000) = 无穷大），所以先减去每行最大值——数学上不改变结果（分子分母同除 exp(max)）：

// softmax_last_dim：沿最后一维逐行独立归一化
for r in 0..rows {
    let mut maxv = f32::NEG_INFINITY;
    for j in 0..d {
        maxv = maxv.max(sd[r * d + j]);                    // 1. 找行最大值
    }
    let mut sum = 0.0;
    for j in 0..d {
        out_data[r * d + j] = (sd[r * d + j] - maxv).exp(); // 2. 减 max 再 exp
        sum += out_data[r * d + j];
    }
    for j in 0..d {
        out_data[r * d + j] /= sum;                        // 3. 归一化
    }
}

稳定版公式：

softmax(x)_i = exp(x_i - max x) / Σ_j exp(x_j - max x)

反向传播公式（softmax_last_dim 的 backward 里已实现）：∂x_i = s_i·(g_i − Σ_j g_j·s_j)，其中 s 是 softmax 输出、g 是上游梯度。这正是第 2 课自动微分的用武之地。

8. 加权求和：out = A·V

把权重矩阵 A 作用到 V 上：

out_i = Σ_j a_ij · v_j

即"每个输出 = 所有位置 V 的凸组合"（权重非负、和为 1）。代码只有两行：

// 6. softmax 得到注意力权重，加权求和
let attn = scores.softmax_last_dim();  // [B*H, T, T_total]   注意力权重
let out = attn.matmul(&v);             // [B*H, T, head_dim]  加权求和

至此，注意力 = 内容相关的可学习加权平均：每个位置根据 query 与各 key 的相似度，从所有 value 里"取"信息。

9. 因果掩码：为什么 GPT 只能看过去

GPT 是自回归语言模型：预测第 t 个 token 时，条件只能是它之前的 token：

P(token_t | token_1, ..., token_{t-1})

• 训练时我们把整段文本并行喂入（这样才快！）。如果不加限制，位置 t 会"偷看"位置 t+1 的答案——信息泄漏，模型退化成"抄答案"
• 解决办法：让位置 i 的查询看不到 j > i 的键——这就是因果掩码（causal mask）

掩码长什么样

GPT::forward 里构造（t_total = t + base，训练时 base = 0）：

let mut mask_data = vec![0.0f32; t * t_total];
for i in 0..t {
    for j in 0..t_total {
        if j > i + base {
            mask_data[i * t_total + j] = f32::NEG_INFINITY;
        }
    }
}
let mask = Tensor::from_vec(mask_data, vec![t, t_total]);

以 t = 4 为例，mask 是这样一个矩阵：

对角线及以下（j ≤ i）是 0，未来位置（j > i）是 -inf。

为什么是 -inf 而不是把列删掉

softmax 里 exp(-inf) = 0，被掩码位置的注意力权重精确为 0——"看不见"等价于"权重为 0"。同时张量形状保持不变，可以继续走统一的批量矩阵运算。

let scores = scores.add(mask);              // [B*H, T, T_total] + 广播 [T, T_total]
let attn = scores.softmax_last_dim();       // -inf -> exp(-inf) = 0 -> 权重为 0

scores.add(mask) 用到了第 3 课的广播：mask 形状 [T, T_total] 自动广播到 [B*H, T, T_total] 的每一个"批"（每个头共享同一张掩码）。 （代码里的 base 是 KV cache 引入的偏移，第 18 课讲解；训练时 base = 0。）

10. 公式 ↔ 代码 总对照

Attention(Q, K, V) = softmax( (Q·Kᵀ / √d_k) + M ) · V

11. 运行与测试

cargo test   # tensor 的 test_softmax / test_matmul_2d / test_matmul_3d 验证本课用到的底层算子
cargo run    # 演示 3 训练小 GPT：每一步前向/反向都跑在本课的注意力计算上

12. 动手练习

1. 手算：取第 5.1 节的 2×2 例子，写出 softmax 后每一行的权重，再算出最终的加权求和结果，对照公式逐项核对。
2. 修改 GPT::forward 里的掩码条件 j > i + base 为 j >= i + base（禁止"看自己"），思考对训练会有什么影响（提示：每个位置的预测将只能依赖更早的 token）。
3. 临时改 softmax_last_dim：去掉"减最大值"那一步，用大数（如 1000）做输入，观察溢出/精度问题，验证数值稳定技巧的必要性。
4. 阅读 softmax_last_dim 的反向代码，对照公式 ∂x_i = s_i·(g_i − Σ_j g_j·s_j) 逐行核对。

13. 本课总结

• 注意力 = 内容相关的加权平均：Q 问、K 答、V 给内容
• scores = Q·Kᵀ/√d_k：点积打分，缩放防止 d_k 大时 softmax 饱和
• softmax：分数 → 权重（非负、和为 1）；先减 max 保证数值稳定
• 因果掩码：未来位置填 -inf → softmax 后权重为 0 → 模型只能看过去
• 下一课：把注意力"切成多份"并行计算——多头注意力！

代码位置：src/tokenizer.rs 演示入口：src/main.rs（demo_bpe） 语料：src/data.rs（CORPUS）

1. 本课要搞懂的问题

1. 模型只能吃数字，文字怎么变成数字？
2. 按词切、按字符切，各有什么问题？有没有更好的方案？
3. BPE 到底是什么算法？训练、编码、解码分别怎么做？

2. 为什么需要分词

大语言模型的输入输出都是数字：输入是一串 token id（usize），输出是对每个 token id 的预测分数。所以第一步要解决的问题是：把人类语言变成一串整数。

核心思想一句话：高频出现的片段合并成一个 token，低频内容用更小的片段表示。

3. 字符级分词 CharTokenizer（对照组）

3.1 结构

pub struct CharTokenizer {
    chars: Vec<char>,            // 词表：语料中出现过的所有字符
    stoi: HashMap<char, usize>,  // 字符 -> id
}

3.2 构建词表

new(text) 扫描语料，按字符首次出现的顺序收集去重：

for c in text.chars() {
    if seen.insert(c) {
        chars.push(c);
    }
}

比如对 "hello world hello"：

• 从左到右第一次遇到的字符依次是 h, e, l, o, ' ', w, r, d → 词表就是这 8 个字符
• stoi = {h:0, e:1, l:2, o:3, ' ':4, w:5, r:6, d:7}，vocab_size() == 8

3.3 编码 / 解码

// 文本 -> id 序列（每个字符查表）
pub fn encode(&self, text: &str) -> Vec<usize> {
    text.chars()
        .map(|c| {
            *self
                .stoi
                .get(&c)
                .unwrap_or_else(|| panic!("词表中没有字符 '{}'", c))
        })
        .collect()
}

// id 序列 -&gt; 文本（按 id 反查字符，越界 id 带下标 panic 提示）
pub fn decode(&self, ids: &[usize]) -> String {
    ids.iter()
        .map(|&i| {
            self.chars
                .get(i)
                .copied()
                .unwrap_or_else(|| {
                    panic!("decode 遇到越界 token id {i}（词表大小 {}）", self.chars.len())
                })
        })
        .collect()
}

编码和解码互为逆运算：decode(encode(text)) == text（单元测试 test_char_tokenizer_roundtrip 验证了这一点）。

注意：字符级编码遇到词表外的字符会直接 panic——这是它最大的短板。 演示程序里 CharTokenizer::new(CORPUS) 得到 35 个字符的词表，"fox" 编码为 [20, 7, 21]。

4. BPE 的直觉：把高频 pair "焊"在一起

BPE（Byte Pair Encoding，字节对编码）源自数据压缩算法，规则很简单：反复找到出现次数最多的相邻符号对，把它们合并成一个新符号。

以单元测试的语料 "low low low low low lowest lowest newest newest newest" 为例：

经过两轮合并，出现 7 次的常用子词 "low" 从 3 个字节压缩成了 1 个 token——这就是"用更少的 token 表示更多文本"。

5. BPE 训练：train()

5.1 初始化：字节级词表

assert!(target_vocab >= 256, "BPE 词表至少 256（字节级）");
// 初始：每个 token 就是一个字节
let mut vocab: Vec<Vec<u8>> = (0u16..=255).map(|b| vec![b as u8]).collect();
let mut merges: Vec<(u16, u16)> = Vec::new();

// 语料 -&gt; 字节 -> id 序列
let mut ids: Vec<u16> = corpus.as_bytes().iter().map(|&b| b as u16).collect();

初始词表 = 256 个字节（0~255），每个 token 恰好是一个字节。为什么用字节而不是字符？

• 任何 UTF-8 文本都可以拆成字节，不存在"词表外"字符（OOV = 0）
• 中文等多语言文本也能直接编码（一个汉字是 3 个字节）
• GPT-2 等真实模型用的就是字节级 BPE

5.2 训练主循环：统计 → 选择 → 合并 → 替换

while vocab.len() < target_vocab {
    // 1. 统计相邻 pair 频率
    let mut pair_freq: HashMap<(u16, u16), usize> = HashMap::new();
    for pair in ids.windows(2) {
        *pair_freq.entry((pair[0], pair[1])).or_insert(0) += 1;
    }
    // 2. 找最高频的 pair（频率相同取 pair 值小者，保证确定性）
    let Some(&best) = pair_freq
        .iter()
        .max_by(|a, b| a.1.cmp(b.1).then_with(|| b.0.cmp(a.0)))
        .map(|(k, _)| k)
    else {
        break; // 没有可合并的 pair 了
    };
    // 3. 合并：新符号 = 两个符号的字节拼接
    let new_id = vocab.len() as u16;
    let mut new_bytes = vocab[best.0 as usize].clone();
    new_bytes.extend_from_slice(&vocab[best.1 as usize]);
    vocab.push(new_bytes);
    merges.push(best);

    // 4. 替换 ids 中所有该 pair
    let mut new_ids: Vec<u16> = Vec::with_capacity(ids.len());
    let mut i = 0;
    while i < ids.len() {
        if i + 1 < ids.len() && ids[i] == best.0 && ids[i + 1] == best.1 {
            new_ids.push(new_id);
            i += 2;               // 一次吞掉两个符号
        } else {
            new_ids.push(ids[i]);
            i += 1;
        }
    }
    ids = new_ids;
}

训练结束后得到两份"产物"：

• vocab: Vec<Vec<u8>>：token id → 它代表的字节序列（解码要用）
• merges: Vec<(u16, u16)>：合并规则表，按下标顺序排列（越早合并优先级越高，编码要用）

词表大小 = 256 + 合并次数。演示里 BPETokenizer::train(CORPUS, 400) 得到 400 = 256 字节 + 144 次合并。

6. BPE 编码：encode()

编码 = 对新文本执行同样的合并。但合并顺序必须和训练时一致：训练时越早合并的规则优先级越高（它对应的 token id 更小）。

pub fn encode(&self, text: &str) -> Vec<usize> {
    let mut ids: Vec<u16> = text.as_bytes().iter().map(|&b| b as u16).collect();
    for (idx, &(a, b)) in self.merges.iter().enumerate() {
        let new_id = (256 + idx) as u16;
        let mut out: Vec<u16> = Vec::with_capacity(ids.len());
        let mut i = 0;
        while i < ids.len() {
            if i + 1 < ids.len() && ids[i] == a && ids[i + 1] == b {
                out.push(new_id);
                i += 2;
            } else {
                out.push(ids[i]);
                i += 1;
            }
        }
        ids = out;
    }
    ids.into_iter().map(|x| x as usize).collect()
}

要点：

• 按规则优先级单趟扫描（GPT-2 的标准实现）：从 merges[0] 到 merges[最后]，每条规则在序列上扫一遍，能合并就替换成它对应的新 token。复杂度 O(len × 合并数)，大语料也能秒级完成（若"每次只合并一个 pair 并全量重扫"是 O(n²×m)，174KB 语料会卡死）
• new_id = 256 + idx：merge 下标 idx 直接映射成 token id——因为训练时第 idx 次合并恰好产生 id 256 + idx
• 字节 id（0~255）直接复用训练时的字节 → id 映射
• 演示里 "the garden" 编码后只有 2 个 token（"the" 和 " garden" 都被压缩成了单个 token）

以 "lowest" 为例走一遍：字节 [l,o,w,e,s,t] → 应用规则 0（假设是 (l,o)）→ [256,w,e,s,t] → 应用规则 1（(256,w)）→ [257,e,s,t] → 继续应用 (e,s)、(s,t) 对应规则……最终 6 个字节被压成 3 个 token（单元测试要求 "low" 编码后不超过 3 个 token）。

7. BPE 解码：decode()

解码是查表 + 拼接：每个 id 查到它代表的字节序列，拼起来还原文本。

pub fn decode(&self, ids: &[usize]) -> String {
    let mut bytes: Vec<u8> = Vec::new();
    for &id in ids {
        let tok = self
            .vocab
            .get(id)
            .unwrap_or_else(|| panic!("decode 遇到越界 token id {id}（词表大小 {}）", self.vocab.len()));
        bytes.extend_from_slice(tok);
    }
    String::from_utf8_lossy(&bytes).to_string()
}

• vocab[id]：id → 字节序列（0~255 是单字节，256+ 是合并出来的多字节序列）
• 越界 id 用 vocab.get(id) 拦截：带下标信息的 panic 提示（decode 也会因传入非法 id 报错，不再是"永不失败"）
• String::from_utf8_lossy：万一拼出非法 UTF-8，用替换字符 � 顶替而不是 panic

完整闭环：decode(encode("lowest new")) == "lowest new"（单元测试 test_bpe_roundtrip 验证）。

8. 训练 / 编码 / 解码 对照总结

三者关系：编码必须复现训练时的合并顺序，解码只是查表，所以编码、解码天然互逆，decode(encode(x)) == x。

9. 运行与测试

cargo test   # 全部测试通过
cargo run    # 演示 2（BPE）：词表 400（256 + 144 次合并）；"Red" -> [82, 101, 100]；"the garden" -> 2 个 token

10. 动手练习

1. 把语料换成中文句子（如 "机器学习机器学习深度学习"），跑一遍 BPETokenizer::train，观察"机器""学习"是否会被合并成单个 token（提示：中文 UTF-8 每字 3 字节，BPE 依然适用）。
2. 修改 train 的 target_vocab，分别用 256 / 300 / 500，比较 encode("lowest") 的 token 数变化。
3. 思考：encode 为什么必须按"merge 下标最小"合并，而不是按"频率最高"合并？（提示：训练时的合并顺序决定了 id 分配，编码必须复现同一顺序才能保证解码还原）
4. 思考：CharTokenizer 遇到词表外字符会 panic，BPE 为什么永远不会有这个问题？

11. 本课总结

• 模型只认数字，分词器负责"文字 ↔ id"的转换
• 字符级分词：简单直观，但序列长、有词表外字符；词级分词：词表大、仍有 OOV
• BPE：字节级底座（256 个字节起步，零 OOV）+ 反复合并最高频相邻 pair
• 训练产出 merges（合并规则）和 vocab（id → 字节序列）；编码贪心复现合并；解码查表拼接
• 下一课：token 变成向量之后，怎么让它们"互相看"？——注意力机制！

代码位置：src/main.rs（demo_xor 函数） 前置知识：第 5 课 Linear / 激活函数、第 6 课 CrossEntropy / SGD

1. 本课要搞懂的问题

1. XOR 为什么"看着简单却学不会"？什么是线性不可分？
2. 为什么必须加隐藏层和非线性激活？
3. demo_xor 的每一行代码在干什么？
4. 为什么 XOR 训练到 4/4 就能证明反向传播是对的？

2. XOR：简单到"一眼看穿"，难到"一条直线搞不定"

异或（XOR）：两个输入相同输出 0，不同输出 1。

把四个点画在平面上（● = 输出 1，○ = 输出 0）：

x2
1 ─  ●(0,1)   ○(1,1)
   │
0 ─  ○(0,0)   ●(1,0)
   └──────────────── x1

问题来了：能不能画一条直线，把 ● 和 ○ 分开？不能。 一条直线最多把平面分成两半，而两个 ● 在对角线上——无论直线怎么放，总有一侧同时混着 ● 和 ○。这就是线性不可分。

3. 为什么必须非线性：单层线性模型必败

• 第 2 课的线性回归只有一层：pred = xW + b，它学出的边界本质是一条直线
• 对 XOR，直线永远分不开 → 单层模型最多只能对 3 个样本正确（3/4），永远到不了 4/4
• 解法：隐藏层 + 非线性激活（第 5 课学的 Tanh）。非线性可以把输入空间"折叠 / 扭曲"，让原本线性不可分的点在高层变得线性可分

直觉（4 个样本为什么能学会）：

• 隐藏层 4 个神经元 ≈ 学出 4 种"特征组合"
• Tanh 把每个特征压缩到 (-1, 1)，给特征带来非线性
• 第二层再线性组合这些特征，就能画出"折线"边界，把四个点完美分开

4. 数据集与网络结构：2 → 4 → 2

数据集（4 个样本，一次全部喂入，batch = 4）：

// 数据集：4 个样本
let x_data = Tensor::from_vec(
    vec![0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0],
    vec![4, 2],
);
let y_targets = vec![0usize, 1, 1, 0]; // XOR 真值表

• x_data：[4, 2]，四行正好是 (0,0)、(0,1)、(1,0)、(1,1)
• y_targets：对应输出 0、1、1、0

网络（两个 Linear 叠成 MLP）：

let fc1 = Linear::new(2, 4, &mut rng);   // 输入 2 维 → 隐藏层 4 维
let fc2 = Linear::new(4, 2, &mut rng);   // 隐藏层 4 维 → 输出 2 维

结构图：

x [4,2] ─► fc1 (2→4) ─► Tanh ─► h [4,4] ─► fc2 (4→2) ─► logits [4,2]

• 输出层是 2 个神经元：第 0 个是"类别 0 的分数"，第 1 个是"类别 1 的分数"
• 为什么不直接用 1 个输出（0 或 1）？因为配合 CrossEntropy 分类更自然，两个分数取 argmax 就是预测

收集参数 + 创建优化器：

let params: Vec<Tensor> = {
    let mut ps = fc1.parameters();
    ps.extend(fc2.parameters());
    ps
};
let opt = SGD::new(0.5, params);

• fc1.parameters() + fc2.parameters()：共 2×4+4 + 4×2+2 = 22 个参数
• 学习率 0.5：XOR 问题简单，步长可以大一点
• 随机种子固定为 42（Rng::new(42)），每次运行初始化相同、结果可复现

5. 训练循环：1000 步

for step in 0..1000 {
    // 前向：tanh(x @ W1 + b1) @ W2 + b2
    let h = tanh(&fc1.forward(&x_data));
    let logits = fc2.forward(&h);
    let loss = cross_entropy_loss(&logits, &y_targets);

    loss.backward();
    opt.step();
    opt.zero_grad();

    if step % 200 == 0 {
        println!("  step {:>4} | loss {:.4}", step, loss.data()[0]);
    }
}

对应第 6 课的"四步舞"：

观察 loss：每 200 步打印一次，趋势是单调下降并趋近 0（具体数值由种子 42 的初始化决定，不必纠结；刚初始化时接近随机猜，2 类问题约等于 −ln 0.5 ≈ 0.69）。

6. 验证：4/4 才算赢

训练完要"考试"：

// 验证正确率
let h = tanh(&fc1.forward(&x_data));
let logits = fc2.forward(&h);
let data = logits.data();
let mut correct = 0;
for i in 0..4 {
    let pred = if data[i * 2] > data[i * 2 + 1] { 0 } else { 1 };
    if pred == y_targets[i] {
        correct += 1;
    }
}
println!("  训练后正确率：{}/4（100% 说明反向传播正确）\n", correct);

• 对每个样本取 argmax：类别 0 的分数 data[i*2] 大于类别 1 的分数 data[i*2+1] → 预测 0，否则预测 1
• 期望输出：训练后正确率：4/4

7. 为什么 4/4 就能证明反向传播正确

推理链条：

1. 参数是随机初始化的（种子 42），绝不是"碰巧答对"；
2. 梯度完全来自自动微分：任何一个运算（matmul / add / relu / softmax / log / sum ...）的反向实现出错，梯度就会偏一点，参数越走越偏，1000 步后必然累积成大错；
3. XOR 没有捷径：它是"必须两层非线性"的问题，梯度要穿透 Tanh 的导数（1 - t²）、两层 matmul 的转置路径、softmax 的雅可比路径——整条反向传播链路被真实训练"压测"了一遍；
4. 结果 100% 正确 + loss 单调下降，还同时验证了优化器（SGD 的 set_data 写回）和训练循环（backward → step → zero_grad 的顺序）都是对的。

换句话说：XOR 训练收敛到 4/4，等于对"张量库 + 网络层 + 损失 + 优化器"整条链路做了一次端到端集成测试。

8. 运行

cargo run

输出第一段（loss 数值随种子确定，趋势如下）：

=== 演示 1：MLP 学习 XOR（第 7 课）===
  step    0 | loss ...
  step  200 | loss ...
  step  400 | loss ...
  step  600 | loss ...
  step  800 | loss ...
  训练后正确率：4/4（100% 说明反向传播正确）

9. 动手练习

1. 把隐藏层大小从 4 改成 2（Linear::new(2, 2, ...)），观察还能不能到 4/4（提示：2 个隐藏神经元很可能不够）。
2. 把 tanh 换成 relu 或 gelu，其他不变，观察收敛速度差异。
3. 把学习率改成 0.05 和 5.0，对比 loss 下降曲线（太小慢、太大震荡）。
4. 去掉 tanh（h = fc1.forward(&x_data) 直接喂给 fc2），验证"没有非线性就学不会 XOR"，观察正确率卡在 3/4。
5. 把训练步数改成 100（可能还没收敛）和 10000（数据只有 4 个样本，看会不会过拟合），观察 loss 与正确率。

10. 本课总结

• XOR 线性不可分 → 必须隐藏层 + 非线性激活
• demo_xor 完整走通：数据集 → 网络 2→4→2 → 训练 1000 步 → 验证 4/4
• 4/4 = 反向传播正确的强证据（端到端集成测试）
• 从现在起，我们有了一个能训练任意小分类问题的完整工具箱
• 下一课：给模型喂真正的文本——BPE 分词器

代码位置：

• src/loss.rs（MSE / CrossEntropy）
• src/optim.rs（SGD / AdamW）

演示入口：src/main.rs（demo_xor 使用 cross_entropy_loss + SGD）

1. 本课要搞懂的问题

1. 损失函数是干什么的？MSE 和 CrossEntropy 分别用在什么任务？
2. CrossEntropy 为什么是"one-hot + log_softmax + 取均值"？
3. SGD 的一行更新公式 θ = θ - lr·g 为什么能让损失下降？
4. 训练循环的"四步"（前向、反向、更新、清零）分别对应哪几行代码？

2. 损失函数：一杆"错得有多离谱"的秤

训练 = 让模型输出不断接近目标。怎么量化"接近"？用损失函数 L(pred, target)：

• L 越大 → 错得越离谱
• 训练就是不断减小 L

两类任务对应两个损失：

语言模型本质是"预测下一个 token 是哪一类"，所以 LLM 用的是 CrossEntropy。

3. MSE：均方误差（回归）

公式：

MSE = (1/N) · Σᵢ (predᵢ - targetᵢ)²

对应代码（loss.rs，全是一行链式调用）：

/// 均方误差：loss = mean((pred - target)²)
pub fn mse_loss(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub(target)
        .pow(2.0)
        .sum()
        .mul_scalar(1.0 / pred.numel() as f32)
}

为什么平方？ 正负误差不会互相抵消；大误差被放大（对离群点更敏感）。

梯度（不需要手写，自动微分自动算出来）：

∂L/∂predᵢ = 2·(predᵢ - targetᵢ) / N

4. CrossEntropy：交叉熵（分类）

分类任务里模型输出的是 logits（每个类别的"原始分数"），要把它变成"概率"再用交叉熵衡量。

输入：

• logits: [B, D] —— B 个样本，每个 D 个类别的分数
• targets: [B] —— 每个样本真实类别的下标

数学公式：

loss = -(1/B) · Σᵢ log( softmax(logitsᵢ)[targetsᵢ] )

对应代码分三步（loss.rs）：

pub fn cross_entropy_loss(logits: &Tensor, targets: &[usize]) -> Tensor {
    assert_eq!(logits.rank(), 2, "交叉熵的 logits 应为 [B, D]");
    let (b, d) = (logits.shape()[0], logits.shape()[1]);

    // —— 第 1 步：one-hot 编码 ——
    let mut onehot = vec![0.0f32; b * d];
    for (i, &t) in targets.iter().enumerate() {
        assert!(t < d, "目标类别越界：{} >= {}", t, d);
        onehot[i * d + t] = 1.0;
    }
    let oh = Tensor::from_vec(onehot, vec![b, d]);

    // —— 第 2 步：log_softmax（数值稳定版，log-sum-exp 技巧）——
    let log_probs = logits.log_softmax_last_dim();

    // —— 第 3 步：取出正确类别的 log 概率，取负，再取均值 ——
    log_probs
        .mul(&oh)        // 只保留正确类别位置的 log 概率
        .sum_last_dim()  // [B,1]：每行求和 = 正确类别的 log 概率
        .neg()           // 取负：正确概率越高 loss 越小
        .sum()           // 所有样本相加
        .mul_scalar(1.0 / b as f32)  // 除以 B 取平均
}

逐步拆解

第 1 步 one-hot：把"下标 2"变成 [0, 0, 1, 0, ...]，只有正确位置是 1。

第 2 步 log_softmax：softmax 把分数变成概率（和为 1、都为正），log 取对数。

注：这里用的是朴素写法 softmax 之后 log（即 log_softmax）。因为 softmax_last_dim 内部已经做了"先减最大值再 exp"的数值稳定处理，所以直接 log 是安全的。

第 3 步 取负 + 取均值：

• mul(&oh) 后每行只有正确类别位置非 0 → sum_last_dim() 恰好得到每个样本正确类别的 log 概率
• 正确概率越高 → -log p 越小 → loss 越小 ✓
• 除以 B（样本数）取平均

数值例子（2 类、2 个样本，示意）

平均 loss ≈ (0.20 + 0.17) / 2 ≈ 0.18。如果模型对正确类别给出概率接近 1，loss 就接近 0。

关于梯度：softmax + 交叉熵的组合，数学上梯度恰好等于 p − y_onehot（预测概率减去 one-hot）。我们的代码没有手写这个式子——反向传播会自动把它算出来。这正好展示了"公式 → 自动微分"的分工。

5. 梯度下降原理：沿着"最陡的下坡"走

回顾第 2 课的核心公式，对每个参数 θ：

θ = θ - 学习率 × ∂L/∂θ

• ∂L/∂θ > 0：θ 增大 loss 也增大 → 减去它，θ 变小
• ∂L/∂θ < 0：θ 增大 loss 减小 → 减去它，θ 变大
• 学习率（lr）控制每步迈多大：太小收敛慢，太大来回震荡甚至发散

几何直觉：loss 是参数的山坡，梯度指向"上升最快的方向"，我们每次都朝反方向走一小步，就像小球从山坡滚下。

loss
  │        .
  │      .   .
  │    .       .
  │  .           .
  └──────────────────→ θ

6. SGD：把更新规则写成代码

**随机梯度下降（SGD）**是最朴素的优化器（optim.rs）：

/// 随机梯度下降（SGD）
pub struct SGD {
    lr: f32,
    params: Vec<Tensor>,
}

impl SGD {
    pub fn new(lr: f32, params: Vec<Tensor>) -> Self {
        SGD { lr, params }
    }

    /// 更新一步：θ = θ - lr * g
    pub fn step(&self) {
        for p in &self.params {
            let g = p.grad();
            let d = p.data();
            let updated: Vec<f32> = d.iter().zip(&g).map(|(v, g)| v - self.lr * g).collect();
            p.set_data(updated);
        }
    }

    pub fn zero_grad(&self) {
        for p in &self.params {
            p.zero_grad();
        }
    }
}

公式与代码的对应：

两个方法的分工：

• step()：必须在 loss.backward() 之后调用，否则梯度全是 0，参数纹丝不动
• zero_grad()：每轮更新后必须清零梯度。因为反向传播是 += 累加（第 2 课讲过），不清零的话下一轮的梯度会叠加上一轮的旧梯度

顺带一提：optim.rs 里还有 AdamW（第 17 课讲）。它把"固定步长"换成"动量 + 自适应步长"，但接口完全一样（step() / zero_grad()），这正是统一接口的好处。

7. 训练循环：四步舞

一次完整的学习 = 四步（所有深度学习框架训练循环的骨架）：

// 1. 前向：算预测和损失
let logits = model.forward(&x);
let loss = cross_entropy_loss(&logits, &y_targets);
// 2. 反向：自动算出每个参数的梯度
loss.backward();
// 3. 更新：θ = θ - lr·g
opt.step();
// 4. 清零梯度：否则下一轮会累加
opt.zero_grad();

顺序很重要：backward → step → zero_grad，反了或漏了都会出错。

8. 运行与测试

cargo test     # 测试全部通过
cargo run      # 演示 1 中 loss 从初始值一路下降、正确率 4/4

（demo_xor 的每一行会在第 7 课完整讲解。）

9. 动手练习

1. 把 cross_entropy_loss 的取均值从"除以 b"改成"除以 b·d"，观察 loss 数值变化，思考哪个更合理。
2. 给 SGD 加动量：v = 0.9·v + g; θ = θ - lr·v，对比收敛速度。
3. 思考：one-hot 为什么只在正确位置是 1？如果给两个类别各标 0.5（soft label）会怎样？
4. 手算一个 3 分类、2 个样本的 cross_entropy_loss，和程序输出对一下。
5. 尝试用 mse_loss 替代 cross_entropy_loss 训练分类任务（直接把 logits 往 one-hot 上回归），观察效果差异。

10. 本课总结

• MSE：回归任务的 mean((pred - target)²)
• CrossEntropy：one-hot + log_softmax → 取负 → 取均值，分类任务标配（LLM 用它）
• SGD：θ = θ - lr·g，一行代码一个循环
• 训练循环 = 前向 → 反向 → 更新 → 清零，四步缺一不可
• 下一课：把 Linear、ReLU、CrossEntropy、SGD 拼起来，训练第一个 MLP！

代码位置：

• src/layers.rs（Linear 层、ReLU / GELU / Tanh / Softmax）
• src/module.rs（Module trait 参数管理）
• src/rng.rs（xorshift64 随机数生成器，权重初始化用）

演示入口：src/main.rs

1. 本课要搞懂的问题

1. 什么是"层"（layer）？为什么神经网络要由层堆叠而成？
2. Linear 层的数学本质 y = xW + b 是什么？代码怎么写？
3. 权重为什么要初始化？xavier 风格的 ±1/√in 是怎么来的？
4. 激活函数是干什么的？ReLU / GELU / Tanh 的公式和反向导数各是什么？
5. 怎么统一管理所有可训练参数？—— Module trait

2. 从"一个参数"到"一层"

第 2 课我们用手写更新 w、b 的方式学会了 y = 2x + 1。但真实问题（图像分类、语言模型）输入是几千维的向量、输出是几万个类别，不可能靠手写公式。

于是我们把"一批神经元"打包成一个层（layer）：

• 每个神经元：对输入做加权和 + 偏置，再过一层非线性函数
• 一层 = 很多个这样的神经元，它们的权重堆在一起就是矩阵 W

神经网络 = 一层一层地堆叠。这一课实现最基本的积木：Linear 层和激活函数。

3. Linear 层：y = xW + b

Linear 又叫"全连接层 / 稠密层"。它把输入向量线性变换成另一个向量：

数学公式（对每个样本 i）：

y[i, j] = Σ_k x[i, k] · W[k, j] + b[j]

等价写法：y = x @ W + b（这就是第 1 课学的矩阵乘法 + 第 3 课学的广播加法）。

对应的 Rust 结构（layers.rs）：

/// 线性层：y = x @ W + b
///
/// - weight: [in_features, out_features]
/// - bias:   [out_features]
pub struct Linear {
    pub weight: Tensor,
    pub bias: Tensor,
}

前向计算的核心就一行：

// y = x @ W + b（b 是 [out]，与 [B, out] 广播相加）
let y = x.matmul(&self.weight).add(&self.bias);

两个关键点：

• weight 和 bias 都用 Tensor::param(...) 构造（requires_grad = true），反向传播时自动得到梯度；
• bias 形状是 [out]，和结果 [B, out] 相加时按广播规则每一行都加上同一个偏置——这正是第 3 课广播的用武之地。

4. 权重初始化：为什么是 ±1/√in

Linear::new 创建时给权重赋随机初始值：

pub fn new(in_features: usize, out_features: usize, rng: &mut Rng) -> Self {
    let std = (2.0 / (in_features + out_features) as f32).sqrt();
    let w: Vec<f32> = (0..in_features * out_features)
        .map(|_| rng.randn() * std)
        .collect();
    Linear {
        weight: Tensor::param(w, vec![in_features, out_features]),
        bias: Tensor::param(vec![0.0; out_features], vec![out_features]),
    }
}

也就是从正态分布 N(0, std) 采样，std = √(2/(in+out))（Xavier 正态初始化），bias 初始化为全零。

为什么不能随便初始化？

直觉推导：假设输入 x 的方差为 1，输出是 in 个项 wᵢxᵢ 的和。若 wᵢ ~ N(0, σ²)，则 Var(wᵢ) = σ²，于是：

Var(Σ wᵢxᵢ) ≈ in · Var(wᵢ) · Var(x) = in · σ²

想让方差保持不变，需要 σ = √(1/in)；Xavier 正态取 σ = √(2/(in+out))（兼顾前向和反向），核心思想——让每一层的输出量级都保持在 1 附近。

随机数从哪来？标准库没有 RNG，我们自己在 rng.rs 实现了一个 xorshift64（几十行位运算，零依赖、可复现）：

/// 生成 [lo, hi) 的均匀浮点数
pub fn uniform(&mut self, lo: f32, hi: f32) -> f32 {
    lo + (hi - lo) * self.next_f32()
}

用固定种子创建（如第 7 课的 Rng::new(42)），每次运行得到的初始化完全相同，实验结果可复现。

5. 支持 2D 与 3D 输入

一个 batch 的输入通常是 2D（[B, in]）；但语言模型里一批数据是 3D（[B, T, in]，B=batch、T=序列长度）。Linear 的 forward 两者都支持：

pub fn forward(&self, x: &Tensor) -> Tensor {
    // 支持 [B, in] 和 [B, T, in]；3D 输入在内部展平计算，输出保持 3D
    let is_3d = x.rank() == 3;
    let orig_shape = x.shape().to_vec();
    let x = match x.rank() {
        2 => x.clone(),
        3 => x.reshape(vec![x.shape()[0] * x.shape()[1], x.shape()[2]]),
        _ => panic!("Linear 输入必须为 2D 或 3D"),
    };
    let y = x.matmul(&self.weight).add(&self.bias);
    if is_3d {
        // 3D 输入 [B, T, in] -> 输出 [B, T, out]（最后一维换成 out_features）
        let mut out_shape = orig_shape;
        let n = out_shape.len();
        out_shape[n - 1] = self.weight.shape()[1];
        y.reshape(out_shape)
    } else {
        y
    }
}

处理方式：3D 先 reshape 成 [B·T, in] 当 2D 算，算完再把最后一维换回 out 并 reshape 成 [B, T, out]。

reshape 不改变数据顺序，梯度也能 1:1 传回，所以这个"展平再恢复"的技巧在自动微分下完全安全。

6. 参数管理：Module trait

深度学习里一切"可训练的结构"都是模块。我们定义统一的接口（module.rs）：

/// 模块接口：任何可训练结构都实现它
pub trait Module {
    /// 返回模块的所有参数（含嵌套子模块）
    fn parameters(&self) -> Vec<Tensor>;
}

/// 便捷方法：清零所有参数的梯度
pub fn zero_grad_all(module: &dyn Module) {
    for p in module.parameters() {
        p.zero_grad();
    }
}

Linear 的实现（layers.rs）：

impl Module for Linear {
    fn parameters(&self) -> Vec<Tensor> {
        vec![self.weight.clone(), self.bias.clone()]
    }
}

为什么需要这个 trait？ 因为优化器（第 6 课）只认"参数列表"。多个模块拼在一起时：

let params: Vec<Tensor> = {
    let mut ps = fc1.parameters();
    ps.extend(fc2.parameters());
    ps
};
let opt = SGD::new(0.5, params);   // 优化器拿到全部 22 个参数

有了 Module，今后 LayerNorm、Embedding、Transformer Block 都能用同一套方式"一键收集参数"，这正是第 7 课 demo_xor 的做法。

7. 激活函数：没有它，多层等于一层

先想一个问题：如果只有线性层会怎样？

f(x) = (xW1 + b1)·W2 + b2 = x·(W1W2) + (b1W2 + b2)

两个线性变换复合还是一个线性变换！堆多少层都等价于一层。所以必须夹入非线性激活函数，多层才有意义。

layers.rs 提供了三个激活函数（都是对第 4 课 tensor.rs 中张量方法的薄包装）：

/// ReLU：max(0, x)，简单、计算快、缓解梯度消失
pub fn relu(x: &Tensor) -> Tensor { x.relu() }

/// GELU：GPT 系列的默认激活，用 tanh 近似，比 ReLU 更平滑
pub fn gelu(x: &Tensor) -> Tensor { x.gelu() }

/// Tanh：S 型，输出 (-1, 1)
pub fn tanh(x: &Tensor) -> Tensor { x.tanh() }

7.1 ReLU

公式与导数：

c = max(0, x)
dc/dx = 1  (x > 0)
      = 0  (x ≤ 0)

对应代码（tensor.rs 中的 relu，核心是反向用 mask 乘梯度）：

// 前向：c = max(0, x)；同时记录 mask：x > 0 的位置为 1，其余为 0
let mask: Vec<f32> = sd.iter().map(|&a| if a > 0.0 { 1.0 } else { 0.0 }).collect();
// 反向：∂x = g * mask（即 g * (x>0)）
sgm[i] += g[i] * mask[i];

特点：计算最快；x<0 时梯度为 0（优点：缓解梯度消失；缺点：可能"神经元死亡"）。

7.2 GELU

公式（tanh 近似，GPT 系列默认激活）：

c = 0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³)))

反向导数（令 a = √(2/π)(x + 0.044715x³)，t = tanh(a)）：

da/dx = √(2/π) · (1 + 3·0.044715·x²)
dc/dx = 0.5·(1 + t) + 0.5·x·(1 - t²)·da/dx

对应代码（tensor.rs 中的 gelu，前向缓存 t 供反向复用）：

const SQRT_2_PI: f32 = 0.797_884_560_8; // sqrt(2/π)
const COEF: f32 = 0.044_715;
// 前向：t_vals 缓存 tanh(a)
let a = SQRT_2_PI * (x + COEF * x * x * x);
let t = a.tanh();
t_vals[i] = t;
data[i] = 0.5 * x * (1.0 + t);
// 反向
let da_dx = SQRT_2_PI * (1.0 + 3.0 * COEF * x * x);
let dy_dx = 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * da_dx;
sgm[i] += g[i] * dy_dx;

特点：处处可导、比 ReLU 平滑，x=0 附近没有"尖角"。

7.3 Tanh

公式与导数：

c = tanh(x) = (e^x - e^-x) / (e^x + e^-x)
dc/dx = 1 - tanh²(x) = 1 - c²

对应代码（tensor.rs 中的 tanh，反向直接用前向输出值 c 算 1-c²）：

// 前向：c = tanh(x)
// 反向：∂x = g * (1 - c²)
sgm[i] += g[i] * (1.0 - data[i] * data[i]);

特点：S 形，输出落在 (-1, 1)，需要居中输出时好用；两端梯度趋近 0。

7.4 三种激活对比

补充：layers.rs 里还有一个 softmax（把一组分数变成概率分布），它是 CrossEntropy 的一部分，下一课会用到。学习计划表里第 5 课原写的是"Sigmoid"，实际代码中用 GELU 取代了它（GPT 时代更常用），ReLU / Tanh 均已实现。

8. 运行与测试

cargo test     # 测试全部通过（含 relu 梯度、softmax、线性回归收敛等）
cargo run      # 演示 1：MLP 学习 XOR（第 7 课完整讲解）

9. 动手练习

1. 把 Linear::new 里的 std 改成 0.01 或 5.0，跑第 7 课的 demo_xor，观察收敛速度变化，解释原因。
2. 给 Linear 加一个"无 bias"选项（bias 全 0 且不出现在 parameters() 里），想想 parameters() 该怎么改。
3. 手动推导：GELU 在 x=0 处的值（应为 0）和导数；tanh 在 x=0 处的导数（应为 1）。
4. 思考：代码里 ReLU 在 x=0 处的梯度定义为 0（a > 0 才为 1）。如果定义为 1 会怎样？训练结果会有区别吗？
5. （选做）查资料对比 xavier 与 kaiming 初始化，并解释为什么 Embedding 用 0.02 的 bound（见 layers.rs 的 Embedding::new）。

10. 本课总结

• Linear：y = xW + b，weight [in, out]、bias [out]，支持 2D / 3D 输入
• 初始化：均匀分布 ±1/√in（xavier 风格），保证激活量级不随层数放大
• Module trait：parameters() 统一收集参数，优化器直接消费
• 激活函数：ReLU / GELU / Tanh，公式、反向导数与代码一一对应
• 下一课：损失函数与优化器，让模型"知道自己错在哪、怎么改"

代码位置：src/tensor.rs（核心：Tensor 结构体与自动微分）、src/module.rs（Module trait）、src/main.rs（模块声明与演示编排） 参考实现：src/layers.rs（Module 的第一个实现者）

1. 本课要搞懂的问题

1. 一个从零写的项目，文件该怎么切？每个文件各管什么？
2. Tensor 里的 Rc<RefCell<Vec<f32>>> 是什么？为什么张量需要"共享可变数据"？
3. 为什么 parents 要用 Rc<Vec<Tensor>> 而不是直接 Vec<Tensor>？直接存会有什么灾难？
4. backward() 里的 DFS 拓扑排序是怎么保证"梯度从 loss 流回叶子"的？
5. Module trait 是干什么的？它和优化器、训练循环怎么配合？

2. 为什么要"重构"

经过前三课，tensor.rs 已经长成了一个功能完整的张量库：构造、形状、逐元素运算、广播、归约、softmax、matmul、自动微分……如果所有代码继续堆在一个文件里，很快就会失控：

• 职责不分：张量运算、自动微分、模型层、训练逻辑混在一起，改一处可能碰坏另一处
• 无法复用：后续每节课都要加新东西（层、损失、优化器、GPT 模型），没有清晰的模块边界就没法增量开发
• 依赖混乱：谁依赖谁说不清，改个接口要满项目搜索

所以这一课做两件事：

1. 横向：把项目按职责拆成多个 .rs 文件，形成清晰的依赖金字塔
2. 纵向：深入理解 Tensor 的 Rc<RefCell> 架构——它是一切上层建筑的地基，前三课里所有算子都建立在这套设计之上

3. src/ 各文件职责分工

main.rs 顶部用 mod 声明了全部模块：

// src/main.rs
mod attention;
mod autograd;
mod checkpoint;
mod cli;
mod config;
mod data;
mod gpu;
mod layers;
mod loss;
mod model;
mod module;
mod optim;
mod rng;
mod rope;
mod sample;
mod tensor;
mod tokenizer;
mod train;

每个文件的职责（与文件头注释一致）：

依赖关系呈金字塔形——越往下越基础，越往上越具体：

要点：

• tensor.rs 在最底层：所有模块都 use crate::tensor::Tensor，但它不依赖任何其他模块。
• module.rs 只依赖 tensor.rs：它只定义一个 trait，是个"协议层"，供 layers.rs / model.rs 实现。
• 上层模块互相独立：loss.rs 不知道 optim.rs 的存在，改动一个不影响另一个；data.rs、tokenizer.rs、rng.rs 与神经网络计算完全解耦。

4. 核心架构：Rc<RefCell<Vec>> 共享可变数据

先看 Tensor 的结构体（src/tensor.rs）：

#[derive(Clone)]
pub struct Tensor {
    data: Rc<RefCell<Vec<f32>>>,   // 数值（前向结果）
    shape: Vec<usize>,             // 形状
    grad: Rc<RefCell<Vec<f32>>>,   // 梯度（反向累加）
    requires_grad: bool,           // 是否需要梯度（参数=true，数据=false）
    parents: Rc<Vec<Tensor>>,      // 父节点列表（我是由谁算出来的）
    backward: Option<BackwardFn>,  // 反向函数：把梯度传给父节点
}

两个问题：为什么数值和梯度要共享？共享之后怎么才能"可变"？

4.1 Rc：多个句柄指向同一份数据

计算图里一个节点经常被多个运算同时引用。比如 z = x·y + w 里 y 只被用一次，但看下面这种图：

          x ──┐
               mul ──► t ──┐
          x ──┘            add ──► z      （x 被用了两次！）
          w ───────────────┘

x 同时是两次运算的输入。如果 Tensor 直接 pub data: Vec<f32>，克隆一个 Tensor 就要复制整份数据；而 Rc<T>（引用计数）允许无数个 Tensor 句柄共享同一个底层数据，克隆只把引用计数 +1，是 O(1) 的操作。这也是 #[derive(Clone)] 能放心用的前提。

4.2 RefCell：共享的同时允许修改

Rc 只解决"共享"，不解决"可变"——Rc<T> 的内容默认不可变。但我们的张量既被多处引用、又需要更新（前向写 data、反向写 grad）。Rust 的常规借用规则（一个值同时只能有一个可变借用）会拒绝这种写法。

RefCell<T>（内部可变性）把"借用检查"从编译期推迟到运行期：borrow() 只读借用、borrow_mut() 可变借用，同一时刻只能有一个 borrow_mut，违反就 panic。这就是"共享可变数据"的标准解法：

对应到代码里的经典模式：

let sd = self.data.borrow();          // 读：前向计算
// …… 用完立刻释放，避免借用重叠
drop(sd);
let mut g = self.grad.borrow_mut();   // 写：反向累加梯度

注意代码里几乎每个算子的前向都有 drop(sd); drop(od);——因为之后构造结果节点要 borrow_mut 自己的东西，如果 data 的只读借用还挂着，同一 RefCell 上"只读借用未释放就再取可变借用"会在运行时 panic。这是 RefCell 使用者的纪律。

4.3 为什么 data 和 grad 分开两个 Rc

细看：data 和 grad 是两个独立的 Rc<RefCell<Vec<f32>>>。除了"语义上一个是数值、一个是梯度"之外，还有两个实际好处：

1. 可以同时读 data、写 grad：反向闭包里常见 let sd_b = sd.borrow(); let mut sgm = sg.borrow_mut();——两者不是同一个 RefCell，互不冲突。
2. data 的指针充当计算图节点 ID：backward 里用 Rc::as_ptr(&t.data) as usize 唯一标识一个节点（见第 6 节），而 grad 的指针用来判断"两个父节点是不是同一个张量"（Rc::ptr_eq(&sg, &og)，处理 x*x 自乘的梯度合并）。

4.4 为什么不是 Arc

Rc 是单线程引用计数（快，无锁），Arc<Mutex> 是多线程方案。我们的训练是单线程的，用 Rc + RefCell 更轻、更符合"从零实现"的教学目标。等将来要并行训练，再换成 Arc<Mutex> 或 Arc<RwLock> 不迟。

5. parents: Rc<Vec> —— 为什么不能直接用 Vec

这是全项目最重要的一个设计决策。看结构体注释（源码原文）：

/// 父节点列表。
/// 注意：用 `Rc<Vec<_>>` 而不是 `Vec<Tensor&gt;`——
/// 若直接存 Vec，`derive(Clone)` 会递归深拷贝整棵祖先计算图，
/// 深层图上每次建节点都是 O(图深) 的灾难。用 Rc 共享后克隆是 O(1)。
parents: Rc<Vec<Tensor>>,

5.1 深拷贝灾难是怎么发生的

每个运算在构造结果时都要记录父节点：

result.parents = Rc::new(vec![self.clone(), other.clone()]);   // 每个算子都这么干

问题出在 self.clone()。如果 parents 的类型是 Vec<Tensor>，那么 clone 一个 Tensor 时：

• data：Rc 共享 → O(1)
• shape：Vec<usize> → O(维度数)
• grad：Rc 共享 → O(1)
• parents：Vec<Tensor> → 递归克隆所有父节点！

而父节点自己又有 parents……于是 clone 会递归复制整棵祖先子树。设计算图深度为 L：

线性回归的计算图只有几层，L 很小，第 2 课那样写还能跑；但到了 GPT（几十层 Transformer），每步前向要建几百上千个节点，O(L²) 的复制会直接拖垮训练——这就是"性能灾难"。

5.2 第 2 课的旧代码 vs 现在的写法

第 2 课文档里的版本（能跑，但有问题）：

// ❌ 旧版：parents 直接存 Vec<Tensor>
pub struct Tensor {
    data: Rc<RefCell<Vec<f32>>>,
    shape: Vec<usize>,
    grad: Rc<RefCell<Vec<f32>>>,
    requires_grad: bool,
    parents: Vec<Tensor>,          // clone 时递归复制祖先图！
    backward: Option<BackwardFn>,
}

重构后的版本：

// ✅ 现在的版本：parents 用 Rc 共享
pub struct Tensor {
    // ……
    parents: Rc<Vec<Tensor>>,      // clone 时只复制一个 Rc 句柄
    backward: Option<BackwardFn>,
}

对比要点：

• 旧版：parents: Vec<Tensor> 里存的是"真实拥有的子图副本"，克隆 = 深拷贝，图越深越慢。
• 新版：parents: Rc<Vec<Tensor>> 里存的是"共享的子图句柄"，克隆 = 引用计数 +1，O(1)。所有节点的 parents 共享同一棵祖先图，图在内存里只有一份。

这也解释了为什么 backward 里能放心地 for p in t.parents.iter() 递归——遍历的是共享图，不会触发任何复制。

6. backward 的 DFS 拓扑排序

6.1 为什么需要拓扑排序

反向传播要求梯度从输出端向输入端逐层传递：z = (x·y) + w 里，必须先执行 z 的 backward（把梯度分给 t、w），t 的梯度就绪后才能执行 t 的 backward（再分给 x、y）。执行顺序由计算图的依赖关系决定，而 DFS 后序 + 逆序执行恰好给出这个顺序。

6.2 代码：迭代式 DFS + HashSet 去重

注：早期版本使用递归 DFS，后改为迭代式栈 DFS（三色标记法）， 因为深层计算图（4 层 Transformer + batch=16）的递归深度会超过 Windows 默认 1MB 栈。 见 src/autograd.rs。

pub fn backward(&self) {
    assert_eq!(self.rank(), 0, "backward() 只支持标量（0 维）输出，当前形状 {:?}", self.shape);
    {
        let mut g = self.grad.borrow_mut();
        g[0] = 1.0;                     // ① loss 的梯度 = 1（∂loss/∂loss = 1）
    }

    // 迭代式 DFS 拓扑排序（避免递归栈溢出）
    // 用显式栈模拟递归：每个元素 (node, child_index) 表示"该节点的第 child_index 个子节点待访问"
    let mut order: Vec<Tensor> = Vec::new();
    let mut visited: HashSet<usize> = HashSet::new();
    let mut stack: Vec<(Tensor, usize)> = Vec::new();

    let key = Rc::as_ptr(&self.grad) as usize;   // 用 grad 的指针作节点唯一 ID
    if visited.insert(key) {
        stack.push((self.clone(), 0));
    }

    // ② 迭代式后序 DFS：先展开所有子节点，再入队
    while let Some((node, idx)) = stack.last_mut() {
        if *idx < node.parents.len() {
            let child = node.parents[*idx].clone();
            *idx += 1;
            let child_key = Rc::as_ptr(&child.grad) as usize;
            if visited.insert(child_key) {
                stack.push((child, 0));           // ③ 先递归父节点
            }
        } else {
            let (node, _) = stack.pop().unwrap();
            order.push(node);                     // ④ 后序入队：子在前、父在后
        }
    }

    for t in order.iter().rev() {                 // ⑤ 逆序遍历：从根走向叶
        if let Some(b) = &t.backward {
            b();                                  // 执行各节点的反向闭包
        }
    }
}

三个关键设计：

6.3 走一遍例子：z = (x·y) + w

        x ──┐
             mul ──► t ──┐
        y ──┘            add ──► z
        w ───────────────┘

DFS 后序结果（子在前、父在后）：order = [x, y, t, w, z]，逆序执行（z → w → t → y → x）：

和 test_chain_rule 验证的结果一致：∂z/∂x = y = 3、∂z/∂y = x = 2、∂z/∂w = 1 ✓

7. module.rs：Module trait 与 zero_grad_all

7.1 trait 定义

深度学习里一切"可训练的结构"（Linear、LayerNorm、Transformer Block、整个 GPT……）都需要回答一个问题：你有多少参数？ module.rs 就把这个回答统一成接口：

// src/module.rs（全文）
use crate::tensor::Tensor;

/// 模块接口：任何可训练结构都实现它
pub trait Module {
    /// 返回模块的所有参数（含嵌套子模块）
    fn parameters(&self) -> Vec<Tensor>;
}

/// 便捷方法：清零所有参数的梯度
pub fn zero_grad_all(module: &dyn Module) {
    for p in module.parameters() {
        p.zero_grad();
    }
}

Module 目前只要求实现一个方法 parameters()。别小看这一个方法——它是优化器、训练循环、梯度清零与模型结构之间的唯一桥梁。

7.2 三个实现者：叶子模块

layers.rs 里三个最简单的实现：

// Linear：线性层 y = xW + b
impl Module for Linear {
    fn parameters(&self) -> Vec<Tensor> {
        vec![self.weight.clone(), self.bias.clone()]
    }
}

// LayerNorm：层归一化
impl Module for LayerNorm {
    fn parameters(&self) -> Vec<Tensor> {
        vec![self.gamma.clone(), self.beta.clone()]
    }
}

// Embedding：词嵌入查表
impl Module for Embedding {
    fn parameters(&self) -> Vec<Tensor> {
        vec![self.table.clone()]
    }
}

注意 clone() 在这里是安全的：因为 Tensor 内部全是 Rc（第 4、5 节），克隆只复制句柄，不复制数据——返回 Vec<Tensor> 给外部持有也不会有性能问题。

7.3 组合模式：大模块递归收集小模块的参数

model.rs 里的 GPT 由若干 TransformerBlock 组成，TransformerBlock 又由 LayerNorm、MultiHeadAttention、Linear 组成。每个组合模块实现 Module 时，把子模块的参数拼接起来：

// model.rs：GPT
impl Module for GPT {
    fn parameters(&self) -> Vec<Tensor> {
        let mut ps = self.tok_emb.parameters();      // 词嵌入
        for block in &self.blocks {
            ps.extend(block.parameters());           // 每层 Transformer Block
        }
        ps.extend(self.ln_f.parameters());           // 最终 LayerNorm
        ps.extend(self.lm_head.parameters());        // 输出头
        ps
    }
}

因为 parameters() 是 trait 方法，block.parameters() 会自动递归调用 TransformerBlock 的实现，后者又调用 MultiHeadAttention 的实现……最终一棵参数树被拍平成一张参数表。

GPT.parameters()
├── tok_emb.parameters()        → [table]
├── block[0].parameters()       → ln1[γ,β] + attn[c_q,c_k,c_v,c_proj] + ln2[γ,β] + mlp[W1,b1,W2,b2]
├── block[1].parameters()       → …
└── ln_f.parameters()           → [γ, β]
    （lm_head 通过权重绑定复用 tok_emb.table.T，不单独存储）

7.4 与优化器的配合（main.rs 的 demo_xor）

main.rs 里第 7 课的 XOR 演示展示了完整闭环：

// ① 收集所有参数，一次性交给优化器
let params: Vec<Tensor> = {
    let mut ps = fc1.parameters();
    ps.extend(fc2.parameters());
    ps
};
let opt = SGD::new(0.5, params);

// ② 训练循环：前向 → 反向 → 更新 → 清梯度
for step in 0..1000 {
    let h = relu(&fc1.forward(&x_data));
    let logits = fc2.forward(&h);
    let loss = cross_entropy_loss(&logits, &y_targets);
    loss.backward();     // 自动算所有参数的梯度
    opt.step();          // SGD 遍历 params 更新：θ -= lr·g
    opt.zero_grad();     // 清零，防累加
}

分工很清晰：

zero_grad_all 就是"把清梯度也统一成接口"：传入任意 &dyn Module，它调用 parameters() 拿到所有参数逐个清零。实际训练中优化器（如上面 SGD::zero_grad）自带同样的清零逻辑，而 zero_grad_all 的价值在于：不经过优化器也能直接清任意模块的梯度（比如想"只求梯度、手动更新"时）。任何新模块只要实现 Module，就能复用这套参数管理设施。

8. 运行与测试

cargo test     # 全部测试通过
cargo run      # 依次演示：XOR MLP → BPE 分词 → 小 GPT 训练与生成

至此，项目的"地基"（阶段一）全部完成。从下一课开始，所有新代码（层、损失、优化器、Transformer）都将站在 tensor.rs + module.rs 这两块基石之上。

9. 动手练习

1. 给 Linear 之外再写一个你自己的模块（比如 BiasOnly { bias: Tensor }），实现 Module trait，并用 zero_grad_all 清一次梯度，确认能编译通过。
2. 思考：parameters() 返回 Vec<Tensor> 时我们敢 clone，靠的是什么？如果 data 不是 Rc 而是裸 Vec<f32>，clone 一次要花多少钱？
3. 画一画 z = (x·y) + w 的计算图，标出 DFS 后序的 order 数组内容，再写出逆序执行的顺序（参考答案在第 6.3 节）。
4. 挑战题：把 parents 改成 Vec<Tensor>，跑 cargo test，确认功能仍然正确；再构造一个 100 层深的计算图，对比两种写法每次前向/反向的耗时，体会"O(1) vs O(L)"的差别。
5. 阅读 model.rs 里 TransformerBlock 的 impl Module，数一数一个 Block 最终贡献多少个参数张量。

10. 本课总结

• 项目按职责拆成 12 个模块，依赖呈金字塔形，tensor.rs 是最底层地基
• Rc<RefCell<Vec<f32>>> = 共享（Rc）+ 可变（RefCell），是单线程下"共享可变数据"的标准方案
• parents: Rc<Vec<Tensor>> 让克隆从 O(图深) 降到 O(1)，避免深拷贝灾难
• backward 用 Rc::as_ptr 做节点 ID、DFS 后序 + 逆序执行，保证梯度按拓扑序从 loss 流回叶子
• Module trait 统一"收集参数"，配合 zero_grad_all、优化器、训练循环形成完整闭环
• 下一课：在 Module 之上实现第一个神经网络层 Linear 与激活函数，搭建第一个真正的网络

代码位置：src/tensor.rs（本章新增：broadcast_shapes / broadcast_map / binary / sum / sum_last_dim / softmax_last_dim / permute / gather_rows / 3D matmul） 演示入口：src/main.rs

1. 本课要搞懂的问题

1. 两个形状不同的张量怎么直接相加？"广播"到底是什么规则？
2. "求和"这个简单的操作，反向传播时梯度应该怎么分？
3. softmax 为什么要"先减最大值"？不减会出什么问题？
4. 形状重排（permute）、按行取数（gather_rows）的梯度怎么传回去？
5. 一批矩阵同时做乘法（3D 批量 matmul）是怎么实现的？

2. 广播 Broadcasting：形状不同也能做运算

2.1 为什么需要广播

上一课的 add 要求两个张量形状完全相同。但神经网络里处处是"小形状 + 大形状"的运算：

• 给一批数据加偏置：[B, D] + [D]
• 给一批样本各自乘一个标量：[B, D] * []
• 线性层 y = xW + b：[B, D] @ [D, out] + [out]

如果每次都要求形状一致，我们就得手动把 [D] 复制成 [B, D] 再算。广播（broadcasting） 就是自动完成这件事：只要形状"兼容"，Rust 帮你把小张量"展开"到大张量的形状再逐元素运算。

2.2 numpy 广播的三条规则

我们的实现完全遵循 numpy 的广播规则：

几个直观例子（1 表示"这一维可以复制"）：

2.3 broadcast_shapes：把规则变成代码

/// 计算两个形状广播后的形状（numpy 广播规则，从右向左对齐）：
/// - 维度相等或其中一个为 1 即可广播
fn broadcast_shapes(a: &[usize], b: &[usize]) -> Option<Vec<usize>> {
    let n = a.len().max(b.len());
    let mut out = vec![1usize; n];
    for i in 0..n {
        // 从右往左对齐：i 从 0 数到 n-1 时，i+n-1 正好指向"倒数第 n-i 个维度"
        // 如果 i + len 超出本形状的维度范围，就当作"前面补了 1"
        let da = if i + a.len() >= n { a[i + a.len() - n] } else { 1 };
        let db = if i + b.len() >= n { b[i + b.len() - n] } else { 1 };
        if da == db {
            out[i] = da;
        } else if da == 1 {
            out[i] = db;
        } else if db == 1 {
            out[i] = da;
        } else {
            return None;   // 规则 ③：无法广播
        }
    }
    Some(out)
}

要点：

• i + a.len() >= n 判断"对齐后当前位是否越过了 a 的最左端"，越过了说明 a 在这一维没有对应维度，视作补 1（规则 ① 的"右对齐"实现）。
• 三条分支对应规则 ②：相等取它、一方为 1 取另一方、都不满足返回 None。
• 返回 Option 而不是直接 panic，把"能不能广播"的判定权交给调用方。

验证一下 [2,3] 与 [3]（n = 2）：

结果 [2, 3] ✓

2.4 broadcast_map：数据怎么搬

光知道"结果形状"还不够——结果每个位置的值要从源数据的哪里取。broadcast_map 就是干这个的：给定源形状 src 和目标形状 target，返回一张映射表 map，map[t] = 目标第 t 个元素对应的源展平下标。

/// 计算源形状 src 广播到目标形状 target 时，每个目标元素对应的源展平下标。
/// src 的维度必须 <= target 且右对齐；src 中大小为 1 的维度索引固定为 0。
fn broadcast_map(target: &[usize], src: &[usize]) -> Vec<usize> {
    let offset = target.len() - src.len();       // 右对齐：源前面补 offset 个"伪维度"
    let total: usize = target.iter().product();
    let mut map = vec![0usize; total];
    for flat in 0..total {
        // 第一步：把目标的一维下标反解成多维下标（行优先）
        let mut r = flat;
        let mut t_idx = vec![0usize; target.len()];
        for d in (0..target.len()).rev() {
            t_idx[d] = r % target[d];
            r /= target[d];
        }
        // 第二步：多维下标映射回源的多维下标，再展平成源的一维下标
        let mut s_flat = 0usize;
        for d in 0..src.len() {
            let td = t_idx[d + offset];          // 源的第 d 维对应目标的第 d+offset 维
            let sd = if src[d] == 1 { 0 } else { td };  // 源大小为 1 的维度，索引固定为 0
            s_flat = s_flat * src[d] + sd;       // 行优先展平
        }
        map[flat] = s_flat;
    }
    map
}

反解多维下标：行优先存储下，一维下标 flat 的分解是"从最后一个维度开始取余、整除"：

t_idx[d] = r % target[d];  r /= target[d];    // d 从大到小

比如目标形状 [2, 3]，flat = 4 → t_idx[1] = 4%3 = 1，r = 1 → t_idx[0] = 1，即位置 (1, 1)。

映射到源：src 的第 d 维对应目标第 d + offset 维；若源该维大小为 1，则无论目标索引多大，都取下标 0（这就是"复制"）。

例：src = [3] 广播到 target = [2, 3]（相当于每行复制同一个向量）：

2.5 binary：把广播装进"逐元素运算"

broadcast_plan 先判断两个张量是否同形状：同形状就直接逐位对齐（无需映射），不同形状才计算两个 broadcast_map：

fn broadcast_plan(&self, other: &Tensor) -> (Vec<usize>, Option<Vec<usize>>, Option<Vec<usize>>) {
    if self.shape == other.shape {
        (self.shape.clone(), None, None)          // 同形状：没有 map
    } else {
        let target = broadcast_shapes(&self.shape, &other.shape)
            .unwrap_or_else(|| panic!("形状无法广播：{:?} vs {:?}", self.shape, other.shape));
        let map_a = if self.shape == target { None } else { Some(broadcast_map(&target, &self.shape)) };
        let map_b = if other.shape == target { None } else { Some(broadcast_map(&target, &other.shape)) };
        (target, map_a, map_b)
    }
}

add / sub / mul / div 四个运算全部复用同一个通用函数 binary：

fn binary(
    &self,
    other: &Tensor,
    fwd: impl Fn(f32, f32) -> f32 + 'static,          // 前向：c = f(a, b)
    back: impl Fn(f32, f32) -> (f32, f32) + 'static,  // 反向：返回 (∂c/∂a, ∂c/∂b)
) -> Tensor {
    let (target_shape, map_a, map_b) = self.broadcast_plan(other);
    let sa = self.data.borrow();
    let sb = other.data.borrow();
    let total: usize = target_shape.iter().product();
    let mut out_data = vec![0.0f32; total];
    for t in 0..total {
        let ia = match &map_a { Some(m) => m[t], None => t };   // 没有 map 就是一一对应
        let ib = match &map_b { Some(m) => m[t], None => t };
        out_data[t] = fwd(sa[ia], sb[ib]);                     // 按映射取数做运算
    }
    drop(sa);
    drop(sb);
    // …… 构造结果节点，注册反向闭包（见下）……
}

反向传播：广播的梯度遵循"谁被复制了，梯度就累加回谁"。因为一个源元素可能对应多个输出位置，∂loss/∂x = Σ_t (∂loss/∂out_t)·(∂out_t/∂x)，所以必须累加（+=）：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let (ga, gb) = if same_shape && Rc::ptr_eq(&sg, &og) {
        // 同一张量参与运算（如 x*x）：两条路径的梯度合并（第 2 课讲过）
        // ……（合并累加逻辑，返回 (true, true)）
        (true, true)
    } else {
        let mut sgm = sg.borrow_mut();
        let mut ogm = og.borrow_mut();
        for t in 0..g.len() {
            let ia = match &map_a_c { Some(m) => m[t], None => t };
            let ib = match &map_b_c { Some(m) => m[t], None => t };
            let (da, db) = back(sd_b[ia], od_b[ib]);
            sgm[ia] += g[t] * da;      // 梯度累加到"源位置"，而不是输出位置
            ogm[ib] += g[t] * db;
        }
        (false, false)
    };
    let _ = (ga, gb);
}));

用测试验证（test_broadcast_add）：

// [2,3] + [3] 广播：b 的每个元素被两行"共享"
let a = Tensor::param(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]);
let b = Tensor::param(vec![10.0, 20.0, 30.0], vec![3]);
let loss = a.add(&b).sum();
loss.backward();
// b 的每个元素被广播两次，梯度应为 2
assert_eq!(b.grad(), vec![2.0, 2.0, 2.0]);
assert_eq!(a.grad(), vec![1.0; 6]);

b[0] 参与了输出的两个位置 (0,0) 和 (1,0)，两个位置的梯度各 1，累加起来就是 2。✓

3. 归约：sum / sum_last_dim / mean

"归约"（reduction）就是把一堆数合并成一个数。神经网络里损失函数 loss 几乎总是某种归约的结果（loss.backward() 要求标量，所以归约是自动微分的必经之路）。

3.1 sum：全部求和成标量

/// 求和成标量，梯度均匀传给每个元素
pub fn sum(&self) -> Tensor {
    let total = self.data.borrow().iter().sum();
    let mut result = Tensor::new(vec![total], vec![], self.requires_grad);
    if self.requires_grad {
        // ……
        result.backward = Some(Rc::new(move || {
            let g = rg.borrow()[0];
            let mut sgm = sg.borrow_mut();
            for v in sgm.iter_mut() {
                *v += g;          // 每个元素的梯度都是 g（因为 ∂sum/∂x_i = 1）
            }
        }));
    }
    result
}

反向逻辑一句话：求和把每个元素的梯度都变成同一个 g。因为 sum(x) = x₀ + x₁ + …，对每个 x_i 求导都是 1，所以 ∂loss/∂x_i = g·1。

3.2 sum_last_dim：沿最后一维求和，保持维度

scores [B, D] 想按行求和得到 [B, 1]（注意最后一维保留为 1，而不是塌缩成 [B]），就用它：

/// 沿最后一维求和，**保持维度**：[..., D] -> [..., 1]
/// 反向：梯度广播回最后一维
pub fn sum_last_dim(&self) -> Tensor {
    assert!(self.rank() >= 1, "sum_last_dim 需要至少 1 维");
    let (pre, d) = (self.numel() / self.shape[self.rank() - 1], self.shape[self.rank() - 1]);
    // pre = 最后一维之前的所有元素数，d = 最后一维大小
    let sd = self.data.borrow();
    let mut out_data = vec![0.0f32; pre];
    for p in 0..pre {
        let mut s = 0.0;
        for j in 0..d {
            s += sd[p * d + j];   // 每 d 个连续元素求和
        }
        out_data[p] = s;
    }
    drop(sd);
    let mut new_shape = self.shape.clone();
    *new_shape.last_mut().unwrap() = 1;   // 最后一维改成 1，保持维度
    // ……
    result.backward = Some(Rc::new(move || {
        let g = rg.borrow();
        let mut sgm = sg.borrow_mut();
        for p in 0..pre {
            for j in 0..d {
                sgm[p * d + j] += g[p];   // 反向：把 g[p] 广播回这一行的 d 个位置
            }
        }
    }));
}

把 [2, 3] 变成 [2, 1]，前向和反向正好是"一对"：

前向： [1 2 3]   sum -> [6]
       [4 5 6]          [15]

反向： g = [2, 3]  broadcast -> [2 2 2]
                               [3 3 3]

3.3 mean：用 sum + mul_scalar 组合出来

我们的库没有单独的 mean 算子——均值就是"求和再除以元素数"，用现成的算子组合即可。看 loss.rs 里 MSE 的实现：

pub fn mse_loss(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub(target)
        .pow(2.0)
        .sum()                                  // Σ (pred - target)²
        .mul_scalar(1.0 / pred.numel() as f32)  // ÷ N，就是 mean
}

这体现了"先保证运算集合最小、用组合满足需求"的设计哲学：少一个算子 = 少一段前向 + 反向代码 = 少一类 bug。

三种归约对比：

4. softmax_last_dim：数值稳定的 softmax

4.1 公式与直觉

softmax 把一行任意实数分数变成概率分布（每个数在 0~1 之间，和为 1），是分类问题（以及 Transformer 注意力）的核心：

softmax(x)_i = e^{x_i} / Σ_j e^{x_j}      （对每一行独立做）

这就是测试 test_softmax 验证的结果：softmax([1,2,3]) = [0.0900, 0.2447, 0.6652]。

4.2 数值爆炸：不先减 max 会怎样

e^x 增长极快：e^88.7 就已经超过 f32 能表示的最大值（约 3.4×10³⁸），再大就是 inf。而神经网络的 logits 动辄上百，直接 exp 会溢出，除以 inf 得到 NaN，整个训练直接崩掉。

4.3 "减最大值"技巧

数学上 softmax 对"整体平移"不敏感，因为分子分母可以同乘一个常数：

softmax(x)_i = e^{x_i} / Σ_j e^{x_j}
             = e^{-m}·e^{x_i} / (e^{-m}·Σ_j e^{x_j})
             = e^{x_i - m} / Σ_j e^{x_j - m}         其中 m = max_j x_j

取 m = max(x) 后，x_i - m ≤ 0，e^{x_i-m} ∈ (0, 1]，永远不会溢出。代码里就是三步：找每行 max → exp(x - max) → 除以行和：

pub fn softmax_last_dim(&self) -> Tensor {
    assert!(self.rank() >= 1, "softmax_last_dim 需要至少 1 维");
    let (rows, d) = (self.numel() / self.shape[self.rank() - 1], self.shape[self.rank() - 1]);
    let sd = self.data.borrow();
    let mut out_data = vec![0.0f32; rows * d];
    // 先存 softmax 结果（反向需要）
    for r in 0..rows {
        let mut maxv = f32::NEG_INFINITY;
        for j in 0..d {
            maxv = maxv.max(sd[r * d + j]);            // ① 找每行最大值
        }
        let mut sum = 0.0;
        for j in 0..d {
            out_data[r * d + j] = (sd[r * d + j] - maxv).exp();  // ② 减 max 再 exp
            sum += out_data[r * d + j];
        }
        for j in 0..d {
            out_data[r * d + j] /= sum;                // ③ 归一化
        }
    }
    drop(sd);
    // ……（反向见下）……
}

注：f32::NEG_INFINITY 作为"比任何数都小"的初始值，保证第一个元素就能更新 max。

4.4 softmax 的反向传播

softmax 的输出互相耦合（每个 s_i 都依赖所有 x_j），梯度公式（s 是输出，g 是输出梯度）：

∂loss/∂x_i = s_i · (g_i - Σ_j g_j·s_j)

代码和公式一一对应：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let mut sgm = sg.borrow_mut();
    for r in 0..rows {
        // dot = Σ_j g_j * s_j
        let mut dot = 0.0;
        for j in 0..d {
            dot += g[r * d + j] * out_data[r * d + j];
        }
        for i in 0..d {
            sgm[r * d + i] += out_data[r * d + i] * (g[r * d + i] - dot);  // s_i(g_i - dot)
        }
    }
}));

推导速览（供参考）：∂s_i/∂x_j = s_i(δ_ij - s_j)（δ 是克罗内克记号），再用链式法则 ∂loss/∂x_i = Σ_j g_j·∂s_j/∂x_i = Σ_j g_j·s_j(δ_ji - s_i) = g_i·s_i - s_i·Σ_j g_j·s_j。

5. permute：任意维重排

5.1 和 reshape 的区别

• reshape：只改解释方式，不改数据顺序。[1,2,3,4,5,6] 排成 2×3 还是 3×2，内存里还是那 6 个数。
• permute：真的把数据搬家。把某个维度整体挪到另一个位置（transpose 就是 2 维的 permute([1,0])）。

permute(&[0, 2, 1]) 的意思是"新张量的第 0 维 = 旧张量的第 0 维，新第 1 维 = 旧第 2 维，新第 2 维 = 旧第 1 维"。

看测试 test_permute：[2, 3] 用 permute([1, 0]) 变成 [3, 2]，数据从 [1,2,3,4,5,6] 变成 [1,4,2,5,3,6]：

原始 [2,3]:          permute([1,0]) 后 [3,2]:
[1 2 3]              [1 4]
[4 5 6]              [2 5]
                     [3 6]

5.2 用"逆映射"实现

要算 out[of] = src[?]，思路是：先反解 of 的多维坐标 out_idx，再问"源的第 d 维应该取输出的哪一维"——答案就是 inv[d]，其中 inv 是 perm 的逆映射：

// 反解 permute 的逆映射：inv[perm[i]] = i
let mut inv = vec![0usize; self.rank()];
for (i, &d) in dims.iter().enumerate() {
    inv[d] = i;
}
// 前向：out_flat -> src_flat
let mut map = vec![0usize; total];
for out_flat in 0..total {
    // ① 反解输出多维下标（行优先）
    let mut r = out_flat;
    let mut out_idx = vec![0usize; self.rank()];
    for d in (0..self.rank()).rev() {
        out_idx[d] = r % new_shape[d];
        r /= new_shape[d];
    }
    // ② 源的第 d 维来自输出的第 inv[d] 维
    let mut src_flat = 0usize;
    for d in 0..self.rank() {
        let sd = out_idx[inv[d]];
        src_flat = src_flat * self.shape[d] + sd;
    }
    map[out_flat] = src_flat;
}
// ③ 按映射搬运数据
let mut out_data = vec![0.0f32; total];
for (of, &sf) in map.iter().enumerate() {
    out_data[of] = sd[sf];
}

反向：和广播类似——permute 也是"每个源元素只去一个地方"，但方向反过来，梯度要"放回原位"。用同一张 map，把输出的梯度 g[of] 累加到 sgm[sf]：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let mut sgm = sg.borrow_mut();
    for (of, &sf) in map.iter().enumerate() {
        sgm[sf] += g[of];        // 梯度"原路返回"
    }
}));

为什么是 +=？因为 permute 之后可能还有别的运算把梯度送到同一个源位置，或者源张量被多处共享，梯度必须累加（全项目统一约定）。

6. gather_rows：按行取数（Embedding 的基石）

6.1 语义

词嵌入（Embedding）的本质就是"查表"：一张 [V, D] 的表格（V 个词的向量），给定一堆词的下标，把对应行拿出来：

table [V, D] + indices [N]  ->  out [N, D]
out[i, :] = table[indices[i], :]

前向代码：

pub fn gather_rows(&self, indices: &[usize]) -> Tensor {
    assert_eq!(self.rank(), 2, "gather_rows 的 table 必须为 2 维");
    let (v, d) = (self.shape[0], self.shape[1]);
    let n = indices.len();
    let sd = self.data.borrow();
    let mut out_data = vec![0.0f32; n * d];
    for (i, &idx) in indices.iter().enumerate() {
        assert!(idx < v, "gather 索引越界：{} >= {}", idx, v);
        for j in 0..d {
            out_data[i * d + j] = sd[idx * d + j];   // 复制第 idx 行
        }
    }
    // ……
}

测试 test_gather_rows：table[3,2] 取第 0、2 行 → [1,2,5,6] ✓

6.2 反向：scatter-add（散射累加）

一行可能被取多次（两个位置都引用同一行），所以反向时梯度累加回表格对应行：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let mut sgm = sg.borrow_mut();
    for i in 0..idx_vec.len() {
        let row = idx_vec[i];                 // 输出第 i 行来自表格第 row 行
        for j in 0..d2 {
            sgm[row * d2 + j] += g[i * d2 + j];   // 累加，不覆盖！
        }
    }
}));

测试验证：

let t = Tensor::param(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]);
let g = t.gather_rows(&[0, 2]);
let loss = g.sum();
loss.backward();
// 第 0 行梯度 1，第 2 行梯度 1，第 1 行梯度 0
assert_eq!(t.grad(), vec![1.0, 1.0, 0.0, 0.0, 1.0, 1.0]);

如果 indices = [0, 0]，第 0 行的梯度就是 2——这就是"scatter add"里 add 的含义。

7. 3D 批量 matmul

7.1 从 2D 到 3D

第 1 课学过 2D 矩阵乘法 C[m,n] = A[m,k] @ B[k,n]。真实模型里我们通常同时处理一批样本，就变成了 3D 批量版本：前两个维度完全一样（都是 B），第三个、第四个维度才是矩阵的 m、k、n：

2D:   C[m,n]     = A[m,k]     @ B[k,n]
3D:   C[B,m,n]   = A[B,m,k]   @ B[B,k,n]     每个 batch 独立做一次 2D 乘法

批量维度 B 必须一致，内维 k 也必须一致，输出是 [B, m, n]。

7.2 前向代码

就是"对每个 batch 跑一遍第 1 课的三重循环"：

pub fn matmul(&self, other: &Tensor) -> Tensor {
    assert!(
        (self.rank() == 2 && other.rank() == 2) || (self.rank() == 3 && other.rank() == 3),
        "matmul 只支持 2D 或 3D（批量），当前 {}-D x {}-D", self.rank(), other.rank()
    );
    if self.rank() == 2 {
        return self.matmul_2d(other);        // 2D 走老路
    }
    // 3D 批量
    assert_eq!(self.shape[0], other.shape[0], "批量维度必须一致");
    let (b, m, k1) = (self.shape[0], self.shape[1], self.shape[2]);
    let (_, k2, n) = (other.shape[0], other.shape[1], other.shape[2]);
    assert_eq!(k1, k2, "矩阵乘法维度不匹配");

    let sd = self.data.borrow();
    let od = other.data.borrow();
    let mut out_data = vec![0.0f32; b * m * n];
    for bi in 0..b {                      // 多了一层 batch 循环
        for i in 0..m {
            for j in 0..n {
                let mut s = 0.0;
                for k in 0..k1 {
                    s += sd[(bi * m + i) * k1 + k] * od[(bi * k1 + k) * n + j];
                }
                out_data[(bi * m + i) * n + j] = s;
            }
        }
    }
    // ……
}

下标规律：每个 batch 内的下标公式和 2D 完全一样，只是多加了 bi * m * k1（A）和 bi * k1 * n（B）的批内偏移。看测试 test_matmul_3d：

A [2,1,2] = [[1,2],[3,4]]      B [2,2,1] = [[5],[6]] , [[7],[8]]
C [2,1,1] = [[1·5+2·6=17], [3·7+4·8=53]]   →  [17, 53] ✓

7.3 反向公式（对每个 batch 独立成立）

第 2 课的反向公式在批量下逐 batch 套用：

代码里就是三重循环写成 += 累加（sgm[(bi*m+i)*k1+k] += Σ_j g[...]·od[...]），并且保留了 Rc::ptr_eq 判断（x @ x 这种自乘场景两条路径合并）。注意维度细节：∂A 是 [B,m,k]（和 A 同形状），∂B 是 [B,k,n]，所以两个循环的累加下标不同：

// ∂A：sgm 用 (bi, i, k) 定位，对 j 求和
sgm[(bi * m + i) * k1 + k] += s;      // s = Σ_j g[bi,i,j] * B[bi,k,j]
// ∂B：ogm 用 (bi, k, j) 定位，对 i 求和
ogm[(bi * k1 + k) * n + j] += s;      // s = Σ_i A[bi,i,k] * g[bi,i,j]

8. 运行与测试

cargo test     # 全部测试通过（含新增：广播、3D matmul、softmax、permute、gather_rows）
cargo run      # 运行演示

重点测试与验证内容：

9. 动手练习

1. 手推一遍 [2,3] + [1,3] 的广播：画两个张量，标出 b 的每个元素被复制到哪些位置；再画反向时梯度如何累加回 b。
2. 用 sum_last_dim 和 mul_scalar 组合实现"按行求均值"（[B, D] -> [B, 1]），并推导梯度。
3. 设 d=2，手推 softmax 反向公式：s = softmax([x0, x1])，写出 ∂loss/∂x0 和 ∂loss/∂x1 的具体表达式。
4. 思考：permute 之后做 sum()，源张量的梯度是什么？为什么说 permute 只改变数据的"排列"而不改变"信息"？
5. 挑战题：gather_rows 若允许同一行被取两次（如 indices=[0,0]），反向时该行梯度是多少？这提示了 Embedding 的梯度为什么必须用 +=。

10. 本课总结

• 广播：右对齐、维度相等或一方为 1、否则失败；broadcast_shapes 定形状，broadcast_map 定搬数路线
• 广播的反向 = 梯度累加回源位置（谁被复制，谁收多份梯度）
• 归约：sum 变标量、sum_last_dim 保持维度；mean 用 sum + 除法组合
• softmax：减 max 保证数值稳定（数学等价、永不溢出），反向公式 s_i(g_i - dot)
• permute / gather_rows：都是"搬家"运算，反向用同一张映射表把梯度放回原位（累加）
• 3D matmul：批量维度不变，内层逐 batch 做 2D matmul，反向逐 batch 套 g@Bᵀ / Aᵀ@g
• 下一课：把项目拆成独立模块，并深入理解支撑这一切的 Rc<RefCell> 架构

代码位置：src/tensor.rs（grad / parents / 各运算的反向闭包）、src/autograd.rs（backward 拓扑排序） 演示入口：src/main.rs

1. 本课要搞懂的问题

1. 什么是梯度？为什么"梯度下降"能让模型学习？
2. 链式法则是怎么回事？为什么反向传播比正向求导更高效？
3. 怎么用 Rust 把"求导"自动化，让代码自动算出每个参数的梯度？

2. 从"拟合直线"说起

上一课我们实现了矩阵乘法。现在想象一个最简单的"学习"任务： 已知一些点 (x, y)，它们近似满足 y = 2x + 1，但参数 w、b 未知。怎么让程序自己找到 w=2、b=1？

思路：先随便给个初始值（比如 w=0, b=0），然后：

1. 算出预测值和真实值的差距 → 损失 loss
2. 问：如果我调大一点 w，loss 会变大还是变小？ → 这个方向就叫梯度
3. 沿着"让 loss 变小的方向"微调 w、b
4. 重复直到 loss 足够小

第 2 步里的"损失随参数的变化率"，数学上叫偏导数。所有参数偏导数的集合叫梯度。

3. 梯度下降公式

对参数 θ，更新规则只有一行：

θ = θ - 学习率 × ∂loss/∂θ

• ∂loss/∂θ > 0：θ 增大 loss 也增大 → 所以减去它，θ 变小
• ∂loss/∂θ < 0：θ 增大 loss 减小 → 减去它，θ 变大
• 学习率（learning rate）控制每次迈多大步子，太小太慢、太大震荡

我们采用的损失是 MSE（均方误差）：

loss = Σ (pred - y)²        # pred = x·w + b

4. 链式法则：复合函数求导的关键

问题来了：loss 是复合函数，loss = g(pred)，而 pred = f(w)。 要求 ∂loss/∂w，不能直接求，需要一层层拆开：

链式法则：如果 z = g(f(x))，那么

dz/dx = (dg/df) × (df/dx)

类比"流水线"：损失一路流经 pred 再到 w，每一段的"变化率"相乘就是总的"变化率"。

例：z = x*y + w，各梯度是：

∂z/∂x = y      （对 x 求导时把 y 当常数）
∂z/∂y = x      （对 y 求导时把 x 当常数）
∂z/∂w = 1      （加法求导是 1）

程序输出正是 dz/dx=3, dz/dy=2, dz/dw=1 ✓

5. 反向传播：把链式法则批量自动化

手工对每个参数求导太累，模型参数动辄百万个。怎么办？

核心思想：前向传播时，把每一步运算记录下来，形成一张计算图；反向传播时，从 loss 出发，把梯度沿图逐层往回传，每层的梯度都复用上一层的结果。

前向（边算边记）：  x ──┐
                     mul ──► t ──┐
                     y ──┘         add ──► z = 7
                     w ───────────┘

反向（梯度倒流）：  x ◄── grad=3
                     t ◄── grad=1 ──► w ◄── grad=1
                     y ◄── grad=2

关键好处：每个中间结果只算一次梯度，总计算量与前向相当，而不是对每个参数各算一次前向（那样要 O(参数个数) 倍计算）。这就是反向传播比"数值微分"快的原因。

6. Rust 实现：每个张量背一个"反向函数"

我们把张量升级成计算图节点：

pub struct Tensor {
    data: Rc<RefCell<Vec<f32>>>,   // 数值（共享可变）
    shape: Vec<usize>,
    grad: Rc<RefCell<Vec<f32>>>,   // 梯度，初始全 0，反向时累加
    requires_grad: bool,           // 这个张量是否需要梯度（参数=true，数据=false）
    parents: Vec<Tensor>,          // 我是由哪些张量算出来的
    backward: Option<BackwardFn>,  // 我的反向函数：把梯度传给我的父母
}

用到的 Rust 核心特性：

• Rc<T>（引用计数）：计算图里一个节点可能被多个运算共享，Rc 让多份句柄指向同一份数据
• RefCell<T>（内部可变性）：让"共享数据"同时可以"可变借用"，这是 Rust 所有权模型下实现可变共享的标准手法
• 每定义一个运算，就写两个东西：前向（算结果）+ 反向（算梯度传给输入）

7. 各运算的反向规则（背下这几条）

以乘法为例看代码（注意累加 +=，因为一个张量可能被多条路径使用，梯度要累加）：

pub fn mul(&self, other: &Tensor) -> Tensor {
    // —— 前向 ——
    let data = /* self * other 逐元素 */;
    let mut result = Tensor::new(data, shape, self.requires_grad || other.requires_grad);
    if result.requires_grad {
        // 捕获需要的句柄
        let rg = result.grad.clone();       // 读：自己的梯度
        let sg = self.grad.clone();          // 写：给 self 累加梯度
        let og = other.grad.clone();
        let sd = self.data.clone();          // 读：self 的数值（计算 g*b 需要）
        let od = other.data.clone();
        result.backward = Some(Rc::new(move || {
            let g = rg.borrow();             // 拿到"传到我这的梯度"
            let sd_b = sd.borrow();
            let od_b = od.borrow();
            let mut sgm = sg.borrow_mut();
            let mut ogm = og.borrow_mut();
            for i in 0..g.len() {
                sgm[i] += g[i] * od_b[i];    // ∂a = g * b
                ogm[i] += g[i] * sd_b[i];    // ∂b = g * a
            }
        }));
    }
    result
}

8. 反向传播主流程：拓扑排序

图是有向无环图（DAG）。要保证"子节点先算，父节点后算"，需要拓扑排序：

pub fn backward(&self) {
    // 1. loss 自己的梯度 = 1（d loss / d loss = 1）
    self.grad.borrow_mut()[0] = 1.0;
    // 2. 迭代式 DFS 收集拓扑序：子在前，父在后（避免递归栈溢出，见第 4 课 6.2 节）
    //    用显式栈模拟递归，三色标记法保证每个节点只入队一次
    // 3. 逆序遍历，逐个调用 backward 闭包：梯度从子流向父
    for t in order.iter().rev() { t.backward() }
}

用 Rc::as_ptr(&t.grad) 作为节点的唯一编号做去重（一个节点被多个运算共享时不能重复入队）。

9. 训练循环：四步舞

一次"学习"包含四步（也是所有深度学习框架训练循环的骨架）：

// 1. 前向：算预测和损失
let loss = x.matmul(&w).sub(&y).mul(...).sum();
// 2. 反向：自动算每个参数的梯度
loss.backward();
// 3. 更新：w -= lr * 梯度（这里手动实现，展示底层）
w.set_data(vec![w.data()[0] - lr * w.grad()[0], ...]);
// 4. 清零梯度：否则下一轮会把旧梯度累加上去
w.zero_grad();

运行 cargo run 可以看到损失从 285 一路降到 0，w→2.00、b→1.00，模型"学会"了 y=2x+1。

10. 一个重要的细节：梯度累加

为什么用 += 而不是 =？因为一个张量可能被多条路径使用。

例：z = x*x。求梯度时 x 有两条路径（z 对第一个 x 的导数是第二个 x，对第二个 x 的导数是第一个 x），必须相加：∂x = g*x + g*x = 2gx。

代码里我们用 Rc::ptr_eq(&sg, &og) 判断两个父节点是不是同一个张量，是的话就合并累加。这也是测试 test_grad_accumulation_and_zero 验证的内容。

11. 运行与测试

cargo test     # 8 个测试全部通过
cargo run      # 演示：链式法则验证 + 线性回归训练

12. 动手练习

1. 给张量加一个 pow(n) 运算（c = a^n），推导并实现它的反向传播（提示：∂a = g * n * a^(n-1)）。
2. 思考：为什么 loss 必须是标量（0 维）才能调用 backward()？如果不是标量会怎样？
3. 把学习率从 0.001 改成 1.0，观察 loss 曲线，解释发生了什么（振荡/发散）。
4. 这是本项目的"心脏"。试着不看代码，自己画出 loss = (x·w - y)² 的计算图并手算 ∂loss/∂w。

13. 本课总结

• 梯度 = 损失对参数的"变化方向"，梯度下降就是沿着反方向走
• 链式法则 + 计算图 = 反向传播，一次性高效算出所有梯度
• Rust 用 Rc<RefCell> + 闭包实现了可自动求导的张量
• 亲手跑通了第一个"会学习"的程序：线性回归
• 下一课：补齐张量运算（广播、softmax 等），为神经网络打基础