从零手搓AI工程：NumPy实现自动微分与注意力机制

1. 从零手搓AI工程为什么我不建议你直接调包第一次看到ai-engineering-from-scratch这个项目名的时候我正坐在工位上啃一个推荐系统的排序模块。说实话那一瞬间我脑子里冒出来的念头是又来了一个从零实现的教程仓库。毕竟这几年从零手搓XX已经快被玩烂了从手写神经网络到从零实现一个操作系统GitHub上能搜出来一大把。但当我真正把这个仓库clone下来花了一个周末把里面的代码从头到尾跑了一遍之后我改变了看法。这个项目解决的是一个非常具体、也非常痛的问题大部分做AI应用的人其实并不理解自己每天调用的那些框架底层到底发生了什么。你写model.fit()的时候梯度是怎么算的你调attention的时候那个矩阵乘法的维度是怎么对齐的你用tokenizer的时候BPE到底是怎么把一句话切成token的这些问题在业务压力下往往被能跑就行给糊弄过去了但一旦线上出问题——loss不收敛、显存爆炸、推理结果诡异——你就抓瞎了。ai-engineering-from-scratch的核心价值就是把这些黑盒一个个拆开用最朴素的Python和NumPy把AI工程链路里的关键组件重新实现一遍。它适合谁我认为有三类人最该认真看第一类是刚入行做AI应用、只会调API的工程师你需要补上底层这一课第二类是要做模型部署和优化的同学你得知道每个算子的开销在哪第三类是想转行做AI但被各种框架劝退的人从零实现反而是最平滑的入门路径。我写这篇东西不是要给你复述一遍仓库的README而是想把我自己踩过的坑、想明白的原理、以及那些文档里不会写但实际很重要的细节完整地摊开讲一遍。你跟着走完至少能做到看到一个AI组件能说出它的输入输出、计算复杂度、以及为什么这么设计。2. 整体设计思路为什么用NumPy而不是PyTorch2.1 从零实现的边界在哪里在动手之前必须先想清楚一个问题从零到底从哪一层开始算起如果你真的从晶体管开始那这辈子都写不完。ai-engineering-from-scratch的取舍很聪明——它把边界划在了张量运算这一层。也就是说底层的矩阵乘法、广播、索引这些交给NumPy而所有AI特有的逻辑比如自动微分、注意力机制、优化器、分词器全部自己写。这个边界的选择背后有很实际的考量。NumPy的底层是C和Fortran实现的BLAS库矩阵运算速度足够快不会成为你理解算法的障碍。如果你连矩阵乘法都自己用Python循环写那跑一个稍微大点的模型等结果的时间足够你睡一觉学习体验直接崩掉。所以用NumPy是性价比最高的从零——既让你看清AI逻辑又不至于被性能拖死。我实测过一个对比用纯Python循环实现一个 512x512 的矩阵乘法大概需要几十秒用NumPy的np.dot毫秒级。这个差距在调试阶段是致命的因为你可能要反复跑几十次来验证梯度对不对。2.2 模块拆分的逻辑整个项目的模块拆分遵循了一条非常清晰的主线数据怎么进来 → 怎么变成向量 → 怎么算 → 怎么学 → 怎么出去。对应下来就是几个核心模块分词与词表构建文本到token的映射这是所有NLP任务的第一步自动微分引擎整个项目的心脏没有它就没法训练基础层实现线性层、激活函数、归一化注意力机制Transformer的核心也是最容易写错的地方优化器SGD、Adam这些决定了模型能不能收敛训练循环与评估把上面所有东西串起来这个拆分方式的好处是每个模块都可以独立测试。你可以先单独验证自动微分对不对再去写层层验证通过了再去搭注意力。这种自底向上、逐层验证的思路是我认为这个项目最值得学习的地方——它把一个大问题拆成了若干个可以独立验证的小问题。2.3 为什么不用现成框架做对比有人可能会问既然PyTorch已经这么好用了为什么还要费劲从零写我的理解是从零实现和用框架是两种完全不同的学习目标。用框架是在学怎么用工具从零实现是在学工具为什么这么设计。前者让你能快速交付后者让你在遇到边界问题时知道往哪个方向调。举个我自己的例子。之前做一个文本分类任务模型死活不收敛loss在0.69附近震荡二分类的随机水平。我调了学习率、换了优化器、加了正则都没用。后来我把从零实现的交叉熵和梯度检查代码翻出来手动算了一遍发现是label的维度搞错了——框架自动广播帮我兜住了错误但结果就是错的。这种问题只有你理解底层才排查得出来。3. 核心细节解析自动微分引擎到底怎么写的3.1 计算图与反向传播的本质自动微分听起来很玄但本质就一句话把前向计算过程记录下来反向的时候沿着记录倒着走一遍用链式法则算梯度。这里的关键是记录——你需要一个数据结构来存前向的每一步操作这个结构就是计算图。在ai-engineering-from-scratch里每个张量它叫Tensor类都维护了两个关键属性data存数值grad存梯度还有一个_backward函数指针指向怎么把梯度传给我的输入。前向的时候每个操作都会生成一个新的Tensor并把它的_backward挂上去反向的时候从loss开始调用_backward一层层把梯度传回去。我用一个最简单的例子说明。假设c a * b那么c._backward就是a.grad b.data * c.gradb.grad a.data * c.grad。这就是乘法求导的链式法则展开。加法更简单梯度直接原样传回去。你看所谓自动微分就是把微积分课本上的求导规则用代码一条条写出来而已。3.2 一个容易踩的坑梯度累加这里有个细节我第一次写的时候栽了跟头。梯度必须累加不能覆盖。为什么因为一个张量可能被用在多个地方比如y x * x x这里的x被用了两次反向的时候两条路径的梯度都要加到x.grad上。如果你写成x.grad ...那后算的那条路径就会把前面的覆盖掉梯度就错了。正确的做法是x.grad ...。而且每次反向传播之前必须把所有参数的梯度清零否则上一轮的梯度会累积到这一轮。这个清零操作在PyTorch里是optimizer.zero_grad()在从零实现里你得自己写。我见过太多新手忘了清零结果loss越跑越大还以为是学习率的问题。提示判断梯度是否正确的黄金标准是数值梯度检验。用(f(xh) - f(x-h)) / (2h)算出来的数值梯度和你反向传播算出来的解析梯度对比相对误差应该在 1e-6 量级。这个检验我建议每写一个新算子都跑一遍。3.3 广播机制的处理NumPy的广播是个双刃剑。前向的时候很方便(3,1) (1,4)自动变成(3,4)但反向的时候你得把广播撤销回去把梯度沿着被广播的维度求和。这一步如果处理不好梯度形状对不上直接报错。具体怎么做假设前向是a(3,1) b(1,4) c(3,4)反向时c.grad是(3,4)那么a.grad需要把c.grad沿着第1维size为4的那维求和变成(3,1)b.grad沿着第0维求和变成(1,4)。这个逻辑要写成一个通用的函数根据前向时的形状自动判断哪些维度需要求和。我当时的实现思路是先看两个输入的维度差把维度少的那个在前面补1然后逐维比较size为1的维度就求和并保持维度。4. 实操过程从分词到训练一个迷你GPT4.1 分词器BPE的朴素实现分词是NLP的第一道关。ai-engineering-from-scratch里实现的是BPEByte Pair Encoding这个算法现在是大模型分词的主流。它的核心思想特别朴素从字符开始不断把出现频率最高的相邻字符对合并成一个新token重复直到词表达到目标大小。我拿一句话举例low lower lowest。初始词表是l, o, w, e, r, s, t这些字符。统计相邻对发现lo出现3次ow出现3次er出现2次……合并频率最高的lo成lo然后继续。几轮之后low就变成了一个token。这个过程本质上是在用统计方法发现语言里的常见片段。实操的时候有几个细节要注意。第一词表大小要设合理。太小了一个词被切成很多碎片序列变长计算量爆炸太大了词表稀疏embedding矩阵浪费。一般英文任务3万到5万中文因为字符本身信息密度高可以小一些。第二特殊token要预留比如pad、unk、bos、eos这些在训练和推理时都要用。第三合并规则要存下来推理的时候要用同样的规则切分否则训练和推理不一致结果直接崩。4.2 注意力机制维度对齐是最大的坑注意力机制是Transformer的灵魂也是从零实现时最容易写错的地方。公式本身不复杂Attention(Q,K,V) softmax(QK^T / sqrt(d_k)) V。但真正写起来维度对齐能让你debug到怀疑人生。我建议你按这个顺序来。首先明确输入形状假设batch size是B序列长度是T每个头的维度是d。那么Q、K、V的形状都是(B, T, d)。第一步算Q K.transpose(-2, -1)得到(B, T, T)的注意力分数矩阵。这里transpose(-2, -1)是把最后两维交换千万别写成K.T那会把batch维也转了形状全乱。第二步除以sqrt(d)。为什么要除因为当d很大时QK的点积结果方差会很大softmax之后会变得非常尖锐梯度接近0训练不动。除以sqrt(d)相当于把方差拉回到1附近让softmax的输出更平滑。这个细节是原论文里的关键很多人抄代码的时候漏掉结果模型就是训不起来。第三步是mask。如果是自回归生成任务比如GPT需要把未来位置的注意力分数设成负无穷这样softmax之后这些位置的权重就是0保证每个位置只能看到自己前面的token。mask的形状是(T, T)的下三角矩阵用masked_fill把上三角部分填成-inf。第四步softmax然后和V相乘得到(B, T, d)的输出。如果是多头注意力还要把多个头的结果拼接起来再过一层线性变换。注意softmax的数值稳定性。直接算exp(x) / sum(exp(x))在x很大时会溢出。标准做法是先减去每行的最大值exp(x - max(x)) / sum(exp(x - max(x)))。这个技巧在从零实现里必须自己加上否则训练到一半就出NaN。4.3 训练循环把零件组装起来有了分词器、自动微分、层、注意力、优化器最后就是训练循环把它们串起来。一个标准的训练循环长这样for epoch in range(num_epochs): for batch in dataloader: # 1. 前向 logits model(batch.input_ids) # 2. 算loss loss cross_entropy(logits, batch.labels) # 3. 清零梯度 model.zero_grad() # 4. 反向 loss.backward() # 5. 更新参数 optimizer.step()看起来简单但每一步都有讲究。前向的时候要注意模型是不是在训练模式影响dropout和batchnorm的行为算loss的时候要注意ignore_indexpadding的位置不应该参与loss计算反向之前必须清零梯度更新参数的时候要注意学习率的调度。我实测下来从零实现一个两层的小Transformer在字符级语言模型任务上比如用莎士比亚文本训练大概几千步就能生成看起来像那么回事的文本。这个过程虽然慢但每一步你都能看到loss在降梯度在变这种看得见的感觉是调框架给不了的。5. 常见问题与排查技巧实录5.1 loss不降反升的排查顺序这是最常见的问题我整理了一个排查顺序按这个走基本能定位到原因排查项现象可能原因解决方法学习率loss震荡或爆炸学习率太大降到1e-4或更低试试梯度清零loss持续上升忘了zero_grad每步反向之前清零数据标签loss不降标签和输入错位检查对齐打印几个样本初始化loss一开始就NaN权重初始化太大用Xavier或He初始化数值稳定训练中途NaNsoftmax/log溢出加max减法、log加eps梯度消失loss降得很慢激活函数饱和换ReLU或加残差连接我自己的经验是先怀疑学习率和梯度清零这两个占了80%的问题。剩下的20%里数据问题又占一大半。所以排查的时候先打印几个batch的输入和标签肉眼看看对不对往往比调参快得多。5.2 显存不够怎么办从零实现虽然用的是NumPy但如果你把它搬到GPU上用CuPy或者自己写CUDA显存问题就来了。几个实用的技巧梯度累积batch size设小一点累积几个batch的梯度再更新一次等效于大batch混合精度用float16存激活值float32存参数显存能省一半激活重计算前向的时候不存中间激活反向的时候重新算一遍用时间换空间检查点只存部分层的激活中间的重算这些技巧在从零实现里理解起来特别直观因为你能看到每一块显存到底被谁占了。用框架的时候这些是黑盒从零写的时候你能精确控制。5.3 训练和推理结果不一致这个问题很隐蔽但很常见。原因通常有几个一是dropout训练时开着推理时要关掉二是batchnorm训练时用batch统计量推理时用全局统计量三是分词训练和推理用了不同的分词规则。从零实现的好处是这些模式切换你都得自己写写的时候就会想明白为什么要有这个切换。提示我习惯在推理前显式调用model.eval()训练前调用model.train()哪怕是从零实现的版本也保留这个习惯。这样代码可读性好也不容易忘。6. 这套东西学完能干什么说实话从零实现一遍AI工程链路不会让你立刻变成算法专家但它会给你一种底气。以前看到一个新模型你只能看论文里的公式现在你能在脑子里把它拆成一个个算子估算它的计算量和显存占用。以前调参靠玄学现在你知道每个超参数在数学上影响什么。我自己最大的收获是在做模型优化的时候。以前遇到推理慢只能瞎试各种框架的加速选项现在我能profile出时间花在哪个算子上知道是矩阵乘法的问题还是softmax的问题优化起来有的放矢。这种能力是调包调不出来的。如果你打算认真走一遍这个项目我的建议是别只看一定要自己敲一遍。看代码的时候觉得都懂自己写的时候才发现到处都是坑。尤其是自动微分和注意力这两块亲手写一遍你对深度学习的理解会上一个台阶。遇到卡住的地方先别急着搜答案自己打印形状、打印梯度、做数值检验这个debug的过程本身就是最好的学习。