把注意力逐项算清楚

本章任务：让一次读取既能前向复算，也能反向复算

前两章已经定义 token、嵌入矩阵、softmax 与交叉熵梯度，但每个位置还没有读取其他位置。现在构造单头自注意力：每个位置依据查询与键的匹配，混合各位置的值向量。最终要从一个标量损失回溯到 Q、K、V，再回到生成它们的投影和输入。这里只计算两个位置的完整参考例子，不加因果遮罩；第四章才限制可见信息，并加入多头和位置。

先修为矩阵乘法、第二章的标量链式法则和共享梯度求和。使用行向量约定：$X\in\mathbb R^{T\times D}$，$T$ 是位置数、$D$ 是输入宽度；$W_Q,W_K\in\mathbb R^{D\times d}$，$W_V\in\mathbb R^{D\times r}$。查询 $Q=XW_Q$、键 $K=XW_K$ 均为 $T\times d$，值 $V=XW_V$ 为 $T\times r$。$d$ 必须相同才能做查询与键的点积，值宽度 $r$ 可以不同。此处 $V$ 表示值矩阵，词表大小不参与本章公式，避免与上一章同字母的另一含义混用。

一、先给三个投影分工，再写读取规则

查询规定当前行以什么特征寻找信息，键规定每个候选位置以什么特征被匹配，值规定被读取的内容。这是三种线性映射的计算职责，不是三份文本或人为规定的三种语义。自注意力表示三者由同一个 X 生成，但权重矩阵彼此独立。若意外将三个矩阵绑定成同一个参数，模型族已经改变，不能只因为形状相同就认为等价。

对查询位置 $i$ 与候选位置 $j$，计算点积分数 $S_{ij}=\sum_a Q_{ia}K_{ja}/\sqrt d$。每个查询要单独在全部候选位置上归一化：$A_{ij}=\exp(S_{ij})/\sum_k\exp(S_{ik})$。最后 $O_{i\ell}=\sum_jA_{ij}V_{j\ell}$。先有匹配分数，再有按行权重，最后是加权读取，三步的索引含义不同。若沿列归一化，得到的是另一种资源分配，通常不能保证每个查询的读取权重和为一。

注意力权重与词表预测概率的样本空间不同。A 的一行对“读取哪个位置”归一化；第二章的 p 对“下一个 token 是哪个类别”归一化。O 是值向量的混合，不必非负，也不必各维和为一。即便把 A 画成概率图，也不能直接拿最大权重位置当作预测 token，或把图中的亮块当成回答正确的证明。

还要把“转置”从形状记号变成实际索引。若 Q 两行为 [1,2]、[0,1]，K 两行为 [1,0]、[2,1]，则查询第一行与两把键的点积分别是一和四，第二行为零和一。正确 QK转置得到 [[1,4],[0,1]]；漏掉转置却得到 [[5,2],[2,1]]。两者都是二乘二，单靠形状断言发现不了。主算例选择对角矩阵方便手算，却会掩盖某些转置错误，因此独立检查必须增加这种非对称输入。

二、缩放来自哪些假设

用一个明确的随机模型解释缩放动机：假设查询和键每个坐标均值零、方差一，查询与键相互独立，不同坐标对也相互独立。对一项乘积 $q_ak_a$，均值为零，二阶矩为 $\mathbb E(q_a^2)\mathbb E(k_a^2)=1$，所以方差为一。不同坐标项的协方差为零，$d$ 项相加的方差为 $d$。除以 $\sqrt d$ 后方差回到一，这就是常用尺度的一个计算依据。

这些独立和方差条件不是训练网络始终满足的定律。例如查询与键来自相同输入，甚至二者相等时，点积包含平方项，均值不再是上面假设的零。缩放仍是我们固定的架构约定，但“方差等于一”的论证只能在声明的随机模型里精确成立。它不保证 softmax 永不集中，也不保证训练不会发散；初始分布、归一化和学习率还会改变数值行为。

算例一：两位置、二维值的完整前向

取 $X$ 为二阶单位矩阵，$d=r=2$。设 $a=\sqrt2\ln3$，$W_Q=aI$，$W_K=I$，$W_V=\operatorname{diag}(2,4)$。因为 $X$ 是单位矩阵，$Q=aI$，$K=I$，$V$ 的两行为 $[2,0]$、$[0,4]$。$QK^{\mathsf T}$ 仍为 $aI$，再除以 $\sqrt2$，$S$ 两行为 $[\ln3,0]$ 与 $[0,\ln3]$。单位矩阵是刻意选择的手算输入，不是实际语言模型必须采用的嵌入。

