P1：向量化与 MLP

对应代码：src/tensor/tensor.py / src/nn/module.py / src/optim/optimizer.py / src/data/dataset.py 对应测试：tests/test_tensor.py / tests/test_mlp.py / tests/difftest_torch.py / tests/test_mnist.py 训练脚本：train/train_mlp.py 参考：Karpathy Zero to Hero（makemore 系列）；PyTorch autograd 源码结构

---

本章目标

读完能回答：

1. 为什么把 P0 的 Value(float) 换成 Tensor(ndarray) 能加速两个数量级？数学原理一字没变，换的是什么？
2. matmul 的反向为什么是 dL/dA = dL/dY @ B.T？能在脑子里”看见”形状匹配吗？
3. softmax 为什么要减最大值？不减会发生什么？
4. softmax + cross_entropy 合起来后梯度为什么化简成 p - y？这个化简带来的不仅是速度，还有数值稳定——为什么？
5. Broadcasting 在反向传播里容易错在哪？”维度被广播出来”意味着反向要做什么？

能动手实现：

1. Tensor 类 ~200 行，支持所有必要算子 + backward()
2. Module / Linear / ReLU / Sequential + SGD / Adam（都不用 torch.nn / torch.optim）
3. MNIST 数据加载 + 训练到 test accuracy ≥ 95%

本章不涉及（defer）

• Conv / BatchNorm / Dropout（P3+ 按需）
• GPU 加速、混合精度
• 学习率调度、梯度裁剪（P5 训 GPT 时引入）
• 高阶导数、JIT

---

1. 问题动机：P0 为什么不够

1.1 量化 P0 的瓶颈

MNIST 一个 batch（64 样本）前向 784 → 128 → 10 MLP：

• 64 × 784 × 128 + 64 × 128 × 10 ≈ 650 万次乘加
• 用 P0 的 Value：每次乘加都要新建一个 Value 对象 + 一个 Python 闭包
• Python 对象创建 + 属性访问 + 闭包调用 ≈ 每次 500 ns → 一步前向以秒计

1.2 numpy 的角色：是”有类型的数学”，不是框架

• numpy 只负责数值计算，不懂微分
• 但它把同一组数学操作压到了 C / BLAS 层—— matmul 一次 784 × 128 × 64 在 CPU 上只要 ~100 μs

⚠️ 常见误解：numpy 不是”另一个 PyTorch”。它没有 autograd，反向传播代码要我们自己写。 本章的内容，就是”把 P0 的 autograd 架构按图索骥地升级成基于 ndarray 的版本”。

1.3 P1 的最小可行目标

x  = Tensor(np.random.randn(64, 784), requires_grad=False)
W  = Tensor(np.random.randn(784, 128) * 0.01)
b  = Tensor(np.zeros(128))
h  = (x @ W + b).relu()            # 一次 forward = 一次 matmul + mask
logits = h @ W2 + b2
loss = logits.cross_entropy(y)
loss.backward()                    # 反向也全压到 numpy 层

---

2. 核心概念

2.1 Tensor 数据结构

字段	类型	作用
data	np.ndarray	前向数值
grad	np.ndarray	dL/d(self)，shape 与 data 一致
_prev	set[Tensor]	父节点集合
_op	str	算子名（调试用）
_backward	闭包	和 P0 同义：把 grad 回传到 _prev
requires_grad	bool	Optimizer 过滤用；不影响计算

与 P0 的关键区别：grad 不再是一个 float，而是一个与 data 同形的 ndarray。所有 += 从标量加变成 ndarray 就地加。

2.2 Matmul 的反向传播

前向：Y = A @ B，形状 (m,k) @ (k,n) = (m,n)

反向（形状匹配就等于正确性）：

• dL/dA = dL/dY @ B.T → (m,n) @ (n,k) = (m,k) ✓
• dL/dB = A.T @ dL/dY → (k,m) @ (m,n) = (k,n) ✓

直觉（和 P0 __mul__ 完全同构）：

• P0 标量：self.grad += other.data * out.grad
• P1 矩阵：A.grad += out.grad @ B.T

“乘以对方的值”在矩阵场景变成”乘以对方的转置”。

2.3 激活函数

ReLU：out = max(0, data)；反向 grad *= (data > 0) 是一个 mask。

softmax 的数值稳定（log-sum-exp trick）：

x = x - x.max(axis=-1, keepdims=True)   # 减最大值
exp_x = np.exp(x)                        # 最大元素现在是 exp(0) = 1
p = exp_x / exp_x.sum(axis=-1, keepdims=True)

