Meta 手写Self-Attention：形状、稳定Softmax与因果遮罩

题目重点：从输入投影到QK缩放、遮罩、逐行归一化和加权V，给出可运行实现及对PyTorch的独立核对。

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

训练是单样本单head、无dropout的scaled dot-product attention；Q[n,d]、K[m,d]、V[m,dv]，输出[n,dv]。self-attention由同X分别乘Wq/Wk/Wv产生。
allowed布尔mask的True表示可看；causal训练要求位置对齐且n=m。padding或cross-attention可传mask，但不暗示某框架mask的True语义必相同。
每query至少一个可看key，全遮罩明确抛错。所有输入有限，使用float64；不支持batched/multihead和KV cache，扩展需明确轴。

解题路线

先标形状：QKᵀ得到[n,m]，除sqrt(d)抑制点积随key维度增加的方差；softmax沿key轴m而非query轴n，最后weights[n,m]乘V[m,dv]。
遮罩在softmax前将不可见logit置−inf，因果矩阵允许keyPosition<=queryPosition。若先softmax再把权重清零却不重新归一，输出权重和不等于1。
稳定softmax逐行减最大logit再exp，整体加常数不改概率；全遮罩行会产生−inf−−inf的NaN，因此训练接口显式拒绝，而不是默默返回错误值。
Self-Attention投影检查X feature维与每W第一维一致，Q/K输出dim需相等，V的dim可不同。position encoding、residual、LayerNorm、FFN是Transformer block额外部分，不混入本函数。
独立验证用实际torch SDPA对照随机mask/形状、概率行和、不可见权重0与future扰动；无需与自身另写一份相同NumPy公式互证。

手动走一遍

Q/K全0时所有合法logit相同，某query可见两个key则各权重1/2，输出是对应V平均。因果第一位置只能看第一个key，输出等于V[0]；改最后一个token不应改变之前位置的causal输出。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

import numpy as np


def attention(query, key, value, allowed=None, causal=False):
    """Unbatched Q[n,d], K[m,d], V[m,dv]. Boolean allowed[n,m]; no dropout.

    causal training mode requires n==m and position-aligned self-attention.
    Fully masked query rows are rejected; exact behavior is part of the contract.
    """
    q,k,v=(np.asarray(a,dtype=np.float64) for a in (query,key,value))
    if any(a.ndim!=2 for a in (q,k,v)) or q.shape[1]==0 or q.shape[0]==0 or k.shape[0]==0:
        raise ValueError('nonempty two-dimensional Q/K and positive key dimension required')
    n,d=q.shape;m=k.shape[0]
    if k.shape[1]!=d or v.shape[0]!=m or v.shape[1]==0:
        raise ValueError('incompatible dimensions')
    if any(not np.isfinite(a).all() for a in (q,k,v)):
        raise ValueError('finite inputs required')
    mask=np.ones((n,m),dtype=bool)
    if allowed is not None:
        supplied=np.asarray(allowed)
        if supplied.shape!=(n,m) or supplied.dtype!=np.bool_:
            raise ValueError('boolean n by m mask required')
        mask &= supplied
    if causal:
        if n!=m:raise ValueError('causal self-attention requires aligned positions')
        mask &= np.tril(np.ones((n,n),dtype=bool))
    if not mask.any(axis=1).all():
        raise ValueError('every query must have an allowed key')
    scores=(q @ k.T)/np.sqrt(d)
    if not np.isfinite(scores).all():
        raise ValueError('dot-product overflow; rescale inputs')
    scores=np.where(mask,scores,-np.inf)
    shifted=scores-scores.max(axis=1,keepdims=True)
    weights=np.exp(shifted)
    weights/=weights.sum(axis=1,keepdims=True)
    return weights @ v,weights


def self_attention(inputs,wq,wk,wv,allowed=None,causal=False):
    x=np.asarray(inputs,dtype=np.float64)
    if x.ndim!=2:raise ValueError('two-dimensional input required')
    matrices=[np.asarray(w,dtype=np.float64) for w in (wq,wk,wv)]
    if any(w.ndim!=2 or w.shape[0]!=x.shape[1] for w in matrices):
        raise ValueError('projection dimension mismatch')
    return attention(x@matrices[0],x@matrices[1],x@matrices[2],allowed,causal)

复杂度

QKᵀ O(nmd)，乘V O(nmdv)，attention matrix O(nm)空间，self-attention n=m时平方；投影额外O(n·dModel·projectionDims)。稳定实现仍会存完整weights，不能宣称线性内存或FlashAttention效率。

易错点与追问

容易失分的边界

沿query轴softmax会把不同query互相归一，不是每query对keys分配权重。
解码cache里query/key位置不对齐时简单tril(n,m)可能错，需真实positionOffset；本实现拒绝这类causal用法。
浮点16数值稳定与kernel性能不是float64基线测试已证明的结论。

继续深挖

Multihead先投影再reshape为[heads,n,dHead]分别注意力，concat后Wo；不要把heads当key序列维。
加入padding mask与causal mask用逻辑AND，保证至少有效key；训练loss也须排除padding输出。
若只问attention为何缩放，独立同分布分量点积方差约d，sqrt(d)令尺度稳定，非随意除序列长度。