第 6 课：损失函数与优化器 —— 让模型知道"错在哪、怎么改"

代码位置：

• 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

两类任务对应两个损失：

任务	预测什么	损失函数
回归	连续数值（房价、温度）	MSE（均方误差）
分类	属于哪个类别（猫/狗、下一个 token）	CrossEntropy（交叉熵）

语言模型本质是"预测下一个 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)
}

步骤	代码	作用
相减	pred.sub(target)	误差 eᵢ = predᵢ − targetᵢ
平方	.pow(2.0)	eᵢ²：消除正负抵消、放大大误差
求和	.sum()	Σ eᵢ²
取平均	.mul_scalar(1/N)	除以元素总数 N（pred.numel()）

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

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

∂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 个样本，示意）

样本	logits	softmax	正确类别	−log p
第 1 个	[2.0, 0.5]	[0.82, 0.18]	0	−ln 0.82 ≈ 0.20
第 2 个	[0.1, 1.8]	[0.15, 0.85]	1	−ln 0.85 ≈ 0.17

平均 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();
        }
    }
}

公式与代码的对应：

公式	代码
θ（当前参数值）	p.data()
g（梯度）	p.grad()
θ − lr·g	d.iter().zip(&g).map(|(v, g)| v - self.lr * g)
写回	p.set_data(updated)

两个方法的分工：

• 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！