完整 decoder 的梯度与资源总账

任务与读者起点：复用已证明的组件

本章把数据管道接到完整 decoder，并证明“全参数训练”在本项目中有可检验含义。读者应已完成 Transformer 实战中注意力、因果性、前馈层、归一化与反向传播的推导；这里不以调一个模型接口替代那些证明，而把已学组件整合成数据到梯度的工程。本项目 lm_lab/decoder.py 是该课自写实现的本地副本，来源路径与哈希单独记录；压缩包包含副本，不运行跨目录导入，也没有下载权重。需要重新学注意力的读者应先补足依赖，再来核对本章总账。

项目固定一层、四头、隐藏宽度二十四、前馈宽度四十八、窗口上限二十四；BPE 实际词表二百九十。注意力的 Q、K、V、输出投影都无 bias；前馈两个线性层有 bias；每层有两个带仿射参数的 LayerNorm，最后还有一次 LayerNorm。位置使用 learned absolute embedding，每个独立窗口从零开始；前馈非线性是精确 GELU；dropout 在此实验设零。输出头与输入 embedding 不共享权重。只有把这些选择列明，参数公式才能精确对齐文件。

从整数到词表 logits 的形状链

输入 x 形状 B×T，整数只能取零至 V−1。查表 E 为 V×d，位置表 P 为 C×d，初始隐藏状态为 $H_{b,t}=E_{x_{b,t}}+P_t$，形状 B×T×d。每个 block 先归一化后做因果多头注意力，并加回原输入，再归一化后送前馈并加回。最后归一化后乘 d×V 输出矩阵得到 B×T×V logits。整条链没有把目标喂入同一位置的隐藏状态；目标只在最后用于交叉熵。

对一个有效目标，令 logits 为 z、softmax 概率为 p、真实类别的 one-hot 为 y。交叉熵对 logits 的梯度是 p−y；多个有效位置再除以批内有效数。若最后隐藏行是 h、输出矩阵是 W，则输出矩阵梯度是外积 $h^T(p-y)$，隐藏行梯度是 $(p-y)W^T$。后者继续经过最后归一化与所有 block，不能在输出头之前 detach。查表会把同一编号多个出现位置的梯度累加到同一行；未出现输入编号的 embedding 行可暂为零，但整个参数张量依然可训练。

残差 $u=h+f(h)$ 的梯度为上游梯度乘 $I+J_f$，其中 $J_f$ 是支路的雅可比。因此残差提供一条恒等路径，同时还必须沿支路传播。把注意力或前馈包在 no_grad 中会删除支路梯度；仅检查总 loss 是否下降无法发现，因为输出头仍可适配固定特征。全参数训练要求所有预期参数 requires_grad 为真、纳入 optimizer、得到有限梯度，并在合理测试输入下有实际更新；这些是不同断言，不应只检查其中一个。

算例一：把一万九千三百二十个参数数出来

用 V 表示词表、d 表示宽度、f 表示前馈宽度、C 表示最大位置、L 表示层数。输入和输出两张矩阵共 $2Vd$，位置表 Cd。每层注意力四张 d×d 矩阵给 $4d^2$；前馈矩阵共 $2df$，bias 共 f+d；两次 LayerNorm 各有缩放与平移，共4d。合计每层 $4d^2+2df+f+5d$；最后归一化再加2d。没有把激活或梯度计成参数，也没有把四头误当四套完整 d×d 矩阵。

代入 V=290、d=24、f=48、C=24、L=1：两词表矩阵是13920，位置表576，注意力2304，前馈矩阵2304，前馈bias72，两次归一化96，最后归一化48，总数19320。独立检查一是遍历实际 model.parameters 的 numel 总和必须相等；独立检查二是按名字分组核验，防止恰好多一处又少一处抵消。头数四只拆分 d 的坐标，不额外乘四；若 d 不能被头数整除，配置应直接拒绝。

算例二：输出能学，但底层已被冻结

考虑一个两特征、两分类的简化输出：h=[1,2]，输出矩阵两列分别 [0.1,0.3] 与 [−0.2,0.4]。logits 是 [0.7,0.6]，真实类别为第一类。softmax 第一项约0.524979，梯度对 logits 是 [−0.475021,0.475021]；对 h 的梯度第一项约−0.142506，第二项约0.047502。这些非零值意味着如果 h 来自更早可训练参数，链式法则应继续传播。若代码把 h 当常量，输出矩阵仍有非零梯度，loss 仍可能下降，却不能说明底层学习。

独立数值检查用中心差分：分别把 h 的一个坐标加减很小的量，重新算交叉熵，两次损失之差除以两倍步长，应接近解析梯度。这个标准库例子只核验最后投影的局部链式法则，不是完整 Transformer 的训练实现。完整实现的验收另在项目测试中运行：十七个命名参数张量均得到有限非零范数，注意力 Q/K/V/O、位置表、两层前馈和归一化均在其中。某个具体参数元素为零不一定是故障，检查不能错误地要求所有元素每步都非零。

import math
W=[[.1,-.2],[.3,.4]]; h=[1.,2.]
def loss(h):
    z=[sum(h[i]*W[i][j] for i in range(2)) for j in range(2)]
    a=max(z); return a+math.log(sum(math.exp(v-a) for v in z))-z[0]
z=[sum(h[i]*W[i][j] for i in range(2)) for j in range(2)]
a=max(z); probs=[math.exp(v-a) for v in z]; total=sum(probs); probs=[v/total for v in probs]
gz=[probs[0]-1,probs[1]]
gh=[sum(gz[j]*W[i][j] for j in range(2)) for i in range(2)]
fd=[]; eps=1e-5
for i in range(2):
    plus=h.copy(); minus=h.copy(); plus[i]+=eps; minus[i]-=eps
    fd.append((loss(plus)-loss(minus))/(2*eps))