不减 max：exp(1000) = inf，接着 inf/inf = nan，loss 一步变 nan。 减 max 之后：任何分量最大不超过 exp(0) = 1，永远安全。

log-softmax 同理：

log_softmax(z) = z - z.max - log(Σ exp(z - z.max))

2.4 cross_entropy：softmax + NLL 的联合梯度

对 logits z（shape (B, C)）和整数标签 y（shape (B,)）：

\[L = -\frac{1}{B}\sum_{i=1}^{B} \log \text{softmax}(z_i)[y_i]\]

核心化简：

\[\frac{\partial L}{\partial z_{i,j}} = \frac{1}{B} (p_{i,j} - \mathbb{1}[j = y_i])\]

其中 $p = \text{softmax}(z)$。

这个化简带来两件事：

1. 便宜：不需要显式构造 softmax 的 (C, C) 雅可比矩阵
2. 稳定：内部不做 log(p)，避免 p ≈ 0 时 log(0) = -inf

每个主流框架（PyTorch、JAX、TF）都做了这个合并——叫 cross_entropy_with_logits 或类似。

2.5 Broadcasting 的反向规则

前向：y = a + b，a.shape = (B, N)、b.shape = (N,) → y.shape = (B, N) 直觉：b 被”复制”了 B 份发给每一行。

反向：b.grad 必须 sum 掉被广播的维度（这里是 axis=0，大小 B）：

b.grad += out.grad.sum(axis=0)

一般规则（封装在 _unbroadcast 里）：

• 若 grad.ndim > target.ndim：左侧多出的维 sum 掉
• 若 target.shape[i] == 1 而 grad.shape[i] > 1：那一维 sum(keepdims=True)

忘了做这件事是 P1 最常见的 silent bug：numpy 在 (N,) += (B, N) 时会触发广播写入，数值不报错但错的离谱。

2.6 Kaiming (He) 初始化

ReLU 网络的方差分析：一个 Linear(nin, nout) 后接 ReLU，如果希望输出方差和输入一致，W 的每个元素标准差要取 $\sqrt{2/\text{nin}}$。

std = np.sqrt(2.0 / nin)
W = rng.standard_normal((nin, nout)) * std

为什么需要：深度网络如果方差一层层衰减（或爆炸），激活会趋近 0 或饱和。Kaiming 让训练从一开始就在合理量级。

2.7 数学 ↔ Python 对齐卡片

操作	数学	Python（numpy）
线性	y = xW + b	y = x @ W + b（广播 b）
softmax	p_i = e^{z_i}/Σe^{z_j}	exp(z - z.max()) / sum()
cross-entropy	L = -mean(log p_y)	合并版：梯度 (p - onehot)/B
matmul 反向	dA = dY B^T	dA = dY @ B.T
sum 反向	常数向量	grad 广播回原 shape
broadcast + 反向	“复制 B 份” → sum	grad.sum(axis=被广播的维)

---

3. 手写实现

完整代码在 src/tensor/tensor.py（~200 行）。下面按层次拆解。

3.1 _unbroadcast 工具

def _unbroadcast(grad, target_shape):
    """把 grad 形状 sum 回 target_shape"""
    while grad.ndim > len(target_shape):
        grad = grad.sum(axis=0)
    for i, dim in enumerate(target_shape):
        if dim == 1 and grad.shape[i] != 1:
            grad = grad.sum(axis=i, keepdims=True)
    return grad

两条规则覆盖所有广播场景：去维、sum 到 1。

3.2 Tensor.__init__

class Tensor:
    def __init__(self, data, _children=(), _op="", requires_grad=True):
        if isinstance(data, Tensor):
            data = data.data
        self.data = np.asarray(data, dtype=np.float64)
        self.grad = np.zeros_like(self.data)   # 立刻分配 buffer
        self._prev = set(_children)
        self._op = _op
        self._backward = lambda: None
        self.requires_grad = requires_grad

和 P0 的差别：grad 在 __init__ 里就预分配成 zeros_like(data)——省掉每个算子首次写入时的 None check，代价是中间 Tensor 多一份 shape 相同的 buffer。MNIST 级别的内存无所谓。

3.3 逐元素算子（含 broadcasting 反向）

def __add__(self, other):
    other = other if isinstance(other, Tensor) else Tensor(other, requires_grad=False)
    out = Tensor(self.data + other.data, (self, other), "+")
    def _backward():
        self.grad  += _unbroadcast(out.grad, self.data.shape)
        other.grad += _unbroadcast(out.grad, other.data.shape)
    out._backward = _backward
    return out

