TikTok Self Attention：Q/K/V 形状、Mask 与多头拼接

题意与假设

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

• 训练单头核心Q为n×d、K为m×d、V为m×dv，True mask表示可访问；每query至少一个可用key。
• 多头由独立投影与输出投影组成；推荐cross-attention的具体语义依据业务另定。

解题路线

1. 输入投影为Q=XWq、K=XWk、V=XWv，scores=QKᵀ/√d，行softmax后A乘V。Q/K内维必须一致，K/V条目数同m；softmax沿key轴，不能沿query轴。
2. self时Q/K/V来自同一序列不同投影，cross时query与被检索key/value来自不同来源。推荐场景可用用户/请求表示作Q、候选历史或上下文作K/V，但不是唯一安排，需讲预测目标。
3. mask在softmax前排除非法key，因果遮未来，padding遮无效位置；相加mask和布尔mask真值含义依接口明确。只在输出后乘零不会阻止非法key占分母。
4. 整行无可用key时softmax无合法归一化，本训练拒绝；生产可规定零输出等特殊分支。不要把全−∞行减max后NaN当正常结果。
5. 多头将投影分成h份独立算注意力，再concat并乘Wo，head维的缩放不能错用总model维。residual、norm、FFN属于完整block，不能遗漏后称已实现整个Transformer。
6. attention权重只做token混合，位置/顺序信息通常另加入，纯无位置信息机制对排列具有对称性。权重也不必然是可靠因果解释，不能只展示热力图就宣称解释模型。
7. 示例使用稳定行max并只遍历可用key，测试query/key长度不同、mask、均匀权重与单可用key。完整训练反向还需投影/softmax梯度，本代码明确只前向。

手动走一遍

Q全零时score全零，两个允许key输出V平均；V=[2,6]输出4。只允许第一个key则输出2，与被遮key的score再大也无关。cross的n与m可不同，不该强制方阵。

Python 参考实现

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

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

from math import exp, isfinite, sqrt


def scaled_dot_attention(q, k, v, mask=None):
    """Single-head Q[n,d], K[m,d], V[m,dv]; mask[n,m] True=allowed."""
    if not q or not k or not v or len(k) != len(v):
        raise ValueError('nonempty Q/K/V and same key/value length required')
    d, dv = len(q[0]), len(v[0])
    if d == 0 or dv == 0 or any(len(row) != d for row in q + k) or any(len(row) != dv for row in v):
        raise ValueError('inconsistent dimensions')
    if mask is not None and (len(mask) != len(q) or any(len(row) != len(k) for row in mask)):
        raise ValueError('invalid mask shape')
    if any(not isfinite(x) for row in q + k + v for x in row):
        raise ValueError('nonfinite input')
    result = []
    scale = sqrt(d)
    for i, query in enumerate(q):
        allowed = [j for j in range(len(k)) if mask is None or mask[i][j]]
        if not allowed:
            raise ValueError('each query needs an allowed key')
        scores = [sum(a * b for a, b in zip(query, k[j])) / scale for j in allowed]
        offset = max(scores)
        weights = [exp(score - offset) for score in scores]
        denominator = sum(weights)
        result.append([sum(weights[t] * v[j][axis] for t, j in enumerate(allowed)) / denominator
                       for axis in range(dv)])
    return result

复杂度

核心前向O(nmd+nmdv)，显式scores/weights一般O(nm)，本参考逐query算权重工作O(m)加输出O(ndv)，输入和投影另计。多头h合并各头成本。

易错点与追问

容易失分的边界

• softmax轴错。
• mask输出后才应用。
• 缩放用错误维度。
• 只有attention就说完整Transformer。

继续深挖

• causal与padding如何组合？
• 不同head维如何计复杂度？
• 全遮行怎么定义？