Keras/TensorFlow中MultiHeadAttention层attention_mask使用示例求助

问题解决方案 #

1. 掩码形状说明 #

自注意力场景下query和key来自同一个输入序列，因此T=S=sequence_length，最终掩码形状为(batch_size, sequence_length, sequence_length)，含义是：对于单个样本的第i个query位置，掩码第i行的第j位为0时，该query在计算注意力时不会关注第j个key位置。

2. 动态掩码实现方案 #

不需要提前用numpy生成全局掩码，直接在层内部用TF算子动态生成，可适配任意batch size，不会占用额外内存。修改后的完整TransformerBlock代码如下：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = keras.Sequential(
            [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs, training):
        # 动态生成padding掩码：取每个时间步的第0个特征判断是否为padding位
        # 得到形状 (batch_size, sequence_length) 的padding掩码，True为有效位置，False为padding
        padding_mask = tf.not_equal(inputs[..., 0], 0)
        # 扩展为自注意力需要的 (batch_size, sequence_length, sequence_length) 形状
        # 相当于每个query位置都使用相同的padding掩码过滤key位置
        attention_mask = padding_mask[:, tf.newaxis, :] & padding_mask[:, :, tf.newaxis]
        
        attn_output = self.att(inputs, inputs, attention_mask=attention_mask)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

如果输入经过嵌入层、CNN处理后padding位置的特征不再为0，可修改call方法支持接收前置层传递的掩码，Keras中嵌入层设置mask_zero=True后会自动传递掩码到后续层：

def call(self, inputs, training, mask=None):
    if mask is not None:
        attention_mask = mask[:, tf.newaxis, :] & mask[:, :, tf.newaxis]
        attn_output = self.att(inputs, inputs, attention_mask=attention_mask)
    else:
        attn_output = self.att(inputs, inputs)
    # 后续逻辑和原有代码一致

3. 掩码的必要性说明 #

• 精度层面：如果当前是序列级分类任务，padding位的信息不会对最终全局池化/平均结果产生明显影响，所以未加掩码时模型表现尚可；如果是token级任务（比如序列标注、翻译），不加掩码会直接导致结果错误。
• 效率层面：加掩码后MultiHeadAttention会跳过掩码为0的位置的注意力计算，确实可以减少计算量，序列padding比例越高加速效果越明显。

内容的提问来源于stack exchange，提问作者R. Giskard