def __mul__(self, other):
    other = other if isinstance(other, Tensor) else Tensor(other, requires_grad=False)
    out = Tensor(self.data * other.data, (self, other), "*")
    def _backward():
        self.grad  += _unbroadcast(other.data * out.grad, self.data.shape)
        other.grad += _unbroadcast(self.data  * out.grad, other.data.shape)
    out._backward = _backward
    return out

def __pow__(self, other):
    assert isinstance(other, (int, float))
    out = Tensor(self.data ** other, (self,), f"**{other}")
    def _backward():
        self.grad += other * self.data ** (other - 1) * out.grad
    out._backward = _backward
    return out

减法、除法、neg 都走组合路径（同 P0），不需要自己的 _backward。

3.4 matmul

def __matmul__(self, other):
    other = other if isinstance(other, Tensor) else Tensor(other, requires_grad=False)
    assert self.data.ndim == 2 and other.data.ndim == 2
    out = Tensor(self.data @ other.data, (self, other), "@")
    def _backward():
        self.grad  += out.grad @ other.data.T
        other.grad += self.data.T @ out.grad
    out._backward = _backward
    return out

本项目只实现 2D matmul。batched matmul（如 P4 attention 的 (B, H, T, D) @ (B, H, D, T)）到时再按需扩展。

3.5 Reduction：sum / mean

def sum(self, axis=None, keepdims=False):
    out = Tensor(self.data.sum(axis=axis, keepdims=keepdims), (self,), "sum")
    def _backward():
        g = out.grad
        if axis is None:
            g = np.broadcast_to(g, self.data.shape)
        else:
            if not keepdims:
                g = np.expand_dims(g, axis=axis)
            g = np.broadcast_to(g, self.data.shape)
        self.grad += g
    out._backward = _backward
    return out

def mean(self, axis=None, keepdims=False):
    n = self.data.size if axis is None else self.data.shape[axis]
    return self.sum(axis=axis, keepdims=keepdims) * (1.0 / n)

sum 的反向是”广播”：sum 把每行加成一个数；反向就把这个数”复制”回每行。

3.6 非线性

def relu(self):
    out = Tensor(np.maximum(0.0, self.data), (self,), "relu")
    def _backward():
        self.grad += (self.data > 0).astype(self.data.dtype) * out.grad
    out._backward = _backward
    return out

def exp(self):
    out = Tensor(np.exp(self.data), (self,), "exp")
    def _backward():
        self.grad += out.data * out.grad
    out._backward = _backward
    return out

def log(self):
    out = Tensor(np.log(self.data), (self,), "log")
    def _backward():
        self.grad += (1.0 / self.data) * out.grad
    out._backward = _backward
    return out

3.7 log_softmax + cross_entropy（合并版）

def log_softmax(self, axis=-1):
    x = self.data
    x_max = x.max(axis=axis, keepdims=True)
    shifted = x - x_max
    log_sum_exp = np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
    result = shifted - log_sum_exp
    out = Tensor(result, (self,), "log_softmax")
    def _backward():
        p = np.exp(result)
        sum_grad = out.grad.sum(axis=axis, keepdims=True)
        self.grad += out.grad - p * sum_grad
    out._backward = _backward
    return out

def cross_entropy(self, targets):
    """self: (B, C) logits；targets: (B,) int；返回标量 loss"""
    assert self.data.ndim == 2
    targets = np.asarray(targets).astype(np.int64)
    B, C = self.data.shape

    # 数值稳定 softmax
    x = self.data
    shifted = x - x.max(axis=-1, keepdims=True)
    exp_x = np.exp(shifted)
    softmax = exp_x / exp_x.sum(axis=-1, keepdims=True)

    log_p = shifted - np.log(exp_x.sum(axis=-1, keepdims=True))
    loss_value = -log_p[np.arange(B), targets].mean()

    out = Tensor(loss_value, (self,), "cross_entropy")
    def _backward():
        # (p - onehot) / B —— 联合梯度的全部奥秘就这两行
        grad = softmax.copy()
        grad[np.arange(B), targets] -= 1.0
        grad /= B
        self.grad += grad * out.grad
    out._backward = _backward
    return out

3.8 backward()

def backward(self):
    assert self.data.ndim == 0 or self.data.size == 1
    topo = []
    visited = set()
    def build(v):
        if v in visited:
            return
        visited.add(v)
        for child in v._prev:
            build(child)
        topo.append(v)
    build(self)
    self.grad = np.ones_like(self.data)
    for v in reversed(topo):
        v._backward()

和 P0 一字不差，只是 grad 从 float 换成 ndarray。这是 P0 → P1 最直接的”同构升级”证据。

3.9 Module / Linear / ReLU / Sequential

