第 5 课：线性层与激活函数 —— 神经网络的"积木"

代码位置：

• 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 又叫"全连接层 / 稠密层"。它把输入向量线性变换成另一个向量：

名字	形状	含义
x	[B, in]（或 [B, T, in]）	输入
W（weight）	[in, out]	权重
b（bias）	[out]	偏置
y	[B, out]（或 [B, T, out]）	输出

数学公式（对每个样本 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 初始化为全零。

为什么不能随便初始化？

初始化方式	后果
全 0	所有神经元输出相同、梯度相同，永远学不出差异
太大（如 ±5）	激活值迅速膨胀 → 数值溢出 / 梯度爆炸
太小（如 ±1e-4）	梯度极小 → 收敛慢如蜗牛
√(2/(in+out))	激活值的方差随层数基本不放大，训练稳定

直觉推导：假设输入 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]。

输入形状	内部计算形状	输出形状
[B, in]	[B, in]（原样）	[B, out]
[B, T, in]	[B·T, in]（展平）	[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 三种激活对比

激活	输出范围	公式	反向导数	特点
ReLU	[0, +∞)	max(0, x)	1（x>0）否则 0	最快，隐藏层最常用
GELU	约 (-0.17, +∞)	0.5x(1+tanh(√(2/π)(x+0.044715x³)))	见 7.2	平滑，GPT 系列默认
Tanh	(-1, 1)	(ex−e-x)/(ex+e-x)	1−c²	输出居中，两端饱和

补充：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，公式、反向导数与代码一一对应
• 下一课：损失函数与优化器，让模型"知道自己错在哪、怎么改"