Nvidia Decaying Attention：距离偏置正负号与稳定 Softmax

题意与假设

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

• 训练版按单头自注意力理解：Q/K 各为 N×D，V 为 N×Dᵥ，q@k 解释为 QKᵀ；行方向对 key 做 softmax。题目矩阵形状未公开。
• 代码默认严格执行所写的 +|i−j|；distance_sign=-1 是用于比较的衰减训练变体，不能声称原题规定了负号。
• 原公式未写 1/√D、causal mask、padding mask 或可学习系数；代码不额外加入。所有输入为有限数且维度一致。

解题路线

1. 先在面试中指出符号歧义：若 bᵢⱼ=+|i−j|，相同 QK 相似度时远处 key 的权重更大；若要距离衰减，通常应为 −λ|i−j| 且 λ>0。请面试官确认最终公式。
2. 对每个 query 位置 i，计算全部 key 位置 j 的 dot(Qᵢ,Kⱼ)+sign×|i−j|。沿 j 取最大 logit 后相减，再 exp 并归一化，避免直接指数溢出。
3. 用归一化权重对 V 的每个通道加权求和。输出形状为 N×Dᵥ；将正负号参数保留在参考实现中，让两种解释有可复现的结果。

手动走一遍

若 Q、K 都为零向量，V=[0,1] 且 N=2，对第 0 个位置，正号的 logits 为 [0,1]，输出 e/(1+e)≈0.731；负号为 [0,−1]，输出 1/(1+e)≈0.269。这一最小例子直接暴露“decaying”与题目正号的冲突。

Python 参考实现

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

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

"""Single-head attention with an explicit distance-bias sign convention."""

from math import exp, isfinite


def distance_attention(q, k, v, *, distance_sign=1.0):
    """Return softmax(Q K^T + sign*|i-j|) V.

    sign=+1 follows the published expression literally; sign=-1 is a decay
    training variant. The published expression has no scaling factor.
    """
    n = len(q)
    if not n or len(k) != n or len(v) != n:
        raise ValueError("Q, K, and V must have the same nonzero sequence length")
    d, value_dim = len(q[0]), len(v[0])
    if not d or not value_dim or not isfinite(distance_sign):
        raise ValueError("dimensions and distance_sign must be valid")
    for row in q:
        if len(row) != d or not all(isfinite(x) for x in row):
            raise ValueError("invalid Q")
    for row in k:
        if len(row) != d or not all(isfinite(x) for x in row):
            raise ValueError("invalid K")
    for row in v:
        if len(row) != value_dim or not all(isfinite(x) for x in row):
            raise ValueError("invalid V")

    output = []
    for i in range(n):
        scores = [
            sum(q[i][t] * k[j][t] for t in range(d))
            + distance_sign * abs(i - j)
            for j in range(n)
        ]
        if not all(isfinite(score) for score in scores):
            raise ValueError("attention score overflow")
        peak = max(scores)
        weights = [exp(score - peak) for score in scores]
        total = sum(weights)
        output.append([
            sum(weights[j] * v[j][c] for j in range(n)) / total
            for c in range(value_dim)
        ])
    return output

复杂度

N 个位置、键维 D、值维 Dᵥ：逐行参考实现时间 O(N²(D+Dᵥ))；除输出外，每行临时 O(N)，总额外空间 O(N+NDᵥ)。若保存完整 attention 矩阵则需 O(N²) 空间；生产实现还要考虑矩阵乘法与 IO 优化。

易错点与追问

容易失分的边界

• 把 b=|i−j| 直接解释成距离衰减；正号实际提高远距离 logit。
• 把 softmax 作用在整个矩阵或 value 维度，而不是每个 query 的 key 维度。
• 按题目之外的习惯自动加缩放、mask 或负号，却未向面试官声明公式已改变。
• 直接 exp(raw logit)，大 dot product 时会溢出；减去行最大值后再归一化。

继续深挖

• 若确认要衰减，请明确 λ、是否学习、是否各 head 独立，以及绝对距离偏置是否截断。
• 若引入 causal 或 padding mask，应在 softmax 前把不允许的位置排除；全被 mask 的行要定义行为。
• 批量、多头与 GPU 实现需要处理张量广播形状、数值精度和大序列的显存/带宽成本。