第一行指数为 [3,1]，总和四，权重 [3/4,1/4]；第二行为 [1/4,3/4]。读取第一行时，四分之三乘 [2,0]，加四分之一乘 [0,4]，得 [3/2,1]；第二行得 [1/2,3]。选择标量探针 J=O第一行第一维+2×O第二行第二维，得到 J=3/2+6=15/2。探针的作用是让后续反向有一个明确目标，不把这个人为线性目标当成语言建模损失。

前向有两种不同检查。先逐行核对 A 的非负与行和，确认归一化方向。再用凸组合边界：由于每行权重非负且和为一，O 第一维必须在两个值向量第一维零和二之间，第二维在零和四之间。输出 [1.5,1] 和 [0.5,3] 都满足。这个坐标范围检查不同于重新相乘，能揭示漏归一化或把分数直接当权重的错误。

凸组合的边界只约束被读取的值及这一层读取输出。值本身已经由输入经过 W_V 投影，不一定与原始字符向量同方向；下一章的输出投影又会再次线性变换。因此不能把“某输入位置权重较大”直接解释成它对最终预测具有相同比例的因果贡献。要检验影响，可以固定其他输入后做明确扰动，并追踪所有后续分支；注意力权重可作为中间观察，却不能跳过完整计算链。

三、从 AV 开始，一步一步反向

用 $dV$、$dA$ 等简写表示 $J$ 对相应矩阵的梯度，并不表示时间微分。先看 $O_{i\ell}=\sum_jA_{ij}V_{j\ell}$。固定一个值坐标 $V_{j\ell}$，它会被所有查询 $i$ 读取，故 $dV_{j\ell}=\sum_iA_{ij}G_{i\ell}$。固定权重 $A_{ij}$，它作用于该值向量的所有维度，故 $dA_{ij}=\sum_\ell G_{i\ell}V_{j\ell}$。写成矩阵就是 $dV=A^{\mathsf T}G$，$dA=GV^{\mathsf T}$。每个求和对应计算图中所有经过该变量的路径。

接下来不能只对 softmax 的分子求导。对同一行省略查询索引，$A_j=e^{S_j}/Z$，$Z=\sum_ke^{S_k}$。固定 $S_k$，用商法则，分子项导数为 $\delta_{jk}e^{S_j}$，分母导数为 $e^{S_k}$。于是导数为 $[\delta_{jk}e^{S_j}Z-e^{S_j}e^{S_k}]/Z^2$，整理成 $A_j(\delta_{jk}-A_k)$。$\delta_{jk}$ 在两索引相同时为一，否则为零。这张所有输出对所有输入的导数表叫雅可比矩阵。

把数代入：第一行的加权梯度均值为 (3/4)×2+(1/4)×0=3/2，所以 dS 第一行是 [(3/4)(2−3/2),(1/4)(0−3/2)]=[3/8,−3/8]。第二行均值为 (1/4)×0+(3/4)×8=6，所以第二行为 [−3/2,3/2]。每行梯度和为零，符合整行分数同加常量不改权重的平移不变性。只保留 A⊙dA 会漏掉分母路径，行和通常不再为零。

再展开分数 $S_{ij}=\sum_aQ_{ia}K_{ja}/\sqrt d$。固定 $Q_{ia}$，对所有 $j$ 求和得到 $dQ_{ia}=\sum_jdS_{ij}K_{ja}/\sqrt d$；固定 $K_{ja}$，对所有查询 $i$ 求和得到 $dK_{ja}=\sum_idS_{ij}Q_{ia}/\sqrt d$。因此 $dQ=dS\,K/\sqrt d$，$dK=dS^{\mathsf T}Q/\sqrt d$。本例 $K=I$，故 $dQ=dS/\sqrt2$；$Q=\sqrt2\ln3\,I$，故 $dK=\ln3\,dS^{\mathsf T}$。数值程序将对全部坐标做差分，而不是只抽最容易的一项。

最后回到投影。根据第二章的线性链式关系，$dW_Q=X^{\mathsf T}dQ$，其余两条对应得到 $dW_K$ 与 $dW_V$。同一个 $X$ 分别产生三条分支，故 $dX=dQW_Q^{\mathsf T}+dKW_K^{\mathsf T}+dVW_V^{\mathsf T}$。它们必须相加，不能以最后计算的值支路覆盖前面两条。此时损失已通过注意力回到输入向量；只有以后把这些梯度与真实交叉熵、所有层和优化器接起来，才构成完整训练。

