手写 ML 模块：7 个值得练习的实现

这七个实现练什么

这 7 个小实现适合练习从公式到代码：输入 shape 对不对、哪些数值会溢出、mask 到底挡住了什么。先写出能运行的版本，再用小例子检查边界，比只记住公式更有用。

参考实现在 code/interview_kit.py，只依赖 NumPy，自带自检，包括反向传播的数值梯度校验。先跑一遍：

python3 00-foundations/code/interview_kit.py

七个实现，先检查这些地方

采样代码还有个很容易漏掉的小细节：[2, 2, 2, 1] 取 top-k、k=2 时，到底留哪两个？本实现按 token ID 打破平局，保留 ID 0、1；如果只用“分数不低于第 k 名”来筛，就会留下前三个。两种约定都可以用，但不能一边写“固定 k 个”，一边在代码里保留所有并列项。已有 mask 不会被重新打开；top_p=1 表示不再截断。

softmax：先把数字缩到能算的范围

比如三个 logits 是 [1000, 1001, 1002]。直接算指数会溢出，但我们只关心它们之间的相对大小。每项都减去 1002，得到 [-2, -1, 0]，再算指数和归一化，概率约为 [0.0900, 0.2447, 0.6652]。

$$ p_i=\frac{e^{x_i-m}}{\sum_j e^{x_j-m}},\qquad m=\max_j x_j. $$

分子分母同时除以 $e^m$，数学结果没变。最大的指数变成 1，其余不超过 1，避免了指数溢出。但减 max 不是对所有数值问题的保证：

这里说的是配套 NumPy 教学实现，不是所有框架的默认行为。整数输入还会检查转换是否精确：像 [2**60, 2**60+1]，直接转成 float64 会丢掉差值，因此报错，不默默给出均匀概率。更高精度也不能保住任意小的概率，或修复已经在上游溢出的 logits。

要的是 log-probability，就直接算它

对 [0, -1000]，第二个概率在常见浮点格式里会下溢到 0。先 softmax 再取 log，得到 -inf；直接计算 log-softmax，却还能得到约 -1000。损失需要的是 log-probability，没必要先把它变成一个存不下的小概率。

import numpy as np

logits = np.array([0.0, -1000.0], dtype=np.float64)
shifted = logits - logits.max()
weights = np.exp(shifted)
probabilities = weights / weights.sum()
with np.errstate(divide="ignore"):
    indirect = np.log(probabilities)
direct = shifted - np.log(weights.sum())
assert np.isneginf(indirect[1])
assert np.allclose(direct, [0.0, -1000.0])

PyTorch 的 log_softmax（2.14，2026-10 核对）也明确区分直接计算与 log(softmax(x))。如果想接着理解怎样分块计算、为什么 attention 不必存整行概率，可以看 online softmax 的例子。这三个问题分别关心数值范围、输出形式和中间量存储，不是同一件事。

把前向和反向连起来

前面的模块可以逐个测试；反向传播则要把它们连起来。建议从一个隐层 ReLU、一个输出 logit 的二分类 MLP 开始，先画出前向过程，再反过来检查每一步的 shape 和梯度。

设 $N$ 个样本的 BCE 取平均，$z_1$ 是隐层输入，$a_1=\max(z_1,0)$，$z_2$ 是输出 logit。关键几步是：

$$\frac{\partial \mathcal{L}}{\partial z_2} = \frac{\sigma(z_2)-y}{N} \quad\longrightarrow\quad \frac{\partial \mathcal{L}}{\partial W_2} = a_1^\top \frac{\partial \mathcal{L}}{\partial z_2} \quad\longrightarrow\quad \frac{\partial \mathcal{L}}{\partial z_1} = \left(\frac{\partial \mathcal{L}}{\partial z_2} W_2^\top\right)\odot \mathbb{1}[z_1 > 0]$$

注意分母：整个 $p-y$ 都要除以 $N$。如果换成 $\frac{1}{N}\sum_i(p_i-y_i)^2$，梯度是 $2(p-y)p(1-p)/N$；若 loss 还带 $1/2$，系数 2 才会消失。例子：$N=2,p=0.8,y=1$，单个 logit 的 mean-BCE 梯度是 -0.1。不要把单样本公式和 batch 平均混在一起，见面试基础题。

写完解析梯度，可以选一个参数，用中心差分对照：

$$\frac{\partial \mathcal{L}}{\partial \theta} \approx \frac{\mathcal{L}(\theta + \epsilon) - \mathcal{L}(\theta - \epsilon)}{2\epsilon}$$

函数足够光滑时，中心差分的截断误差是 $O(\epsilon^2)$，前向差分是 $O(\epsilon)$。但 $\epsilon$ 太小会放大浮点舍入误差，跨过 ReLU 的不可导点也不适用这个误差结论。用 float64，检查几档步长，并避开激活值（activations）恰好为 0 的位置。

怎么练

可以先选一个模块，不必一次写完七个。比如先写 softmax，再检查每行概率和是否接近 1、被 mask 的位置是否为 0，以及加上同一个常数后结果是否在容差内一致。卡住时回来看对应的小节就好。

熟悉以后，试着先写测试再写实现。除了正常输入，也放进很大的 logits、部分 mask 和整行 mask。这样检查的不只是公式有没有写对，还有这段代码遇到不支持的输入时会怎么做。

继续阅读

• 面试基础题：大半在问同一件事：概念那一半
• 多头注意力 · 归一化 · 解码策略
• 参考实现：interview_kit.py 是这页的配套，attention_numpy.py 是完整的多头版本

快速学习：手写公式时先回答哪五件事