class Module:
    def parameters(self):
        params, seen = [], set()
        def add(p):
            if isinstance(p, Tensor) and p.requires_grad and id(p) not in seen:
                params.append(p); seen.add(id(p))
        def walk(v):
            if isinstance(v, Tensor): add(v)
            elif isinstance(v, Module):
                for p in v.parameters(): add(p)
            elif isinstance(v, (list, tuple)):
                for item in v: walk(item)
        for _, v in self.__dict__.items():
            walk(v)
        return params

    def zero_grad(self):
        for p in self.parameters():
            p.grad = np.zeros_like(p.data)


class Linear(Module):
    def __init__(self, nin, nout, rng=None):
        rng = rng or np.random.default_rng()
        std = np.sqrt(2.0 / nin)                      # Kaiming
        self.W = Tensor(rng.standard_normal((nin, nout)) * std)
        self.b = Tensor(np.zeros(nout))
    def __call__(self, x):
        return x @ self.W + self.b


class ReLU(Module):
    def __call__(self, x):
        return x.relu()


class Sequential(Module):
    def __init__(self, *layers):
        self.layers = list(layers)
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

parameters() 的 id 去重：避免”一个 Tensor 在两处被引用”（如 self.also_W = self.W）时被 Optimizer 更新两次、步长翻倍。

3.10 SGD 与 Adam

class SGD(Optimizer):
    def __init__(self, params, lr=1e-2, momentum=0.0):
        super().__init__(params)
        self.lr, self.momentum = lr, momentum
        self.v = [np.zeros_like(p.data) for p in self.params]
    def step(self):
        for p, v in zip(self.params, self.v):
            v *= self.momentum
            v += p.grad
            p.data -= self.lr * v


class Adam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
        super().__init__(params)
        self.lr = lr
        self.b1, self.b2 = betas
        self.eps = eps
        self.t = 0
        self.m = [np.zeros_like(p.data) for p in self.params]
        self.v = [np.zeros_like(p.data) for p in self.params]
    def step(self):
        self.t += 1
        bc1 = 1.0 - self.b1 ** self.t
        bc2 = 1.0 - self.b2 ** self.t
        for i, p in enumerate(self.params):
            g = p.grad
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * g * g
            p.data -= self.lr * (self.m[i] / bc1) / (np.sqrt(self.v[i] / bc2) + self.eps)

Adam 的 bias correction 为什么关键：m, v 初始都是 0，前几步 m ≈ (1-β1)g（量级被压到 0.1 量级）。 除以 bc1 = 1 - β1^t（前几步接近 0.1）就把它”放大回”合理量级。 不做 bias correction 的观察结果：前 100 步 loss 几乎不动。

3.11 MNIST idx 解析

def _load_idx(path):
    with gzip.open(path, "rb") as f:
        magic = struct.unpack(">I", f.read(4))[0]
        ndim = magic & 0xFF
        shape = tuple(struct.unpack(">I", f.read(4))[0] for _ in range(ndim))
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(shape)

MNIST 的”idx 格式”：4 字节 magic number（低字节 = 维数）+ 每维 4 字节大小 + raw bytes。一个 5 行的解析器足够。

---

4. 梯度验证：对拍 PyTorch

4.1 两道防线

第一道：tests/test_tensor.py —— 数值梯度对拍。构造随机输入，解析 grad vs finite difference，< 1e-5。 第二道：tests/difftest_torch.py —— PyTorch oracle 对拍。torch.Tensor 算一遍作为”标准答案”，精度 < 1e-6。

和 TEMU 的 difftest with Spike 完全同构：异构实现跑同样输入，靠 oracle 证明正确性。

4.2 一个对拍用例

def test_dt_cross_entropy():
    targets = np.array([2, 0, 4])
    t_targets = torch.tensor(targets, dtype=torch.long)
    _run_pair(
        lambda a: a.cross_entropy(targets),
        lambda a: torch.nn.functional.cross_entropy(a, t_targets),
        [rng.standard_normal((3, 5))],
    )

_run_pair 内部做三件事：

1. 同一输入在两边做 forward；比较前向数值
2. 两边 backward()；比较 grad
3. np.testing.assert_allclose 带 atol=1e-6

49 个测试全绿 是 P1 正确性的硬证据。

---

5. 训练与结果

5.1 架构

model = Sequential(
    Linear(784, 128, rng=rng),
    ReLU(),
    Linear(128, 10, rng=rng),
)
opt = Adam(model.parameters(), lr=1e-3)

参数量：784 × 128 + 128 + 128 × 10 + 10 = 101,770 ≈ 10 万。

5.2 训练循环

loader = DataLoader(X_tr, y_tr, batch_size=64, shuffle=True, seed=0)

