Transformer语言理解任务中scaled_dot_product_attention函数的掩码工作原理及数学逻辑问询

嗨，我来帮你把这个掩码的逻辑讲透，其实核心就是利用softmax函数的特性来实现忽略填充值的目的，咱们一步步拆解你的两个疑问：

一、为什么要给填充位置加上-1e9？ #

首先得回忆Transformer注意力机制的关键步骤：计算完scaled_attention_logits后，会通过tf.nn.softmax把这些对数概率转换成注意力权重——权重总和为1，代表模型把多少注意力分配到对应位置的value上。

填充元素（你例子里掩码矩阵值为1的位置）是无意义的，我们绝对不希望模型把注意力浪费在这些位置上。那怎么让softmax自动忽略它们呢？

先看softmax的公式：
$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}}$$

你会发现，当$x_i$是一个极小的负数时，$e^{{x_i}$会趋近于0。这里选-1e9是因为这个数值足够小，小到$e}{-1e9}$在数值计算里几乎等于0，但又不会触发严重的数值下溢问题（如果用-1e18这类极端值，可能会出现计算异常）。

所以给填充位置的logits加上-1e9，本质就是把这些位置的logits拉到极低，让它们在softmax计算中几乎没有权重占比。

二、数学层面如何实现忽略填充值？ #

咱们结合你给的示例掩码和代码来推导：

假设你的掩码矩阵是：

[[0, 0, 1],
 [1, 0, 0],
 [0, 1, 0]]

值为1的位置是填充元素，执行scaled_attention_logits += (mask * -1e9)后，填充位置的logits会变成原logit值 + (-1e9)，非填充位置保持不变。

举个具体例子：假设某个填充位置的原logit是3，计算后变成3 - 1e9，这个值几乎等同于-1e9。代入softmax公式：

• 分子：$e^{3 - 1e9} = e^3 * e^{-1e9} ≈ e^3 * 0 = 0$
• 分母：所有位置的$e^{{x_j}$之和，其中填充位置的项几乎为0，相当于只计算非填充位置的$e}{x_j}$之和

这样一来，填充位置对应的注意力权重就会趋近于0。接下来执行output = tf.matmul(attention_weights, v)时，权重为0的位置会把对应v的值完全“抵消”，相当于模型直接忽略了这些填充元素的信息。

结合代码看完整流程 #

你提供的scaled_dot_product_attention函数里，掩码操作的位置非常关键——必须在softmax之前执行，才能在权重计算阶段就把填充位置的权重“清零”：

def scaled_dot_product_attention(q, k, v, mask):
    """计算注意力权重。
    q、k、v必须具有匹配的前置维度。
    k、v必须具有匹配的倒数第二个维度，即：seq_len_k = seq_len_v。
    掩码的形状取决于其类型（填充掩码或前瞻掩码），但必须可广播以进行加法操作。
    参数：
        q: query的形状 == (..., seq_len_q, depth)
        k: key的形状 == (..., seq_len_k, depth)
        v: value的形状 == (..., seq_len_v, depth_v)
        mask: 浮点张量，形状可广播至(..., seq_len_q, seq_len_k)，默认为None。
    返回：
        output, attention_weights
    """
    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
    # 缩放matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
    # 将掩码添加到缩放后的张量上
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)
    # softmax在最后一个轴（seq_len_k）上归一化，使得分数之和为1
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)
    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)
    return output, attention_weights

内容的提问来源于stack exchange，提问作者mousa alsulaimi