局部反向中把 Q、K、V 分别扰动，并不是假设它们在统计上相互独立。微积分允许先研究函数在三个自变量坐标上的变化，再通过链式法则接回共同输入。统计独立是缩放动机中的随机模型条件，坐标偏导则是计算图分解，两者不是同一个概念。若差分检查时本想只改 Q，却由 X 同时重算三项，再与局部 dQ 比较，就会比较两个不同函数，自然得不到一致结果。

算例二：相同键让查询失去区分能力

保留查询和值，只把两行键都改成 [1,1]。对任何固定查询，它与两行相同键的点积相等，所以 softmax 权重必为 [1/2,1/2]，两行输出均为 [1,2]，探针 J=1+2×2=5。查询怎样变化都无法区分这两个键，因此在这个点附近仅改变 Q，输出不变。这个退化例子计算的是匹配机制的限制，不能以“模型一定会自己学到不同键”跳过。

反向也能核对上述结论。dA 仍为 [[2,0],[0,8]]，均值分别为一和四，因此 dS 两行为 [1/2,−1/2]、[−2,2]。每行与两行相同 K 相乘都抵消，dQ 全零。但 dK 未必为零，因为分别扰动某一行键会打破相同条件。零查询梯度和非零键梯度并不矛盾：它们对应不同的独立扰动。把所有梯度都设零反而是错误的。

也可以在二位置特例里独立验证雅可比。设一行权重为 [a,1−a]，上游权重梯度为 [u,v]。第一个分数导数应为 a(1−a)(u−v)，第二项恰为它的相反数。两项权重的竞争来自共同分母：提高第一个分数不仅提高第一权重，也降低第二权重。当 u=v 时，总读取目标对两者的分配不敏感，两项导数全零。若只看 softmax 分子，会错误地产生非零更新。

四、独立有限差分与故障对照

中心差分在冻结其他输入时，分别将一个坐标加减 h，用损失差除以 2h。程序把 Q、K、V 都当作可独立扰动的实数矩阵，检验刚才推导的局部函数；再将三个投影接回共同的 X，另外检查 dX，检测遗漏分支。采用双精度浮点和两个不同步长，误差门槛事先固定。差分仅提供本次数值证据，解析推导仍承担一般公式的理由。

import math
import copy

def transpose(A): return [list(row) for row in zip(*A)]
def mm(A,B):
    return [[sum(a*b for a,b in zip(row,col)) for col in zip(*B)] for row in A]
def scaled(A,c): return [[c*x for x in row] for row in A]
def softmax(z):
    m=max(z); e=[math.exp(v-m) for v in z]
    return [v/sum(e) for v in e]
def forward(Q,K,V):
    S=scaled(mm(Q,transpose(K)),1/math.sqrt(len(Q[0])))
    A=[softmax(row) for row in S]
    return mm(A,V),A
Q_test=[[1.,2.],[0.,1.]]; K_test=[[1.,0.],[2.,1.]]
assert mm(Q_test,transpose(K_test))==[[1.,4.],[0.,1.]]
assert mm(Q_test,K_test)!=mm(Q_test,transpose(K_test))
G=[[1.,0.],[0.,2.]]
def objective(Q,K,V):
    O,_=forward(Q,K,V)
    return sum(x*g for row,gr in zip(O,G) for x,g in zip(row,gr))
def backward(Q,K,V):
    O,A=forward(Q,K,V)
    dV=mm(transpose(A),G); dA=mm(G,transpose(V)); dS=[]
    for i,row in enumerate(A):
        mean=sum(a*g for a,g in zip(row,dA[i]))
        dS.append([A[i][j]*(dA[i][j]-mean) for j in range(len(row))])
    assert all(abs(sum(row))<1e-12 for row in dS)
    dQ=scaled(mm(dS,K),1/math.sqrt(len(Q[0])))
    dK=scaled(mm(transpose(dS),Q),1/math.sqrt(len(Q[0])))
    return dQ,dK,dV

a=math.sqrt(2)*math.log(3)
Q=[[a,0.],[0.,a]]; K=[[1.,0.],[0.,1.]]; V=[[2.,0.],[0.,4.]]
O,A=forward(Q,K,V)
assert max(abs(A[i][j]-[[.75,.25],[.25,.75]][i][j]) for i in range(2) for j in range(2))<1e-12
assert abs(objective(Q,K,V)-7.5)<1e-12
analytic=backward(Q,K,V); fd_errors=[]
for h in (1e-5,1e-6):
    for field in range(3):
        error=0.
        for i in range(2):
            for j in range(2):
                plus=copy.deepcopy([Q,K,V]); minus=copy.deepcopy([Q,K,V])
                plus[field][i][j]+=h; minus[field][i][j]-=h
                numeric=(objective(*plus)-objective(*minus))/(2*h)
                error=max(error,abs(numeric-analytic[field][i][j]))
        fd_errors.append(error)