for epoch in range(5):
    for xb, yb in loader:
        opt.zero_grad()
        logits = model(Tensor(xb, requires_grad=False))
        loss = logits.cross_entropy(yb)
        loss.backward()
        opt.step()

5.3 实测结果（macOS M 芯片，CPU）

loading MNIST ...
  train: (60000, 784), test: (10000, 784)
model params: 101,770
epoch 1  avg_loss 0.2974  test_acc 0.9508  (0.5s)
epoch 2  avg_loss 0.1310  test_acc 0.9664  (0.6s)
epoch 3  avg_loss 0.0911  test_acc 0.9727  (0.7s)
epoch 4  avg_loss 0.0690  test_acc 0.9728  (0.6s)
epoch 5  avg_loss 0.0548  test_acc 0.9760  (0.6s)

final test accuracy: 0.9760   ← 超出 95% 验收线

每 epoch ≤ 1 秒——938 step × (64 × 784 × 128) 计算量在 numpy 上是 sub-second 级别的事。

5.4 性能对照（P0 → P1 加速比）

环节	P0（Value）	P1（Tensor）
一个 batch 前向	~1 秒（估）	~0.5 ms
一个 epoch	不现实	0.6 s
加速比	1×	~10³–10⁴×

变的只有粒度——数学完全不变。

---

踩坑清单

1. matmul 形状搞反：A.T @ dY 和 dY @ A.T 混了。方阵情形会静默算错，其他情形会立刻 raise——所以 test 一定要用非方阵
2. Broadcasting 反向忘了 sum：b.shape=(N,), out.grad.shape=(B,N)，直接 b.grad += out.grad 触发 numpy in-place 广播写入，数值全错但不报错
3. softmax 忘了减 max：exp(1000) = inf → inf/inf = nan，loss 一步变 nan，训练直接死
4. cross_entropy 天真写 -log(softmax(z))[y]：正常输入能跑，但某个 p ≈ 0 时 log(0) = -inf；要用合并版（梯度 p - y）
5. ReLU 的 mask 依赖 self.data：如果后续对 self 做 in-place 修改（我们不做，但要警惕），_backward 会拿到改过的 data——所以要么不做 in-place、要么在闭包里提前把 mask 捕获
6. Adam 忘了 bias correction：前几步 m/v ≈ 0，除出来的 update 几乎为零——表现是”前 100 步 loss 不动”；debug 半天才定位
7. MNIST 没归一化：像素 0-255 直接喂进去，第一次 matmul 输出就 ~10⁴ 量级，ReLU 后激活爆炸
8. 参数共享被 parameters() 重复收集：共享 Tensor 被 Optimizer 更新两次，步长翻倍；用 id 去重
9. backward() 前忘了 zero_grad()：P0 的坑在 P1 重演，训练指数发散
10. backward() 在非标量上调用：梯度种子 ones_like(data) 意义模糊；强制 assert 只在标量 loss 上调用
11. log_softmax 的 logsumexp 实现：直接 log(exp(z).sum()) 对极大 logits 会 overflow；必须 z - z.max 之后再做

---

本章小结

• 实现：Tensor 类（~200 行）+ matmul / softmax / cross_entropy + backward()
• 实现：Module / Linear / ReLU / Sequential
• 实现：SGD / Adam（手写 bias correction）
• 实现：MNIST idx 解析 + DataLoader
• 实现：MLP 训练到 test acc = 97.60%（≥ 95% 验收）
• 验证：49/49 tests 绿，包含对拍 PyTorch
• 能解释：matmul 反向 dL/dA = dL/dY @ B.T 的形状直觉
• 能解释：softmax+CE 合并梯度 (p - y)/B 既便宜又稳定
• 能解释：broadcasting 反向必须 sum 掉被广播的维

---

与 P2 的连接

P1 处理静态向量输入：MNIST 的每个样本都是固定的 784 维。 P2 要引入序列：每个样本是一串可变长度的字符。

从最简单的序列模型（bigram）开始——它其实就是一个 Embedding(vocab → vocab) 加 softmax。 你会发现 P2 的模型本质就是 P1 的 MLP 换了输入输出语义——autograd 机制一字不改。

P3 再用 MLP 处理一个”固定长度窗口”的字符序列（Bengio 2003 的 N-gram 神经语言模型）。 P4 引入 attention，让模型动态关注序列不同位置。 P5 把这些组合成完整的 mini-GPT。

主轴从这里开始从”分类器”切换到”语言模型”——但 autograd 基础设施就是 P1 构建的这套，不再改动。

---

P1 完成。下一章见 textbook/P2-语言模型基础.md。