attention_gradients.py

解析反向传播。对应第31章。依赖：NumPy。

原始Python文件 · 65行

"""单头注意力的完整解析反向；NumPy float64；不是GPU算子。"""from __future__ import annotationsimport numpy as np  def forward(q, k, v, allowed=None):    q, k, v = [np.asarray(x, dtype=np.float64) for x in (q, k, v)]    if any(x.ndim != 2 for x in (q,k,v)) or q.shape[1] != k.shape[1] or k.shape[0] != v.shape[0]:        raise ValueError('Q/K/V形状不兼容')    if min(q.shape + k.shape + v.shape) < 1 or not all(np.isfinite(x).all() for x in (q,k,v)):        raise ValueError('只支持有限非空张量')    scores = q @ k.T / np.sqrt(q.shape[1])    if allowed is None:        allowed = np.ones_like(scores, dtype=bool)    allowed = np.asarray(allowed, dtype=bool)    if allowed.shape != scores.shape or not allowed.any(axis=1).all():        raise ValueError('掩码形状错误或存在全掩码行')    scores = np.where(allowed, scores, -np.inf)    exp = np.exp(scores - scores.max(axis=1, keepdims=True))    p = exp / exp.sum(axis=1, keepdims=True)    return p @ v, (q, k, v, p, allowed)  def backward(upstream, cache):    q, k, v, p, allowed = cache    upstream = np.asarray(upstream, dtype=np.float64)    if upstream.shape != (q.shape[0], v.shape[1]) or not np.isfinite(upstream).all():        raise ValueError('上游梯度形状或数值错误')    dp = upstream @ v.T    ds = p * (dp - (dp * p).sum(axis=1, keepdims=True))    ds = np.where(allowed, ds, 0.0)    scale = np.sqrt(q.shape[1])    return ds @ k / scale, ds.T @ q / scale, p.T @ upstream  def finite_difference(q, k, v, upstream, allowed=None, eps=1e-6):    originals = [np.array(x, dtype=np.float64, copy=True) for x in (q,k,v)]    grads=[]    for variable in range(3):        grad = np.zeros_like(originals[variable])        for index in np.ndindex(grad.shape):            plus = [x.copy() for x in originals]            minus = [x.copy() for x in originals]            plus[variable][index] += eps            minus[variable][index] -= eps            fp = float(np.sum(forward(*plus, allowed)[0] * upstream))            fm = float(np.sum(forward(*minus, allowed)[0] * upstream))            grad[index] = (fp - fm) / (2*eps)        grads.append(grad)    return tuple(grads)  def run():    rng=np.random.default_rng(42)    q,k,v,g=rng.normal(size=(3,3)),rng.normal(size=(4,3)),rng.normal(size=(4,2)),rng.normal(size=(3,2))    allowed=np.arange(4)[None,:] <= np.arange(3)[:,None]+1    _, cache=forward(q,k,v,allowed)    analytic=backward(g,cache)    numerical=finite_difference(q,k,v,g,allowed)    return {'kind':'解析反向与有限差分；非训练能力/非GPU性能评测',            'max_absolute_errors':{name:float(np.max(np.abs(a-n))) for name,a,n in zip(('dQ','dK','dV'),analytic,numerical)}} if __name__=='__main__':    import json    print(json.dumps(run(),ensure_ascii=False,indent=2))