assert max(fd_errors)<1e-7

X=[[1.,0.],[0.,1.]]; Wq=copy.deepcopy(Q); Wk=copy.deepcopy(K); Wv=copy.deepcopy(V)
def projected_loss(x): return objective(mm(x,Wq),mm(x,Wk),mm(x,Wv))
dQ,dK,dV=analytic
branches=[mm(dQ,transpose(Wq)),mm(dK,transpose(Wk)),mm(dV,transpose(Wv))]
dX=[[sum(branch[i][j] for branch in branches) for j in range(2)] for i in range(2)]
x_error=0.; h=1e-6
for i in range(2):
    for j in range(2):
        plus=copy.deepcopy(X); minus=copy.deepcopy(X)
        plus[i][j]+=h; minus[i][j]-=h
        numeric=(projected_loss(plus)-projected_loss(minus))/(2*h)
        x_error=max(x_error,abs(numeric-dX[i][j]))
assert x_error<1e-7
same_keys=[[1.,1.],[1.,1.]]
O_same,A_same=forward(Q,same_keys,V)
dQ_same,_,_=backward(Q,same_keys,V)
assert all(abs(x)<1e-12 for row in dQ_same for x in row)
assert O_same==[[1.,2.],[1.,2.]]
print("attention_weights",A)
print("output",O,"probe_loss",objective(Q,K,V))
print("dQ",dQ,"dK",dK,"dV",dV)
print("QKV_finite_difference_max_error",max(fd_errors))
print("shared_X_finite_difference_max_error",x_error)
print("same_key_output",O_same,"same_key_dQ",dQ_same)
print("all_attention_checks_passed")

在以标量乘加计费的普通稠密实现中，QK转置约需 T²d 量级运算，AV 约需 T²r，存储完整权重需 T² 个数。这是明确计算模型下的量级，不是某台设备的耗时预测；投影本身还要另外计费。小参考函数刻意保存中间矩阵以便检查，日后换加速实现时，可复用同一数值、梯度和因果测试，而不能因速度更快就省略等价验证。

探针给出的上游梯度在本例中非常稀疏，是为了让每步乘法可复算。真实语言模型里，词表交叉熵经过输出矩阵、归一化和其他层后会形成另一张上游梯度，注意力的局部反向公式不因此变化。必须从真实损失计算这张梯度，不能把本章固定 G 塞进模型训练。局部算子的可复用性来自输入输出与导数定义一致，不是来自人为探针恰好能让目标下降。

差分输入还应避开数值上几乎完全饱和的权重。如果某个指数已经下溢，微小扰动后的输出可能在机器上完全不变，解析与数值都接近零，看似吻合却很难检出错误。先用本章有明显混合权重的例子验证，再另设极端分数检查稳定性。一个检查针对导数链，一个针对有限精度边界，不能因为前者通过就省略后者。

五、递进练习与解答

# causal_attention
def causal_attention(query, key, value, query_offset=0, dropout=None):
    raise NotImplementedError('attention')

# 第3章仅实现无缓存注意力；第10章在同一函数增加offset。
def causal_attention(query,key,value,query_offset=0,dropout=None):
    if query_offset != 0:
        raise NotImplementedError('第10章再处理历史位置')
    scores = query @ key.transpose(-2,-1) / math.sqrt(query.shape[-1])
    future = torch.arange(key.shape[-2],device=key.device)[None,:] > torch.arange(query.shape[-2],device=query.device)[:,None]
    probabilities = torch.softmax(scores.masked_fill(future,float('-inf')),dim=-1)
    weights = probabilities if dropout is None else dropout(probabilities)
    return weights @ value, probabilities

本章项目增量与下一章依赖

现在具备单头注意力的完整前向、解析反向和独立差分参考，包含三个投影共享输入时的梯度累加。本章所有位置互相可见，所以不能用于宣称下一 token 模型已经因果。第四章将在相同计算链中引入因果与填充许可、多头拆合和输出投影，以及主项目采用的可学习绝对位置，检查整个前缀接口。后面的全参数训练会使用这些实际投影，不能把它们冻结后仍称完整训练。