assert max(abs(a-b) for a,b in zip(gh,fd))<1e-9
V,d,f,C,L=290,24,48,24,1
params=2*V*d+C*d+L*(4*d*d+2*d*f+f+5*d)+2*d
assert params==19320
print('logits',z,'hidden gradient',gh,'finite difference',fd,'parameters',params)

from lm_lab.data import read_corpus
from lm_lab.engine import build,update,parameter_count
split,_=read_corpus()
model,opt,tok,windows,sampler=build(split)
receipt=update(model,opt,windows,sampler)
assert sum(p.numel() for p in model.parameters())==parameter_count(model.config)==19320
assert set(receipt['gradient_tensor_norms'])==set(dict(model.named_parameters()))
print('trainable tensors',len(receipt['gradient_tensor_norms']),'valid targets',receipt['targets'])

内存与算力估算的适用范围

若参数、梯度与 Adam 的一阶二阶状态均用 float32，稳态仅这四份张量约占十六乘参数数目的字节，本例是309120字节，约三百零二 KiB。还未计激活、临时工作区、优化器步数标量、Python对象、框架库和分配器。不能把它当进程内存测量。注意力分数约 B×头数×T×T，随窗口平方增长；词表 logits 约 B×T×V，词表变大也增加峰值。这个项目小到框架开销可能压过矩阵计算，因此短墙钟时间不能外推 GPU 或大模型效率。

常见训练计算近似 $6PD$ 把每个训练token的前后向投影成本压成参数数目的常数倍；它忽略部分注意力二次项、激活、归一化和实际内核，所以本课只把它作为预算量级说明，不当实测 FLOPs。对词表占比高、层少、窗口变化大的极小模型，误差尤其可能显著。第八章会同时报告参数、更新数、有效目标数和墙钟时间；若没有硬件计数器测量，就不会把近似数标成真实吞吐。

本章项目增量是可完整反传的模型、参数公式和梯度收据。模型架构来源明确，所有参数从随机初始化学习，没有“输出层训练代替 Transformer 训练”的简化。下一章固定训练循环与恢复状态，把一次有效更新扩大为可重放的八十次更新；我们需要证明中断前后的下一批、下一步 loss、参数和优化器状态都一致，而不是只证明保存文件能重新打开。

归一化本身也有参数与条件。对每个位置的d维向量a，均值是各坐标和除以d，方差是各坐标到均值的平方偏差再除以d；本实现不用无偏样本方差。归一化输出为缩放参数逐元素乘以中心化向量除以方差加epsilon的平方根，再加平移参数。epsilon固定为十万分之一。缩放与平移各有d个元素，这解释每个LayerNorm为何贡献二倍d参数，也解释bias不是被归一化“自动消失”后就无需记录。

前馈层先把每个位置独立地从d维投影到f维，再用GELU，最后回到d维。精确GELU是输入a乘标准正态累积分布函数；它不像ReLU在负半轴完全截成零，而是连续调节。所有位置共享同一前馈权重，但不混合位置；跨位置的信息混合来自注意力，位置身份来自位置表。把这些职责说清楚，才能理解为什么只在前馈层加宽改变参数与非线性容量，却没有直接扩大上下文上限。

局部输出梯度还可以从对数似然完整推导。单个目标k的loss是log各logit指数之和减去z_k；对z_j求导，第一项给exp(z_j)除以指数和，第二项给当j等于k时的一，否则零，合起来就是p_j减指示值。稳定实现先减最大logit不会改变概率，因为分子分母同乘一个因子；最大值在精确等价表达中只是数值基准。批内忽略位置完全不进入总和，平均梯度才除以实际有效数。

以h来自一个更早线性映射为例，h=xA。如果刚才求得隐藏梯度gh，则A的梯度是x的转置与gh的外积，输入梯度为gh乘A的转置。这一步说明输出投影的有限差分只是链上的局部验证；注意力、归一化和GELU仍有各自雅可比，完整自动微分需要沿全部节点累计。项目没有手写全部反传，而使用框架autograd执行这些已定义的操作，测试检查真实链路是否连通，不能把autograd这个调用本身当成推导。

检查更新也要考虑优化器。一个张量有梯度但没被放进optimizer参数组，仍不会被step修改；一个张量梯度恰好为零，AdamW衰减也可能使它改变。因此“参数变了”不能单独证明当前数据的梯度到达。合理验收组合包括参数对象集合与模型命名集合一致、梯度数值有限、所选受控批次各预期张量范数非零，以及实际更新差异。当前十三项测试分别覆盖配置总数和梯度，不把一个标量总范数代替十七个张量的细查。

完整模型前向的一个受控形状例子是B为二、T为三、d为二十四。每个头宽六，Q与K拆头后为二乘四乘三乘六，分数为二乘四乘三乘三；拼头后的向量回到二乘三乘二十四，再经输出头成为二乘三乘二百九十。四个头只是对同一二十四维投影结果做坐标分块，不能把输出宽度写成九十六。有效目标若只有五个，交叉熵分母五，尽管logits仍有六行。这个例子让形状总账与目标总账对接，也能定位把head轴、time轴或vocab轴放错的错误：softmax注意力沿key位置，输出分类softmax沿词表，两个归一化轴不可互换。

输入查表索引、注意力归一化轴和输出分类轴都有独立约定，形状相同不代表概率含义相同。

python -m unittest -v tests.test_pipeline.PipelineTests.test_parameter_formula_and_all_gradients tests.test_pipeline.PipelineTests.test_causality_and_padding_invariance